模型量化论文阅读#3----Up or Down? Adaptive Rounding for Post-Training Quantization

前言

在量化中，常用的量化函数一般的四舍五入函数，即在量化时以0.5为阈值，当超过0.5时，结果为1；不超过0.5时，结果为0.这就是所谓的四舍五入函数。很多时候，直接进行四舍五入量化，结果往往不是最优的。为什么不是最优的，可以通过下面一个简单的例子来进行说明。
对于一个训练好的网络来说，其权重为www,当进行量化时，可以把量化带来的权重的变化Δw\Delta wΔw视为一种扰动，那么由于扰动而带来的任务损失的变化表述为下式：

这里主要用到了二阶泰勒展开。因为网络权重已经训练收敛，一阶项一阶接近0，所以忽略一阶项，只保留二阶项。二阶项中的海森矩阵H(w)H^{(w)}H(w)表征了不同权重之间的相互作用。

以下面的一个简单的例子说明：

对于四舍五入函数来说，其忽略了H(w)H^{(w)}H(w)中的非对角项，即忽略了不同权重之间的相互影响，只考虑每个权重在量化前后的变化。很明显，如果不考虑对角项，只需要考虑Δw\Delta wΔw的幅值即可，因为其取平方之后，一定为正值。但是由于存在非对角项Δw1Δw2\Delta w_1 \Delta w_2Δw1Δw2 ,此时就要考虑Δw1\Delta w_1Δw1和Δw2\Delta w_2Δw2 的正负符号，如果它们的符号相反，会进一步的改善该损失。因此为了最小化该量化损失，就需要在对角项和非对角项之间做一个折衷权衡。四舍五入因为没有考虑非对角项，其解自然是次优的。

Method

上面介绍了四舍五入函数的缺陷，下面该论文提出一种自适应的舍入方法。

量化后的权重表示为：
w^i(ℓ)∈{wi(ℓ),floor ,wi(ℓ),ceil }\widehat{\mathbf{w}}_{i}^{(\ell)} \in\left\{\mathbf{w}_{i}^{(\ell), \text { floor }}, \mathbf{w}_{i}^{(\ell), \text { ceil }}\right\} wi(ℓ)∈{wi(ℓ), floor ,wi(ℓ), ceil }
其中:
wi(ℓ),floor =s(ℓ)⋅clip⁡(⌊wi(ℓ)s(ℓ)⌋,n,p),wi(ℓ),ceil=s(ℓ)⋅clip⁡(⌈wi(ℓ)s(ℓ)⌉,n,p)\mathbf{w}_{i}^{(\ell), \text { floor }}=\mathrm{s}^{(\ell)} \cdot \operatorname{clip}\left(\left\lfloor\frac{\mathbf{w}_{i}^{(\ell)}}{\mathrm{s}^{(\ell)}}\right\rfloor, \mathrm{n}, \mathrm{p}\right), \mathbf{w}_{i}^{(\ell), \text { ceil}}=\mathrm{s}^{(\ell)} \cdot \operatorname{clip}\left(\left\lceil\frac{\mathbf{w}_{i}^{(\ell)}}{\mathrm{s}^{(\ell)}}\right\rceil, \mathrm{n}, \mathrm{p}\right)wi(ℓ), floor =s(ℓ)⋅clip(⌊s(ℓ)wi(ℓ)⌋,n,p),wi(ℓ), ceil=s(ℓ)⋅clip(⌈s(ℓ)wi(ℓ)⌉,n,p),

那么，量化带来的扰动就表示为Δwi(ℓ)=wi(ℓ)−w^i(ℓ)\Delta \mathbf{w}_{i}^{(\ell)}=\mathbf{w}_{i}^{(\ell)}-\widehat{\mathbf{w}}_{i}^{(\ell)}Δwi(ℓ)=wi(ℓ)−wi(ℓ)。

需要注意，在优化舍入函数前，尺缩系数就已经确定下来了。而关于如何确定尺缩系数S，这篇论文后面也进行了相关研究。在此先不对s的确定做太多讨论，还是主要集中在舍入函数的优化。

