NSNQuant: A Double Normalization Approach for Calibration-Free Low-Bit Vector Quantization of KV Cache

发表时间: 2025-05 · arXiv:2505.18231 (NeurIPS 2025)

原文: https://arxiv.org/abs/2505.18231

速读

一句话结论
提出了一种免校准的 KV Cache 低比特向量量化方法 NSNQuant,通过将特征分布强制对齐到标准正态分布,解决了依赖校准数据导致的泛化失效问题,在 1-bit 和 2-bit 压缩下均实现了最优效果并带来最高 3 倍的吞吐量提升。

要解决什么问题
大语言模型在处理长上下文时,KV Cache 的显存占用会随序列长度线性增长,导致推理严重受限于显存带宽。近期引入的向量量化(将一组数值联合映射到密码本索引)方法如 CQ,通过在校准数据集上聚类来学习密码本,取得了极高的压缩率。但原有做法卡在一个致命缺陷上:极易受到分布偏移的影响。当推理时的输入数据分布与校准集不一致时,量化误差会急剧放大。具体到机制上,由于不同数据集的词元分布差异极大(例如 WikiText-2 中的逗号带有空格,而 C4 数据集中的逗号不带空格),仅在小规模校准集上学习到的聚类中心,在遇到未见过的标点符号特征时会产生巨大的量化误差。这种误差会导致注意力权重严重扭曲,使得模型在跨数据集测试时困惑度大幅恶化,泛化能力失效。

怎么做的
核心思路是逆向思维:既然让密码本去适配多变的 KV 分布容易过拟合,不如通过数学变换,把所有输入的 KV 分布都强制拉平到一个已知的先验分布(标准正态分布),然后用一个通用的、在纯随机高斯数据上训练好的密码本进行量化。这样就能彻底绕开对校准数据的依赖。关键设计由三个核心部件构成:首先是 Normalize-Shift-Normalize (NSN) 变换。它包含三个步骤:第一步是 Token 级别的归一化,将每个 Token 向量除以其范数,防止极个别异常值主导后续计算;第二步是通道级别的平移,减去每个通道的均值使其零中心化;第三步是再次进行 Token 级别的归一化。整个过程定义为:

$$ s_1 \leftarrow \text{norm}(v, \text{dim=token}) / \sqrt{d}, \quad v_n \leftarrow v / s_1 $$

$$ o \leftarrow \text{mean}(v_n, \text{dim=channel}), \quad v_{ns} \leftarrow v_n - o $$
$$ s_2 \leftarrow \text{norm}(v_{ns}, \text{dim=token}) / \sqrt{d}, \quad v_{nsn} \leftarrow v_{ns} / s_2 $$
经过 NSN 处理后,再配合 Hadamard 变换(一种能将数据特征打散的正交矩阵乘法),通道分布在理论和实际上都被严格约束到了标准正态分布。其次是自适应缩放。在从密码本中查表还原出量化向量 $v_Q$ 后,作者并没有直接使用它,而是对其进行缩放,使得量化误差与原向量正交:
$$ v_Q \leftarrow \frac{\|v\|_2^2}{v \cdot v_Q} v_Q $$
这一步的职责是最大程度保留原向量的平行分量,确保每个 Token 的特征依然具备区分度。最后是密码本构建与双重量化。因为输入已经被标准化,密码本只需在合成的标准正态分布数据(如 torch.randn)上微调即可。同时,NSN 过程中产生的均值 $o$ 和缩放因子 $s_1$ 会被进一步做 4-bit 的双重量化,从而将额外显存开销压缩到极致。

效果如何
实验在 LLaMA2(7B/13B)、LLaMA3/3.1(8B)和 Mistral-7B 模型上展开。对比基线点名了代表通道/Token独立量化路线的 KIVI 和 KVQuant,以及代表向量量化路线的强基线 CQ。量化结果显示,在 WikiText-2 和 C4 数据集的困惑度测试中,NSNQuant 在 1-bit 和 2-bit 设置下均展现出极强的泛化能力。特别是在 C4 数据集上,CQ 在 LLaMA3 模型上因标点符号分布偏移导致困惑度崩溃(飙升至 16.16),而 NSNQuant 依然稳定在 9.03。在 LongBench 长文本理解、GSM8K 数学推理和 HumanEval 代码生成等任务中,NSNQuant-1b 均大幅超越其他 1-bit 基线,NSNQuant-2b 也展现出最能忠实保留全精度模型输出的能力。在硬件效率上,得益于定制的 CUDA 算子,NSNQuant 相比 FP16 基线支持了 4 倍的批处理大小,并实现了 3 倍的吞吐量提升。代价与局限性在于:首先,NSN 和 Hadamard 变换引入了额外的计算开销,导致 Prefill(预填充)阶段的延迟有所增加,因此该方法更适合 Decode(解码)密集的任务。其次,作者承认在模型极少数的早期层中,由于存在方差极大的异常通道,NSN 无法完美将其标准化,尽管这在当前并未造成严重的量化误差,但仍是未来需要处理的边界情况。

