FlexRound：基于逐元素除法的可学习舍入后训练量化方法
-
- 摘要
- [1. 引言](#1. 引言)
- [2. 相关工作](#2. 相关工作)
- [3. 方法](#3. 方法)
- [4. 实验](#4. 实验)
-
- [4.1 消融实验](#4.1 消融实验)
- [4.2 ResNet和MobileNetV2在ImageNet数据集上的实验](#4.2 ResNet和MobileNetV2在ImageNet数据集上的实验)
- [4.3 语言模型实验](#4.3 语言模型实验)
- [5. 结论](#5. 结论)
摘要
后训练量化（PTQ）在资源受限设备上部署深度神经网络方面越来越受欢迎，因为与量化感知训练不同，它完全不需要完整的训练数据集，也不需要端到端的训练。实践表明，基于重构每个层或块输出的后训练量化方案在提升量化模型性能方面效果显著，因此，近期的研究工作开发了相关算法，设计并学习新的权重舍入方案，以便更好地重构每个层或块的输出。在本文中，我们提出了一种简单且高效的后训练量化权重舍入新机制，名为FlexRound。该机制基于逐元素除法而非典型的逐元素加法，使得FlexRound能够联合学习一个通用的量化网格大小以及每个预训练权重的不同缩放因子。得益于逐元素除法所带来的导数倒数法则，FlexRound在更新预训练权重对应的缩放因子时，能够自然地利用预训练权重的信息，从而根据预训练权重的数值大小对其进行灵活量化。我们在多种模型和任务上通过实验验证了FlexRound的有效性。据我们所知，本文首次开展了全面的实验，实验不仅涵盖图像分类和自然语言理解任务，还包括自然语言生成任务。此外，我们首次证明，通过逐块重构输出的方式，大型语言模型能够被高效量化，且相较于半精度基准模型，其性能损失微乎其微。相关代码可在以下链接获取：https://github.com/onliwad101/FlexRound_LRQ 。
1. 引言
近年来，深度神经网络在计算机视觉、自然语言处理、自动语音识别等多个领域取得了前所未有的成功。然而，随着这些网络不断改进并超越人类水平的性能，其架构日益复杂，所需的计算资源和内存占用也随之增加。为了减小模型规模并加快推理运算速度，研究人员尝试了多种压缩技术，如网络量化（Courbariaux等人，2016）和网络剪枝（Han等人，2016）。本文主要关注网络量化，因为4位整数（INT4）或8位整数（INT8）量化具有显著优势，借助英伟达A100张量核心图形处理器（Wu等人，2020）或ARM Cortex微控制器（Kim等人，2021）等现成的加速器，就能实现量化神经网络的加速。
网络量化技术大致可分为两类：量化感知训练（QAT）和后训练量化（PTQ）。量化感知训练是在训练过程中融入网络量化的一种方法，Jung等人（2019）、Jain等人（2019）、Zhao等人（2020）、Esser等人（2020）、Lee等人（2021）等众多研究都提出了相关方案。值得注意的是，采用量化感知训练时，神经网络的全精度版本与量化版本在性能上的差异很小。但量化感知训练需要在完整的训练数据集上进行端到端的重新训练或微调，这往往要耗费大量时间和资源，才能得到性能具有竞争力的量化神经网络。此外，出于数据隐私问题，或者由于需要使用遗留模型，完整的训练数据集可能无法获取。量化感知训练的这些缺点，使得研究人员近年来更多地将目光投向了后训练量化（Zhao等人，2019；Wang等人，2020；Nahshan等人，2021），因为后训练量化完全不需要完整的训练数据集，也无需端到端的学习过程。
后训练量化最初是通过"就近舍入"的方式，在参数空间中最小化量化误差来实现的。但这种方法会导致模型性能严重下降。有研究指出，若将量化权重视为受扰动的权重，量化所导致的性能损失可以通过泰勒展开式近似为二阶误差。基于这一观点，Nagel等人（2020）和Li等人（2021）证实，在某些假设条件下，重构每个层或块的输出，等同于最小化量化导致的性能损失近似值。因此，近期的研究工作（Nagel等人，2020；Li等人，2021；Hubara等人，2021；Wei等人，2022）提出，不再采用"就近舍入"的方式，而是通过设计和学习新的权重舍入方案来重构每个层或块的输出，以确保后训练量化后的模型仍能保持全精度模型的性能。然而，现有研究中设计的所有这些新舍入方案，在对预训练权重进行自适应舍入或量化时，均采用逐元素加法的方式。
本文提出了一种新颖的后训练权重量化方法------FlexRound。该方法不同于以往基于逐元素加法的方法，而是采用逐元素除法的思路。通过联合学习通用的量化网格大小和预训练权重的除法因子，FlexRound为后训练量化提供了一种新的思路。有趣的是，借助逐元素除法所带来的导数倒数法则，FlexRound在为每个预训练权重学习对应的单独缩放因子时，能够自然地利用预训练权重的信息。具体而言，我们证实，在对数值较大的预训练权重进行量化时，需要探索更广泛的离散值范围。这种方法背后的原理是，权重的数值大小可以被解读为其在网络中的相对重要性。正如Han等人（2016）的研究所示，数值较大的权重对网络性能的影响比数值较小的权重大。因此，为了在量化后仍保持预训练模型的性能，对于绝对值较大的权重，相较于绝对值较小的权重，放宽其量化约束条件就显得尤为重要（也就是说，重要的权重不仅可以被量化到与其最接近的两个离散值之一，还可以被量化到距离其更远的离散值）。基于此，FlexRound能够根据预训练权重自身的数值大小对其进行灵活量化，进而实现更优的性能。
本文的贡献主要体现在以下三个方面：
- 提出了FlexRound作为后训练权重量化的一种新舍入方案。该方案基于逐元素除法原理，不仅能够为每个预训练权重学习单独的缩放因子，还能为一组（如一个通道或一个层）权重学习通用的量化网格大小。
- 从理论和实验两方面证明，这种基于逐元素除法的新舍入方案在更新预训练权重对应的缩放因子时，会考虑预训练权重的数值大小。因此，与仅能向上或向下舍入的方式相比，FlexRound对数值较大的预训练权重（即重要的预训练权重）的量化更具灵活性。
- 据我们所知，本文首次在逐张量均匀后训练量化设置下，针对图像分类、自然语言理解以及自然语言生成任务开展了广泛实验，实验所使用的模型包括ResNet、MobileNetV2、BERT、GPT-Neo、OPT和GPT-2等多种模型。同时，我们也首次在常识推理和因果语言建模任务上，对LLaMA等大型语言模型进行了均匀后训练量化重构实验。
2. 相关工作
近年来，许多研究人员尝试在不进行任何（重新）训练的情况下，对适用于计算机视觉、自然语言理解/生成等各类任务的多种模型进行量化。异常通道拆分（OCS）（Zhao等人，2019）方法会复制包含异常值的通道，然后将这些通道中的异常值减半。尽管异常通道拆分方法能明确处理异常值，但当权重和激活值都被量化到低比特时，模型仍会出现严重的精度下降。为解决这一问题，Wang等人（2020）提出了位拆分（Bit-Split）方法，该方法将一个整数拆分为多个比特位，并分别对这些比特位进行优化。在低比特设置下，位拆分方法的性能可与全精度模型相媲美，但对于MobileNetV2等特定架构，其效果可能并不理想。
为克服上述局限性，Nagel等人（2020）和Hubara等人（2021）分别提出了名为AdaRound和AdaQuant的新权重舍入机制，并通过逐层最小化全精度层输出与量化层输出之间的均方误差，实现了权重的学习。这种逐层重构误差最小化的方法为4位后训练量化领域开辟了新路径。在此基础上，Li等人（2021）提出了基于块重构的BRECQ方法，该方法不仅考虑了层间依赖关系，还实现了MobileNetV2模型的4位全量化。除了块重构方法外，Wei等人（2022）提出了QDrop方法，在了你后训练量化重构过程中随机取消激活值的量化，使激活值量化与权重量化保持同步。然而，BRECQ和QDrop均基于AdaRound方法，而AdaRound方法在量化时，最多只能在"固定"的量化网格大小下对权重进行向上或向下舍入。AdaQuant方法虽然能够同时学习量化网格大小并对权重进行自适应量化，但在对MobileNetV2模型进行低比特量化时，会导致严重的性能下降。
在后训练量化研究的另一个方向上，一些后训练量化技术专门用于对BERT、GPT类等语言模型进行量化。Bondarenko等人（2021）首次将后训练量化应用于BERT模型，提出了一种逐嵌入组的激活值量化方案，以应对激活值范围高度动态变化的问题。Bai等人（2021）研究了BERT模型后训练量化重构的并行实现方式。Yao等人（2022）提出了ZeroQuant方法，该方法通过逐层知识蒸馏驱动的逐标记激活值量化，以逐组权重量化的方式对BERT和GPT-3模型进行量化（不过ZeroQuant方法需要专用的CUDA内核）。Dettmers等人（2022）采用逐向量权重量化和混合精度分解（激活值采用16位浮点数）的方式，对OPT等大型语言模型（LLMs）进行量化。为避免使用16位浮点数激活值，Xiao等人（2022）提出了SmoothQuant方法，该方法将激活值量化的难度转移到权重量化上，从而实现了大型语言模型中权重和激活值的8位整数量化。遗憾的是，Dettmers等人（2022）和Xiao等人（2022）的方法都假设激活值中的异常值会以特定模式出现。
上述大多数后训练量化研究要么只针对视觉模型，要么只针对语言模型，很少有研究能同时适用于两类模型。据我们所知，本文首次在逐张量均匀后训练量化设置下，针对从图像分类、自然语言理解到自然语言生成的多种任务开展了广泛实验。此外，我们也首次证明，通过逐块重构输出，无需假设激活值异常值会以特定模式出现，就能对大型语言模型进行高效量化，且相较于半精度基准模型，其精度损失极小。
3. 方法
本节首先介绍全文使用的符号以及后训练量化（PTQ）的背景知识，然后详细阐述用于均匀后训练量化重构方法的FlexRound的概念和设计，最后深入分析在FlexRound中采用逐元素除法原理的优势。
3.1 预备知识
符号说明
标量、向量和矩阵（或张量）分别用非粗体字母、小写粗体字母和大写粗体字母表示（例如， s s s、 s \boldsymbol{s} s和 S \boldsymbol{S} S）。 W ^ \widehat{\boldsymbol{W}} W 表示 W \boldsymbol{W} W的量化版本。若前面所有层均保持完整，二维卷积层或线性层的输入用 X \boldsymbol{X} X表示；若前面所有层都经过量化，则用 X ~ \tilde{\boldsymbol{X}} X~表示。矩阵 A \boldsymbol{A} A的元素表示为 A ( i , j ) A_{(i,j)} A(i,j)，四维张量 A \boldsymbol{A} A的元素表示为 A ( i , j , k , l ) A_{(i,j,k,l)} A(i,j,k,l)。 ⊙ \odot ⊙和 ⊘ \oslash ⊘分别表示逐元素乘法和逐元素除法，其运算方式与Python NumPy中的广播机制类似。 ⌊ ⋅ ⌉ \lfloor\cdot\rceil ⌊⋅⌉和 ⌊ ⋅ ⌋ \lfloor\cdot\rfloor ⌊⋅⌋分别表示舍入函数和下取整函数。 ∥ ⋅ ∥ F \|\cdot\|_F ∥⋅∥F表示弗罗贝尼乌斯范数。
后训练量化背景
传统的均匀后训练量化方法是通过"就近舍入"将预训练权重 W \boldsymbol{W} W量化为 W ^ = s 1 ⌊ W s 1 ⌉ \widehat{\boldsymbol{W}} = s_1 \left\lfloor \frac{\boldsymbol{W}}{s_1} \right\rceil W =s1⌊s1W⌉，其中量化网格大小 s 1 ∈ R > 0 s_1 \in \mathbb{R}_{>0} s1∈R>0的设置需满足最小化 ∥ W − W ^ ∥ F 2 \|\boldsymbol{W} - \widehat{\boldsymbol{W}}\|_F^2 ∥W−W ∥F2。但在参数空间中最小化量化误差，并不能等同于最小化最终任务的损失。正如Li等人（2021）所证明的，量化导致的性能损失可以近似表示为网络输出与其海森矩阵的二次形式。因此，有多项研究致力于利用少量数据，逐层或逐块地最小化 ∥ W X − W ^ X ~ ∥ F 2 \|\boldsymbol{WX} - \widehat{\boldsymbol{W}}\tilde{\boldsymbol{X}}\|_F^2 ∥WX−W X~∥F2（其中 W ^ \widehat{\boldsymbol{W}} W 要么是带有特定函数 h ( ⋅ ) h(\cdot) h(⋅)的 s 1 ( ⌊ W s 1 ⌋ + h ( V ) ) s_1(\lfloor \frac{\boldsymbol{W}}{s_1} \rfloor + h(\boldsymbol{V})) s1(⌊s1W⌋+h(V))（Nagel等人，2020），要么是 s 1 ⌊ W + V s 1 ⌉ s_1 \left\lfloor \frac{\boldsymbol{W} + \boldsymbol{V}}{s_1} \right\rceil s1⌊s1W+V⌉（Hubara等人，2021））。然而，所有这些舍入机制都是基于逐元素加法实现的。
3.2 FlexRound
与以往基于逐元素加法的研究不同，我们采用逐元素除法对预训练权重进行量化。基于逐元素除法，我们提出的权重舍入方案可表示为：
其中， S \boldsymbol{S} S是 W \boldsymbol{W} W的除法因子，其维度与 W \boldsymbol{W} W相同，且 S \boldsymbol{S} S的所有元素以及 s 1 s_1 s1均为正数且可学习。与之前的研究类似，更新 s 1 s_1 s1和 S \boldsymbol{S} S的目的都是为了最小化 ∥ W X − W ^ X ~ ∥ F 2 \|\boldsymbol{WX} - \widehat{\boldsymbol{W}}\tilde{\boldsymbol{X}}\|_F^2 ∥WX−W X~∥F2。
公式（1）表明，FlexRound的基本公式支持逐张量均匀后训练量化。尽管FlexRound也可以通过简单地将标量 s 1 s_1 s1替换为向量 s 1 \boldsymbol{s}_1 s1来实现逐通道权重量化，但正如后续实验所展示的，仅通过FlexRound实现的逐张量均匀后训练量化就足以使模型达到全精度模型的性能。因此，除非另有说明，本文主要关注逐张量均匀后训练量化重构。图1展示了在逐张量均匀后训练量化设置下FlexRound的整体流程。
接下来详细介绍 S \boldsymbol{S} S的设计。首先，我们将 S \boldsymbol{S} S表示为 S = s 1 ⊙ S 2 \boldsymbol{S} = s_1 \odot \boldsymbol{S}_2 S=s1⊙S2，其中 S 2 \boldsymbol{S}_2 S2是与 W \boldsymbol{W} W维度相同的矩阵或张量缩放因子，且 S 2 \boldsymbol{S}2 S2的每个元素均为正数且可学习。后续实验将证明，当 S = s 1 ⊙ S 2 \boldsymbol{S} = s_1 \odot \boldsymbol{S}2 S=s1⊙S2时，在逐张量均匀后训练量化设置下，公式（1）所表示的方法与现有的基于逐元素加法的权重舍入方案相比，已能取得更优的性能。但为了进一步提升基于逐元素除法的新舍入方案的性能，我们对 S 2 \boldsymbol{S}2 S2进行了补充。对于线性层 W ∈ R C o u t × C i n \boldsymbol{W} \in \mathbb{R}^{C{out} \times C{in}} W∈RCout×Cin，我们为 S 2 \boldsymbol{S}2 S2补充了一个额外的可学习张量 s 3 ∈ R > 0 C o u t × 1 \boldsymbol{s}3 \in \mathbb{R}{>0}^{C{out} \times 1} s3∈R>0Cout×1。已有大量研究（Nagel等人，2019；Lou等人，2020）证实，输出通道的统计特性差异显著，因此，我们通过为 S 2 \boldsymbol{S}2 S2补充 s 3 \boldsymbol{s}3 s3来考虑输出通道统计特性的变化。对于二维卷积层 W ∈ R C o u t × C i n × H × W \boldsymbol{W} \in \mathbb{R}^{C{out} \times C{in} \times H \times W} W∈RCout×Cin×H×W，我们特别为 S 2 \boldsymbol{S}2 S2补充了两个额外的可学习张量 s 3 ∈ R > 0 C o u t × 1 × 1 × 1 \boldsymbol{s}3 \in \mathbb{R}{>0}^{C{out} \times 1 \times 1 \times 1} s3∈R>0Cout×1×1×1和 s 4 ∈ R > 0 1 × C i n × 1 × 1 \boldsymbol{s}4 \in \mathbb{R}{>0}^{1 \times C{in} \times 1 \times 1} s4∈R>01×Cin×1×1。
由此，对于线性层， S \boldsymbol{S} S可表示为 S = s 1 ⊙ S 2 ⊙ s 3 \boldsymbol{S} = s_1 \odot \boldsymbol{S}_2 \odot \boldsymbol{s}_3 S=s1⊙S2⊙s3（如图2所示）；对于二维卷积层， S \boldsymbol{S} S可表示为 S = s 1 ⊙ S 2 ⊙ s 3 ⊙ s 4 \boldsymbol{S} = s_1 \odot \boldsymbol{S}_2 \odot \boldsymbol{s}_3 \odot \boldsymbol{s}_4 S=s1⊙S2⊙s3⊙s4。此时，公式（1）可转化为：
我们将公式（2）所表示的方法称为"FlexRound"。其中， S 2 \boldsymbol{S}_2 S2、 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4的所有元素均初始化为1，这样做是为了便于从传统的"就近舍入"方法（即 s 1 ⌊ W s 1 ⌉ s_1 \left\lfloor \frac{\boldsymbol{W}}{s_1} \right\rceil s1⌊s1W⌉）过渡到FlexRound方法进行学习。在更新所有参数（ s 1 s_1 s1、 S 2 \boldsymbol{S}_2 S2、 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4）以最小化 ∥ W X − W ^ X ~ ∥ F 2 \|\boldsymbol{WX} - \widehat{\boldsymbol{W}}\tilde{\boldsymbol{X}}\|_F^2 ∥WX−W X~∥F2的过程中，需满足 s 1 s_1 s1、 S 2 \boldsymbol{S}_2 S2、 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4的所有元素均为正数这一约束条件。
在公式（2）中，逐元素除法与逐元素加法的作用相似，都能实现比"就近舍入"更有效的舍入方案。通过逐元素除法实现这种新的舍入策略，我们可以让 s 1 s_1 s1、 S 2 \boldsymbol{S}_2 S2、 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4均成为可学习参数。这使得FlexRound能够在舍入过程（例如FlexRound中的 S 2 ⊙ s 3 \boldsymbol{S}_2 \odot \boldsymbol{s}_3 S2⊙s3或 S 2 ⊙ s 3 ⊙ s 4 \boldsymbol{S}_2 \odot \boldsymbol{s}_3 \odot \boldsymbol{s}_4 S2⊙s3⊙s4）中，联合学习通用的量化网格大小（即 s 1 s_1 s1）。此外，正如以下命题从理论和实验两方面所证明的，逐元素除法所带来的导数倒数法则，使得FlexRound在学习相应缩放因子时能够利用预训练权重的信息。
命题3.1 ：设 c c c为根据公式（2）计算得到的重构误差， S ′ \boldsymbol{S}' S′为公式（2）中用于缩放预训练权重 W \boldsymbol{W} W的矩阵（或张量），即对于线性层， S ′ = S 2 ⊙ s 3 \boldsymbol{S}' = \boldsymbol{S}_2 \odot \boldsymbol{s}3 S′=S2⊙s3；对于二维卷积层， S ′ = S 2 ⊙ s 3 ⊙ s 4 \boldsymbol{S}' = \boldsymbol{S}2 \odot \boldsymbol{s}3 \odot \boldsymbol{s}4 S′=S2⊙s3⊙s4。当使用直通估计器（Bengio等人，2013）时，损失 L L L对 S ′ \boldsymbol{S}' S′中某个元素的梯度 ∂ L ∂ S ( i , j ) ′ \frac{\partial L}{\partial S'{(i,j)}} ∂S(i,j)′∂L（或 ∂ L ∂ S ( i , j , k , l ) ′ \frac{\partial L}{\partial S'{(i,j,k,l)}} ∂S(i,j,k,l)′∂L）与对应的预训练权重 W ( i , j ) W{(i,j)} W(i,j)（或 W ( i , j , k , l ) W{(i,j,k,l)} W(i,j,k,l)）成正比。
命题3.1表明，对于线性层，元素 S ( i , j ) ′ S'{(i,j)} S(i,j)′会（部分地）受到 W ( i , j ) W{(i,j)} W(i,j)的影响，因此 W ˉ ( i , j ) = ⌊ W ( i , j ) s 1 ⊙ S ( i , j ) ′ ⌉ \bar{W}{(i,j)} = \left\lfloor \frac{W{(i,j)}}{s_1 \odot S'{(i,j)}} \right\rceil Wˉ(i,j)=⌊s1⊙S(i,j)′W(i,j)⌉的更新也会受到 W ( i , j ) W{(i,j)} W(i,j)的影响。也就是说，预训练权重 W ( i , j ) W_{(i,j)} W(i,j)的数值越大，在后续训练量化重构过程中， W ˉ ( i , j ) \bar{W}_{(i,j)} Wˉ(i,j)获得更大更新的可能性就越高。权重的数值大小可以作为衡量权重重要性的指标，用于剪枝不重要的权重（Han等人，2015）。因此，数值较大的权重比数值较小的权重更为重要（Han等人，2016）。为了减小预训练全精度模型与其量化版本之间的性能差距，合理的做法是放宽对数值较大的预训练权重（即潜在重要的预训练权重）的量化约束，也就是说，允许这些权重不仅有机会被量化到与其最接近的两个量化网格值，还能被量化到距离更远的量化网格值，而对于数值较小的预训练权重则不给予这种灵活性。上述结论同样适用于二维卷积层。
图3展示了通过FlexRound对MobileNetV2和ResNet-18模型进行权重更新的情况。图3的左侧和中间部分是直方图，分别展示了数值较小的预训练权重（ ∣ W ∣ < 1 |W| < 1 ∣W∣<1，左侧）和数值较大的预训练权重（ ∣ W ∣ > 1 |W| > 1 ∣W∣>1，中间）对应的 W ˉ ( i , j , k , l ) \bar{W}{(i,j,k,l)} Wˉ(i,j,k,l)的变化情况。右侧的散点图则展示了权重相对于"就近舍入"（RTN）网格的偏移量。可以看出，借助FlexRound，MobileNetV2和ResNet-18模型的量化结果存在明显差异。例如，在图3（a）所示的MobileNetV2模型中，通过最小化重构误差 c c c，当 W ( i , j , k , l ) W{(i,j,k,l)} W(i,j,k,l)的绝对值大于1时， W ˉ ( i , j , k , l ) \bar{W}_{(i,j,k,l)} Wˉ(i,j,k,l)的变化更为显著（即舍入时可能会偏离不止一个步长的向上或向下舍入），这表明FlexRound对数值较大的预训练权重的量化更具灵活性，如图3（a）中红色虚线方框所示。在MobileNetV2模型第一个块的第一个二维卷积层中，约有12.8%的权重在量化时发生了显著的舍入偏移。而对于ResNet-18模型，如图3（b）所示，不存在数值大于1的预训练权重，因此大多数预训练权重仅进行向上或向下舍入（例如，在ResNet-18模型第一个块的第一个二维卷积层中，仅有约1.5%的权重在量化时发生了显著的舍入偏移）。附录A中直观地对比了AdaRound、AdaQuant和FlexRound三种方法的舍入结果。
尽管FlexRound在更新预训练权重对应的缩放因子时会考虑权重的数值大小，但有人可能会提出疑问：从图3右侧部分可以看出，FlexRound对数值中等的预训练权重的量化似乎比数值较大的预训练权重更灵活。而我们提出FlexRound的初衷是，强调与数值相对较小的预训练权重相比，数值相对较大的预训练权重更有可能被灵活地量化。如附录B所述， ∣ ∂ L ∂ S ( i , j ) ′ ∣ \left| \frac{\partial L}{\partial S'{(i,j)}} \right| ∂S(i,j)′∂L 与 ∣ W ( i , j ) ∂ L ∂ W ^ ( i , j ) ∣ \left| W{(i,j)} \frac{\partial L}{\partial \widehat{W}{(i,j)}} \right| W(i,j)∂W (i,j)∂L 直接成正比。无论 W ( i , j ) W{(i,j)} W(i,j)的数值多大，只要 ∣ ∂ L ∂ W ^ ( i , j ) ∣ \left| \frac{\partial L}{\partial \widehat{W}{(i,j)}} \right| ∂W (i,j)∂L 趋近于零， ∣ ∂ L ∂ S ( i , j ) ′ ∣ \left| \frac{\partial L}{\partial S'{(i,j)}} \right| ∂S(i,j)′∂L 也会趋近于零。在这种情况下，数值较大的预训练权重可能会被量化到"就近舍入"所能得到的网格值。而如果 ∣ ∂ L ∂ W ^ ( i , j ) ∣ \left| \frac{\partial L}{\partial \widehat{W}_{(i,j)}} \right| ∂W (i,j)∂L （显著）大于零，那么数值较大的预训练权重就可能被量化到距离其最近两个网格值较远的网格值，如图4所示。总之，虽然预训练权重的数值大小会影响FlexRound中对应缩放因子的更新，但这并不意味着数值越大的权重就一定比数值较小的权重被量化得更灵活。
需要注意的是，随着比特宽度的增加，FlexRound对权重的量化灵活性也会提高。对比图3右侧部分和图5可以发现，比特宽度越大，FlexRound使得权重相对于"就近舍入"网格的偏移量就越大。与AdaRound不同（无论使用多少比特，AdaRound都只能对权重进行向上或向下舍入），FlexRound随着比特宽度的增加，能够实现更灵活的权重量化，因此相较于AdaRound，FlexRound更适合对需要更高比特宽度的模型（如大型语言模型）进行量化。
4. 实验
本节首先通过实验验证联合学习量化网格大小 s 1 s_1 s1与舍入过程的重要性，以及额外张量 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4对FlexRound的独特贡献。随后，在逐张量均匀后训练量化设置下，将FlexRound与最先进的后训练量化方法在以下场景中进行性能对比：在ImageNet数据集（Russakovsky等人，2015）上，使用ResNet（He等人，2016）和MobileNetV2（Sandler等人，2018）模型进行图像分类任务（4.2节）；在GLUE数据集（Wang等人，2018）上，使用BERT（Devlin等人，2018）和GPT-Neo（Black等人，2021）模型进行自然语言理解任务（NLU）（4.3节）；在WikiText2数据集（Merity等人，2016）、Penn Treebank（PTB）数据集（Marcus等人，1993）上，使用GPT-Neo和OPT（Zhang等人，2022）模型进行自然语言生成任务（NLG），以及在WebNLG数据集（Gardent等人，2017）上，使用GPT-2（Radford等人，2019）模型进行自然语言生成任务（4.3节）。最后，我们验证了通过逐块重构输出，无需假设激活值异常值会以特定模式出现，就能对大型语言模型（LLMs）进行量化，且相较于半精度基准模型，其精度损失微乎其微。我们以LLaMA（Touvron等人，2023）模型为研究对象，采用逐通道权重量化和逐张量激活值量化方法，在六个常识推理基准数据集（BoolQ（Clark等人，2019）、PIQA（Bisk等人，2020）、HellaSwag（Zellers等人，2019）、WinoGrande（Sakaguchi等人，2021）、ARC easy和challenge（Clark等人，2018）、OpenBookQA（Mihaylov等人，2018））以及WikiText2数据集的因果语言建模任务上进行了实验（4.3节）。
为简洁起见，我们用"B + X"和"Q + X"分别表示在BRECQ（Li等人，2021）或QDrop（Wei等人，2022）所描述的实验设置（实验设置包括用于最小化重构误差的块单元定义，或取消激活值量化的概率大小等）下，采用某种舍入方案"X"。与BRECQ和QDrop中介绍的一致，在更新激活值量化的步长时，我们也采用了LSQ技术（Esser等人，2020）。所有实验结果均基于开源代码通过自主实现获得。
4.1 消融实验
消融实验1
如第2节所述，尽管AdaRound在现有的后训练量化方法中表现出最先进的性能，但它无法将量化网格大小 s 1 s_1 s1与舍入过程进行联合学习。为了明确联合学习 s 1 s_1 s1与舍入过程的重要性，我们在ImageNet数据集上，将权重量化为4位（激活值不进行量化），评估了固定 s 1 s_1 s1时FlexRound的性能（表1中的消融实验1）。从表1可以看出，当 s 1 s_1 s1固定时，除ResNet-18模型外，FlexRound与AdaRound的性能相近。这表明，无论采用AdaRound还是FlexRound量化方法，固定 s 1 s_1 s1都会阻碍量化模型性能的进一步提升。然而，当将 s 1 s_1 s1与舍入过程进行联合学习时，FlexRound在所有模型上的性能都超过了AdaRound。联合学习 s 1 s_1 s1与舍入过程是缩小全精度模型与其量化版本性能差距的关键因素。正如第3.2节所提及的，由于FlexRound基于逐元素除法实现，因此具备联合学习的能力，而AdaRound则不具备这一能力。
消融实验2
为了验证在FlexRound中加入额外张量 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4的合理性，我们在ImageNet数据集上进行了消融实验：将权重量化为4位（激活值不进行量化），对比FlexRound在包含和不包含 s 3 \boldsymbol{s}_3 s3与 s 4 \boldsymbol{s}_4 s4情况下的性能（表1中的消融实验2）。从表1最后两行的结果可以看出，加入 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4后，所有模型的top-1精度均有所提高。值得注意的是，即使不加入 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4，FlexRound的性能仍然优于AdaQuant和AdaRound。这一结果支持了我们的观点：将舍入方案从基于逐元素加法转变为基于逐元素除法，是显著提升量化质量的关键所在。
表1 ：在ImageNet数据集上仅对权重进行4位量化时的top-1/top-5精度（%）。"B + X"表示在BRECQ的设置下实现方案X。" s 1 s_1 s1"列表示在后训练量化重构过程中， s 1 s_1 s1是固定的还是可学习的。" S 2 \boldsymbol{S}_2 S2"和" s 3 \boldsymbol{s}_3 s3、 s 4 \boldsymbol{s}_4 s4"列分别表示在FlexRound中是否包含 S 2 \boldsymbol{S}_2 S2以及 s 3 \boldsymbol{s}_3 s3、 s 4 \boldsymbol{s}_4 s4（"O"表示包含，"X"表示不包含）。例如，FlexRound（本文提出）和消融实验1所采用的公式为公式（2），而消融实验2所采用的公式为 W ^ = s 1 ⌊ W s 1 ⊙ S 2 ⌉ \widehat{\boldsymbol{W}} = s_1 \left\lfloor \frac{\boldsymbol{W}}{s_1 \odot \boldsymbol{S}_2} \right\rceil W =s1⌊s1⊙S2W⌉。
| 方法 | s 1 s_1 s1 | S 2 \boldsymbol{S}_2 S2 | s 3 \boldsymbol{s}_3 s3、 s 4 \boldsymbol{s}_4 s4 | ResNet-18 | ResNet-50 | MobileNetV2 |
|---|---|---|---|---|---|---|
| 全精度 | 不适用 | 不适用 | 不适用 | 71.00 / 89.97 | 76.63 / 93.04 | 72.62 / 90.67 |
| B + AdaQuant | 可学习 | 不适用 | 不适用 | 67.50 / 87.75 | 72.79 / 90.77 | 15.17 / 32.89 |
| B + AdaRound | 固定 | 不适用 | 不适用 | 70.18 / 89.38 | 75.86 / 92.62 | 69.46 / 88.85 |
| B + FlexRound（本文提出） | 可学习 | O | O | 70.28 / 89.44 | 75.95 / 92.68 | 70.82 / 89.67 |
| → 消融实验1 | 固定 | O | O | 70.09 / 89.43 | 75.88 / 92.61 | 69.47 / 88.85 |
| → 消融实验2 | 可学习 | O | X | 70.22 / 89.45 | 75.92 / 92.63 | 70.51 / 89.49 |
4.2 ResNet和MobileNetV2在ImageNet数据集上的实验
我们从ImageNet数据集中随机抽取1024张图像，在低比特后训练量化重构中对ResNet-18、ResNet-50和MobileNetV2模型进行量化。权重和/或激活值的量化采用线性对称逐张量量化格式，而Li等人（2021）和Wei等人（2022）采用的是线性非对称逐通道量化格式，这导致我们自主实现的BRECQ和QDrop与Li等人（2021）和Wei等人（2022）所报告的结果存在差异。对于FlexRound，在5000次迭代过程中对每个层或块的输出进行重构，同时使用一个学习率（例如，对于采用3位或4位量化的ResNet模型，学习率为4e-4；对于采用2位量化的ResNet模型和MobileNetV2模型，学习率为1e-3）更新所有可学习参数（即 s 1 s_1 s1、 S 2 \boldsymbol{S}_2 S2、 s 3 \boldsymbol{s}_3 s3和 s 4 \boldsymbol{s}_4 s4）。与Li等人（2021）的做法一致，将模型的第一层和最后一层量化为8位，并将批量归一化层折叠到卷积层中。除非另有说明，实验均基于BRECQ github仓库¹提供的全精度预训练模型进行。基于PyTorch官方仓库提供的全精度预训练模型的实验结果详见附录C。实验结果均为五次随机试验的中位数。
在仅对权重进行量化的情况下，我们将FlexRound与同样基于逐元素加法原理的AdaRound和AdaQuant进行了性能对比。表2的结果显示，FlexRound的性能始终优于这两种基于加法的舍入策略。需要注意的是，在表2中，AdaQuant的性能不如AdaRound。因此，为节省篇幅，后续实验中FlexRound仅与AdaRound进行对比。表3展示了同时对权重和激活值进行量化时，AdaRound和FlexRound在BRECQ或QDrop设置下的模型精度。值得注意的是，在表3中，正如第3.2节所解释的原因，FlexRound对MobileNetV2模型（包含数值较大的权重）的量化效果尤为显著。此外，从表3还可以观察到，即使在逐张量均匀后训练量化设置下，将ResNet模型的权重和激活值均量化为4位，其性能相较于全精度预训练模型的下降幅度也很小（不到2%）。
¹https://github.com/yhhhli/BRECQ
表2：仅对权重进行量化时在ImageNet数据集上的top-1/top-5精度（%）。"B + X"表示在BRECQ的设置下实现方案X。
| 方法 | 比特数（权重/激活值） | ResNet-18 | ResNet-50 | MobileNetV2 |
|---|---|---|---|---|
| 全精度 | 32 / 32 | 71.00 / 89.97 | 76.63 / 93.04 | 72.62 / 90.67 |
| B + AdaQuant | 4 / 32 | 67.50 / 87.75 | 72.79 / 90.77 | 15.17 / 32.89 |
| B + AdaRound | 4 / 32 | 70.18 / 89.38 | 75.86 / 92.62 | 69.46 / 88.85 |
| B + FlexRound（本文提出） | 4 / 32 | 70.28 / 89.44 | 75.95 / 92.68 | 70.82 / 89.67 |
| B + AdaQuant | 3 / 32 | 57.09 / 80.82 | 52.13 / 75.22 | 0.20 / 0.79 |
| B + AdaRound | 3 / 32 | 68.79 / 88.62 | 74.31 / 91.81 | 62.51 / 84.52 |
| B + FlexRound（本文提出） | 3 / 32 | 68.65 / 88.54 | 74.38 / 91.81 | 66.87 / 87.56 |
| B + AdaQuant | 2 / 32 | 0.23 / 0.92 | 0.10 / 0.50 | 0.10 / 0.50 |
| B + AdaRound | 2 / 32 | 61.99 / 84.81 | 48.47 / 77.09 | 39.57 / 66.18 |
| B + FlexRound（本文提出） | 2 / 32 | 62.57 / 84.84 | 63.67 / 85.72 | 46.04 / 72.48 |
表3：同时对权重和激活值进行量化时在ImageNet数据集上的top-1/top-5精度（%）。"B + X"和"Q + Y"分别表示在BRECQ和QDrop的设置下实现方案X和Y。
| 方法 | 比特数（权重/激活值） | ResNet-18 | ResNet-50 | MobileNetV2 |
|---|---|---|---|---|
| 全精度 | 32 / 32 | 71.00 / 89.97 | 76.63 / 93.04 | 72.62 / 90.67 |
| B + AdaRound | 4 / 4 | 69.18 / 88.85 | 74.44 / 91.80 | 61.05 / 83.30 |
| B + FlexRound（本文提出） | 4 / 4 | 69.32 / 88.83 | 74.56 / 91.87 | 63.74 / 85.01 |
| Q + AdaRound | 4 / 4 | 69.20 / 88.96 | 74.90 / 92.15 | 65.42 / 86.23 |
| Q + FlexRound（本文提出） | 4 / 4 | 69.26 / 88.81 | 75.08 / 92.20 | 66.66 / 87.21 |
| B + AdaRound | 3 / 3 | 64.83 / 86.12 | 67.01 / 87.28 | 3.74 / 11.54 |
| B + FlexRound（本文提出） | 3 / 3 | 64.99 / 85.93 | 68.29 / 87.89 | 25.43 / 48.28 |
| Q + AdaRound | 3 / 3 | 65.71 / 86.96 | 70.49 / 89.93 | 39.86 / 66.00 |
| Q + FlexRound（本文提出） | 3 / 3 | 65.43 / 86.60 | 70.74 / 89.78 | 51.49 / 76.90 |
4.3 语言模型实验
本文中的所有语言模型均基于Transformer（Vaswani等人，2017）架构。除非另有说明，为将这些模型的精度降至8位，权重和激活值均采用线性非对称逐张量量化方案。后训练量化的重构步骤应用于每个Transformer层，包括注意力子层和前馈子层。注意力子层和前馈子层中的所有权重均量化为8位。激活值在进入每个线性层之前即时量化为8位，而softmax层和归一化层的输入则保持全精度，这与Zafrir等人（2019）和Zhang等人（2020）的建议一致。除GPT-2和LLaMA模型的实验外，预训练语言模型（PLMs）和数据集均来自HuggingFace（Wolf等人，2020）仓库。在SQuADv1（Rajpurkar等人，2016）数据集上对微调后的BERT模型进行问答任务的实验结果详见附录G。
BERT和GPT-Neo在GLUE数据集上的实验
我们在GLUE基准数据集上对多种模型（包括BERT-Base、BERT-Large、GPT-Neo-125M、GPT-Neo-1.3B和GPT-Neo-2.7B）进行微调，以评估FlexRound在自然语言理解（NLU）任务上的性能。由于篇幅限制，仅报告在MNLI、QQP和MRPC数据集上的实验结果，完整实验结果详见附录H。为证明"Q + FlexRound"无需花费大量精力为每个任务选择最优学习率，就能在整体性能上超越"Q + AdaRound"，对于BERT模型，所有可学习参数（ s 1 s_1 s1、 S 2 \boldsymbol{S}_2 S2和 s 3 \boldsymbol{s}_3 s3）的学习率均设为2e-4；对于GPT-Neo模型，学习率均设为3e-4，不考虑具体任务差异。重构过程使用1024个随机样本，迭代20000次。最后一个随机初始化的层保持全精度。更多实验细节详见附录H。表4报告了"Q + AdaRound"和"Q + FlexRound"的性能，这两种方案在表3中已显示出较好的潜力。可以发现，在所有模型和自然语言理解任务中，"Q + FlexRound"的自然语言理解得分均高于"Q + AdaRound"。特别是在MNLI和QQP数据集上，除GPT-Neo-125M模型外，"Q + FlexRound"在逐张量均匀后训练量化设置下的性能能够与全精度模型相当，甚至更优。
GPT-Neo和OPT在WikiText2和PTB数据集上的实验
我们在WikiText2数据集和PTB数据集上对微调后的预训练语言模型（包括GPT-Neo-125M、GPT-Neo-1.3B、GPT-Neo-2.7B、OPT-125M、OPT-1.3B和OPT-2.7B）进行测试，以评估FlexRound在自然语言生成（NLG）任务上的性能。从下游任务训练数据中随机抽取128个样本，采用逐张量量化方式，分别通过AdaRound和FlexRound对微调后的预训练语言模型（用于自然语言生成任务）进行量化。更多实验设置细节详见附录I。表5的结果表明，在所有模型和数据集上，"Q + FlexRound"的性能均优于"Q + AdaRound"，这意味着FlexRound不仅在图像分类和自然语言理解任务中有效，在自然语言生成任务中同样具有良好的效果。值得注意的是，即使对于OPT模型，"Q + FlexRound"的性能也接近全精度模型。
GPT-2在WebNLG数据集上的实验
截至目前，我们对BERT、GPT-Neo和OPT模型在下游任务中均采用了全量微调的方式。然而，对于语言模型，存在多种微调技术（Houlsby等人，2019；Liu等人，2022；Hu等人，2022），这些技术所需的可训练参数比全量微调更少，且能实现更优的性能。为评估FlexRound与其他微调方法的兼容性，我们进行了将FlexRound与LoRA（Hu等人，2022）（一种最先进的微调方法）融合的GPT-2模型量化实验。重构过程从WebNLG训练集中随机选择128个样本。更多实验细节详见附录J。表6的结果显示，"Q + FlexRound"的性能优于"Q + AdaRound"，且与采用LoRA的全精度模型性能相当，甚至更优。因此，除全量微调外，FlexRound还与其他最先进的微调技术具有良好的兼容性。
LLaMA在常识推理和WikiText2数据集上的实验
最后，我们评估了LLaMA-33B模型在六个常识推理基准数据集和WikiText2数据集的因果语言建模任务上的零样本性能。旨在验证：通过逐块重构输出，无需假设激活值异常值会以特定模式出现，就能对大型语言模型进行高效量化，且相较于半精度基准模型，其精度损失微乎其微。在表7中，重构过程从C4（Raffel等人，2020）训练数据集中随机选择512个样本。权重采用线性非对称逐通道量化，激活值采用线性非对称逐张量量化。LLaMA-7B、LLaMA-13B和LLaMA-33B模型的零样本和五样本性能以及相关实验细节详见附录K。表7的结果显示，"Q + FlexRound"能够保持半精度基准模型的精度，且性能超过"Q + AdaRound"。在不对大型语言模型中激活值异常值的模式做任何假设的情况下，FlexRound仍能对大型语言模型进行量化，同时保持半精度基准模型的性能。
表4：在GLUE数据集上的性能。对于评估指标，MNLI数据集报告匹配准确率和不匹配准确率，QQP数据集报告F1分数和准确率，MRPC数据集报告准确率。"Q + X"表示在QDrop的设置下实现方案X。注意力子层和前馈子层的权重及输入激活值均采用逐张量非对称方案量化为8位。
| 数据集 | 方法 | BERT-Base | BERT-Large | GPT-Neo-125M | GPT-Neo-1.3B | GPT-Neo-2.7B |
|---|---|---|---|---|---|---|
| MNLI | 全精度 | 84.49 / 85.20 | 86.05 / 85.98 | 79.11 / 79.63 | 85.12 / 86.04 | 86.36 / 87.02 |
| Q + AdaRound | 83.69 / 84.61 | 85.75 / 85.86 | 72.67 / 74.11 | 84.90 / 85.82 | 86.33 / 86.75 | |
| Q + FlexRound（本文提出） | 84.53 / 84.98 | 85.93 / 85.99 | 72.94 / 74.24 | 85.56 / 86.14 | 86.41 / 86.89 | |
| QQP | 全精度 | 88.06 / 91.08 | 88.66 / 91.59 | 85.20 / 88.99 | 88.26 / 91.28 | 88.62 / 91.50 |
| Q + AdaRound | 87.65 / 90.58 | 87.48 / 90.62 | 72.97 / 79.35 | 87.98 / 91.04 | 88.38 / 91.27 | |
| Q + FlexRound（本文提出） | 87.81 / 90.83 | 88.38 / 91.31 | 73.75 / 80.65 | 88.27 / 91.18 | 88.60 / 91.39 | |
| MRPC | 全精度 | 85.05 | 85.54 | 80.15 | 85.05 | 87.99 |
| Q + AdaRound | 81.62 | 82.35 | 75.25 | 84.80 | 85.78 | |
| Q + FlexRound（本文提出） | 84.07 | 84.31 | 75.49 | 85.05 | 86.76 |
表5：GPT-Neo和OPT在WikiText2和PTB数据集上微调后的性能。采用困惑度（PPL）作为性能指标，困惑度越低，性能越好。"Q + X"表示在QDrop的设置下实现方案X。注意力子层和前馈子层的权重及输入激活值均采用逐张量非对称方案量化为8位。
| 数据集 | 方法 | GPT-Neo-125M | GPT-Neo-1.3B | GPT-Neo-2.7B | OPT-125M | OPT-1.3B | OPT-2.7B |
|---|---|---|---|---|---|---|---|
| WikiText2 | 全精度 | 21.96 | 12.09 | 10.78 | 19.85 | 11.52 | 10.27 |
| Q + AdaRound | 30.52 | 12.47 | 14.09 | 27.96 | 12.66 | 10.97 | |
| Q + FlexRound（本文提出） | 24.30 | 12.37 | 12.43 | 21.43 | 12.02 | 10.63 | |
| PTB | 全精度 | 24.20 | 16.09 | 14.70 | 16.50 | 11.62 | 10.80 |
| Q + AdaRound | 31.40 | 16.63 | 19.80 | 20.28 | 13.00 | 12.02 | |
| Q + FlexRound（本文提出） | 26.03 | 16.32 | 16.87 | 17.68 | 12.22 | 11.29 |
表6：GPT-2中型（M）和大型（L）模型通过LoRA在WebNLG数据集上微调后的性能。"Unseen""Seen"和"All"分别表示在WebNLG测试集中未见过、见过和所有类别的BLEU分数。BLEU分数越高，性能越好。"Q + X"表示在QDrop的设置下实现方案X。注意力子层和前馈子层的权重及输入激活值均采用逐张量非对称方案量化为8位。
| 模型 | 方法 | Unseen | Seen | All |
|---|---|---|---|---|
| GPT-2 M | 全精度（LoRA） | 47.16 | 62.31 | 55.43 |
| Q + AdaRound | 45.70 | 60.92 | 54.05 | |
| Q + FlexRound（本文提出） | 46.85 | 61.83 | 55.06 | |
| GPT-2 L | 全精度（LoRA） | 48.06 | 64.39 | 56.97 |
| Q + AdaRound | 48.09 | 63.98 | 56.75 | |
| Q + FlexRound（本文提出） | 48.42 | 64.47 | 57.16 |
表7：LLaMA-33B模型在6个常识推理任务（BoolQ、PIQA、HellaSwag、WinoGrande、ARC easy和challenge、OBQA）和WikiText2数据集的因果语言建模任务上的零样本性能。常识推理任务报告准确率（%），因果语言建模任务报告困惑度（PPL）。困惑度越低，性能越好。"Q + X"表示在QDrop的设置下实现方案X。注意力子层和前馈子层的权重采用逐通道非对称格式量化为8位，而这些子层的输入激活值采用逐张量非对称方案量化为8位。
| 模型 | 方法 | BoolQ | PIQA | HellaSwag | WinoGrande | ARC-e | ARC-c | OBQA | WikiText2 |
|---|---|---|---|---|---|---|---|---|---|
| LLaMA-33B | 半精度 | 68.38 | 80.09 | 79.21 | 72.93 | 58.92 | 45.48 | 42.00 | 6.35 |
| Q + AdaRound | 64.86 | 74.65 | 68.64 | 57.93 | 49.28 | 36.95 | 41.00 | 10.39 | |
| Q + FlexRound（本文提出） | 69.08 | 79.16 | 77.43 | 72.53 | 56.61 | 44.97 | 44.00 | 6.82 |
5. 结论
本文提出了一种新的舍入方案FlexRound，用于后训练权重量化。该方案基于逐元素除法原理，能够联合学习通用的量化网格大小和每个预训练权重的单独缩放因子。我们通过实验验证，FlexRound能够根据预训练权重自身的数值大小更新对应的缩放因子，从而对预训练权重进行灵活量化。因此，FlexRound可应用于多种模型，甚至包括大型语言模型，且能确保这些模型在量化后精度损失微乎其微。