本文介绍我们针对大语言模型量化的工作 IntactKV，可以作为插件有效提升 GPTQ、AWQ、QuaRot 等现有主流量化方法效果。论文作者来自清华大学、华为诺亚、中科院自动化所和香港中文大学。论文代码已经开源，欢迎大家使用！

• 论文链接：IntactKV: Improving Large Language Model Quantization by Keeping Pivot Tokens Intact
• 开源代码：https://github.com/ruikangliu/IntactKV

摘要

• 大语言模型（LLM）在自然语言处理的各项任务上都取得了出色的表现，但其庞大的参数量也给实际部署带来了新的挑战。模型量化是一种常用的模型压缩方法，可以通过将模型权重和激活值压缩到低比特来有效降低模型推理时的算力和内存开销，然而量化往往伴随着模型的精度损失。对此，我们提出了一种新的量化方法 IntactKV。我们发现 LLM 中存在关键词元（pivot tokens），这些关键词元的表征上存在极大的离群值，这会导致模型将大量自注意力得分分配到这些词元上，而模型量化会使得这些关键词元表征受损，严重损害了量化模型精度。IntactKV 通过预先缓存关键词元的无损 KV cache 来保证量化模型中关键词元表征无损，理论推导证明这能有效降低模型的量化误差上界。此外，缓存的无损 KV cache 还可以作为模型的额外参数做进一步校准，进一步弥补量化误差
• IntactKV 实现简单并且与 GPTQ[1]、AWQ[2]、QuaRot[3] 等当前主流的 LLM 量化方法正交，实验表明，IntactKV 可以作为插件有效提升现有量化模型精度，同时不会带来额外的推理开销，在权重量化和权重激活值量化上均取得了当前的 SOTA 精度

关键词元（Pivot Tokens）

• 现有研究[3][4][5]发现，LLM 激活值的特定通道上存在离群值，这些离群值往往出现在所有词元的激活值上，使得 LLM 的激活值难以量化。而我们发现 LLM 中还存在另一种特殊的离群值，如图 2 所示，它只存在于 LLM 输入的首词元 [BOS] 以及开头的一些特定词元上（如 “,”、“.” 等标点符号），并且这类离群值比特定通道上的离群值还要大得多。更为重要的是，这类特定词元上的离群值会导致 LLM 将大量的自注意力得分分配到这些词元上，形成 Attention Sink[6]，因此我们将这些特定词元称为关键词元（pivot tokens）。由于汇聚了大量自注意力得分，关键词元的表征对 LLM 来说至关重要。类似的现象也在同期工作 Massive Activations[7] 中被发现并得到了更为细致的研究。然而，还没有其他研究针对这一现象对于量化模型的影响做深入探索

方法概述

• 我们发现，模型量化会不可避免地导致关键词元表征受损，进而严重影响模型精度，因此如果能保证关键词元的表征，也就是其 KV cache 无损，就能有效降低量化误差。为此，我们提出使用 IntactKV — 先使用全精度模型生成关键词元的无损 KV cache 并将其缓存下来（i.e., IntactKV），量化模型在推理时就能直接使用无损的关键词元表征，从而有效提升量化模型精度。如图 3 所示，IntactKV 包含了从首词元 [BOS] 开始的若干词元，随着 IntactKV 包含词元数的增加，量化误差也不断降低，尤其是当 IntactKV 包含关键词元时，量化误差会得到明显下降，这也进一步验证了关键词元的重要性
  
• 然而，模型在每次推理时的输入序列各不相同，这意味着我们并不能使用任意长度的 IntactKV，否则在每次量化模型推理前都必须调用全精度模型来生成对应输入的 IntactKV，这会带来额外的推理开销。对此，我们利用了 LLM 推理时的公共前缀，如果 IntactKV 只包含公共前缀，那么全精度模型在生成 IntactKV 后就可以被释放，不会给之后量化模型的推理带来任何额外开销。对于预训练 LLM，输入序列前一般都会加上 [BOS] 词元作为公共前缀，因此可以使用只包含 [BOS] 词元的长度为 1 的 IntactKV，称为 IntactKV_[B]。如图 3 所示，由于首词元总为关键词元，因此使用 IntactKV_[B] 就可以明显地降低量化误差。而对于 SFT 模型，输入序列前一般会包含一段系统提示词作为公共前缀，因此可以使用长度为系统提示词的 IntactKV，称为 IntactKV_[P]。由于使用了更长的 IntactKV，IntactKV_[P] 能覆盖更多关键词元，进一步降低量化误差。我们在实验中发现，对于 Vicuna 模型，长度为 34 的系统提示词就足以覆盖全部关键词元
• 此外，缓存的 IntactKV 还可以作为模型的额外参数做校准来进一步弥补量化误差。具体来说，我们在大小为128的校准集上做了轻量化训练，可学习参数仅为缓存下来的 IntactKV，损失函数采用量化模型和全精度模型输出的 MSE loss：
  $$\mathcal{L}(\Theta )=\frac{1}{2}\sum _{l=1}^L\Vert f_l(\mathbf{w},\mathbf{x})-f_l(\hat{\mathbf{w}},\mathbf{x};\Theta ){\Vert }_2^2$$其中，$\Theta$ 为 IntactKV，$f_l$ 为第 $l$ 个 Transformer 层的映射函数，$L$ 为 LLM 中 Transformer 层的层数， $\mathbf{x}$ 为输入激活值， $\mathbf{w},\hat{\mathbf{w}}$ 分别为全精度权重和量化权重