作者/机构:Donghyun Son, Euntae Choi, Sungjoo Yoo (Seoul National University)

主要贡献

大型语言模型(LLM)的推理通常是内存密集型的,尤其是在处理大批量和长序列时,这主要归因于键值(KV)缓存的巨大尺寸。虽然最近采用了向量量化(VQ)来缓解这一问题,但现有的方法由于依赖于校准数据集,容易受到分布偏移(Distribution Shift)的影响。为了解决这一限制,本文提出了NSNQuant,这是一种专为KV缓存低比特压缩设计的无校准向量量化(VQ)技术。

本文的主要创新点如下:
1. 通过实证表明,键值分布对输入数据集的强烈依赖可能导致现有基于校准的VQ方法产生严重误差。特别观察到,CQ-4c9b学习到的质心在C4数据集上无法准确量化LLaMA3-8B和LLaMA3.1-8B中重要的标点符号token,从而导致性能退化。
2. 提出了NSNQuant,这是一种用于KV缓存的无校准VQ方法。NSNQuant通过“归一化-中心化-归一化”(Normalize-Shift-Normalize, NSN)过程和Hadamard变换,有效地将键和值的分布与标准正态分布对齐。由于键值的近似分布在推理前是已知的,因此无需任何外部数据即可构建用于VQ的有效码本。
3. 进行了全面的实验和分析,证明了NSNQuant在不同任务和模型中的有效性。在LLaMA和Mistral模型上的结果清楚地表明,NSNQuant在1-bit和2-bit量化中均优于其他基线方法。此外,为低比特计算实现了高效的CUDA内核,提高了吞吐量并减少了内存使用。

图1:(a) 使用LLaMA3.1-8B的PPL评估结果。尽管CQ在WikiText-2(分布内)中实现了较低的PPL,但在C4(分布外)上表现较差。(b) LLaMA3.1-8B键和值的t-SNE可视化。聚类模式表明键和值的分布强烈依赖于输入数据。更多可视化见图11
图1:(a) 使用LLaMA3.1-8B的PPL评估结果。尽管CQ在WikiText-2(分布内)中实现了较低的PPL,但在C4(分布外)上表现较差。(b) LLaMA3.1-8B键和值的t-SNE可视化。聚类模式表明键和值的分布强烈依赖于输入数据。更多可视化见图11

背景知识与动机

LLM推理与KV缓存
LLM推理包含预填充和解码两个主要阶段。在预填充阶段,所有提示词token由Transformer解码器层同时处理。在解码阶段,新token以自回归方式逐个生成。由于掩码自注意力的因果性质,在这两个阶段中,每个token只关注先前的token。为了避免冗余计算,KV缓存存储了来自先前token的键值对。它在预填充期间初始化,并在每次解码步骤中通过附加最新的键值对进行更新。当处理长序列时,其大小随序列长度线性增加,成为主要的瓶颈。

向量量化(VQ)
与每个标量值被单独量化的标量量化(SQ)不同,向量量化(VQ)使用码本联合压缩一组值。在VQ中,一个$d$维向量与码本中最接近的条目匹配,并量化如下:
$ \operatorname{VQ}(v) = \operatorname{argmin}_i \operatorname{D}(v, \mathbb{C}[i]) $
其中$\mathbb{C}$表示码本,$\operatorname{D}(a, b)$是向量$a$和$b$之间的距离函数。

图2:NSNQuant下注意力的整体结构。为简单起见,省略了残差。我们使用上标$k$和$v$分别标记与键和值相关的值。由于NSN在RoPE之前应用于键,因此需要两个分支来正确计算注意力分数。图中显示的注意力计算细节见附录B。
图2:NSNQuant下注意力的整体结构。为简单起见,省略了残差。我们使用上标$k$和$v$分别标记与键和值相关的值。由于NSN在RoPE之前应用于键,因此需要两个分支来正确计算注意力分数。图中显示的注意力计算细节见附录B。

