论文链接:
https://openreview.net/pdf?id=POWv6hDd9XH
代码开源地址:
https://github.com/yhhhli/BRECQ
动机与背景
模型量化将浮点输入和参数映射为定点数,利用硬件的整型指令加速算子推理。作为一种通用的模型压缩和加速方式,已经广泛应用于各种计算机视觉任务。目前的工业级量化生产主要需要解决三个问题:
真的能快:量化点对齐硬件的量化位置,考虑部署常用的折叠BN等操作。保持精度:在一些极限要求和特殊场景下仍能保持与浮点模型效果的一致性。高效生产:不要在浮点模型生产的基础之上额外增添过多时间和过于复杂的流程。这三个问题自然的引出了对生产量化模型方式的辩证思考,选择离线量化还是在线量化?
表1 离线量化(PTQ)与在线量化(QAT)的对比
算法位置
所需数据
时间
上手复杂度
能力上限
离线量化
贴近部署
少量校准集
10-100张
几分钟-10分钟
低,几乎一个命令行
能解决8bit的绝大部分问题
在线量化
贴近训练
大量训练数据
完整训练集合
几小时甚至几天
高,需要一定程度侵入训练代码
解决8bit的hard-case和更低比特
由上表可以看出,离线量化在更贴合硬件部署的位置,更容易做硬件的对齐,不需要耦合训练过程,只需要很少的数据、很短的时间和很简单的命令行调用,即可完成量化模型的生产,从“天性上”就更适合量化模型的工业级落地,而其唯一的缺陷就是面临更大的精度挑战,面对一些8bit的难点问题或者更低比特问题时,有更大的准确度损失。那么离线量化是否能尽可能接近在线量化的效果呢?
方法
针对这个问题,本文重新审视了离线量化的优化粒度,并尝试从理论视角分析已有方式存在的问题。
最早期的量化工作主要关注在量化本身的误差上,并通过如下方式建模量化误差:
其中代表浮点权重,代表量化权重。然而直观上来说仅考虑量化本身的误差其实是不够的,因为这个误差会随着神经网络的计算流程逐渐累积,这也是Bias Correction[1]一类的方法能够有效的原因。因此很直观的是想到以更大的优化粒度寻找更好的量化参数。Intel的LAPQ[2]即是这样做的,不过它采用了一些坐标下降系列的优化方法。直到ICML2020上,高通的AdaRound[3]一鸣惊人,提出除了量化参数,权重也是可以在离线量化场景中被优化的,优化的空间主要在于上下取整。作者在文中提出了一种估计损失函数的方法,即使用二阶泰勒展开去计算损失函数的近似:
这里代表量化权重后产生的扰动,g和H分别代表梯度向量和Hessian矩阵(二阶梯度)。因为在离线量化中,全精度模型已经训练到收敛,因此梯度可以近似等于0。但是估计全部Hessian矩阵是非常困难的,因为矩阵的维度是,N是参数量,存下一个Hessian矩阵将需要TB级别的容量空间。为了估计Hessian矩阵,[3]假设H是层间独立的,即矩阵是层-对角化的(任何两层之间的二阶导数均为0)。之后经过一些推导,得到的结论是只需要对每一层的输出进行重建:
其中是全精度的第l层输出,则是量化网络的第l层输出。在重建过程中,只需要几百张训练图片(无需label标注)就能显著提升量化模型的效果。
AdaRound的方法确实取得了很好的效果,但是重新思考其分析过程会发现,它一方面想考虑整网的任务损失,一方面又推导到了逐层优化,这里存在一个优化粒度选择的问题,逐层优化忽略了层间依赖关系,显然是损失了信息的。因此本文对这个问题进行了重新思考:
离线量化的难点主要在于两个方面(1)如何正确的获取到Task loss信息,并且使其降低(2)如何利用有限的数据集在短时间内校准出一个量化模型。对于问题(1)本文认为用层对角化的Hessian估计会丢失过多的信息,因此需要寻找到一个更精确地估计方案。定义 为所有层权重,那么二阶导数可以计算为:
其中,为网络输出。当网络已经收敛时,假设Hessian矩阵是半正定的,因此可以使用高斯-牛顿矩阵来估计Hessian,即忽略掉上式中的第一项。接下来,可以用矩阵表达式来表示该矩阵:
将上式中的H替换到中,并且使用泰勒一节逼近,可以将替换为 。这意味着我们可以重建网络的最终输出层的误差(和网络蒸馏非常相似)来校准量化模型。
但是当使用整个网络重建之后,发现效果并不如预期。这是因为在网络蒸馏中可以用整个数据集来优化模型,但是离线量化中,只能使用几百张图片做校准,在整网重建过程中,校准集的误差的确变得比逐层重建更低,但是更容易出现过拟合现象。而在逐层重建中,量化模型的每一层输出都强制与全精度模型保持相同输出,相当于引入了一个逐层的正则化,从而具有更高的泛化能力。但其问题是无法从更全局的视角评估量化误差,不能保证task loss的一致性。
因此本文希望尝试找到一个折中的方法来重建量化模型,定义了4个重建粒度,分别是逐层重建,逐块重建,逐阶段重建,整网重建。其中,块(Block)定义为组成模型的一个基本模块,通常只有三四层。比如在残差网络ResNet-50中,块被定义为BottleneckBlock,由两个1x1和一个3x3的卷积层加上一个残差分支(residual branch)组成。而阶段(Stage)则是由几个块组成,在每个阶段,卷积网络的感受野将会被下采样,从而提取更高维度的特征。在ResNet中,每个网络都有4个阶段,例如ResNet-34每个阶段含有的块数量分别为3,4,6,3。不同优化粒度之间的关系见下图。
图1 优化粒度
本文在这四个重建粒度之间进行了选择,发现块重建在大量网络结构以及图像分类和检测任务中都是最优的选择。
同时,为了进一步提升离线量化上限,本文也探索了离线量化的混合精度问题。在以往的文献中,混合精度往往都是层间独立的,即用不同的比特数求出每层的敏感度,再将敏感度加起来得到整个网络的敏感度。可以用如下公式来表示,c是比特向量,包含了每层的比特数,H(c)是硬件评估函数,测量了模型的规模,或者延迟等硬件参数。
延续前文的观点,层与层之间不是独立的,因此使用块敏感度度量,将每个块内所有的比特排列的敏感度测量出来,块之间的敏感度则是认为独立可直接相加。
实验
本文首先做了消融实验,验证四个不同的重建粒度之间的区别。在ImageNet数据集,ResNet-18和MobilenetV2上面做了实验,结果如下:
表2 消融实验:不同优化粒度的离线量化效果
可以发现,使用逐块重建的效果是最好的。
为了更广泛的验证各种网络结构和不同的表现,本文在ImageNet数据集上做了不同比特数下的离线量化实验,且使用了大量的网络结构包括ResNet, RegNet, MobileNetV2, MNasNet进行验证。实验结果如下:
表3 图像分类离线量化效果(只量化权重)
表4 图像分类离线量化效果(权重和激活值都量化)
可以看出,在分类的4比特量化任务上,BRECQ在效果方面实现了非常大的跨越,以往的方法量化到2W4A一般都是崩的,而本文的方法首次实现了可接受的精度,提升了50%-60%的准确率。
同时我们在MS COCO上面做了目标检测的离线量化实验:
表5 目标检测离线量化效果
可以看出,块重建技术在目标检测上也取得了显著的优势。
如前文所说,引入BRECQ之后,离线量化在一些更极限的比特数下首次实现了接近QAT的水平,具体效果见下表:
表6 与QAT方法的对比
在4比特下,BRECQ甚至可以打败一些在线量化方法的准确率,同时享受到大概240-250倍的时间优势。这对于降低量化模型生产成本来说是一个极大的好消息,为离线量化技术的进一步规模应用打开了新的空间!
通过进一步引入混合比特,可以将离线量化的精度持续提升,为了验证这一点,我们在ARM和FPGA上进行了验证,同时考虑硬件延迟/模型体积和准确率,利用遗传算法寻找最优解:
图2 ARM CPU混合比特结果
可以看出,引入混合比特之后,确实可以相比固定比特实现进一步的精度上限拉升。
总结
本文提出了一套块重建方法用于提升更加实用的离线量化的精度上限,通过寻找最优重建粒度,首次将离线量化的精度拉升到在线量化的水平,该方法有极高的实用性,在多种任务和模型结构上取得了稳定一致的效果提升!
作者介绍
论文共同第一作者李雨杭,耶鲁大学准博士生,主要研究方向为模型压缩与加速应用,包括模型量化、模型结构搜索、硬件加速以及神经拟态计算。以第一作者在ICLR、AAAI等会议上发表四篇论文。
论文共同第一作者龚睿昊,商汤科技高级研究员,主要工作方向为模型压缩与AI-System协同优化,包括模型量化稀疏、系统视角的模型结构推荐、神经网络编译等。以第一/二作者在ICLR、CVPR、ICCV、PR等期刊会议上发表6篇论文,服务大量业务线量化模型落地。个人主页:https://xhplus.github.io/。
团队介绍
商汤研究院-Spring工具链团队致力于通过System+AI技术打造顶尖的深度学习核心引擎。开发的模型训练和模型部署工具链已服务于公司多个核心业务。团队在量化模型的在线/离线生产、部署对齐、标准工具等方向有着明确的技术规划。团队目前正在招聘模型量化/稀疏方向的系统算法研究员/实习生,欢迎感兴趣的小伙伴投递简历到 [email protected]
电子科技大学顾实教授,2017年入选国家特聘青年专家。2016年毕业于美国宾夕法尼亚大学,获应用数学与计算科学博士。2016-2017年在宾夕法尼亚大学医学院做博士后研究。主要研究领域为网络神经科学、神经影像分析、类脑人工智能等,曾在Nature Communication, PNAS等顶尖期刊发表高水平论文。其团队在相关领域有广泛深入的研究,常年招硕士/博士生。
实验室主页:https://guslab.org/
参考文献:
[1] Markus Nagel, Mart van Baalen, Tijmen Blankevoort, and Max Welling. Data-free quantization through weight equalization and bias correction. In Proceedings of the IEEE International Conference on Computer Vision, pp. 1325–1334, 2019.
[2] Yury Nahshan, Brian Chmiel, Chaim Baskin, Evgenii Zheltonozhskii, Ron Banner, Alex M Bronstein, and Avi Mendelson. Loss aware post-training quantization. arXiv preprint arXiv:1911.07190, 2019.
[3] Markus Nagel, Rana Ali Amjad, Mart van Baalen, Christos Louizos, and Tijmen Blankevoort. Up or down? adaptive rounding for post-training quantization. arXiv preprint arXiv:2004.10568, 2020.