LCQ:基于低比特量化精度提升的可学习压扩量化方法
本文是日本冲电气集团(Oki Electric Industry)发表在CVPR 2021上的关于非均匀量化的文章,通过提出可学习的压扩函数灵活而非均匀地控制权值和激活的量化级别。结合权重归一化技术达到了SOTA的量化结果。
- 论文题目:Learnable Companding Quantization for Accurate Low-bit Neural Networks
- 论文链接:https://arxiv.org/abs/2103.07156
摘要
量化深度神经网络是一种有效的减少内存消耗和提高推理速度的方法,因此适用于资源受限的设备。然而,极低位模型仍然很难达到与全精度模型相当的精度。为了解决这个问题,本文提出了可学习扩展量化 (LCQ) 作为一种新的非均匀量化方法用于 2-bit、3-bit 和 4-bit 模型量化。LCQ 联合优化模型权重和可学习的压扩函数,这些函数可以灵活而非均匀地控制权值和激活的量化级别。本文还提出了一种新的权重归一化技术,允许更稳定的量化训练。实验结果表明,在图像分类和目标检测任务中,LCQ 优于传统最先进的方法,缩小了量化模型和全精度模型之间的差距。值得注意的是,ImageNet 上的2-bit ResNet-50 模型达到了最高的 75.1% 的精度,并将差距缩小到 1.7% ,使 LCQ 能够进一步挖掘非均匀量化的潜力。
简介
深度神经网络(DNNs)已经成功地应用于基于图像的任务,如图像分类和目标检测,但由于需要大量的乘法累积(MAC)操作和参数,它们在资源受限的移动或边缘设备上的实现仍然困难。为了缓解这一问题,人们提出了各种压缩 DNNs 的技术,同时保持性能,如修剪、知识蒸馏、低秩近似和网络量化。其中,网络量化作为一种有效提高内存消耗和推理速度的重要方法。然而,众所周知,网络量化会随着位宽减少的数量而降低原始模型的性能。
在网络量化中,DNNs的权值或激活通常由一个量化函数离散化。虽然量化函数是不可微的,但一个直通估计器(STE)可以用来近似反向传播的梯度计算。量化函数分为均匀量化和非均匀量化两种类型,其中输入值分别进行线性离散和非线性离散。由于权重或激活分布在经验上与均匀分布不同,通过适当的优化,可以期望非均匀量化比均匀量化进一步减少量化和预测误差。例如,以前关于非均匀量化的工作曾尝试使用固定和对数量化级别或可学习的量化级别来最小化量化错误。然而,精确估计有效的量化水平并不容易,特别是在低位模型中,其精度往往低于均匀量化方法。因此,本文旨在挖掘非均匀量化器的潜力,进一步弥合量化模型与全精度模型之间的精度差距。
本文提出了一种非均匀量化方法,称为可学习扩展量化(LCQ)。图 1 显示了LCQ的概述。本文的方法是基于广泛应用于电信和信号处理的压缩扩展技术,通过非线性限制输入信号的动态范围来减小输入信号的位宽。本文假设扩展对 DNN 的量化在两个方面是有效的:一是利用非线性函数及其逆函数保持输入输出之间的尺度不变,并通过反向传播减小量化误差和稳定训练。二是如果压扩函数是可微的,则可对其参数进行优化,直接使任务损失最小。然后,由于使用舍入前后路径的两个梯度的和来更新参数,因此可以对它们进行具有较大量化影响的训练。具体来说,本文将一个可学习的分段线性函数作为一个非线性压缩函数,允许通过优化对量化水平进行灵活和非均匀的控制。本文还提出了一种新的权值归一化方法,通过将量化权值的标准差恢复到原来的水平来提高精度,并讨论了一种利用查找表进行有效推理的训练技巧。本文的主要贡献总结如下:
- 本文提出了一种 LCQ 方法作为一种新的非均匀量化方法,通过训练可学习的压扩函数来优化非均匀量化水平以最小化任务损失。
- 本文为权重量化器提供了一种简单的归一化方法,称为有限权重归一化(LWN),它可以产生稳定的训练和更好的解决方案。
- 本文评估了用于图像分类和目标检测任务的各种网络上的 LCQ,结果显示在CIFAR-10/100、ImageNet 和 COCO 数据集上有良好的性能。
相关工作
为了量化训练中的权重和激活,通常采用两种操作:输入的 “裁剪” 和 “量化”。量化误差在各个过程中都会发生,因此提出了各种方法来减少量化误差剪裁技术。为了进行量化,首先使用裁剪来约束输入值的取值范围。确定剪切阈值的最简单方法是使用给定的固定值,但是这样做不能适应训练期间输入值的动态范围的变化。为了解决这个问题,Jacob等人提出了一种用指数移动平均跟踪的输入最大值作为阈值的方法。Choi等人提出了一种将阈值视为可学习参数的方法,对其进行优化,使任务损失与权值一起最小化。Zhao等人在每次迭代中使用模拟梯度来估计接近最优阈值。一些先前的著作已经提出了改进的可学习阈值方法的公式,使用由预量化值和量化值之间的残差计算的梯度来更新参数。
均匀量化
均匀量化将剪切值映射到等间距的离散能级之一。虽然这样的映射是通过一个不可微的步进函数来执行的,但是 STE 经常被用于近似梯度计算和基于 backpropagation 的参数更新。Gong等人提出了一种减少梯度逼近误差的方法,将量化函数表示为几个tanh函数的串联,并训练它们的形状参数逐步收敛于阶跃函数。Li等人将均匀量化应用于具有批归一化折叠的目标检测模型。Zhuang等人提出了一种渐进的两阶段量化方法。Jung等人引入了参数化量化区间并对其进行优化以使任务损失最小化。Liu等人使用了一种不应用 STE 的方案,而是使用量化前和量化后值的加权平均值来逐步向量化值转移。Zhuang等人提出了一种训练方案,利用连接到低位网络的辅助模块,为其提供来自其他损失的梯度。
非均匀量化
由于DNN 权重和激活在经验上是非均匀分布的,将输入离散到不相等水平的非均匀量化应该有效工作。Han等人使用k-means聚类作为量化的方法来分享权重。Xu等人采用了相同的聚类策略,但逐层逐步共享权值。Miyashita等人引入了使用二次方 (PoT) 函数的非均匀量化,并表明 DNNs 中的乘法可以被更便宜的位移操作所取代。Polino 等人将量化水平制定为可学习的参数,并使用梯度下降和蒸馏方法训练它们。Zhang 等人提出了量化能级的参数化基,并连续估计了最小化量化误差的解析解。Li等人出了一种可加性幂次量化器来解决幂次函数映射极低的量化水平到更大的输入值(刚性分辨率)的问题。
方法
基础知识
网络量化的目标是用较低的位宽代替DNNs的浮点权值和激活,以减少内存消耗并加快MAC操作。对于统一量化,标准的统一量化器
将输入
量化为:
其中
是一个 clip 参数,
是一个统一量化函数,下标
是位宽。这种裁剪操作通过将量化后的值再乘以
来减小量化误差,使其返回到初始值范围。让剪切的输入为
,
可以表示为:
Learnable Companding Quantization
定义 LCQ 的量化器为:
其中,
为压缩函数;
为均匀量化函数;
为扩张函数。组合到一块的
为压扩函数。
然后本文使用一个可学习的分段线性函数作为压缩函数,训练其参数 θ 以减少任务损失。分段线性函数可以根据断点(或间隔)的数量灵活地改变其斜率,适合于对量化级别的细粒度控制。例如,下图给出了分段线性函数(图a)及其压扩函数(图b)在不同位宽下的演化过程。这些图表明,压扩函数的量化水平和区间可以通过分段线性函数每个区间的斜率变化来精细地确定。通过这种方式,本文通过给模型直接调整量化级别的能力来生成精确的低比特 DNNs。
详细公式
具体来说,这样的分段线性函数需要是单调递增的,并满足 [0,1) 的输入范围的约束,以解释量化函数。在本文的公式中,首先让构成第 k 个间隔(其中
)的断点等间隔,这意味着所有间隔长度都是
。然后本文准备了可学习的参数
,并使用softmax 函数将它们的值范围限制为[0,1],就像在
中一样。我们进一步定义第k 个间隔的斜率为
,第 k个间隔的总长度为
,输出级别的累积和为
,并设置
和
。经过上述准备,分段线性压缩膨胀函数可表示为:
其中,
是一个指示函数,如果
返回1,否则返回0。最后通过
和
,使用梯度下降算法对
进行优化。
压扩的反向传播
根据链式规则,量化器
相对于
的梯度可以写成:
这里,
相对于
和
的梯度应该仔细计算,因为压缩和扩展函数在经过量化函数时可能使用具有不同对应关系的间隔。为简单起见,让
作为量化和压缩函数的输出,即
。那么扩展函数
的梯度可以表示为:
其中,
注意,本文使用STE近似量化函数的导数,
由于舍入可能超过上限值范围[0,1) ,但在这种情况下从
减去一个无穷小值使其保证在范围内。
相对于
的梯度可以写成:
类似地,
是上式
替换
的过程。由于梯度中含裁剪参数α,在
的优化中考虑了裁剪效果。
相对于截断输入v的梯度受到与
类似的量化函数的影响。但是,本文通过经验发现,当
较小而
较大时,梯度可能过大,而且这样的梯度会使训练不稳定,所以本文改用
。然后量化器
相对于输入x的梯度变为
for
,否则为0,就像均匀量化器一样。这种不修改(“直通式”)输入梯度的策略也被用于其他非均匀量化方法(APoT和LQ-Net)。
Clipping 的反向传播
本文估计裁剪参数 α 基于训练,具体地根据量化器的梯度更新α,表示为:
注意 α 是与扩展参数 θ 联合训练的。在参数化压缩函数时,与输入间隔相关的断点被设置为等间距,而不是训练,从而防止裁剪参数的变化对断点位置产生显著影响,降低训练效率。使用等间距不那么灵活,但是可以通过增加断点(或间隔)的数量来弥补。
Limited Weight Normalization
APoT 报告说,在使用量化器之前,用其均值和标准差对权重进行标准化,可以稳定每一层的裁剪参数训练。这有两个主要原因:其一,因为权值分布是以零为中心的,满足有符号量化的对称假设。其二,因为 clip 参数对标准偏差的变化不太敏感。
但是,由于很多量化方法都是用预先训练好的全精度权值对量化模型进行初始化,以获得良好的精度,权值归一化会导致量化前后线性层的输出尺度存在差距,这可能会对训练产生负面影响。因此,本文还提出了一种限定权重归一化(LWN)的方法,将归一化的有效范围限制在权重量化器上。LWN可以写成:
其中
是一个输入权重,
和
分别是权重的样本均值和标准差,
是任意有符号量化器。请注意,
和
的梯度并没有像 APoT 中那样用于更新所有可学习参数。LWN 和 APoT 的唯一区别是量化后的标准差是否相乘。这种简单的方法不仅可以在正向传播中将标准偏差恢复到它的预归一化水平,而且可以使可学习量化器参数的梯度依赖于反向传播中的标准偏差。实验表明,该方法具有较好的稳定性和求解效果。
实验结果及分析
CIFAR-10
ImageNet
COCO
消融实验
Number of intervals
压扩函数中的分段线性函数需要区间 K 的个数作为预定义的超参数。本文对 CIFAR-100 数据集的 ResNet-20 模型中间隔数与预测精度之间的关系进行了测试。上图显示了当间隔数从4增加到16时,对于每个不同位数的比特,相对于
的相对精度差异。
的2、3、4位精度分别为65.2、67.4、67.6。对于比特数,准确性往往会随着间隔数的增加而提高,本文发现,准确性往往会随着比特数的减少而显著提高。由于压扩函数在控制量化水平方面提供了更大的灵活性,因为区间的数目增加了,本文推断这种灵活性与精度有关。
Effect of LWN
本文使用 ResNet-18/20、CIFAR10 和 ImageNet 数据集来研究 LWN 对准确性的有效性。上表显示了不同位宽的结果。注意,在不使用 LWN 时,我们改用传统方法,该方法只对预量化的权重进行标准化。我们观察到,结果显示在 2-bit 上有较大的改进,而在 3-bit 和 4-bit 上有较小的改进。因此,LWN 是一种简单而可靠的方法,有助于提高精度。
Outer bit-widths
上表显示了 3-bit ResNet18/20 模型的最高精度和每层LUTs的内存成本,以及之前描述的外部位宽(权重为
,激活为
)的 CIFAR-10 和 ImageNet 数据集。我们观察到CIFAR-10 的精度在 8-bit 到 6-bit 之间保持不变,但在 4-bit 时降低了。相比之下,与 CIFAR-10 相比,ImageNet 的精度下降得更明显,并且随着外位宽数的减少呈线性下降趋势。然而,这两种情况都有一个优势,特别是在 4-bit 的情况下,因为精度下降仅为 0.2% 点,并且 LUT 的内存成本是 8-bit 情况的一半。
结论
本文提出了 LCQ 作为一种非均匀量化方法,它可以通过一个可学习的压扩函数来优化量化 levels。本文给出了压扩函数,使量化 levels 可以通过训练,灵活而非均匀地控制。本文还发现,通过将归一化的有效范围限制在权重量化器(LWN)上,可以稳定量化训练。此外,本文通过应用重量化技术降低了 LUTs 所需的内存成本,以进行有效的推理。各种涉及图像分类和极低比特模型目标检测任务的实验表明,LCQ 的性能优于或堪比传统的均匀和非均匀量化方法。本文还进行了三项消融研究。结果表明,压扩函数的间隔数与其精度之间可能存在一定的比例关系,LWN 有助于精度的提高,通过减小与 LUT 大小相关的外位宽可以在一定程度上保持精度。虽然本文证明了非均匀量化具有强大的潜力,但在资源受限的设备上的快速推理在实践中需要一个有效的硬件加速器,因此本文计划在未来的工作中解决这个问题。