Up or Down? Adaptive Rounding for Post-Training Quantization个人理解

亮点

提出一个新的无需重训练的量化方法（post-training quantization，与之对应的就是量化感知训练）。利用矩阵傅里叶展开求量化导致的误差，证明同一层的权重误差是会互相影响的，因此量化最优解并非rounding-to-nearest quantization。

核心公式

推导

假设 $\Delta \omega$ 是量化前后的差值，那么，量化对$\mathcal{L}$（loss函数）的影响为：
$$\begin{array}{rl}& \mathbb{E}[\mathcal{L}(x,y,w+\Delta \omega )-\mathcal{L}(x,y,w)]\\ & \stackrel{a}{\approx }\mathbb{E}[\Delta \omega \cdot ({\nabla }_w\mathcal{L}(x,y,w){)}^T+\frac{1}{2}\Delta {\omega }^T\cdot {\nabla }_w^2\mathcal{L}(x,y,w)\cdot \Delta \omega ]\\ & =\Delta {\omega }^T\cdot g^{\left(w\right)}+\frac{1}{2}\Delta {\omega }^T\cdot H^{\left(w\right)}\cdot \Delta \omega \end{array}\text{\hspace{1em}(1)}$$
这是很简单的矩阵傅里叶2阶展开之后得到的化简式，不作详细证明。

易知$g^{\left(w\right)}$是权重的梯度，$H^{\left(w\right)}$是Hessian矩阵（向量的二阶微分公式）。当网络被训练好的时候，$g^{\left(w\right)}$趋向于0，所以可以不考虑。所以量化带来的主要误差是由$H^{\left(w\right)}$引入的。
文中做了一个简单假设：
$$\Delta {\omega }^T=[\Delta {\omega }_1\Delta {\omega }_2]\text{\hspace{1em}(2)}$$
$$H^{\left(w\right)}=\left[\begin{array}{cc}1& 0.5\\ 0.5& 1\end{array}\right]\text{\hspace{1em}(3)}$$
则
$$\Delta {\omega }^T\cdot H^{\left(w\right)}\cdot \Delta \omega =\Delta {\omega }_1^2+\Delta {\omega }_2^2+\Delta {\omega }_1\Delta {\omega }_2\text{\hspace{1em}(4)}$$
由式(4)可知，(1)式的最优化不仅与权重差值的绝对值有关，也与差值的符号有关。文中的描述为：

“为了最小化量化对任务损失的总体影响，我们需要在对角线项和非对角线项的贡献之间进行权衡。舍入到最接近的方法忽略了非对角线的贡献，这使得它常常不是最优的。”

然后作者尝试了100次的随机舍入量化，发现比起rounding-to-nearest，确实有更优的方法存在，并且精度的提升与$\frac{1}{2}\Delta {\omega }^T\cdot H^{\left(w\right)}\cdot \Delta \omega$的下降相关。

接着就提出了他们的量化方法：Adaround。假设量化后的权重为${\hat{w}}_i^{\left(l\right)}$，那么它的取值有两种：
$${\hat{w}}_i^{\left(l\right)}ϵ\left\{{\hat{w}}_i^{\left(l\right),floor},{\hat{w}}_i^{\left(l\right),ceil}\right\}\text{\hspace{1em}(5)}$$
其中
$${\hat{w}}_i^{\left(l\right),floor}=s^{\left(l\right)}\cdot clip\left(\lfloor \frac{w_i^{\left(l\right)}}{s^{\left(l\right)}}\rfloor ,n,p\right)\text{\hspace{1em}(6)}$$
$${\hat{w}}_i^{\left(l\right),ceil}=s^{\left(l\right)}\cdot clip\left(\lceil \frac{w_i^{\left(l\right)}}{s^{\left(l\right)}}\rceil ,n,p\right)\text{\hspace{1em}(7)}$$
clip函数就是让小于$n$的数均为$n$，大于$p$的数均为$p$，中间数为自身。则

$$\Delta w_i^{\left(l\right)}=w^{\left(l\right)}-{\hat{w}}_i^{\left(l\right)}\text{\hspace{1em}(8)}$$
假设$s^{\left(l\right)}$为定值来优化舍入方式，那么当优化$\Delta w_i^{\left(l\right)}$的cost function时，${\hat{w}}_i^{\left(l\right)}$只为式(5)中的两个值。

