IntactKV：保持关键词元无损的大语言模型量化方法

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

论文链接：https://arxiv.org/abs/2403.01241

开源代码：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精度。

图1 IntactKV示意图

关键词元（Pivot Tokens）

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

图2 关键词元示意图

方法概述

我们发现，模型量化会不可避免地导致关键词元表征受损，进而严重影响模型精度，因此如果能保证关键词元的表征，也就是其KV cache无损，就能有效降低量化误差。为此，我们提出使用IntactKV—先使用全精度模型生成关键词元的无损KV cache并将其缓存下来（i.e., IntactKV），量化模型在推理时就能直接使用无损的关键词元表征，从而有效提升量化模型精度。如图3所示，IntactKV包含了从首词元[BOS]开始的若干词元，随着IntactKV包含词元数的增加，量化误差也不断降低，尤其是当IntactKV包含关键词元时，量化误差会得到明显下降，这也进一步验证了关键词元的重要性。

图3 保证关键词元表征无损能有效降低量化误差

然而，模型在每次推理时的输入序列各不相同，这意味着我们并不能使用任意长度的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：

实验验证

1. 权重量化

我们首先在PPL任务上进行了实验。如表1所示，IntactKV能很好地与GPTQ、OmniQuant、AWQ等主流量化方法结合，无需任何训练或额外推理开销即可在LLaMA系列的各种不同模型上有效提升量化模型精度，其中AWQ+IntactKV在各个模型上都达到了最优效果。

表1 权重量化PPL任务实验结果

我们基于Vicuna模型，在MMLU和QA任务上做了进一步实验。如表2所示，实验结论与PPL任务基本一致。

表2 权重量化MMLU和QA任务实验结果

我们对量化模型的生成质量也进行了评估，通过GPT-4对模型生成回答进行打分来衡量量化模型的生成能力。如表3所示，IntactKV_[P]已经能明显提升量化模型的回复得分，对IntactKV的轻量化校准则进一步提升了模型的回复质量。

表3 权重量化MT-bench任务实验结果

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]中。

图4 权重和KV Cache量化MMLU任务实验结果

3. 权重和激活值量化

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

表4 权重和激活值量化PPL任务实验结果