DuQuant:通过正交变换分散 LLM 离群值,实现 SOTA 4bit 量化
一、背景
本文中我们介绍一下最近被 NeurIPS 2024 接收为 Oral Presentation 的 LLM 量化工作 DuQuant。这是一种令旋转变换和排列变换来更有效地缓解 Massive Outlier 和 Normal Outlier 的新方法。其达到了 4bit 量化的新 SOTA。具体来说:
- 作者发现 LLM 的 FFN 模块中,在 down_proj 的输入中存在明显的Massive Outlier;这种 Outlier 不同于以往发现的 Outlier,表现为绝对值大于几百的个别异常值;这些 Outlier 会显著限制 4bit 权重激活量化方法的性能。
- 作者提出了基于先验知识(Outlier Channel)的旋转变换来将Massive Outlier旋转至临近 Channel,考虑效率平衡,作者使用分块旋转矩阵。而分块旋转矩阵会导致部分 Block 的平均激活值的绝对值大于其他 Block,所以作者进一步提出通道置换变换,重新平衡不同 Block 之间的 Outlier,然后使用第二次旋转变换进一步最小化 Outlier。
- DuQuant 的训练非常快速,可以在 50s 内完成 7B 模型的量化过程,即插即用。
对应的 Paper 为:[2406.01721] DuQuant: Distributing Outliers via Dual Transformation Makes Stronger Quantized LLMs [1]
对应的代码库为:https://github.com/Hsu1023/DuQuant [2]
二、 引言
2.1 量化方法
在每个 Transformer Block 中,多头自注意力(MSA)和前馈网络(FFN)基本上都由线性层组成,将其表示为:,其中, 是激活输入, 表示权重矩阵。
模型量化通过将浮点数(FP16, BF16)表示的模型权重或激活值转换为低比特的浮点数或整数(INT8、FP8、INT4)来减少内存占用,权重激活量化可以通过使用低比特乘法的算子进一步加快模型推理的速度。本文中,作者重点关注低比特整型量化,旨在获得更好的硬件支持。具体来说, 位量化过程将 FP16 张量 映射为低位整数:
符号 表示最接近的舍入操作, 是量化步长, 表示零点。遵循主流量化方法,作者对激活 采用逐 Token 量化,对权重 采用逐 Channel 量化,这意味着为 的每个 Token 分配不同的步长(),为 的每个输出 Channel 分配不同的步长()。
2.2 动机
作者首次在实验中发现,LLM FFN 模块中的 down_proj layer 存在明显的 [2402.17762] Massive Activations in Large Language Models [3],表现为绝对值大于几百的异常值并局限于个别的 Token 中。这些 Massvie Outlier 造成 [2211.10438] SmoothQuant: Accurate and Efficient Post-Training Quantization for Large Language Models [4] 和 [2308.13137] OmniQuant: Omnidirectionally Calibrated Quantization for Large Language Models [5] 等算法在 4bit WA 量化中表现较差。在 down_proj layer 中发现 Massive Outlier 现象为该工作首次提出。
如下图 Figure 1 (a)(b) 对比了常见的 Normal Outlier,和在 FFN 中出现的 Massive Outlier。
SmoothQuant 中将激活值除以每 Channel 平滑因子并将其乘回权重矩阵,尝试将量化难度从激活值转移到权重。具体来说,SmoothQuant 使用逐 Channel 的平滑对角矩阵,记为 ,将原始线性层重新表达为:,对角矩阵 中的元素 的计算方式为 ,其中 是一个超参数,表示迁移强度。
然而,作者观察到在输入端进行这种转移可能会导致权重矩阵也出现明显的难以量化的 Outlier(如上图 Figure 1 (d) 所示),这一问题的根源在于 Massive Outlier 使平滑因子 变得异常大。
此外,极大的 Outlier 还可能导致基于优化的方法出现梯度爆炸的问题,所以基于梯度优化的 OmniQuant 和 [2403.12544] AffineQuant: Affine Transformation Quantization for Large Language Models [6] 等工作会直接跳过 down_proj layer,直接退化为 SmoothQuant 的处理方式。这都表明需要一种更新、更好的方式处理两种 Outlier,特别是平滑掉 down_proj 输入侧的 Massive outlier。
三、本文方案
3.1 概览
DuQuant 提出通过学习旋转和置换变换矩阵,在 Activation 矩阵内部将 Outlier 转移到其他通道,最终得到平滑的激活矩阵,从而大幅度降低量化难度。如下图 Figure 2 所示为 DuQuant 算法说明:
- a. 逐步展示了 DuQuant 算法对 Normal Outlier 的处理过程。
- b. DuQuant 显著降低了 Massive Outlier。
- c. 一个 Example 说明 DuQuant 有效降低了激活矩阵的量化难度。
简单来说,DuQuant 算法包含三个步骤:
- 旋转矩阵的构造有效利用了特定 Outlier Channel 的位置索引,作者使用分块对角的旋转矩阵,在每个 Block 内部通过贪心算法将 Outlier 平摊到其他的 Channel 中。
- 由于 Block Size 的限制,可能导致某些 Block 在旋转之后组内的平均值大于其他 Block,因此作者进一步使用通道置换技术重新分配 Activation Channel,利用 zigzag 顺序排列使各组均值的方差大幅减小。
- 进一步执行一次旋转变换达到更均匀的 Activation 分布,从而大幅降低量化难度。
3.2 旋转矩阵
作者希望应用旋转矩阵 进行行或列变换,减轻 Normal Outlier 和 Massive Outlier 的影响。由于 Massive Outlier 通常随机分布在激活空间中,直接找到能够通过单次旋转变换减轻 Outlier 的最优旋转矩阵 是具有挑战性的。为了解决这一问题,作者采用带有先验知识的贪心搜索方法来计算旋转矩阵 ,从而近似理想的旋转矩阵 。具体来说, 的计算包括以下步骤:
- 识别 Outlier 主要集中的特征维度 ,即:, 其中, 表示 中第 行和第 列的元素。
- 基于搜索到的维度 ,构建旋转矩阵如下: 是用于交换激活值的第 1 列和第 列的交换矩阵, 表示一个正交初始化的旋转矩阵,其第 1 行均匀分布。这样做的目的是通过 变换后减轻第 1 列中的 Outlier。为了进一步增加随机性,作者保留了减轻 Outlier 后的第 1 列,并通过与随机正交矩阵 相乘,随机旋转其他列。
- 设 为贪心搜索的步数,则近似的旋转矩阵为 ,其中 。每个 根据上述公式和识别到的特征维度 构建。
通过这种构建方式,可以确保近似的最优旋转矩阵 能够有效减轻具有较大幅度 Outlier,而不仅仅是使用随机选择的正交旋转矩阵。然而,直接构建整个旋转矩阵非常耗时,并且会导致大量的内存开销。为了实现快速矩阵乘法,参考 [2306.11987] Training Transformer with 4ibts [6],作者选择以分块的方式近似旋转矩阵 ,
其中, 表示第 个块的方阵,该矩阵按照上述三步构建。块的数量 通过 计算得出。
3.3 通道置换矩阵
尽管采用了块对角旋转矩阵 来提高时间和存储效率,但其专注于局部信息的特性带来了进一步减少 Outlier 的潜在限制。由于在每个小块内进行的旋转变换无法整合跨不同块的信息,一个块中可能有相对较大的 Outlier,而另一个块中则有较小的 Outlier,导致不同块之间存在较高的方差。因此,作者提出利用通道置换矩阵平衡不同 Block 之间 Outlier 的幅度。
具体来说,在每个小块中,将维度 中最大的 Outlier 记为 。同时, 表示第 个块中所有 的平均值,其中 ,各个块之间激活幅度的方差可以表示为:
作者引入了之字形置换矩阵 。具体来说,通过生成一个之字形序列,首先将激活值最高的通道分配给第一个块,接着,将激活值次高的通道按递减顺序分配到后续的块,直到第 个块。在到达最后一个块后,顺序反转,从下一个激活值最高的通道开始,按递增顺序分配。这个往复模式贯穿所有块,确保没有单个块持续接收最高或最低激活值的通道。通过使用之字形置换,DuQuant 实现了不同块之间 Outlier 的均衡分布,从而能够使用额外的旋转变换进一步平滑 Outlier,如图 Figure 2 所示。
需要说明的是:
- 通道置换其实是非常重要的一步,也很简单快速(对于推理效率的影响很小,可见后面实验部分),即可以避免像 [2405.16406] SpinQuant: LLM quantization with learned rotations [7] 那样复杂的训练流程,也比 [2404.00456] QuaRot: Outlier-Free 4-Bit Inference in Rotated LLMs [8] 的 Hadamard 旋转性能表现更好。
- 旋转矩阵和置换变换矩阵都是正交矩阵,保证了 输出的不变性,作者还通过严谨的理论推导证明了两种变换有效降低了量化误差,具体证明可以阅读 Paper 里面的 Appendix。
四、实验和结果
4.1 主要结果
作者基于提出的方案验证了 LLaMA、Vicuna、Mistral 系列模型,其在 PPL、QA、MMLU 和 MT-Bench 等任务上都明显提升了量化模型的性能。
如下图 Table 1 所示,相比几种常见的 W4A4 量化方案,DuQuant 在 LLaMA-1 和 LLaMA-2 模型上的量化达到了 SOTA 效果:
如下图 Table 4 所示,DuQuant 在 LLaMA3 模型的量化效果,4-bit 下同样达到了 SOTA:
如下图 Table 3 所示,DuQuant 在 Vicuna-v1.5-13B 模型上的量化效果,0-shot 和 5-shot 任务同样都打到了 SOTA:
4.2 推理速度
作者在 NVIDIA RTX 3090 GPU 上进行了推理速度测试,其 W4A4 Kernel 来自 Quarot。如下图 Table 9 所示,结果表明 DuQuant 在 Prefiling 阶段相比 FP16 加速比可以达到 2.08x(LLaMA2-7B)、2.34x(LLaMA2-13B)。同时,如下图 Figure 5 所示,DuQuant 相比 INT4 推理速度会降低 10% 左右,略微高于 QuaRot,但带来更多精度提升。
4.3 与 QuaRot 对比
DuQuant 与使用 Hadamard 旋转矩阵的 QuaRot 主要有以下两点不同:
1. DuQuant 构造的旋转矩阵利用了先验知识(具体的 Outlier Channel 索引),因此 DuQuant的 旋转矩阵可以比 QuaRot 更好地平滑激活空间,如下图 Figure 4 展示了 DuQuant 单次旋转和 Hadamard 旋转对 LLaMA2-7B Attention Key_proj 输入的变换效果。
2. QuaRot 依赖于耗时的 GPTQ 算法来提升性能,而作者引入的通道置换矩阵可以帮助 DuQuant 在极短时间内进一步平衡 Outlier 的分布,从而取得了更好的量化效果,如下图 Table 7 所示。
五、参考链接
1. https://arxiv.org/abs/2406.01721
2. https://github.com/Hsu1023/DuQuant
3. https://arxiv.org/abs/2402.17762
4. https://arxiv.org/abs/2211.10438
5. https://arxiv.org/abs/2308.13137
6. https://arxiv.org/abs/2306.11987
7. https://arxiv.org/abs/2405.16406
8. https://arxiv.org/abs/2404.00456