Additive Powers-of-Two (APoT) Quantization：硬件友好的非均匀量化方法

本文是电子科大&哈佛大学&新加坡国立联合发表在 ICLR2020 上的一篇非均匀量化（APoT）的工作。本文，在非均匀量化中通过采用Additive Powers-of-Two（APoT）加法二次幂量化，综合考虑了计算上有效性，低比特量化导致的模型精度下降问题。并实现了不错的量化效果！

• 论文题目：Additive Powers-of-Two Quantization: A Non-uniform Discretization for Neural Networks
• 论文链接：https://arxiv.org/pdf/1909.13144v2.pdf
• 论文代码：https://github.com/yhhhli/APoT_Quantization

摘要

本文首先提出了Additive Powers-of-Two（APoT）加法二次幂量化，一种针对钟形和长尾分布的神经网络权重，有效的非均匀性量化方案。通过将所有量化数值限制为几个二次幂相加，这APoT量化有利于提高计算效率，并与权重分布良好匹配。其次，本文通过参数化Clipping函数以生成更好的更新最佳饱和阈值的梯度。最后，提出对权重归一化来调整权重的输入分布，使其在量化后更加稳定和一致。实验结果表明，本文提出的方法优于最先进的方法，甚至可以与全精度模型竞争，因此证明了本文提出的APoT量化的有效性。例如，本文在 ImageNe t上的 3bit 量化 ResNet-34 仅下降了 0.3％ 的 Top-1 和 0.2％ Top-5 的准确性。

方法

Additive Powers-of-Two量化 (APoT)

量化表示

$$\hat{\mathbf{W}}={\Pi }_{\mathcal{Q}(\alpha ,b)}\lfloor \mathbf{W},\alpha \rceil$$

• $\mathbf{W}\in {\mathbb{R}}^{C_{\text{out}}\times C_{\text{in}}\times K\times K}$
• $\alpha$ 代表裁剪阈值。
• $\lfloor \cdot ,\alpha \rceil$ 代表Clip函数，将权重裁剪到$[-\alpha ,\alpha ]$。
• $\mathbf{W}$中每个元素通过 $\Pi (\cdot )$ 映射成量化值
• $\mathcal{Q}(\alpha ,b)$代表量化候选数值
• $b$ 代表量化位宽

均匀量化表示

$${\mathcal{Q}}^u(\alpha ,b)=\alpha \times \left\{0,\frac{\pm 1}{2^{b-1}-1},\frac{\pm 2}{2^{b-1}-1},\dots ,\pm 1\right\}$$

Powers-of-Two (PoT) 量化表示

$${\mathcal{Q}}^p(\alpha ,b)=\alpha \times \left\{0,\pm 2^{-2^{b-1}+1},\pm 2^{-2^{b-1}+2},\dots ,\pm 2^{-1},\pm 1\right\}$$

基于Powers-of-Two (PoT) 的非均匀量化模式有一个好处是在计算过程中可以采用移位的方式代替复杂的乘法运算，因此幂次的非均匀量化可以显著提高计算效率。如下公示所示：

$$2^{k}x=\{\begin{array}{lr}x& \text{ if }k=0\\ x<<k& \text{ if }k>0\\ x>>k& \text{ if }k<0\end{array}$$

基于Powers-of-Two (PoT) 的非均匀量化十分适配基于钟型的weights形式，可以实现0附近权重集中的位置量化表示多，长尾部分量化表示少。

Additive Powers-of-Two（APoT）量化表示

PoT 量化虽然十分适配基于钟型的weights形式，但是，对于增加bit数是没有明显增益。比如，我们将位宽从$b$设置为$b+1$时$\left[0,\pm 2^{-2^{b-1}+1}\right]$范围内的间隔（interval）不会发生变化，只是在$\left[-2^{-2^{b-1}+1},2^{-2^{b-1}+1}\right]$范围进一步缩小间隔。这个问题被定义为 rigid resolution（刚性分辨率）问题。为解决此问题，本文提出了APoT量化表示。

$${\mathcal{Q}}^a(\alpha ,kn)=\gamma \times \left\{\sum _{i=0}^{n-1}{p}_i\right\}\text{ where }{p}_i\in \left\{0,\frac{1}{2^i},\frac{1}{2^{i+n}},\dots ,\frac{1}{2^{i+\left(2^k-1\right)n}}\right\}$$

• $\gamma$ 是一个缩放系数，以确保 ${\mathcal{Q}}^a$ 中的最大级别是 $\gamma$。
• $k$ 称为基位宽，即每个加法项的位宽，$n$ 是加法项的数量。
• 当设置了位宽$b$和基位宽$k$时，$n$可由$n=\frac{b}{k}$计算得到。

参数化Clipping函数 (RCF)

传统的STE仅仅对clip函数边界以外的阈值梯度值进行更新，对于边界以内的阈值参数梯度均为零，这不利于寻找最优的clip阈值边界。

$$\frac{\partial \hat{W}}{\partial \alpha }\approx \frac{\partial \lfloor W,\alpha \rceil }{\partial \alpha }=\mathrm{sign}(W)\text{ if }|W|>\alpha \text{ else }0$$

针对传统STE梯度约束不完整的问题，本文对权重内外的阈值边界梯度均进行梯度约束，以便于更快更好的训练得到最优阈值。

$$\hat{\mathbf{W}}=\alpha {\Pi }_{\mathcal{Q}(1,b)}\lfloor \frac{\mathbf{W}}{\alpha },1\rceil$$

$$\frac{\partial \hat{\mathbf{W}}}{\partial \alpha }=\{\begin{array}{ll}\mathrm{sign}(\mathbf{W})& \text{ if }|\mathbf{W}|>\alpha \\ {\Pi }_{\mathcal{Q}(1,b)}\frac{\mathbf{W}}{\alpha }-\frac{\mathbf{W}}{\alpha }& \text{ if }|\mathbf{W}|\le \alpha \end{array}$$

class _pq(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, alpha):
        input.div_(alpha)                          # weights are first divided by alpha
        input_c = input.clamp(min=-1, max=1)       # then clipped to [-1,1]
        sign = input_c.sign()
        input_abs = input_c.abs()
        if power:
            input_q = power_quant(input_abs, grids).mul(sign)  # project to Q^a(alpha, B)
        else:
            input_q = uniform_quant(input_abs, b).mul(sign)
        ctx.save_for_backward(input, input_q)
        input_q = input_q.mul(alpha)               # rescale to the original range
        return input_q

    @staticmethod
    def backward(ctx, grad_output):
        grad_input = grad_output.clone()             # grad for weights will not be clipped
        input, input_q = ctx.saved_tensors
        i = (input.abs()>1.).float()
        sign = input.sign()
        grad_alpha = (grad_output*(sign*i + (input_q-input)*(1-i))).sum()
        return grad_input, grad_alpha

权重归一化

权重归一化为裁剪（Clip）和投影（projection）提供了相对一致且稳定的输入分布，这便于在训练过程中更平滑地优化不同层和迭代。 此外，将权重的平均值设为零可以使得量化更加对称。权重归一化公式如下，主要是通过权重值减均值除方差完成，使得归一化后的权重分布满足均值为0方差为1。

$$\tilde{\mathbf{W}}=\frac{\mathbf{W}-\mu }{\sigma +ϵ},\text{ where }\mu =\frac{1}{I}\sum _{i=1}^I{\mathbf{W}}_i,\sigma =\sqrt{\frac{1}{I}\sum _{i=1}^I{\left({\mathbf{W}}_i-\mu \right)}^2}$$

APoT量化伪代码

实验结果

CIFAR-10

ImageNet

---