根据上述讨论，如何获得最佳的舍入函数即可等价于下述的优化目标函数：
arg⁡min⁡ΔwE[L(x,y,w+Δw)−L(x,y,w)](12)\underset{\Delta \mathbf{w}}{\arg \min } \mathbb{E}[\mathcal{L}(\mathbf{x}, \mathbf{y}, \mathbf{w}+\Delta \mathbf{w})-\mathcal{L}(\mathbf{x}, \mathbf{y}, \mathbf{w})] (12) ΔwargminE[L(x,y,w+Δw)−L(x,y,w)](12)
为了降低计算负担，将上述函数利用泰勒展开为下式：
arg⁡min⁡ΔwE[Δw(ℓ)TH(w(ℓ))Δw(ℓ)](13)\underset{\Delta \mathbf{w}}{\arg \min } \mathbb{E}\left[\Delta \mathbf{w}^{(\ell)^{T}} \mathbf{H}^{\left(\mathbf{w}^{(\ell)}\right)} \Delta \mathbf{w}^{(\ell)}\right] (13) ΔwargminE[Δw(ℓ)TH(w(ℓ))Δw(ℓ)](13)
上述问题是一个二次无约束二值优化问题（QUBO），因为Δw\Delta wΔw是一个只能取值为{0,1}的二值变量。

至此，就将第一个针对任务损失函数的优化函数转换为第二个针对权重汉森矩阵的优化函数，这样就可以不用从头训练损失函数，从而减少了很多的计算。但是后面的目标函数是否可以很好的替代第一个目标函数呢？
这篇论文中基于resnet18，使用了100个随机的舍入函数对模型进行量化，然后看量化后的Δw(ℓ)TH(w(ℓ))Δw(ℓ)\Delta \mathbf{w}^{(\ell)^{T}} \mathbf{H}^{\left(\mathbf{w}^{(\ell)}\right)} \Delta \mathbf{w}^{(\ell)}Δw(ℓ)TH(w(ℓ))Δw(ℓ) 和量化后模型精度的关系，如下图，基本可以看出两点，第一，Δw(ℓ)TH(w(ℓ))Δw(ℓ)\Delta \mathbf{w}^{(\ell)^{T}} \mathbf{H}^{\left(\mathbf{w}^{(\ell)}\right)} \Delta \mathbf{w}^{(\ell)}Δw(ℓ)TH(w(ℓ))Δw(ℓ) 与任务损失有很强的相关性，可以用来代替任务损失作为量化的目标函数；第二，四舍五入量化不是一种最优的量化方法。

但是上述优化目标（13）还存在两个问题：

H(w(ℓ))\mathbf{H}^{\left(\mathbf{w}^{(\ell)}\right)}H(w(ℓ)) 在参数量较大时，很难计算，对目前的网络来说，去计算基本不可能；
（13）仍然是一个NP-hard问题，一旦Δw\Delta wΔw的维度稍微大一点，就很难解决。

下面针对上述两个问题，分别给出相应的解决方法。

从泰勒展开到局部损失

以全连接层为例，损失函数关于其权重的海森矩阵可以表示为：
∂2L∂Wi,j(ℓ)∂Wm,o(ℓ)=∂∂Wm,o(ℓ)[∂L∂zi(ℓ)⋅xj(ℓ−1)]=∂2L∂zi(ℓ)∂zm(ℓ)⋅xj(ℓ−1)xo(ℓ−1)（15）\begin{aligned} \frac{\partial^{2} \mathcal{L}}{\partial \mathbf{W}_{i, j}^{(\ell)} \partial \mathbf{W}_{m, o}^{(\ell)}} &=\frac{\partial}{\partial \mathbf{W}_{m, o}^{(\ell)}}\left[\frac{\partial \mathcal{L}}{\partial \mathbf{z}_{i}^{(\ell)}} \cdot \mathbf{x}_{j}^{(\ell-1)}\right] \\ &=\frac{\partial^{2} \mathcal{L}}{\partial \mathbf{z}_{i}^{(\ell)} \partial \mathbf{z}_{m}^{(\ell)}} \cdot \mathbf{x}_{j}^{(\ell-1)} \mathbf{x}_{o}^{(\ell-1)} \end{aligned} （15） ∂Wi,j(ℓ)∂Wm,o(ℓ)∂2L=∂Wm,o(ℓ)∂[∂zi(ℓ)∂L⋅xj(ℓ−1)]=∂zi(ℓ)∂zm(ℓ)∂2L⋅xj(ℓ−1)xo(ℓ−1)（15）
其中，z(ℓ)=W(ℓ)x(ℓ−1)\mathbf{z}^{(\ell)}=\mathbf{W}^{(\ell)} \mathbf{x}^{(\ell-1)}z(ℓ)=W(ℓ)x(ℓ−1)，是l层的输出（激活前）。

