Lin J, Tang J, Tang H, et al. Awq: Activation-aware weight quantization for llm compression and acceleration[j]. Machine Learning System. Best Paper. https://arxiv.org/abs/2306.00978
1. 背景和动机
- 直接在FP16精度上Round成INT3/INT4会造成极大的性能损失
- 基于activation distribution对重要的weight做精度保留则可以很大程度上提高模型性能。
- 但是混合存储FP16和INT3/4，在推理系统实现的时候过于复杂且对于硬件非常的不友好。
2. 算法
2.1 原理和假设
原有的Round方法(图a)：
$$ Q(\mathbf{w})=\Delta\cdot\mathrm{Round}(\frac{\mathbf{w}}\Delta),\quad\Delta=\frac{\max(|\mathbf{w}|)}{2^{N-1}} $$
其中$\mathbf{w}$表示一组参数，$Q(\mathbf{w})$表示量化函数，$N$表示量化位数。
改进后的量化方法(图c)：
$$ Q(w\cdot s)\cdot\frac xs=\Delta^{’}\cdot\mathrm{Round}(\frac{ws}{\Delta^{’}})\cdot x\cdot\frac1s $$
其中$w \in \mathbf{W}$。即先对特定的$w$做Scaling然后再Scaling回去。这样做的理由是，误差可以成倍的减小，如下面的公式和观察出来的现象：
$$ \begin{aligned}\operatorname{Err}(Q(w)x)&=\Delta\cdot\operatorname{RoundErr}(\frac w\Delta)\cdot x \newline \operatorname{Err}(Q(w\cdot s)(\frac xs))&=\Delta^{’}\cdot\operatorname{RoundErr}(\frac{ws}{\Delta^{’}})\cdot x\cdot\frac1s\end{aligned} $$
其中由于$\operatorname{Round}$函数是四舍五入，所以误差$\operatorname{RoundErr}\in [0,0.5]$且是一个均匀分布。平均在0.25。不管是否被缩放了，这个分布是不变的。
- 由于一组权重$\mathbf{w}$的最大值在缩放一个$w$后是基本不变的，所以我们可以认为$\Delta^{’} \approx \Delta$。
在此基础上，我们可以看出使用了Scaling以后得误差变小了，将上述提到的误差做个比值可以看出，$k=\frac{\Delta^{’}}{\Delta} \times \frac{1}{s}$。
2.2 优化：如何找到最优的Scaling值呢？
$$ \mathbf{s}^{*} = \arg \mathop{\min}_{s}\mathcal{L}(\mathbf{s}) $$
$$ \mathcal{L}(s) = \Vert Q(\mathbf{W} \cdot \text{diag}(s))(\text{diag}(s)^{-1}\cdot \mathbf{X}) - \mathbf{WX} \Vert $$
实际上就是用L2 Norm搜索出来一组最为合适的Scaling参数，注意，Scaling是对于每个权值都做的。 作者团队为了简化搜索的空间，对$\mathbf{s}$做了一定的约束，如下：
$$ \mathbf{s} = \mathbf{s_X}^{\alpha}，\alpha^{*}=\arg \mathop{\min}_{\alpha}\mathcal{L}(\mathbf{s_X}^{\alpha}) $$
其中，$\alpha \in [0, 1]$为缩放系数，$\mathbf{s_{X}}$是activation perchannel 的值。当$\alpha = 0$时，意味着不缩放；反之，意味着完全缩放。
3. 总结
非常有用，简单粗暴，Best Paper实至名归？是一个非常符合直觉的工程性质的工作，但是工程的实现复杂性比较高，在多数场景上没有被使用。 问题：由于一组权重$\mathbf{w}$的最大值在缩放一个$w$后是基本不变的，所以我们可以认为$\Delta^{’} \approx \Delta$。这个假设并不是很成立啊？