LLM旋转量化笔记

LLM 旋转量化笔记

背景：随着 LLM 尺寸的爆发，4-bit 量化（W4A4）成为了离线与端侧推理的刚需。然而由于激活中离群点（Outliers）的存在，传统量化方法在 4-bit 激活量化时会产生极大的精度退化。基于正交旋转的方法（如 QuaRot、SpinQuant、ConvRot）是攻克 4-bit 量化的核心技术路线。

1. 为什么 LLM 量化这么难？

1.1 离群点（Outliers）是头号敌人

在 LLaMA、OPT 等主流 LLM 的激活（Activations）中，少数几个特定的通道（Channel）数值会系统性地远大于其他通道（幅度差达到几十甚至上百倍）。这种现象特别集中在注意力机制的输出、FFN 的 gate/up 投影，以及 down_proj 层的前一激活层。

1.2 为什么离群点导致 4-bit 量化崩溃？

以最常用的对称均匀量化为例，其量化公式为：

$$Q(x) = \text{clip}\left(\left\lfloor \frac{x}{s} \right\rceil, -q_{\max}, q_{\max}\right)$$

其中 $s$ 为缩放因子（Scaling Factor），$\lfloor \cdot \rceil$ 代表四舍五入取整。

  • 量化误差在 $\frac{s}{2}$ 左右。
  • 为降低硬件乘加指令的复杂度，矩阵乘法在进行累加（K 维度）时，通常需要在一个 Tensor 内共享同一个 $s$（Per-tensor 量化），或者在一行（Per-token）内共享 $s$。
  • 致命问题：若某一维度出现极大的离群点 $x_{\text{out}}$，为了防止该通道溢出，缩放因子被迫取得非常大（$s = \frac{|x_{\text{out}}|}{q_{\max}}$）。这将导致占绝大多数的正常通道值 $x_{\text{normal}}$ 在被 $s$ 除之后落入 $[-0.5, 0.5]$ 的区间，四舍五入后全部被舍入为 0，有效信息彻底丢失。

1.3 历史应对路线的局限性

方法核心思路局限性
LLM.int8() (2022)动态检测离群通道，保留 FP16 运行，其余通道用 INT8 计算引入混合精度分支，硬件执行不友好，不便于量化加速
SmoothQuant (2023)利用等价缩放，将激活部分的量化难度等比迁移到权重上本质是启发式搜索参数，在 4-bit 量化下容易使权重崩溃
AWQ (2023)保护 1% 重要的权重（保留高精度），其余进行低比特量化主要是 Weight-Only 量化，对激活离群点（Activation Outliers）无能为力
GPTQ (2023)基于 Hessian 矩阵引导权重更新，补偿量化误差同样偏向 Weight-Only 优化，激活离群问题依然存在
Rotation-based (2024+)引入正交旋转，在不改变物理意义的前提下，从源头打散激活离群点运行时有少量 online 旋转开销（如 Q/K 的在线旋转）

2. 旋转量化的核心思想（等价变换）

2.1 数学原理

对于一个线性层 $Y = X W^T$，假设引入一个正交变换矩阵 $R$（满足 $R R^T = I$），则有：

$$Y = X W^T = (X R) (R^T W^T) = \tilde{X} \tilde{W}^T$$
  1. 激活值旋转：$\tilde{X} = X R$（在线计算，负责将离群点“平摊”到所有通道中）。
  2. 权重旋转：$\tilde{W} = W R$（离线计算，不带任何运行时计算开销）。

由于 $R R^T = I$，在数学上这个等价变换是完全恒等的。我们可以在不修改任何权重训练结果的前提下，通过对激活和权重做对应的旋转，消除激活中的高振幅离群通道。

2.2 为什么正交旋转能平滑分布

正交变换具有能量守恒的特性（即保持向量的 $L_2$ 范数不变）：

$$\|\tilde{x}\|_2 = \|x R\|_2 = \|x\|_2$$