此时优化问题转化为：
$$\underset{\Delta w}{\mathrm{argmin}}\mathbb{E}[\mathcal{L}(x,y,w+\Delta \omega )-\mathcal{L}(x,y,w)]\text{\hspace{1em}(9)}$$
经过化简并舍去梯度项，最终为：
$$\underset{\Delta w^{\left(l\right)}}{\mathrm{argmin}}\mathbb{E}[\Delta {\omega }^T\cdot H^{\left(w\right)}\cdot \Delta \omega ]\text{\hspace{1em}(10)}$$
上式中下标$i$代表着单个权重，没有下标则是当前层所有权重的优化，随着层大小的提升，式(10)的求解会越来越复杂，文中描述其为NP-hard问题。

接着就是一个有意思的化简了，

根据
$$z^{\left(l\right)}=W^{\left(l\right)}{x}^{\left(l-1\right)}\text{\hspace{1em}(11)}$$
其中，$W^{\left(l\right)}$是权重矩阵，$x^{\left(l-1\right)}$是输入向量，$z^{\left(l\right)}$是当前层预激活结果，即激活函数之前的输出。

上图中，$W_{i,j}^{\left(l\right)}$我的理解是当前层$l$，第$i$个卷积核的第$j$个通道，所以式(11)就变为了
$$z_i^{\left(l\right)}=W_{i,j}^{\left(l\right)}{x}_j^{\left(l-1\right)}\text{\hspace{1em}(12)}$$
也就是说，第$i$个卷积核的第$j$个通道的矩阵与输入图像的第$j$个通道的输入相乘就可得到第$i$个通道的输出。也就是说，$i,j$并不相关，想要得到所有输出，需要在第$i$个卷积核的情况下遍历所有通道。即，$\forall iϵ\left[1,I\right]$，遍历所有$j$, $jϵ[1,J]$。

所以$H^{\left(w\right)}$的矩阵形式为：
$$H^{\left(W^{\left(l\right)}\right)}=\mathbb{E}[{\nabla }_{W^{\left(l\right)}}^2\mathcal{L}]\text{\hspace{1em}(13)}$$
而
$${\nabla }_{W^{\left(l\right)}}^2\mathcal{L}=\left[\begin{array}{cccc}\frac{{\partial }^2\mathcal{L}}{\partial W_1^2}& \frac{{\partial }^2\mathcal{L}}{\partial W_1\partial W_2}& \cdots & \\ \frac{{\partial }^2\mathcal{L}}{\partial W_2\partial W_1}& \ddots & & \\ ⋮& & \ddots & \\ & & & \frac{{\partial }^2\mathcal{L}}{\partial W_n^2}\end{array}\right]\text{\hspace{1em}(14)}$$
这是Hessian矩阵的展开。
而
$$\frac{{\partial }^2\mathcal{L}}{\partial W_i\partial W_m}=\left[\begin{array}{cccc}\frac{{\partial }^2\mathcal{L}}{\partial W_{i,1}\partial W_{m,1}}& \frac{{\partial }^2\mathcal{L}}{\partial W_{i,1}\partial W_{m,2}}& \cdots & \\ \frac{{\partial }^2\mathcal{L}}{\partial W_{i,2}\partial W_{m,1}}& \ddots & & \\ ⋮& & \ddots & \\ & & & \frac{{\partial }^2\mathcal{L}}{\partial W_{i,n}\partial W_{m,n}}\end{array}\right]\text{\hspace{1em}(15)}$$
根据上面的推导，易知：
$$\frac{{\partial }^2\mathcal{L}}{\partial W_i\partial W_m}=\left[\begin{array}{cccc}{x}_1^2& x_1{x}_2& \cdots & \\ x_2{x}_1& \ddots & & \\ ⋮& & \ddots & \\ & & & x_n{x}_n\end{array}\right]\cdot \frac{{\partial }^2\mathcal{L}}{\partial z_i\partial z_m}\text{\hspace{1em}(16)}$$
由此，易知：
$$H^{\left(W^{\left(l\right)}\right)}=\mathbb{E}[{\nabla }_{W^{\left(l\right)}}^2\mathcal{L}]=\mathbb{E}[x^{\left(l-1\right)}{x^{\left(l-1\right)}}^T\otimes {\nabla }_{z^{\left(l\right)}}^2\mathcal{L}]\text{\hspace{1em}(17)}$$

这是第二个核心公式。从(17)中可以看出，${\nabla }_{z^{\left(l\right)}}^2\mathcal{L}$带来了主要的计算复杂度，所以作者做出了第二个假设：${\nabla }_{z^{\left(l\right)}}^2\mathcal{L}$是一个对角矩阵，即：