Hadamard变换
Hadamard矩阵是一个正交矩阵,其所有条目具有相同的量级。它可以通过Sylvester构造递归地将其大小翻倍,该构造通过以特定模式组合较小Hadamard矩阵的副本来构建更大的矩阵。这种递归定义构成了Walsh-Hadamard变换的基础,允许在$O(d \log d)$时间内进行矩阵-向量乘法。QuIP#【37,QuIP#: Even better LLM quantization with hadamard incoherence and lattice codebooks+2024+ICML】采用了随机Hadamard变换(RHT),其中每行和每列的符号以1/2的概率独立翻转。本文也采用了RHT来计算理论界限。

分布不匹配导致的量化失效
现有的VQ方法,即耦合量化(CQ)【43,Kv cache is 1 bit per channel: Efficient large language model inference with coupled quantization+2024+NeurIPS】,在分布外(OOD)数据集上测试时遭遇了严重的性能退化。这归因于数据集之间的分布不匹配。这种不匹配可能在某些数据集中引起严重误差,一个显著的例子是LLaMA3-8B和LLaMA3.1-8B中标点符号token观察到的大量化误差。CQ的校准数据集WikiText-2仅包含带有空格的" , " token,而C4数据集包含没有空格的"," token。这导致LLaMA3模型第一层中对应于标点符号","的键产生巨大误差,因为CQ的质心仅从WikiText-2获得。这导致注意力权重严重扭曲,因为在某些注意力头中,"," token占据了超过$90\%$的注意力权重。通过在第一层中保留这些token对应的键,CQ-4c9b在C4上的困惑度在LLaMA3-8B中从13.97改善到9.15,在LLaMA3.1-8B中从12.24改善到9.16,这与NSNQuant的结果非常接近。

无校准量化方法的提出
为了避免这些由校准引起的误差,提出了一种无校准的向量量化(VQ)方法NSNQuant,它不依赖于任何外部数据。CQ试图将码本与键值(KV)分布匹配,而本文提出将KV分布与众所周知的先验分布匹配。受近期基于Hadamard的方法在产生与输入无关的类正态输出分布方面取得成功的启发,引入了一种新颖的转换——Normalize-Shift-Normalize (NSN)——将键和值通道对齐到标准正态分布。这使得能够构建一个专门针对标准正态分布的单一可复用码本,使NSNQuant成为无校准且对各种输入具有鲁棒性的方法。

图3:NSN对每个通道值分布影响的视觉说明。当与Hadamard变换一起使用时,我们的Normalize-Shift-Normalize (NSN)过程将分布与标准正态分布对齐。
图3:NSN对每个通道值分布影响的视觉说明。当与Hadamard变换一起使用时,我们的Normalize-Shift-Normalize (NSN)过程将分布与标准正态分布对齐。

方法细节

Normalize-Shift-Normalize (NSN)的处理过程
NSN包含三个步骤:1) token级别的归一化(Normalize),2) channel级别的中心化(Shift),3) 第二次token级别的归一化(Normalize)。令$v \in \mathbb{R}^{l \times d}$为张量,其中$l$是序列长度,$d$是每个头的隐藏层维度。首先在第一个Normalize步骤中,将每个token归一化为具有$\sqrt{d}$的范数,其中$d$是每个头的token维度。这防止了异常值token【7,Prefixquant: Static quantization beats dynamic through prefixed outliers in llms+2024+arXiv】主导后续步骤并对幅度较小的token产生巨大误差。接着在Shift步骤中,计算并减去通道级均值,以便得到的分布是零中心化的。最后在第二个Normalize步骤中,再次将每个token归一化为具有$\sqrt{d}$的范数。整个过程公式化如下:
1. Normalize: $s_1 \gets \mathrm{norm}(v, \mathrm{dim}{=}\mathrm{token}) / \sqrt{d}, \quad v_{\mathrm{n}} \gets v / s_1$
2. Shift: $o \gets \mathrm{mean}(v_{\mathrm{n}}, \mathrm{dim=channel}), \quad v_{\mathrm{ns}} \gets v_{\mathrm{n}} - o$
3. Normalize: $s_2 \gets \mathrm{norm}(v_{\mathrm{ns}}, \mathrm{dim=token}) / \sqrt{d}, \quad v_{\mathrm{nsn}} \gets v_{\mathrm{ns}} / s_2$
每个步骤都会产生一个副产品——表示为$s_1, o$和$s_2$——用于通过$v = s_1 (s_2 v_{\mathrm{nsn}} + o)$恢复原始张量。尽管最后一步可能使通道级均值偏离零,但其影响可以忽略不计。

