LSQ:基于可训练Step_size的低比特量化
from ICLR2020
IBM Research
Abstract
神经网络以低精度在运行推理时具有低功耗和更易于存储的优势，但却牺牲了精度，因此本文就致力于解决这一问题。本文提出的LSQ量化方案应用在不同架构时，在ImageNet数据集上均达到了SOTA，其中权重和激活值量化到了2、3、4bit，以及量化位数为3bit时达到了全精度模型的精度。LSQ在现有的学习权重量化的基础上改进了量化器（quantizer)本身的配置方式;值得一提的是引入了一种新的梯度损失，用于估计和缩放每一层权重和激活值的量化步长，这样就可以与网络中的其他参数一同训练。这种量化策略可以根据给定系统的不同需求进行设定，只需要对原有的代码进行修改即可。
Section I Introduction
深度学习在一些革命性技术，如图像识别、语音识别、自动驾驶等领域大放异彩，但充分发挥深度学习的潜能的同时还需要考虑实际性能、吞吐率、能效、模型紧凑型等指标，这就需要算法和硬件部署的协同设计。当前众多学者致力于在保持深层次网络高精度的同时，尽可能压缩权重和激活值所需的精度，从而降低模型的计算量和所需内存，目前按照这一思路设计的网络已经成功应用到了一些系统中。
量化就是将高精度的权重、激活值映射到一个由较低精度组成的离散数值集合中，通常是8bit甚至更少。本文希望能够对每一层提出单独的量化方案，从而最大化模型性能，这仍然是一个悬而未决的问题。
目前大多数网络量化都采用了一致的量化方案，主要确定一个单一的量化步长参数；随后进一步考虑了非一致的量化映射，这更为复杂。如Hubara等人设计了一种混合配置的量化器，Rastegari等人则聚焦于利用数据分布的特点完成数据的量化映射，Cai等人则聚焦于如何在训练过程中最小化量化误差。
近期人们的目光聚焦于如何借助反向传播随机梯度下降来优化量化产生的误差。因为前文提及的固定量化方案虽然简单，但不能保证网络性能；最小化量化误差的思路虽然理论上可以获得最小的量化误差，但可能也不是最优的。
通过最小化损失函数来优化量化器的映射方案看起来十分有吸引力，因为它直接优化最关键/最感兴趣的部分，但是由于量化器本身就是不连续的，如何计算梯度是一个大问题，目前的近似处理已经是一种很粗糙的方式了，已经忽略了不同状态之间的跳变了。
因此，本文提出一种在深层网络中学习每一层映射的LSQ策略-Learned Step Size Quantization(步长学习策略）。主要由两大亮点：
首先提出了一种简单的逼近量化梯度的方法，将量化步长作为参数可以在更细粒度层面上优化量化器；
其次提出一种简单易行的启发式算法，能够更好地平衡步长幅度与权重更新的幅度，这样有助于提升模型的收敛。
LSQ既可用于权重也可用于激活值的量化，并且与现有的反向传播、随机梯度下降兼容。
已经在ImageNet数据集上测试了量化后的多种网络架构，均比原有的量化方案具有更佳的精度；甚至里程碑式的实现了3bit量化网络与全精度网络达到一样的准确率。
Section II Methods
我们假定网络在推断时使用低精度的整形数完成卷积层和全连接层的计算，这就需要度权重和激活值进行量化
。
对于待量化的数据v，s为需要学习的步长参数-step size，Qp和Qn分别代表正负两个量化级别，而量化器就负责计算v量化后的结果vbar,vbar乘以缩放因子s就会得到v^。
clip(z,r1,r2)返回值z,如果小于r1则shezhiweir1,如果高于r2则设置为r2，[z]操作负责将z取整到最接近的整数；对于无符号数采用bbit量化，则量化范围（Qn,Qp）=（0,2^b-1),
对于符号数则量化范围（Qn,Qp）=（2b-1,2b-1-1)。
因此进行推断时wbar和xbar就可以作为输入进行低精度的整数乘法运算，而卷积层和全连接层的输出就可以按照步长重新缩放，以较低的成本进行高精度的标量乘法，甚至可以和BN等归一化操作合并进行。Fig1展示了具体的计算流程。
Part A Step Size Gradient
Step size梯度计算如下：
通过Benigo的STE算法（straight through estimator）得到梯度的近似估计，STE核心的思想就是我们的参数初始化的时候就是float这样的连续值，当我们forward的时候就将原来的连续的参数映射到{-1, 1}带入到网络进行计算，这样就可以计算网络的输出。然后backward的时候直接对原来float的参数进行更新，而不是对二值化的参数更新。
以Fig2为例，（A）是常规量化（B）则是LSQ,QIL,PACT不同量化策略的输出，可以看到LSQ的梯度对跳变点周围更敏感。相比之下QIL只对距离clip point的远近有关，PACT更不用说了梯度恒0，因此我们有理由认为使用LSQ计算得到的梯度训练网络可以达到更高的精度。因为在QIL和PACT方法中，v与状态跳变点的距离并不会影响量化参数的梯度。
也可以这样认为，给定的输入数据v越靠近跳变点，就越有可能因为网络参数的更新而改变其量化的bin，从而引起v^一个较大的变化。因此我们自然而然期望∂ˆv/∂s能够随着v到跳变点的距离越近而增加，也确实在LSQ中看到了这种趋势，这比采用固定的量化方案或者STE近似计算要更实用。
本文中网络的每一层权重和激活值都有自己单独的step size，用fp32表示，初始化为
，分别在初始权重值或第一批激活值上计算。
Part B Step Size Gradient Scale
步长梯度的缩放
有研究表明，网络中各层权重的更新，其更新幅度与平均参数的幅度大致相近时可以使得训练有较好的收敛；只要设置合适恰当的学习率就可以保证更新的步子既不太大也不会在local minima附近重复跳而导致收敛时间过长。
因此本文也因为步长的更新也应该遵循这一规律，因此通过下面这个式子：
来计算调整的比重。
当量化位数较大是，采用较小的step-size；当量化位数较少时采用较大的step-size。因此需要给梯度乘以一个缩放因子g，对于权重和激活值g的计算分别为：
其中Nw代表某层的权重数目；Nf代表某层的特征数目。
Appendix Section A进一步说明了这样有助于提升训练精度。
Part C Training
LSQ通过赋予量化步长可学习的参数来优化量化器，为了验证LSQ的效果，本文训练了一个BinaryConnect的网络，在BinaryConnect中会将权重和激活值二值化从而减少前向传播和反向传播中对乘法的需求，因此梯度计算如下：
训练时为了简单起见，还用v^作为矩阵乘法的输入，这就相当于在进行推理任务。随后本文将除first和last层的其它层的权重和激活值进行2、3、4、8bit量化，其他的参数则均使用32位浮点数。所有参与对比的量化网络均从一个全精度的网络初始化而来。
训练使用交叉熵损失函数，momentum = 0.9，cosine衰减，理论上8bit应该接近全精度网络，因此8bit网络仅训练1epoch，其他网络训练90epoch。
网络框架有ResNet,VGG,SqueezeNet。输入图像为256x256，裁剪至224x224大小，框架使用的是Pytorch.。
Section III Results
Part A Weight Decay
本文认为降低模型精度可以减少过拟合的强项，也可以避免为了实现良好性能使用正则化，为此本文对ResNet-18不同量化位数权重衰减进行了参数扫描，从Table II可以看出确实量化位数较低的网络可以通过设置较小的权重衰减参数来达到更高的精度，比如3BIT量化时把权重衰减减半，2bit时权重衰减为原来的1/4时取得最佳精度。
Part B与其他量化方案的对比
Table I展示了LSQ与其他量化方案的对比。为了公平起见，本文假定模型中所有卷积层和全连接层的权重、激活值都被量化到特定位数，只有第一层和最后一层保持高精度，可以发现在一些框架下LSQ量化甚至比full-precision精度更高，而且不管是2、3、4bit量化LSQ均取得了最佳的精度；此外还可以看到2bit精度时ResNet-18的精度只掉了2.9，但是SqueezeNet精度却掉了14.0，一种可能的解释就是SqueezeNet设计时就力图用尽可能少的参量充分发挥模型性能，导致一些点对精度十分的敏感。
Part C Accuracy vs Model Size
对于对模型规模有限制的一些应用，就需要在满足内存限制的前提下尽可能选择性能更佳的模型。因此本文对不同Model size的模型Top1精度进行了对比，可以看出规模和精度下ResNet都是不二之选。
Part D Step Size Gradient Scale Impact
本文还对缩放指标R进行了分析，基于ResNet-18一个epoch内500次迭代的平均值；如果不进行缩放计算出的步长会比权重每次更新的步幅高2-3个量级，通过对step size进行缩放，这种步长与权重更新的不平衡也在很大程度上被消除。
TableIII展示了2bitResNet-18缩放不同尺度的对比结果，可以看到如果不缩放训练甚至无法收敛；通过调整学习率使得网路收敛的精度也比baseline低3.4；这一问题通过对gradient进行缩放至1/√N即可解决，精度也只下降了0.3.
Part E衰减学习率的作用
本文使用余弦衰减学习率，这样就避免了设置学习率衰减的超参数，这在大多数框架内都适用且还不增加训练时间。比如2bit ResNet18使用LSQ训练90个epoch达到了67.2的精度，虽然比等效的模型精度低了0.4，但仍然比其他训练方法平均提高了1.5.
Part F 量化误差
下面来探究，LSQ是否是在最小化量化误差，虽然在损失函数中没有明显的表明这一点。经分析得知LSQ实际上没有最小化量化误差，而是直接寻求精度更佳的量化方案，也表明了简单用optimizer转换到相应的数据分布可能不是最佳的量化方案。
Section IV Conclusions
本文提出的低比特量化方案LSQ，在诸多网络框架上超过了之前其他量化方案的性能，LSQ的核心思想是学习每一层的权重和激活值各自的量化步长，从而保证量化后网络的最佳性能，而方法本身只需要在每一层额外增加一个参数即可。
尽管提出的LSQ的目的是提升低比特网络的精度，但还不确定量化至2bit效果如何，现在看起来和全精度网络还存在一定差距；但对于一些较为流行的框架已经被证实食欲用的，比如当模型限制为8M大小时，采用2bit ResNet50的结果比4bit ResNet-34的精度更佳。
低比特量化的思路还会继续探索下去，这可以增强人工神经网络与生物神经网络之间的类比，比如生物神经网络中的突触可用单比特的spike表示；另一方面降低模型的比特数同时还能保持模型精度在压缩模型、提升吞吐量等方面都大有可为。