实验验证

1. 权重量化

• 我们首先在 PPL 任务上进行了实验。如表 1 所示，IntactKV 能很好地与 GPTQ、OmniQuant、AWQ 等主流量化方法结合，无需任何训练或额外推理开销即可在 LLaMA 系列的各种不同模型上有效提升量化模型精度，其中 AWQ+IntactKV 在各个模型上都达到了最优效果
  
• 我们基于 Vicuna 模型，在 MMLU 和 QA 任务上做了进一步实验。如表 2 所示，实验结论与 PPL 任务基本一致
  
• 我们对量化模型的生成质量也进行了评估，通过 GPT-4 对模型生成回答进行打分来衡量量化模型的生成能力。如表 3 所示，IntactKV_[P] 已经能明显提升量化模型的回复得分，对 IntactKV 的轻量化校准则进一步提升了模型的回复质量
  

2. KV Cache 量化

• 为了将 IntactKV 用于 KV cache 量化，我们采用混合精度量化策略，将 IntactKV 保存为全精度，其余 KV cache 量化到低比特。对于 KV cache 量化，由于模型仍然需要先将量化的 KV cache 反量化为全精度才能进行自注意力计算，因此上述混合精度方案在推理时只需要将 IntactKV 和反量化后的 KV cache 重新拼接到一起即可，并不会带来额外的推理开销。如图 4 所示，IntactKV 在不同的量化比特设置下均可以有效提升 AWQ 量化模型精度，尤其是在低比特时提升更加明显。目前，IntactKV 已经被集成到了 SOTA KV cache 量化方法 KVQuant[8] 中
  

3. 权重和激活值量化

• 对于权重激活值量化，KV cache 需要全部量化到低比特才能有效利用低比特矩阵乘进行运算加速。我们发现，关键词元的 KV cache 表征明显比其余 KV cache 表征更加平坦，就算直接将 IntactKV 量化到低比特也只会带来极少的精度损失，因此对于激活值量化，我们选择将 IntactKV 进一步压缩到低比特以免引入额外的推理开销。如表 4 所示，IntactKV 能有效提升当前SOTA量化方法精度，进一步缩小了 W4A4 量化模型与原始模型间的精度损失
  

参考文献

• [1] Frantar, Elias, et al. “OPTQ: Accurate quantization for generative pre-trained transformers.” The Eleventh International Conference on Learning Representations. 2022.
• [2] Lin, Ji, et al. “Awq: Activation-aware weight quantization for llm compression and acceleration.” arXiv preprint arXiv:2306.00978 (2023).
• [3] Ashkboos, Saleh, et al. “Quarot: Outlier-free 4-bit inference in rotated llms.” arXiv preprint arXiv:2404.00456 (2024).
• [4] Dettmers, Tim, et al. “Gpt3. int8 (): 8-bit matrix multiplication for transformers at scale.” Advances in Neural Information Processing Systems 35 (2022): 30318-30332.
• [5] Xiao, Guangxuan, et al. “Smoothquant: Accurate and efficient post-training quantization for large language models.” International Conference on Machine Learning. PMLR, 2023.
• [6] Xiao, Guangxuan, et al. “Efficient Streaming Language Models with Attention Sinks.” The Twelfth International Conference on Learning Representations. 2023.
• [7] Sun, Mingjie, et al. “Massive Activations in Large Language Models.” arXiv preprint arXiv:2402.17762 (2024).
• [8] Hooper, Coleman, et al. “KVQuant: Towards 10 Million Context Length LLM Inference with KV Cache Quantization.” arXiv preprint arXiv:2401.18079 (2024).