以矩阵形式表示海森矩阵：
H(w(ℓ))=E[x(ℓ−1)x(ℓ−1)T⊗∇z(ℓ)2L]（16）\mathbf{H}^{\left(\mathbf{w}^{(\ell)}\right)}=\mathbb{E}\left[\mathbf{x}^{(\ell-1)} \mathbf{x}^{(\ell-1)^{T}} \otimes \nabla_{\mathbf{z}^{(\ell)}}^{2} \mathcal{L}\right] （16） H(w(ℓ))=E[x(ℓ−1)x(ℓ−1)T⊗∇z(ℓ)2L]（16）

所以这里的复杂度主要来源于∇z(ℓ)2L\nabla_{\mathbf{z}^{(\ell)}}^{2} \mathcal{L}∇z(ℓ)2L，其需要对网络一系列的层的输出求二阶导数。
为了降低复杂度，论文提出一个假设：∇z(ℓ)2L\nabla_{\mathbf{z}^{(\ell)}}^{2} \mathcal{L}∇z(ℓ)2L是一个对角矩阵，即假设一层中不同的输出之间没有相互影响，该对角矩阵表示为diag⁡(∇z(ℓ)2Li,i)\operatorname{diag}\left(\nabla_{\mathbf{z}^{(\ell)}}^{2} \mathcal{L}_{i, i}\right)diag(∇z(ℓ)2Li,i)，那么，上式子简化为：
H(w(ℓ))=E[x(ℓ−1)x(ℓ−1)T⊗diag⁡(∇z(ℓ)2Li,i)]（17）\mathbf{H}^{\left(\mathbf{w}^{(\ell)}\right)}=\mathbb{E}\left[\mathbf{x}^{(\ell-1)} \mathbf{x}^{(\ell-1)^{T}} \otimes \operatorname{diag}\left(\nabla_{\mathbf{z}^{(\ell)}}^{2} \mathcal{L}_{i, i}\right)\right] （17） H(w(ℓ))=E[x(ℓ−1)x(ℓ−1)T⊗diag(∇z(ℓ)2Li,i)]（17）
将(17)带入到（13）得到转换后的目标函数为：
arg⁡min⁡ΔWk,:(ℓ)E[∇z(ℓ)2Lk,k⋅ΔWk,:(ℓ)x(ℓ−1)x(ℓ−1)TΔWk,:(ℓ)T]=(a)arg⁡min⁡ΔWk,:(ℓ)ΔWk,:(ℓ)E[x(ℓ−1)x(ℓ−1)T]ΔWk,:(ℓ)T=arg⁡min⁡ΔWk,:(ℓ)E[(ΔWk,:(ℓ)x(ℓ−1))2](20)\begin{array}{l} \underset{\Delta \mathbf{W}_{k,:}^{(\ell)}}{\arg \min } \mathbb{E}\left[\nabla_{\mathbf{z}^{(\ell)}}^{2} \mathcal{L}_{k, k} \cdot \Delta \mathbf{W}_{k,:}^{(\ell)} \mathbf{x}^{(\ell-1)} \mathbf{x}^{(\ell-1)^{T}} \Delta \mathbf{W}_{k,:}^{(\ell)^{T}}\right]\\ \stackrel{(a)}{=} \underset{\Delta \mathbf{W}_{k,:}^{(\ell)}}{\arg \min } \Delta \mathbf{W}_{k,:}^{(\ell)} \mathbb{E}\left[\mathbf{x}^{(\ell-1)} \mathbf{x}^{(\ell-1)^{T}}\right] \Delta \mathbf{W}_{k,:}^{(\ell)^{T}}\\ =\underset{\Delta \mathbf{W}_{k,:}^{(\ell)}}{\arg \min } \mathbb{E}\left[\left(\Delta \mathbf{W}_{k,:}^{(\ell)} \mathbf{x}^{(\ell-1)}\right)^{2}\right] \end{array} (20) ΔWk,:(ℓ)argminE[∇z(ℓ)2Lk,k⋅ΔWk,:(ℓ)x(ℓ−1)x(ℓ−1)TΔWk,:(ℓ)T]=(a)ΔWk,:(ℓ)argminΔWk,:(ℓ)E[x(ℓ−1)x(ℓ−1)T]ΔWk,:(ℓ)T=ΔWk,:(ℓ)argminE[(ΔWk,:(ℓ)x(ℓ−1))2](20)
其中，(a)代表进一步做了一个假设：∇z(ℓ)2Li,i=ci\nabla_{\mathbf{z}^{(\ell)}}^{2} \mathcal{L}_{i, i}=\mathrm{c}_{i}∇z(ℓ)2Li,i=ci ，ci\mathrm{c}_{i}ci表示常数。
论文做了很强的假设，但论文表示这些假设并没有消极影响。