如果输入向量为 $x = [0.1, 0.1, \dots, 0.1, 10.0]$（最后一个通道存在大小为 10.0 的离群点，其余为 0.1）。 通过归一化 Hadamard 矩阵 $R = \frac{1}{\sqrt{K}} H_K$ 旋转后，每一个输出位置都会被均匀地乘以 $\pm \frac{1}{\sqrt{K}}$。 旋转后的每个通道值将收敛在 $\approx \frac{10.0}{\sqrt{K}}$。 若通道维度 $K = 4096$，则最大值被摊薄到 $\frac{10.0}{64} \approx 0.156$。原本陡峭的通道能量被平缓摊开，从而极大地降低了量化难度。

3. QuaRot（NeurIPS 2024）详解

第一个在 LLaMA-2 70B 等模型上做到 W4A4 KV4 量化基本无损的工作。

3.1 旋转在 Transformer 块中的平摊与吸收（Absorb）

为了不增加运行时的计算开销，我们希望绝大多数的旋转矩阵 $R$ 能够直接被相邻线性层的权重“吸收”：

上一层输出 X_rotated = X · H
       ↓
RMSNorm (保持旋转透明，见 3.4)
       ↓
Q/K/V 投影 (权重吸收旋转：W_new = H^T · W)
       ↓
计算 Q, K, V
       ↓
Attention 算子 (Q @ K^T, 旋转自然抵消，见 3.2)
       ↓
O (Output) 投影 (权重吸收旋转：W_new = W · H) → 输出重新被旋转：Y · H
       ↓
FFN up/gate 投影 (权重吸收旋转：W_new = H^T · W)
       ↓
FFN down 投影 (权重吸收旋转：W_new = W · H)

通过这一传导链，相邻线性层之间形成了完美的对接：上一个线性层输出端产生的旋转，直接被下一个线性层的输入端吸收。在整个推理链条上，只有极少数地方必须在运行时（online）实时进行旋转计算。

3.2 Attention 模块的在线旋转与抵消机制

在 Attention 结构中，由于有 RoPE（旋转位置编码）以及 Key-Value 缓存，旋转不能被轻易离线吸收。QuaRot 采用了一种极其精妙的设计：

3.2.1 Query 与 Key 的在线旋转

  1. 基本原理：对于 Q 与 K，旋转发生在 RoPE 之后，在存入 KV Cache 前在线乘以每个 Head 尺度的 Hadamard 矩阵 $H_{\text{head}}$： $$\tilde{Q} = \text{RoPE}(Q) \cdot H_{\text{head}}, \quad \tilde{K} = \text{RoPE}(K) \cdot H_{\text{head}}$$
  2. 抵消机制：当进行 QK 矩阵乘法计算注意力权重时： $$\tilde{Q} \tilde{K}^T = \left(\text{RoPE}(Q) \cdot H_{\text{head}}\right) \left(\text{RoPE}(K) \cdot H_{\text{head}}\right)^T = \text{RoPE}(Q) \left(H_{\text{head}} H_{\text{head}}^T\right) \text{RoPE}(K)^T$$ 由于 $H_{\text{head}}$ 的正交性（即 $H_{\text{head}} H_{\text{head}}^T = I$），旋转在计算矩阵乘时自动抵消，完全不影响 Attention Matrix 的数值！但保存在 KV Cache 中的 $\tilde{K}$ 通道却被彻底平滑，得以便于 4-bit 量化。

3.2.2 Value 的在线旋转

  1. 对于 Value，在线旋转作用于 Value 本身：$\tilde{V} = V \cdot H_{\text{head}}$，从而消除了 Value Cache 里的离群值。
  2. 注意力机制输出为 $A \tilde{V} = A (V \cdot H_{\text{head}}) = (A V) \cdot H_{\text{head}}$。
  3. 这个附带在 Attention 输出上的 $H_{\text{head}}$，可以直接通过转置并离线融合进后续的 Output Projection 权重中（即 $W_O \leftarrow W_O H_{\text{head}}$）。在运行时完全不需要对 Value 的输出再作逆旋转。