NSN与Hadamard变换的对齐效果
当与随后的Hadamard变换一起使用时,NSN过程有效地将通道分布与标准正态分布对齐。正如先前研究【37,QuIP#: Even better LLM quantization with hadamard incoherence and lattice codebooks+2024+ICML】、【39,Bitnet v2: Native 4-bit activations with hadamard transformation for 1-bit llms+2025+arXiv】、【9,Sdr: Efficient neural reranking using succinct document representation+2021+arXiv】、【20,Ostquant: Refining large language model quantization with orthogonal and scaling transformations for better distribution fitting+2025+arXiv】所指出的,Hadamard变换导致类正态分布,这由中心极限定理支持。NSN过程在与随后的Hadamard变换一起使用时,大致标准化了每个通道的分布。将这两点结合起来,得到的通道分布与标准正态分布对齐。这可以通过引理1来证明,它为方差提供了理论界限。由于NSN收紧了界限,并且对于大多数层来说协方差较小,观察到得到的方差通常接近于1。然而,在早期层的某些头中这并不成立,这是由于第一层中存在具有巨大方差的异常值通道。尽管存在这一限制,量化误差在这些层中仍然保持在较低水平。

残差策略的引入
为了在解码阶段实现第二步(Shift),引入了来自KIVI【30,Kivi: A tuning-free asymmetric 2bit quantization for kv cache+2024+arXiv】的残差概念。遵循KIVI的做法,将KV缓存分为两部分:一部分是量化的KV缓存,另一部分是全精度KV缓存(残差)。如果残差的大小达到其最大容量,则将残差中的KV缓存刷新、量化并附加到量化部分。引入了一个称为残差大小(residual size)的超参数来控制其大小。为了确保一致性,在预填充阶段也以残差大小的块应用NSN。在实验中,将残差大小设置为64。

键与值的差异化应用
NSN在键和值上的应用略有不同。对于键,NSN在投影层之后立即应用,并且$v_{\mathrm{nsn}}$在RoPE和Hadamard变换之后被量化。虽然在RoPE之后应用NSN似乎更直观(因为RoPE可能影响通道级均值),但发现这种排序产生了更好的量化质量。由于RoPE尚未应用于$o$,因此在计算注意力分数时在自定义内核中应用它。对于值,Hadamard变换被融合到投影层中,并且NSN在之后立即应用。由于Hadamard变换等效于乘以旋转矩阵,它与相邻NSN的顺序不会改变输出。

自适应尺度调整
令$v \in \mathbb{R}^d$为经过NSN和Hadamard变换处理的token向量。然后将其分为8维子向量,并使用码本进行量化。令$v_Q$为通过查找码本恢复的向量,即$v_Q \doteq \mathbb{C}[\mathrm{VQ}(v)]$。发现与其原样使用$v_Q$,不如自适应地缩放$v_Q$以提高性能。具体来说,将$v_Q$缩放如下是有益的(这与缩放$s_2$相同,因为恢复时$v_Q$乘以$s_2$):
$ v_Q \gets \frac{\|v\|_2^2}{v \cdot v_Q} v_Q \quad (\mathrm{i.e., } s_2 \gets \frac{\|v\|_2^2}{v \cdot v_Q} s_2) $
这是一种使$v_Q - v$正交于$v$的缩放策略。换句话说,它保留了平行于$v$的成分,同时允许一些正交误差。将$o$解释为局部上下文,将$v$解释为独特的token特征,这种策略可以解释为使每个token具有独特性,考虑到KV缓存的选择性特性,这对于KV缓存是必不可少的。

码本微调
构建了一个用于将8维向量压缩为整数索引的单一全局码本,遵循QuIP#【37,QuIP#: Even better LLM quantization with hadamard incoherence and lattice codebooks+2024+ICML】的做法。NSNQuant-2b使用8位表示符号,8位表示码本索引,而NSNQuant-1b仅使用8位表示索引。可以通过在标准正态数据上使用K-Means构建一个简单的基线,但其局部最优性限制了性能。通过在合成的标准正态数据(torch.randn)上进行微调来改进这一点,以最小化原始向量和量化向量之间的余弦距离,因为尺度调整的误差仅取决于它们之间的角度。由于查找是不可微的,梯度仅通过查找后操作传播。这个过程的PyTorch实现在RTX 3090上运行不到5分钟,不像CQ或KVQuant的校准过程需要通过模型权重进行反向传播。

双重量化策略
为了进一步减少内存开销,采用了QLoRA【10,Qlora: Efficient finetuning of quantized llms+2023+arXiv】中提出的双重量化(DQ),该方法对用于量化的参数进行量化。具体来说,使用4位最近舍入(RTN)量化,分别以32和残差大小为组大小对$o$和$s_1$进行量化。DQ显著降低了平均位宽。结果,当残差大小设置为64时,NSNQuant在NSN过程中平均仅额外消耗0.23位。此外,为了减少码本所需的共享内存量,还对码本条目应用了4位量化。DQ几乎不影响NSNQuant的有效性。