$$H^{\left(W^{\left(l\right)}\right)}=\mathbb{E}[x^{\left(l-1\right)}{x^{\left(l-1\right)}}^T\otimes diag\left({\nabla }_{z^{\left(l\right)}}^2{\mathcal{L}}_{i,i}\right)]\text{\hspace{1em}(18)}$$
将(18)代入(10)，可得：

这步推导是利用了对角矩阵的特点，(a)是第三个假设，${\nabla }_{z^{\left(l\right)}}^2{\mathcal{L}}_{k,k}=c_i$是一个输入数据的独立参数，不影响最优化问题的后续求解，所以在第二步舍去。

$\Delta W_{K,:}$代表矩阵的一行，此时可以发现，最优化问题已经和loss以及后续层没有关系了，只与当前层的权重以及输入有关。在：
$$\underset{\Delta W_{k,:}^{\left(l\right)}}{\mathrm{argmin}}\mathbb{E}\left[{\left(\Delta W_{K,:}{x}^{\left(l-1\right)}\right)}^2\right]\text{\hspace{1em}(19)}$$
中，做的是最小化Mean Squared Error(MSE)问题。作者表示虽然其他研究者之前也有这么优化的，但他是第一个用理论推导出来的，是在不了解正在优化的那一层网络其他部分的情况下所能做到的最好的结果。

式(19)中的问题可以通过预先计算$\mathbb{E}[x^{\left(l-1\right)}{x^{\left(l-1\right)}}^T]$，然后优化$\Delta W_{K,:}$，或者在优化过程中，对每个$\Delta W_{K,:}$进行单层前向传递。

核心推导

式(19)对于大规模优化任务还是太复杂了，所以作者用软量化变量将其继续relax。
$$\underset{V}{\mathrm{argmin}}\Vert Wx-{\tilde{W}x\Vert }_F^2+\lambda f_{reg}\left(V\right)\text{\hspace{1em}(20)}$$
这是弗罗宾尼斯范数（矩阵中元素的平方和），$\tilde{W}$ 是软量化权重：
$$\tilde{W}=s\cdot clip\left(\lfloor \frac{W}{s}\rfloor +h\left(V\right),n,p\right)\text{\hspace{1em}(21)}$$
矩阵乘法可以替换为卷积

此时$\Delta W_{K,:}$的优化问题就变为对$V_{i,j}$的优化问题。$h\left(V_{i,j}\right)$是一个连续变量并且$h\left(V_{i,j}\right)ϵ[0,1]$，附加项 $\lambda f_{reg}\left(V\right)$则是为了让$h\left(V_{i,j}\right)$收敛于0或1，即 h ( V i , j ) ϵ { 0 , 1 } h\left(V_{i,j}\right)\epsilon\left\{ 0,1 \right\} h(Vi,j​)ϵ{0,1}。

于是作者利用sigmoid构造了一个函数：

$\zeta ,\gamma$都是弹性参数，此时固定为1.1和-0.1，这个式子是从[1]里借过来的，不会梯度消失，利于优化到极值。

而
$$\lambda f_{reg}\left(V\right)=\sum _{i.j}1-|2h\left(V_{i,j}\right)-{1|}^{\beta }\text{\hspace{1em}(22)}$$
$\beta$是退火参数，当其值较大时，$h\left(V_{i,j}\right)$可以自由变化使得MSE最小，在之后的优化中，$\beta$变小以让$h\left(V_{i,j}\right)$收敛于0或1。

为了解决深度带来的误差（之前的优化都只是对当前层的优化），作者提出：
$$\underset{V}{\mathrm{argmin}}\Vert f_a\left(Wx\right)-{f_a\left(\tilde{W}\hat{x}\right)\Vert }_F^2+\lambda f_{reg}\left(V\right)\text{\hspace{1em}(23)}$$
$\hat{x}$是前面所有层都量化之后的得到的当前层的输入，$f_a$则是激活函数。(23)是最终的优化公式，即AdaRound。

实验结果

不做赘述，直接贴图。图中公式序号与本文序号不对应。

[1]: Louizos, C., Welling, M., and Kingma, D. P. Learning sparse neural networks through l0 regularization. International Conference on Learning Representations (ICLR), 2018.
[2]: Nagel M, Amjad R A, Van Baalen M, et al. Up or down? adaptive rounding for post-training quantization[C]//International Conference on Machine Learning. PMLR, 2020: 7197-7206.