[ICLR 2025] SpinQuant: LLM Quantization with Learned Rotations

Introduction

• 作者提出 SpinQuant，在 QuaRot 基础上训练正交矩阵以达到更优精度

Method

• Rotation parameterization. 作者 follow 了 QuaRot，利用 Hadamard 变换缓解 outliers. 另外有一个小区别：attn 层里 out proj 的输入激活值没有用大小为 hidden dim 的 Hadamard 矩阵去做变换，而是只使用的大小为 head dim 的 Hadamard 矩阵 $R_2$，这样 $R_2$ 可以融到权重里，不会带来额外的推理开销；但是代价是 out proj 输入处的量化必须用 head dim 的 group-wise 量化，相当于是省去一次在线变换，但 out proj 输入处的量化必须用 group-wise 量化，也会引入额外开销，否则也会掉精度；不过这样的设计对 TP 的兼容性比较好，如果是 QuaRot + TP，那 o_proj 前的在线变换就是 ${\mathbf{H}}_{n_h/tp}\otimes \mathbf{I}$，当 TP = 8 时，$n_h/tp$ 很可能找不到对应维数的 Hadamard 矩阵，而如果像 SpinQuant 这样都做 head-wise 的变换就不会有这种问题
  
• (2) Cayley-optimized rotation. 作者对 $R_1,R_2$ 做了训练，$R_3,R_4$ 由于无法融到权重里，因此这些矩阵保持不变，这样可以使用快速 Hadamard 变换，推理开销比较小。训练方式上采用 Cayley transform 用于保证训练过程中 $R$ 正交
  其中，$Y$ 为反对称矩阵，$G$ 为可学习参数 (Cayley trick 可参考 Orthogonal Finetuning)；损失函数采用 CE loss，在 800 个样本的 Wiki 校准集上训 100 epochs，8 卡训 7B 模型大约需要 1.25h，训练开销还是比较大的；注意到，在 QuaRot 框架中，$R_1$ 影响着所有 Transformer Layer，因此如果要优化 $R_1$，就必须做端到端优化，训练开销一定是很大的
  

Experiments

• Quantization Settings. follow QuaRot，不一样的是激活值量化采用非对称量化，另外也和 QuaRot 一样用了 weight clipping，但没有用 activation clipping，因为作者发现没啥收益；作者在论文里没说，但是代码里实际和 QuaRot 一样也是用的自注意力层混精量化 (Q 不量化，只量化 KV)

---

• Main results.
  
• Ablation studies. (1) Compatibility with GPTQ. 在结合 GPTQ 时，作者在只量化激活值的模型上优化 rotation matrices，把权重量化的工作完全交给 GPTQ，也让正交变换更专注于缓解激活值 outliers
  (2) Impact of each rotation.
  (3) Rotation type. 训练后不管初始化是随机正交矩阵还是 Hadamard 矩阵，最终精度都差不多
  

References

• Liu, Zechun, et al. “SpinQuant: LLM Quantization with Learned Rotations.” arXiv preprint arXiv:2405.16406 (2024).
• code: https://github.com/facebookresearch/SpinQuant