高效内核实现
为NSNQuant-2b和NSNQuant-1b的高效执行实现了CUDA内核。对于码本匹配,在流多处理器(SMs)中动态计算和管理距离,同时将码本块加载到共享内存中。对于反量化和矩阵-向量乘法,内核将码本加载到共享内存以最小化对DRAM的访问。为$qK^T$和$Wv$($W$:注意力权重)实现了两种不同的矩阵-向量乘法内核,因为它们使用不同的轴进行归约。还融合了量化部分和残差部分的计算,以在小批量场景下最大化GPU利用率。

实验环境

  • 数据集名称、规模及用途

    • 语言建模困惑度(PPL)评估:WikiText-2 和 C4(序列长度4096)。
    • LongBench子集:用于长文本理解评估,包括Qasper(单文档问答)、QMSum和MultiNews(摘要)、TREC、TriviaQA和SAMSum(少样本学习)、LCC和RepoBench-P(代码补全)。
    • 综合生成能力评估:GSM8K(数学推理,8-shot CoT)、HumanEval(代码生成)、CoQA(对话问答)、MMLU(多任务语言理解,4-shot CoT)。
    • AIME-2024(附录):用于评估长上下文推理能力。
  • 模型架构关键参数

    • 测试模型包括:LLaMA2-7B, LLaMA2-13B, LLaMA3-8B, LLaMA3.1-8B, Mistral-7B-v0.3, LLaMA2-13B-Chat, LLaMA2-7B-Chat, LLaMA3-8B-Instruct, DeepSeek-R1-Distill-Llama-8B。
  • 硬件配置

    • PPL评估在配备2个RTX Titan GPU的Linux服务器上执行。
    • LongBench, GSM8K, HumanEval, CoQA, MMLU评估在配备8个RTX 3090 GPU的Linux服务器上进行。
    • 效率分析(内存和吞吐量)在配备单个A100-80GB GPU的Linux服务器上进行。
  • 软件配置

    • 基于HuggingFace Transformers库使用PyTorch框架实现。
    • 比较的基线方法包括:KIVI、KIVI + Hadamard、KVQuant、CQ。所有方法统一采用NSNQuant的残差策略(残差大小设为64),并在预填充阶段统一使用全精度缓存(PPL评估除外)。