3.3 必须 online 计算的旋转

由于以上精妙的吸收和抵消设计，QuaRot 的运行时额外开销极低：

  • FFN、Projection 层：运行时 0 额外开销（全部被离线融合进权重中）。
  • Attention 模块：仅需对 Query 和 Key 进行在线 FWHT（Fast Walsh-Hadamard Transform），维度仅为 head_dim（通常为 128），计算开销微乎其微。

3.4 RMSNorm 为何不影响旋转传递？

能够实现全流程等价融合的关键在于 RMSNorm 是尺度不变且对正交变换“透明”的。

RMSNorm 计算公式为：

$$\text{RMSNorm}(x) = \frac{x}{\text{RMS}(x)} \odot \gamma$$

其中，$\text{RMS}(x) = \sqrt{\frac{1}{d} \sum_{i=1}^d x_i^2} = \frac{1}{\sqrt{d}} \|x\|_2$。 当我们输入带旋转的向量 $\tilde{x} = x R$ 时（由于 $R$ 是正交矩阵，$\|\tilde{x}\|_2 = \|x\|_2$）：

$$\text{RMS}(\tilde{x}) = \frac{1}{\sqrt{d}} \|x R\|_2 = \frac{1}{\sqrt{d}} \|x\|_2 = \text{RMS}(x)$$

因此：

$$\text{RMSNorm}(x R) = \frac{x R}{\text{RMS}(x R)} \odot \gamma = \left(\frac{x}{\text{RMS}(x)}\right) R \odot \gamma$$
  • 注：由于可学习参数 $\gamma$ 是逐元素（Element-wise）相乘的，如果不加处理，它会破坏与正交矩阵 $R$ 的乘法交换律。
  • QuaRot 的工程解法：在量化转换时，将 RMSNorm 的 $\gamma$ 参数直接吸收并折算到后一个线性层的权重中，即令原本 Norm 内部的 $\gamma = \mathbf{1}$（全 1 向量）。这样上述公式退化为： $$\text{RMSNorm}(x R) = \text{RMSNorm}(x) R$$ 旋转完美地从 RMSNorm 中穿透了过去！

4. SpinQuant（ICLR 2025）

可学习正交旋转（Learned Rotations）的代表作。

4.1 与 QuaRot 的核心区别

  • QuaRot 采用的是固定、通用的随机或标准 Hadamard 矩阵（RHT）。
  • SpinQuant 认为：Hadamard 虽然数值友好，但对于特定的 LLM 并不是最优解。应该针对具体模型的权重分布，通过学习的方式得到最优的正交旋转矩阵。

4.2 基于 Cayley 变换的参数化（如何优化正交矩阵）

在反向传播中直接优化正交矩阵非常困难，因为普通的梯度更新会破坏矩阵的正交性。SpinQuant 引入了 Cayley 变换： 若 $A$ 是一个反对称矩阵（即 $A = -A^T$），那么如下定义的 $R$ 恒为正交矩阵：

$$R = (I - A)(I + A)^{-1}$$

为了让反对称矩阵 $A$ 可学习，SpinQuant 引入一个无约束的自由矩阵 $G$：

$$A = \frac{G - G^T}{2}$$

在训练时，直接通过梯度下降优化参数矩阵 $G$，其经过 Cayley 变换生成的 $R$ 将始终严格保持正交性。这使得正交矩阵的训练变得简单且稳定。

4.3 SpinQuant 的工程优化

  • 训练开销小：只在有限的校准集上优化 $G$，7B 模型仅需 8 卡约 1.25 小时。
  • TP 兼容性好：由于大模型常采用 Tensor Parallelism（张量并行），如果使用全局 Hadamard 旋转，需要确保 hidden_dim / TP_size 依然是 2 的幂。SpinQuant 将优化设计为 Head-wise 旋转（在每个 Attention Head 维度内做旋转，例如 128 维），天然兼容任何 TP 分片配置。

5. 旋转量化的其他前沿进展（2024 - 2025+）

