ICML2020AlgoAdaRound: Up or Down? Adaptive Rounding for Post-Training Quantization

从二阶视角推导权重的自适应舍入,取代四舍五入。

总结:本篇文章作者首先从数学角度证明了在模型量化过程中,直接将浮点数进行四舍五入round到最近定点数的方法并不是精度最优的。并且通过了一个简单的实验验证了猜想,随后基于此作者进行一系列的数学推导和数学近似推导除了最终的优化目标:最小化由于量化在预激活值中引入的均方误差,从而提出了自适应的Round方法:AdaRound.这种方法在进行量化时,自适应地决定将浮点值转到最近右定点还是左定点值。AdaRound可以在不需要QAT or finetune的情况下仅使用少量无标签的校准数据在精度上达到SOTA,甚至4bit量化也可以保留较好的精度。

个人评价:这篇文章的行文流惯,公式推导顶级,从完全理论的方式推导出了大部分量化论文中量化目标函数。

首先作者将量化过程定义为了一个对预训练模型权重w的微小扰动$\Delta w$.我们的目标为最小化这个扰动对损失函数$L(w)$造成的影响,即最小化$E[L(w+\Delta w)-L(w)]$,为了近似这个损失,采用了二阶泰勒展开有:

$$ L(w + \Delta w)\approx L(w) + \nabla L(w)^\top \Delta w + \frac{1}{2}\Delta w^{\top} H(w)\Delta w\\ \Delta L \approx \nabla L^{\top}\Delta w + \frac{1}{2}\Delta w^\top H\Delta w $$

由于模型经过预训练损失函数的梯度很小,可以忽略,而高阶项只要扰动$\Delta w$不是特别大,二阶近似往往就是准确的。对于4-bit或更高精度而言这个是成立的。

因此我们可以认为影响模型精度的主要是$\Delta w$以及损失函数的曲率$H(w)$相关。

令$\Delta w^\top = [\Delta w_1,\Delta w_2]$,$H^{(w)} = \begin{bmatrix}1 & 0.5\\0.5 &1\end{bmatrix} $,那么由此我们可以计算出,量化导致的损失为:

$$ \Delta w^\top H^{(w)}\Delta w = \Delta w_1^2 + \Delta w_2^2 + \Delta w_1\Delta w_2 $$

对于对角线项$\Delta w_1^2,\Delta w_2^2$而言四舍五入是最优的,最小化了误差,但是对于非对角线项$\Delta w_1\Delta w_2$采用四舍五入就不一定最优了。例如若二者符号取反乘积为负就可以抵消一部分损失的增量。

因此从理论上分析出了四舍五入方法的局限性。后续也从实验上进行了论证,作者采用四舍五入,全部向上,全部向下,随机舍入进行比较,发现在随机舍入中存在比四舍五入高出10%的取舍法,说明在取舍办法中,存在更优的方法。

这个取舍办法的选取可以通过如下问题描述。

假设每层权重量化,量化后的权重为$\hat w_i^{(l)}$

$$ \hat w_i^{(l)}\in \{w_i^{(l),floor},w_i^{(l),ceil}\} $$

$\Delta w_i^{(l)} = w^{(l)} - \hat w_i^{(l)}$,由此,最优的舍入过程可以描述为以下二元优化问题:

$$ \arg \min_{\Delta w} \mathbb{E}[L(x,y,w+\Delta w) - L(x,y,w)] $$

直接对这个式子进行优化并不现实,因为,每次调整$\Delta w$都需要进行一次前向传播,计算成本太高,我们采用前面理论分析时的泰勒展开近似。此外,忽略属于不同层之间权重的交互。优化目标近似为:

$$ \arg \min_{\Delta w^{(l)}} \mathbb{E}[\Delta w^{(l)}H^{(w^{(l)})}\Delta w^{(l)}] $$

但是这个优化过程受限于Hessian矩阵的计算困难以及问题本身是一个NP-Hard问题。因此无法将这个作为最终的优化目标。

我们从Hessian矩阵计算的复杂性来分析

$$ \frac{\partial^2 L}{\partial W^{(l)}_{i,j}\partial W^{(l)}_{m,o}} = \frac{\partial}{\partial W_{m,o}^{(l)}}[\frac{\partial L}{\partial z_i^{(l)}}\cdot x_j^{(l-1)}] = \frac{\partial^2 L}{\partial z^{(l)}_i\partial z_m^{(l)}}\cdot x_j^{(l-1)}x_i^{(l-1)} $$

写作矩阵的形式

$$ H(w^{(l)}) = \mathbb{E}[x^{(l-1)} x^{(l-1)\top}⊗ \nabla^2_{z^{(l)}}L] $$

其中⊗为Kronecker积。由此看出Hessian矩阵的复杂性主要来于二阶导的求取,它需要通过网路的后续层反向传播二阶导数(见对角近似)。

为了解决这个问题,我们采用Hessian矩阵的对角近似,即将其近似为对角矩阵,记作$diag(\Delta^2_{z^{(l)}}L)$。

$$ H(w^{(l)}) = \mathbb{E}[x^{(l-1)} x^{(l-1)\top}⊗ diag(\nabla^2_{z^{(l)}}L)] $$

将这个近似带入优化方程中有:

$$ \arg \min_{\Delta W_{k,:}^{(l)}} \mathbb{E}[\nabla^2_{z^{(l)}}L_{k,k}\cdot \Delta W_{k,:}^{(l)}x^{(l-1)}x^{(l-1)\top}\Delta W_{k,:}^{(l)\top}]\\ =\arg\min_{\Delta W_{k,:}^{(l)}} \Delta W_{k,:}^{(l)}\mathbb{E}[x^{(l-1)}x^{(l-1)\top}]\Delta W_{k,:}^{(l)\top}\\ =\arg \min_{\Delta W_{k,:}^{(l)}} \mathbb{E}[(\Delta W_{k,:}^{(l)}x^{(l-1)})^2] $$

这里是认为$\nabla^2_{z^{(l)}}L_{i,i}$是一个与输入样本数据无关的常量结果。

由此我们推导出,我们只要最小化由于量化而在激活函数$z^{(l)}$中引入的均方误差。这与大部分量化的论文中的结论一致(如AdaQuant)

想要通过直接求解上面的优化方程仍然是一件困难的事情,因为它是NP-Hard的,因此作者将优化目标放宽为如下形式

$$ \arg \min_{V}||Wx-\hat Wx||^2_{F}+\lambda f_{reg}(V) $$

其中$||\cdot||^2_F$为F范数,$\hat W$为优化的软量化权重

$$ \hat W = s\cdot clip([\frac{W}{s}] + h(V),n,p) $$

$h(V_{i,j})$可以是任何在0和1之间取值的可微函数,$f_{reg}(V)$是一个可微正则项,用于鼓励$h(V_{i,j})$收敛到0或1.

但是这个方法存在一个缺陷,无法避免量化误差的不断积累且没有考虑到激活函数,所以做了进一步优化

$$ \arg \min_{V}||f_a(Wx)-f_a(\hat W\hat x)||_F^2 + \lambda f_{reg}(V) $$

其中$fa(\cdot)$为激活函数$\hat x$为当前层的反量化输入,x为当前层的浮点输入