实验结果

  1. 困惑度(PPL)评估
    - 实验内容:在WikiText-2和C4数据集上评估语言建模的量化误差。
    - 实验结果:在WikiText-2上,CQ和NSNQuant表现相当(2-bit下接近,1-bit下CQ略优)。但在分布外数据集C4上,NSNQuant在1-bit和2-bit设置下均一致优于CQ。例如,CQ在LLaMA3-8B和LLaMA3.1-8B上遭遇严重性能退化,而NSNQuant保持强劲性能(如LLaMA3.1-8B上NSNQuant-2b PPL为9.15,而CQ-4c9b为12.24)。
    - 分析结论:NSNQuant跨数据集的泛化能力更好,而CQ在偏离校准数据集的分布偏移下挣扎。
    - 图表引用:Table 2。

  2. LongBench长文本理解评估
    - 实验内容:在LongBench的多个子任务上评估模型性能。
    - 实验结果:在1-bit量化中,NSNQuant-1b大幅优于其他基线。在2-bit量化中,所有方法表现相似。通过额外测量与FP16输出的ROUGE-L分数来评估模型保留原始输出的能力,NSNQuant取得了最高分。
    - 分析结论:在极端低比特(1-bit)下,NSNQuant优势显著。2-bit下指标相似是因为某些任务(如代码生成)的评估指标存在噪声,ROUGE-L分数明确证实了NSNQuant最能忠实地保留原始模型的输出。
    - 图表引用:Table 3, Table 19(附录)。

  3. 多样化生成任务评估
    - 实验内容:在GSM8K, HumanEval, CoQA, MMLU上评估生成质量。
    - 实验结果:NSNQuant在大多数设置下优于其他基线。特别是在GSM8K和MMLU的少样本CoT(思维链)设置下,NSNQuant表现卓越。
    - 分析结论:NSNQuant对于需要生成精确且严格推理路径的任务具有显著优势。
    - 图表引用:Table 4。

  4. 内存使用与吞吐量评估
    - 实验内容:模拟ShareGPT场景(平均输入长度161,平均生成长度338),测量不同批次大小下的峰值内存和吞吐量。
    - 实验结果:FP16基线在大批次时遭遇OOM,而NSNQuant-2b和NSNQuant-1b能高效扩展,实现了4倍更大的批次大小和3倍的吞吐量加速。
    - 分析结论:NSNQuant显著提升了运行时效率和内存效率。
    - 图表引用:Figure 4。

  5. 延迟分解评估
    - 实验内容:测量前向传递中每个操作的延迟(批次大小32,输入长度512,生成64个token)。
    - 实验结果:NSNQuant在预填充阶段因额外的Hadamard变换、VQ和NSN计算导致延迟较高。但在解码阶段,NSNQuant实现了更低的延迟(如NSNQuant-2b每步36.49ms对比FP16的50.80ms)。
    - 分析结论:NSNQuant通过将KV缓存压缩为低比特,缓解了注意力计算的内存受限问题,特别有利于解码繁重的任务(如推理或代码生成)。
    - 图表引用:Table 5。

  6. 消融实验 - NSN与Hadamard的作用
    - 实验内容:在WikiText-2上单独移除NSNQuant-2b的各个组件并测量PPL。
    - 实验结果:跳过任何步骤都会导致更高的PPL。第一次token归一化影响最大(PPL升至6.293),第二次token归一化影响最小。
    - 分析结论:NSN和Hadamard变换共同作用导致了最佳的对齐效果,第一次归一化在中心化之前抑制异常值token起到了关键作用。
    - 图表引用:Table 6。

  7. 消融实验 - 码本微调的效果
    - 实验内容:对比微调码本和K-Means码本在量化前后的余弦相似度。
    - 实验结果:微调版本提供了更高的余弦相似度。使用模型键值token测量的相似度与使用合成标准正态数据测量的相似度非常接近。码本微调将NSNQuant-2b的PPL从5.294降至5.285,将1b的PPL从6.910降至6.703。
    - 分析结论:NSN成功地将输出分布与标准正态分布对齐,仅在合成数据上训练的码本能有效地量化这些结果。
    - 图表引用:Figure 5。

图4:随批次大小变化的峰值内存使用量(左)和吞吐量(右)。残差大小设置为64。不同残差大小的结果见图8。
图4:随批次大小变化的峰值内存使用量(左)和吞吐量(右)。残差大小设置为64。不同残差大小的结果见图8。
图5:将VQ应用于经NSN处理的LLaMA2-7B键和值时,原始向量和重建向量之间的余弦相似度。相似度是在WikiText-2数据集上使用不同码本测量的。KM表示K-Means码本,FT表示微调码本。带有标记的线显示使用模型键值测量的余弦相似度,而虚线显示使用合成标准正态数据测量的结果。
图5:将VQ应用于经NSN处理的LLaMA2-7B键和值时,原始向量和重建向量之间的余弦相似度。相似度是在WikiText-2数据集上使用不同码本测量的。KM表示K-Means码本,FT表示微调码本。带有标记的线显示使用模型键值测量的余弦相似度,而虚线显示使用合成标准正态数据测量的结果。

结论

本文提出了NSNQuant,一种用于压缩LLM KV缓存的无校准向量量化(VQ)方法。NSNQuant通过NSN(归一化-中心化-归一化)过程有效地将token分布与标准正态分布对齐,从而允许使用专为标准正态分布设计的专门码本。通过广泛的实验表明,与基于校准的VQ方法不同,NSNQuant在不同的任务和数据集上具有良好的泛化能力。特别是在1-bit量化方面,NSNQuant表现卓越,以巨大优势超越了先前的最先进方法。此外,为NSNQuant实现了高效的CUDA内核,并验证了相较于FP16基线,NSNQuant实现了3倍的加速。未来的工作可以像混合策略那样整合token的重要性。

附录

引理1的证明与协方差分析
为了获得关于引理1界限的见解,测量了协方差矩阵的层级平均非对角线Frobenius范数。对于键,在Hadamard变换之前计算协方差。对于值,移除值投影矩阵中融合的Hadamard变换,并在NSN之后立即计算协方差。结果表明,前几层的Frobenius范数很大,而在后面的层中保持较低。这与在早期层中标准化失败的观察一致。前几层Frobenius范数很大的原因在于,键和值的协方差矩阵中存在异常值,导致了巨大的Frobenius范数。