随着旋转量化路线的成功，学术界与工业界在 2025 年前后涌现出了一批更具针对性的优化工作：

5.1 DuQuant（NeurIPS 2024 Oral）

  • 针对痛点：发现 LLaMA-3 中某些特定层（如 down_proj）存在超出普通范畴的超级离群点（Massive Outliers）。
  • 解决方案：提出 Rotation + Permutation（旋转 + 通道排列） 协同设计。在旋转前，先通过排列（Permutation）操作将包含极大离群值的维度移动到指定位置，然后再实施针对性的局部正交变换，避免为了平摊个别超级离群点而过度模糊正常通道。

5.2 FlatQuant

  • 核心发现：量化精度并不单单取决于“是否存在离群点”，而是取决于通道分布的整体平坦度（Flatness）。
  • 方法：放弃了全局正交旋转这种开销较大的操作，转而为每个线性层学习极轻量级的仿射变换（Learnable Affine Transform），从而以极低开销在 LLaMA-3 70B 实现了 W4A4 近无损的 SOTA 效果。

5.3 MambaQuant (ICLR 2025)

  • 背景：在量化 Mamba / State Space Models 时，由于其特有的 Parallel Scan（并行扫描）操作，导致传统的 QuaRot Hadamard 旋转无法保持等价性，从而导致精度大跌。
  • 创新：使用 KLT 变换（Karhunen-Loève Transform） 替代固定 Hadamard，自动适应 Mamba 通道的协方差分布，实现了非 Transformer 架构的 4-bit 量化。

5.4 OSTQuant (ICLR 2025)

  • 引入了 QSUR（Quantization Space Utilization Ratio，量化空间利用率） 作为度量指标，在数学上更系统地定义了旋转变换如何使张量更贴合低比特网格。

6. 为什么在 Diffusion Transformer (DiT) 上会失效？（ConvRot 的核心立论点）

ConvRot 是第一个成功将“旋转量化”搬到 DiT 模型上的工作。它深入剖析了为什么 LLM 的 QuaRot 等方法不能直接用在 DiT 上：

6.1 离群点特征不同：行向（Row-wise）离群点的威胁

  • LLM 的离群点主要集中在列向（通道维度）。
  • DiT（如 PixArt, Sora 等）因为输入是 Patch 化的图像空间数据，在某些层（如 proj_out）会出现明显的行向（Token / 空间维度）离群点。
  • 致命后果：传统的 Sylvester 型 Hadamard 矩阵第一列全部为 1，作用于带有系统性偏置的行向向量时，会把所有的行偏置无差别地求和并放大 $\sqrt{n}$ 倍。这会导致旋转后反而催生出极其严重的全新离群通道。
  • ConvRot 解法：弃用 Sylvester 构造，改用 Regular Hadamard 矩阵（其列偏差为理论最低值 $\sqrt{n}$），有效避免了行向离群点的求和放大。

6.2 归一化限制：AdaLN 破坏旋转等价性

  • LLM 使用的 RMSNorm 天生对正交旋转“透明”。
  • DiT 使用的是 AdaLN（自适应层归一化），其在层归一化后要乘以并加上由 Time/Label Embedding 动态生成的调整系数。这破坏了正交矩阵的交换律，导致旋转无法“穿透” AdaLN。
  • ConvRot 解法：在 AdaLN 之前引入反向旋转进行对冲，并限制在小组（Group-wise，如 256 或 1024 组大小）内完成局部 RHT 旋转，极大地缓解了无法全局融合权重的计算痛点。

7. 核心结论与一句话总结

LLM 旋转量化的本质： 利用正交矩阵（通常是 Hadamard）的能量守恒与平铺机制，在数学等价的前提下，将激活中高耸的离群点平摊到所有通道。QuaRot 实现了将旋转离线融合进模型权重的极简设计，开启了 W4A4 推理的新时代；而 SpinQuant、DuQuant、FlatQuant 等后续工作则让这种旋转变得更为精准和自适应。