Adaround

(20)求解的复杂度虽然降低，但是由于Δw\Delta wΔw是一个离散变量，求解起来仍然比较困难。因此，论文提出一种松弛化的方案，如下：
arg⁡min⁡V∥Wx−W~x∥F2+λfreg(V),（21）\underset{\mathbf{V}}{\arg \min }\|\mathbf{W} \mathbf{x}-\widetilde{\mathbf{W}} \mathbf{x}\|_{F}^{2}+\lambda f_{r e g}(\mathbf{V}), （21） Vargmin∥Wx−Wx∥F2+λfreg(V),（21）

其中，W~\widetilde{\mathbf{W}}W 是软量化的权重：
W~=s⋅clip⁡(⌊Ws⌋+h(V),n,p)（22）\widetilde{\mathbf{W}}=\mathrm{s} \cdot \operatorname{clip}\left(\left\lfloor\frac{\mathbf{W}}{\mathrm{s}}\right\rfloor+h(\mathbf{V}), \mathrm{n}, \mathrm{p}\right) （22） W=s⋅clip(⌊sW⌋+h(V),n,p)（22）

其中，V\mathbf{V}V是连续的可训练的变量，是我们要优化的对象。h(V)h(\mathbf{V})h(V)是一个可微分的函数，其输出值的范围是[0,1]。freg(V)f_{r e g}(\mathbf{V})freg(V)是一个正则化项，可以使得h(V)h(\mathbf{V})h(V)向两个端点0，1收敛。

很明显，在真正的量化中，h(V)h(\mathbf{V})h(V)只能取{0,1}中的一个值，但是这里在训练过程中引入软量化，使得h(V)h(\mathbf{V})h(V)可以在[0,1]之间变化，然后随着训练的进行，通过调节某个超参数使得正则化项$f_{r e g}(\mathbf{V})最终收敛向0或1.

这个论文里，h(V)h(\mathbf{V})h(V)使用经过整流的sigmoid函数：
h(Vi,j)=clip⁡(σ(Vi,j)(ζ−γ)+γ,0,1)h\left(\mathbf{V}_{i, j}\right)=\operatorname{clip}\left(\sigma\left(\mathbf{V}_{i, j}\right)(\zeta-\gamma)+\gamma, 0,1\right) h(Vi,j)=clip(σ(Vi,j)(ζ−γ)+γ,0,1)
之所以不直接使用sigmoid函数的原因是 sigmoid函数在接近0或1时，其梯度基本趋近于0了。而我们这里的需求又是要训练V\mathbf{V}V最终收敛到0或1，如果梯度接近0的话，那就很难再更新V\mathbf{V}V了，因此，这里通过对sigmoid函数做一些小的调整，使得其在接近0或1时，仍能有较大的梯度值。

正则项freg(V)f_{r e g}(\mathbf{V})freg(V)如下：
freg (V)=∑i,j1−∣2h(Vi,j)−1∣βf_{\text {reg }}(\mathbf{V})=\sum_{i, j} 1-\left|2 h\left(\mathbf{V}_{i, j}\right)-1\right|^{\beta} freg (V)=i,j∑1−∣2h(Vi,j)−1∣β
β\betaβ就是前面提到的用来控制freg(V)f_{r e g}(\mathbf{V})freg(V)逐渐收敛到1的超参数。

下面给出freg(V)f_{r e g}(\mathbf{V})freg(V)和h(V)h(\mathbf{V})h(V)这两个函数随着β\betaβ从20逐渐退货减小到0的变化图（红色代表h(V)h(\mathbf{V})h(V)，蓝色代表freg(V)f_{r e g}(\mathbf{V})freg(V)）：

随着β\betaβ的增大，freg(V)f_{r e g}(\mathbf{V})freg(V)越来越接近一个标准的方波函数。

其他细节请见论文。

@inproceedings{nagel2020up, title={Up or down? adaptive rounding for
post-training quantization}, author={Nagel, Markus and Amjad, Rana
Ali and Van Baalen, Mart and Louizos, Christos and Blankevoort,
Tijmen}, booktitle={International Conference on Machine Learning},
pages={7197–7206}, year={2020}, organization={PMLR} }