图6:协方差矩阵的平均非对角线Frobenius范数。结果使用LLaMA2-7B在WikiText-2上测量。
图6:协方差矩阵的平均非对角线Frobenius范数。结果使用LLaMA2-7B在WikiText-2上测量。
图7:LLaMA2-7B第一层中键(上)和值(下)的非对角线协方差的最小-最大范围。对于键和值,某些通道在协方差矩阵中受到异常值的影响。我们使用WikiText-2的第一个样本进行可视化。
图7:LLaMA2-7B第一层中键(上)和值(下)的非对角线协方差的最小-最大范围。对于键和值,某些通道在协方差矩阵中受到异常值的影响。我们使用WikiText-2的第一个样本进行可视化。

NSN中的注意力计算
NSNQuant使用NSN过程的副产品来计算注意力权重和输出。首先,查询和键之间的点积计算如下:
$ qK^T = q(\mathrm{RoPE}(K_{\mathrm{pre-RoPE}}))^T = q\left(\mathrm{RoPE}\left(s_1^k (s_2^k v_{\mathrm{nsn}}^k + o^k)\right)\right)^T $
$ \quad = q\left(s_1^k s_2^k \mathrm{RoPE}(v_{\mathrm{nsn}}^k) + s_1^k \mathrm{RoPE}(o^k)\right)^T $
$ \quad = s_1^k s_2^k \mathrm{HT}(q)(\mathrm{HT}(\mathrm{RoPE}(v_{\mathrm{nsn}}^k)))^T + s_1^k q(\mathrm{RoPE}(o^k))^T $
$ \quad \simeq s_1^k s_2^k q_{\mathrm{Had}} v_Q^{k^T} + s_1^k q(\mathrm{RoPE}(o^k))^T $
这里,$\operatorname{HT}(\cdot)$是Hadamard变换,$v_Q^k := \mathbb{C}[\mathrm{VQ}(\mathrm{HT}(\mathrm{RoPE}(v_{\mathrm{nsn}}^k)))]$,$q_{\mathrm{Had}} := \mathrm{HT}(q)$。由于$o^k$的长度与$K$不同,在内核中扩展它以匹配形状。通过$W := \mathrm{softmax}(qK^T)$获得注意力权重,然后计算输出如下:
$ Wv = W s_1^v (s_2^v v_{\mathrm{nsn}}^v + o^v) \simeq W s_1^v (s_2^v v_{\mathrm{Q}}^v + o^v), \quad v_Q^v := \mathbb{C}[\mathrm{VQ}(v_{\mathrm{nsn}}^v)] $

随机Hadamard变换(RHT)的影响
从引理1中发现,在NSN过程之后采用随机Hadamard变换(RHT)为每个通道的方差提供了理论界限。然而,由于通道之间的协方差倾向于具有统一的符号,发现使用朴素的Hadamard变换已经足够好。两种变换给出了相似的结果,由于RHT需要更多的参数和计算,因此在最终设计中使用朴素的Hadamard变换。

残差大小的影响
残差大小是决定全精度缓存数量的重要超参数。使用LLaMA3.1-8B-Instruct评估了3种不同的残差大小(32, 64, 128)。对于NSNQuant-1b,观察到性能对残差大小的选择高度敏感。128的残差大小在大多数基准测试中产生了最佳结果,而32表现最差。相反,NSNQuant-2b在不同残差大小下显示出相当的性能。这表明NSNQuant-2b产生了更高质量的量化,使其对残差大小的变化更具鲁棒性。测量结果还显示,具有较小残差大小的方法需要更少的内存并实现更高的吞吐量。

图8:在不同批次大小和残差大小下测量的峰值内存使用量(左)和吞吐量(右)。
图8:在不同批次大小和残差大小下测量的峰值内存使用量(左)和吞吐量(右)。

尺度调整的影响
在不进行任何缩放的情况下使用量化向量$v_Q$会导致次优近似。为了确定更有效的缩放策略,评估了三种可能的方法。第一种方法是缩放$v_Q$以最小化$v$和$v_Q$之间的L2误差,第二种方法是缩放$v_Q$以匹配$v$的大小,第三种方法是缩放$v_Q$以保留平行于$v$的成分。对于键和值,策略3显示出最佳的量化质量。通过将其应用于NSNQuant-2b中的键和值,LLaMA2-7B在WikiText-2上的困惑度从5.395降低到5.285。

图9:三种测试的尺度调整策略的可视化。
图9:三种测试的尺度调整策略的可视化。

码本微调的影响
由于尺度调整的误差仅取决于$v$和$v_Q$之间的角度,将微调目标设置为最小化它们之间的余弦距离。结果表明,微调在NSNQuant-1b和NSNQuant-2b中都提高了量化质量,并在2-bit量化中匹配了E8P【37,QuIP#: Even better LLM quantization with hadamard incoherence and lattice codebooks+2024+ICML】的性能。

