【Bit-level量化】BSQ: Exploring Bit-Level Sparsity for Mixed-Precision Neural Network Quantization

论文题目：
[ICLR 2021] BSQ: Exploring Bit-Level Sparsity for Mixed-Precision Neural Network Quantization

为什么研究这个（以前工作哪里不好）？

• 不同层的混合精度难以被定义为一个可微分的目标
• 由于混合精度量化巨大的搜索空间，NAS搜索方法面临巨大的搜索成本
• Hessian方法可以对不同层的重要度排序，但是需要手工决定每层的精度选择

主要贡献

• 提出了一种基于梯度的bit-level量化训练算法。该算法将每一位量化的权重作为一个独立的可训练变量，允许使用基于梯度的优化以及直通估计器（STE）
• 提出了一种bit-level的group Lasso正则器，以动态地降低每层的权重精度，从而产生混合精度的量化方案
• BSQ只使用一个超参数，即正则器强度，来权衡模型的性能和大小

核心方法

量化训练的前向传播和反向传播

其中，$w$是浮点表示，$w_q$是对应的$n$-bit定点表示。反向传播时，由于Round函数不可微分，所以求导时用浮点$w$代替$w_q$进行梯度计算。

前向传播使用$w_q$计算模型输出和损失函数，反向传播使用浮点$w$计算梯度，并且$w$在整个训练过程中都保持浮点表示！

Bit-level量化的训练策略

第1步：提取W的动态范围
$$W=s\cdot W_s$$
其中，$s=max(|W|)$是scaling factor，$W_s$是缩放后的权重，于是$W_s$的元素绝对值都在$[0,1]$范围内。这一步本质就是将浮点数除以scaling factor。

第2步：将$W_s$量化到$n$-bit表示
现在对$W_s$的元素进行量化：
$$w_q=\frac{Round(|w_s|\cdot (2^n-1))}{2^n-1}$$

第3步：将$W_q$写成$n$-bit二进制数表示
$$w_q=\frac{{\sum }_{b=0}^{n-1}{w}_s^{(b)}\cdot 2^b}{2^n-1}$$

以上3个步骤如下图(a)所示，其中$s$和$W_s^{(b)}$是可训练的独立参数。

于是，浮点$W$可以表示如下：

但是，$W_s^{(b)},sign(W)$是离散函数，不可优化，因此为了松弛$W_s^{(b)}$的二值限制，采用如下的STE训练。

训练过程如下图中的例子所示，整个过程和上述公式一致。

为了使$sign(W)$也可以动态更新，本文将$W_s$拆分成了两个部分，一部分是$W_p$代表$W_s$中的正数，另一部分是$W_n$代表$W_s$中的负数绝对值。于是，
$$W_s^{(b)}=W_p^{(b)}-W_n^{(b)}$$

Bit-level Group Lasso

为了在训练过程中引入混合精度量化方案，定义了如下正则器：

其中，$W_p^{(b)},W_n^{(b)}$是从$W^g$得到的bit表示。
$B_{GL}$的目的是使$W_p^{(b)},W_n^{(b)}$中某一位$b$中所有元素同时为0。因此，可以安全地去除该位，以降低精度。

训练流程

给损失函数加了memory consumption-aware reweighing，考虑不同层的内存占用（计算参数量）。

重量化

由于BSQ用浮点变量训练模型的位表示，本文执行重量化，将$W_p^{(b)},W_n^{(b)}$转换为精确的二进制值，并识别可以去除的全零位。
重量化过程从$W_p^{(b)},W_n^{(b)}$中重构quantized scaled权重$W_q^{\prime }$：
$$W_q^{\prime }=Round(\sum _{b=0}^{n-1}{W}_p^{(b)}\cdot 2^b-\sum _{b=0}^{n-1}{W}_n^{(b)}\cdot 2^b)$$
于是，$W_q^{\prime }$被转换为了$(n+1)$-bit二进制数，其中每位用$W_q^{(b)}$表示。

精度调整

重量化后，需要调整每一层的精度。具体操作如下【我觉得这部分是本文核心，怎么丢比特？】：
首先，从MSB到LSB检查$W_q^{(b)}$，移除所有具有零元素的位，直到第一个非零位。这个过程中scaling factor不变。
然后，从LSB到MSB检查$W_q^{(b)}$，当从LSB端被移除一位时，$s$需要翻倍（因为LSB移除一位0后，相当于整个数右移了一位）。
假设精度调整使每一层的精度$n$变到了$n^{\prime }$，scaling factor则需要更新为$s^{\prime }=s\cdot \frac{2^{n^{\prime }-1}}{2^n-1}$。精度调整前后的$W$是等价的：

重量化和精度调整在训练过程中以固定的间隔周期性地执行。

激活值量化

BSQ只修改权重的精度，不影响激活。因此，激活的精度是固定的，本文取>=4比特。

主要实验

对比不同的$\alpha$
越大的$\alpha$使整体精度越低。

对比不同量化方法

相比于常用的HAWQ和PACT，BSQ的压缩率和准确率都有所提升。

BTW，这篇文章的实验组织还挺不同于别的文章，ablation study在主实验的前面，第一次见。