文章目录
- Abstract
- 1 INTRODUCTION
- 2 RELATED WORK
- 3 METHODOLOGY
- 3.1 AFFINEQUANT
- 3.2 REVERSIBILITYANDGRADUALMASK
- 3.3 EFFICIENCY
- 4 EXPERIMENTS
- 4.1 实现细节
- 4.2 评估实验
- 4.3 消融研究
- 5 CONCLUSION
Abstract
大规模语言模型（LLMs）所需的显著资源引发了对压缩和加速神经网络技术的广泛兴趣。在这些技术中，训练后量化（Post-Training Quantization，PTQ）因其在训练环境中的显著压缩效率和成本效益而备受关注。现有的LLM PTQ方法将优化范围限制在量化前后权重之间的缩放变换。这种限制在低比特配置下导致了显著的量化误差。
本文提倡在PTQ中使用等效的仿射变换（AffineQuant）进行直接优化。该方法扩展了优化范围，从而显著减少了量化误差。此外，通过使用对应的逆矩阵，我们可以确保PTQ量化前后输出的等效性，从而保持其效率和泛化能力。为了在优化过程中确保变换的可逆性，我们进一步引入了一种渐进掩码优化方法。该方法最初侧重于优化对角元素，然后逐步扩展到其他元素。这种方法符合Levy-Desplanques定理，从理论上确保了变换的可逆性。
结果表明，在不同的LLMs和多样化的数据集上都显现出了显著的性能提升。特别是在使用非常低比特量化时，这些改进尤为显著，允许在边缘设备上部署大型模型。例如，在LLaMA2-7B模型的W4A4量化配置下，我们在C4数据集上取得了15.76的困惑度（相比OmniQuant的18.02下降了2.26）。在零样本任务中，AffineQuant在使用4/4比特量化的LLaMA-30B模型时达到了58.61%的平均准确率（相比OmniQuant的56.63提高了1.98%），并为LLMs的PTQ设立了新的标杆。
1 INTRODUCTION
基于Transformer的大型语言模型（LLMs）在自然语言处理任务中占据主导地位，如语言翻译和文本摘要（Zhang等，2023；Touvron等，2023；Zhang等，2022b）。视觉Transformer（VITs）则将这一理念应用于计算机视觉领域。然而，由于训练LLMs所需的巨大计算资源和高质量数据（Zhang等，2023a；c），通过模型微调实现量化具有挑战性。因此，研究界越来越重视无需训练的算法，即训练后量化（Post-Training Quantization，PTQ）（Wei等，2022c；Xiao等，2023；Wei等，2023；Lin等，2023；Shao等，2023；Yuan等，2023）。训练后量化允许使用较少的校准数据进行高效优化。然而，这一过程可能导致显著的性能下降，特别是在小型模型或低比特情况下。
等效变换在大多数PTQ方法中被广泛采用。如图1所示，AWQ（Lin等，2023）通过优化统计量来增强尺度计算，并首次引入均方误差损失作为LLMs中量化前后特征图的优化指标。最近，OmniQuant（Shao等，2023）引入了块级可学习的尺度和偏移参数，以增强优化。在更高维度上，等效量化的概念也引起了关注。RPTQ（Yuan等，2023）通过对激活值列进行排序，每个簇内进行激活量化。重新排序可以通过将尺度从向量转换为矩阵形式来数学表示，其中每行和每列对应于单个尺度值。这一变换有效地重新排列了激活列和权重行。
总之，等效量化从手动设计进化到梯度优化，从低维变换到高维变换，从单尺度合并到多操作组合，包括平移和重新排序。等效量化在两个主要方面具有优势。首先，通过确保量化前后输出的一致性，可通过优化等效变换参数有效减轻引入的量化噪声。这与训练后量化的概念一致，等效变换作为噪声改进的中介。其次，不同类型的等效变换彼此正交。直观上，引入每种新类型的等效变换都会扩展参数优化空间，从而提升性能。
因此，我们提出了一种用于等效仿射变换的算法。具体而言，我们将仿射变换矩阵左乘线性层的权重，并将其逆矩阵右乘于激活值。在均方误差损失的指导下，我们优化仿射变换矩阵，优化过程中损失始终低于其他算法。此外，我们在优化过程中探索了矩阵的可逆性。Levy-Desplanques定理（Naimark & Zeheb，1997）表明，严格对角占优矩阵是可逆的。为确保仿射变换矩阵严格对角占优，我们采用对角初始化和渐进掩码方法。通过这种方式，我们在有限校准数据下优化高维矩阵。
在推理效率方面，我们的方法在矩阵合并后与其他方法一致。最终，我们的方法在LLMs量化中取得了最先进的性能，特别是在小型模型或低比特配置的场景中。总体而言，我们的贡献总结如下：
- 我们提出了一种新颖的PTQ仿射变换方法，保留了PTQ的优势，确保了效率和泛化能力，显著减少了量化误差，特别是在低比特量化下，并支持在边缘设备上部署LLMs。
- 我们提出了一种保证过程可逆性的新颖优化算法，利用Levy-Desplanques定理，同时降低了计算成本。
- 我们的方法在大型语言模型量化中取得了最先进的性能，特别是在低比特或小型模型上。在LLaMA2-7B的w4a4配置下，我们在C4数据集上的困惑度为15.76（相比OmniQuant的18.02下降了2.26）。同样，在LLaMA-30B的w4a4配置下，我们在6个零样本任务中的平均准确率为58.61%（相比OmniQuant的56.63提高了1.98%）。
2 RELATED WORK
量化可以根据算法效率和数据需求分为两大类：量化感知训练（Quantization-Aware Training，QAT）和训练后量化（Post-Training Quantization，PTQ）。QAT（Bondarenko等，2023；Choi等，2018；Yao等，2021；Gong等，2019；Wang等，2019；Esser等，2019；Sun等，2020；Lee等，2021）需要大量数据来微调模型权重，这使得支持LLMs（Zhang等，2023b）变得具有挑战性。相比之下，我们在本研究中重点关注PTQ（Nagel等，2020；Li等，2021；Wei等，2022b；Ma等，2023；Liu等，2021；Cai等，2020；Cheng等，2023），因为其算法效率较高。
在卷积神经网络（CNN）领域，PTQ主要集中在优化权重舍入策略。Adaround（Nagel等，2020）通过优化权重舍入策略改进量化模型，考虑向上或向下舍入。BRECQ（Li等，2021）引入了块级优化过程，并结合了平方梯度信息。QDROP（Wei等，2022b）通过随机丢弃量化激活值来提升量化模型的性能。
大型语言模型量化。为了应对计算资源限制，研究社区集中于大型语言模型（LLMs）的高效量化算法。LLMs量化可以根据是否量化激活值分类为仅权重量化（Frantar & Alistarh，2022；Lin等，2023；Shao等，2023；Kim等，2023）和权重-激活量化（Xiao等，2023；Wei等，2023；Yao等，2022；Dettmers等，2022；Shao等，2023）。鉴于LLMs的巨大规模，内存访问效率成为加速的主要瓶颈。仅权重量化通过将模型权重压缩到较低比特精度，有效缓解了内存壁垒问题（Kim等，2023）。
3 METHODOLOGY
在本节中，我们介绍了AffineQuant，这是一种利用等效仿射变换进行量化的方法。与其他方法相比，AffineQuant在优化过程中始终保持最佳均方误差。我们还探讨了仿射变换矩阵在优化过程中的可逆性。为了确保稳定性，我们提出了一种基于Levy-Desplanques定理（Naimark & Zeheb，1997）的渐进掩码方法，以保持仿射变换矩阵为严格对角占优矩阵。最后，我们分析了经AffineQuant优化后的大型语言模型（LLMs）的推理效率。
3.1 AFFINEQUANT
从物理角度考虑等效变换的概念，我们可以类比某些操作。例如，在SmoothQuant（Xiao等，2023）中，我们可以将尺度类比为向量的缩放操作；在Outlier Suppression+（Wei等，2023）中，我们可以将平移类比为向量的平移操作。同样，向量的旋转也可以归类为等效变换。
我们定义伪量化函数如下：
[ Q(x) = \Delta \cdot \left\lfloor \text{clamp}\left(\frac{x}{\Delta} + z_p, 0, 2^n - 1\right) \right\rfloor - z_p, ]
其中 (\Delta)、(z_p) 和 (n) 分别是量化步长、零点和位数。 (\lfloor \cdot \rfloor) 是舍入操作。 如图1所示，AffineQuant通过将仿射变换矩阵 (A) 左乘于权重矩阵 (W) 来更好地对齐权重分布与量化函数 (Q(\cdot))。扩展的优化空间使得在变换后的权重中实现更小的量化误差，从而减少困惑度。同时，我们将仿射变换矩阵 (A) 的逆矩阵右乘于激活值 (X) ，以保持激活值和权重之间矩阵乘积输出的不变性。对于单个线性层，AffineQuant构建了以下优化问题：
[ \underset{A}{\text{argmin}} | XW - XA^{-1}Q(AW) |_F^2 ]
AffineQuant结合了AWQ（Lin等，2023）和SmoothQuant（Xiao等，2023）的精髓，当矩阵 (A) 的主对角元素从权重和激活统计量计算得出时。它与OmniQuant（Shao等，2023）一致，仅更新 (A) 的对角元素。RPTQ（Yuan等，2023）使用的重新排序矩阵是仿射变换矩阵 (A) 的子集，当 (A) 的每一行和每一列仅包含一个元素1时。总而言之，AffineQuant包含了各种以前的等效量化算法，从而扩展了权重分布 (W) 的优化可能性。
在图1中，假设权重矩阵 (W \in \mathbb{R}^{2 \times 2}) 有2个输出通道和输入通道。缩放因子、平移因子和仿射变换矩阵分别表示为 (s)、(b) 和 (A)。我们根据输出通道将权重矩阵划分为2个向量 ({v_1, v_2})。缩放变换 (s_i \cdot v_i) 均匀缩放 (v_i) 的每个元素。平移变换 (v_i + b_i) 沿不同轴平移 (v_i)。仿射变换 (A v_i) 允许 (v_i) 的任意重新定位。然而，缩放和平移变换在将 (v_i) 的维度映射到相邻的量化固定点方面有限。相比之下，仿射变换保证了向量中所有维度与量化函数 (Q(x)) 引入的噪声对齐，从而减少量化误差。值得注意的是，通过对仿射变换矩阵按行归一化 (A \rightarrow s’ A’)，其中矩阵 (A’) 的每一行的范数为1， (A’) 转化为标准旋转矩阵。这个旋转矩阵旋转权重的输出通道，同时保持其幅度。缩放因子 (s’) 对旋转后的向量进行缩放。因此，仿射变换矩阵 (A) 结合了缩放和旋转等效变换，并且与平移变换正交。
困惑度（ppl）与交叉熵（CE）损失呈指数关系，后者与量化前后输出激活的均方误差呈正相关（Nagel等，2020；Li等，2021）。因此，通过优化量化前后的均方误差可以实现困惑度的优化。具体而言，
[ \text{PPL} \propto L_\text{CE} \propto | XW - XA^{-1}Q(AW) |_F^2 ]
在大型语言模型量化中，AffineQuant的优化目标如下，
[ \underset{A, \delta}{\text{argmin}} | f_i (X,W) - f_i \left((X - \delta)A^{-1}, Q(AW), b + \delta W\right) |_F^2, ]
其中 (f_i) 是第 (i) 个Transformer块。 ((X - \delta)A^{-1})、(Q(AW)) 和 (b + \delta W) 分别是等效变换后的激活、权重和偏置。我们结合了仿射和平移变换，并将量化前后Transformer块输出的均方误差作为优化目标。
值得注意的是，AffineQuant在初始损失方面表现优于OmniQuant（Shao等，2023），这是由于前几个块中仿射变换矩阵的优异性能。此外，AffineQuant展示了在最后一个块中的更快的损失收敛和更优的总体优化性能。这些结果再次确认了可逆矩阵优化的巨大潜力。附录A.4中的图5和图6展示了多个稳定性因子（alpha）的随机采样，这些因子对量化损失收敛产生影响。数据揭示了最后一个Transformer块的量化损失与量化模型性能之间的显著联系，这表明AffineQuant在减少量化误差方面的有效性。
3.2 REVERSIBILITYANDGRADUALMASK
在优化过程中，有必要对仿射变换矩阵进行求逆。然而，我们在目标函数（公式4）中并未包含任何约束，以确保矩阵保持满秩或良好条件。因此，如何在优化过程中保持矩阵的可逆性？首先，我们定义严格对角占优矩阵如下：
定义1（严格对角占优矩阵）：当一个矩阵 (A) 的每个对角元素的绝对值大于对应行中其他元素绝对值之和时，该矩阵被认为是严格对角占优的。具体地，
[ |a_{ii}| > \sum_{i \neq j} |a_{ij}|, \quad \forall i. ]
Levy-Desplanques定理（Naimark & Zeheb, 1997）表明，所有严格对角占优的矩阵都是可逆的。通过初始化仿射变换矩阵的对角元素，我们确保其初始为严格对角占优的。虽然在优化器中使用二阶动量和较低的学习率可以帮助满足Levy-Desplanques定理的要求，但随着模型规模的增加，优化大仿射变换矩阵仍面临不稳定性挑战。
为了确保仿射变换矩阵在优化过程中保持严格对角占优，我们引入了一种渐进掩码方法，如图2所示。在每个优化块的开始，我们冻结除主对角线以外的所有元素。随着优化的进行，我们逐步解冻接近主对角线的元素。最终，所有矩阵元素都可以进行优化。这种冻结机制，称为渐进掩码（Gradual Mask, GM），定义如下：
[ GM_{ij} = \begin{cases}
1 & \text{如果 } i = j, \
\alpha & \text{如果 } 0 < |i-j| \leq e \times \text{hiddensize}, \
0 & \text{否则},
\end{cases} ]
其中，(GM_{ij}) 是掩码矩阵的第 (i) 行、第 (j) 列元素，(t) 是目标训练周期，(e \in [1,t]) 是当前训练周期，“hidden size” 是仿射变换矩阵的维度，(\alpha) 是稳定因子。在注意力模块内，我们在每个注意力头中应用渐进掩码。GM 是一个学习率调节器，通过与矩阵 (A) 元素逐个点积来实现其目的。具体而言，GM矩阵对优化过程的影响可分为两个方面。以下是结合GM后的矩阵 (A) 优化过程：
前向传播：
[ A^*_e = A_e \circ GM_e, ]
反向传播：
[ A_{e+1} = A_e + \eta \frac{\partial L}{\partial A^_e} \cdot \frac{\partial A^_e}{\partial A_e}, ]
[ = A_e + \eta GM_e \cdot \frac{\partial L}{\partial A^*_e}. ]
其中，(\circ) 是Hadamard积，(A_e) 和 (GM_e) 分别是第 (e) 训练周期的矩阵 (A) 和渐进掩码（GM）矩阵，(\eta) 是矩阵 (A) 的学习率，(L) 是优化损失。GM矩阵在前向传播过程中有效地减少了矩阵 (A) 中非主对角线元素的幅度，当稳定因子 (\alpha) 小于1时，这确保了在优化过程中的第 (e) 训练周期内 (A^*) 的稳定逆矩阵的存在，根据Levy-Desplanques定理。在反向传播中，GM影响学习率 (\eta)，从而抑制了矩阵 (A) 中非主对角线元素的更新速率。因此，GM对 (\eta) 的影响确保了第 (e+1) 训练周期内的矩阵 (A) 保持严格对角占优，满足Levy-Desplanques定理。
值得注意的是，当 (\alpha) 接近0时，优化过程稳定收敛，并且变得等效于OmniQuant（Shao等，2023）。此外，附录A.2中包含一个定理，证明了足够小的稳定因子 (\alpha) 确保了优化过程中矩阵 (A) 的严格对角占优。
渐进适应的概念在视觉模型的训练后量化中也存在。在Adaround（Nagel等，2020）中，使用逐渐增加的 (\beta) 值在软量化函数中控制参数的梯度更新。当 (\beta) 足够大时，由于梯度的限制，仅更新接近0或1的值。随着优化的进行，所有舍入值的梯度逐渐释放。然而，需要注意的是，两种方法的目标和方法是不同的。Adaround（Nagel等，2020）采用逐渐增加的 (\beta) 值以防止目标函数的快速收敛，这可能导致次优优化。另一方面，AffineQuant中的渐进掩码确保了仿射变换矩阵的严格对角占优。附录A.6展示了在不同训练周期内的各个块的仿射变换矩阵的热图，证明了渐进掩码方法在保持严格对角占优矩阵方面的有效性。
3.3 EFFICIENCY
优化效率。PyTorch的线性代数库（Paszke等，2019）提供了浮点和双精度的矩阵求逆计算。因此，我们在整个优化过程中保持模型的精度为浮点或双精度。此外，由于计算机数值精度的限制，矩阵求逆的近似计算可能包含错误。因此，我们在第4.3节分析了两种精度类型的内存消耗、优化时间、误差幅度及其对模型性能的影响。
推理效率。与类似算法一致，我们将仿射变换矩阵与其他层集成。随后，我们在网络上执行半精度推理。对于所有线性层，我们将仿射变换矩阵与权重和偏置参数合并。此外，在LayerNorm后的权重-激活量化中，我们仅优化仿射矩阵的对角元素。这使我们能够将仿射矩阵与LayerNorm的权重和偏置合并。因此，AffineQuant可以实现而不对模型推理引入任何额外开销。表2和表3展示了在零样本和困惑度任务中，AffineQuant在使用4/4位量化时，即使没有额外开销，也表现出优于其他方法的性能。
4 EXPERIMENTS
4.1 实现细节
算法细节：在AffineQuant中，随着模型规模的增加、量化位数的减少和分组大小的增加，稳定因子 (\alpha) 逐渐减少。对于OPT-6.7B及更小的模型，我们设置 (\alpha = 1)。随着模型规模的增加，对于权重量化为3位或以上的配置，我们使用 (\alpha = 1e-2)。对于其他配置，我们从集合 ({1e-2, 1e-3, 1e-4}) 中选择 (\alpha)。然后，我们在MLP模块中的两个线性层之间排除仿射变换。这是因为在膨胀维度中优化大变换矩阵具有挑战性。此外，激活函数的存在使得矩阵的等效变换无效。
4.2 评估实验
如表1和表3所示，我们观察到在各种量化配置下，所有模型的一致性能提升。这表明AffineQuant并不依赖于特定的量化配置。值得注意的是，AffineQuant在低位量化或较小模型规模的情况下表现出显著的改进。具体而言，在OPT-125M模型的w3a16g128配置中，我们实现了困惑度（PPL）减少5.10，显著超越了OmniQuant（Shao等，2023）的表现。此外，在LLaMA2-7B模型上的C4和WikiText2数据集的w4a4量化配置下，我们分别实现了困惑度减少2.26和1.57。上述结果强调了在具有挑战性的量化任务中扩展等效因子的优化空间的重要性。有关其他数据集的结果，请参阅附录。
4.3 消融研究
数值精度的影响：我们比较了不同精度方案的各种指标，包括合并错误、困惑度（PPL）、内存使用和优化运行时间。具体而言，在“float-double”方案中，我们将激活值或权重转换为双精度，与变换矩阵相乘，然后截断为浮点精度。对于合并错误，我们定义了输入和输出通道均设置为4,096的两个线性层。我们随机采样仿射变换矩阵 (A \in \mathbb{R}^{4096 \times 4096}) 和输入激活 (X \in \mathbb{R}^{2048 \times 4096})。我们进行1,000次运行，以计算不同精度方案下线性层输出在合并前后的均方误差平均值。结果见表4。值得注意的是，在双精度优化的情况下，矩阵求逆的计算错误被最小化。尽管相比其他方案，时间和内存使用更高，但较小的合并错误导致了困惑度的轻微改进。然而，对于更大规模的模型，这种改进并不显著。
稳定因子的效果：在表5中，我们调整了公式6中的稳定因子 (\alpha)。随着 (\alpha) 接近0，仿射变换理论上收敛到尺度变换。我们观察到，随着 (\alpha) 的减少，OPT-125M和LLaMA-7B模型的性能收敛到OmniQuant（Shao等，2023）。然而，在LLaMA-7B的情况下，较大的稳定因子不能保证变换矩阵的严格对角占优，这可能导致训练崩溃。因此，有必要在确保Levy-Desplanques定理（Naimark & Zeheb, 1997）有效的情况下增大 (\alpha)。
渐进掩码的贡献：在公式6中，我们逐渐释放接近对角线的掩码矩阵元素。在表6中，当我们移除渐进方法时，OPT-125M和LLaMA-7B模型表现不佳或无法完成训练。这表明在一开始就更新仿射变换矩阵的所有参数不利于保持其可逆性或良好条件性。渐进掩码方法提供了一种稳定的方法来优化大矩阵。
5 CONCLUSION
基于等效性的训练后量化显示出显著的潜力。然而，先前的等效方法限制了可优化的权重空间，导致在变换后的权重分布中出现较大的量化误差。这一限制在小模型和低位量化的情况下变得更加明显。仿射变换方法通过显著扩展可优化的权重空间来解决这一问题。先前的变换方法可以看作是仿射变换的一种特殊情况或者与之正交的情况。此外，Levy-Desplanques定理为在高维空间中保持矩阵优化的稳定性提供了理论基础。根据这一定理，我们提出的渐进掩码方法确保了在优化过程中矩阵始终严格对角占优。这保证了矩阵的可逆性或良好条件性，并进一步减少了均方误差目标函数。我们的方法在所有模型的各种量化配置下持续提高了性能。值得注意的是，仿射变换方法在提高性能方面表现出巨大的潜力，特别是对于小模型和低位配置。在未来，更有效地优化仿射变换矩阵值得仔细考虑。