RoPE前NSN的影响
NSNQuant在RoPE之前将NSN应用于键。虽然在RoPE之后应用NSN似乎很自然(因为RoPE可能对通道级均值引入噪声),但RoPE前NSN实现了略低于RoPE后NSN的PPL。这归因于两个原因:首先,RoPE对通道级均值的影响很小,因为残差内的token共享相似的旋转角度;其次,计算$o$的全精度RoPE是有益的,因为RoPE旋转矩阵包含重要的位置信息。

双重量化的影响
应用双重量化时,性能退化极小,而内存节省是巨大的。

扩展CQ的校准集
测试了2个将校准集大小加倍到32的CQ变体。第一个变体使用WikiText-2中的32个样本(W32)。另一个变体使用WikiText-2的16个样本和C4的16个样本(W16C16),这增加了多样性。虽然W32仅显示出边缘改进,但W16C16提供了更明显的好处,表明为了提高CQ的鲁棒性,使校准集多样化很重要。尽管性能有所提高,NSNQuant仍然表现出更好的结果。由于CQ需要梯度计算,随着校准集的增长,校准过程会产生与微调相当的成本,因此NSNQuant仍然是一个高度实用和鲁棒的选择。

NSN的变体 - 使用Weiszfeld算法替换Shift步骤
NSN的最后归一化步骤可能会引入小偏差,使通道级均值略微偏离零。为了实证探索消除这种漂移的影响,将第二步(Shift)重新表述为寻找几何中位数。使用Weiszfeld算法(几何中位数的标准迭代求解器)计算中位数并在NSN的第二步中减去它。由此产生的变体使归一化后的均值基本保持在零。然而,这种改变仅产生了边缘负面影响,尽管增加了迭代更新的额外成本。因此,在最终的NSN设计中保留了原始的均值减法步骤。

NSN的变体 - 使用通道级缩放替换第二次归一化步骤
尝试用通道级缩放替换第三步(token级归一化),其中每个通道除以样本标准差。结果表明,使用通道级缩放不如token级归一化有效。这是因为通道级缩放不能抑制异常值token,尤其是注意力下沉(attention sink)token。考虑到码本的球形特性,这些token的量化误差会很大,因为它们的幅度远离零。相反,token级归一化有效地调节了尺度,使码本有效工作。

局限性与广泛影响
虽然NSN在经验上被证明能将输出分布与标准正态分布对齐,但发现由于存在巨大方差的异常值通道,它在早期层中效果不佳。此外,NSNQuant没有考虑通道间的相关性或依赖性,而CQ则因为过度依赖它们而引入误差。在两种方法之间寻找中间立场将是有益的。在广泛影响方面,本研究没有提供关于安全方面的任何实证分析。应用量化可能会引入意想不到的后果,例如幻觉或可靠性降低,因此在安全关键型应用中部署量化模型时建议保持谨慎。

图10:不同通道之间的层级平均绝对相关性(MAC)。结果是在WikiText-2上使用LLaMA2-7B测量的。
图10:不同通道之间的层级平均绝对相关性(MAC)。结果是在WikiText-2上使用LLaMA2-7B测量的。
图11:LLaMA3.1-8B键和值的t-SNE可视化。
图11:LLaMA3.1-8B键和值的t-SNE可视化。
图12:在WikiText-2(上)和C4(下)上对LLaMA2-7B应用NSN和Hadamard变换后,键的通道级均值和标准差的可视化。每个数据集测试分割的第一个样本用于可视化。虽然NSN总体上很好地执行了标准化,但它在早期层的某些头中遇到了困难。
图12:在WikiText-2(上)和C4(下)上对LLaMA2-7B应用NSN和Hadamard变换后,键的通道级均值和标准差的可视化。每个数据集测试分割的第一个样本用于可视化。虽然NSN总体上很好地执行了标准化,但它在早期层的某些头中遇到了困难。
图13:在WikiText-2(上)和C4(下)上对LLaMA2-7B应用NSN和Hadamard变换后,值的通道级均值和标准差的可视化。每个数据集测试分割的第一个样本用于可视化。虽然NSN总体上很好地执行了标准化,但它在早期层的某些头中遇到了困难。
图13:在WikiText-2(上)和C4(下)上对LLaMA2-7B应用NSN和Hadamard变换后,值的通道级均值和标准差的可视化。每个数据集测试分割的第一个样本用于可视化。虽然NSN总体上很好地执行了标准化,但它在早期层的某些头中遇到了困难。