LLM推理效率革命：KVLatent下采样技术破解内存Bottleneck！

备注好友：方向-学校/公司-姓名/昵称

【AIGC 先锋科技】交流群

基于Transformer解码器的巨型语言模型（LLMs）已成为对话生成式AI的首选方案。尽管解码器架构总体上具有优势，但在推理过程中逐渐增大的键值（KV）缓存已成为主要的效率 Bottleneck ，体现在内存消耗和数据传输带宽限制方面。为应对这些挑战，作者提出了一种名为KVLatent的范式。

通过将键值向量维度下采样到潜在空间，作者可以显著减少KV缓存的占用空间并提升推理速度，仅需少量额外训练，少于预训练1%的时间。

此外，作者通过修改其频率采样机制，增强了应用于低维向量的旋转位置嵌入的稳定性，避免了高频引入的噪声，同时保留了位置衰减特性。

作者的实验，包括具有分组 Query 注意力和不具有分组 Query 注意力的模型，均取得了令人满意的结果。最后，作者进行了对比实验，研究了分别减少 Key和Value 组件对模型性能的影响。

llm-KV-Latent_2507使得构建更高效的语言模型系统成为可能，并为KV缓存节省和高效LLMs开辟了新的可能性。

作者的代码可在https://github.com/ShiLuohe/KVLatent获取。

1 引言

ChatGPT（Brown等人，2020）的发布开启了生成式AI的趋势，而作为这些最先进模型的核心架构，Transformer解码器（Vaswani等人，2017）获得了广泛关注。它们的能力非常显著。不可否认，随着大语言模型（LLMs）越来越融入人们的生活，训练和推理相关的成本越来越难以忽视（Yan等人，2024）。虽然训练成本相对固定且集中，但推理成本随着用户采用量的增加而线性增长，并且通常分布在不同的空间和时间框架内，这使得优化模型推理成本日益紧迫。LLMs采用的Transformer解码器架构作为因果模型运行，避免了在自回归生成过程中重新计算大多数中间状态。然而，它仍然需要保留某些中间状态。具体来说，作为一种基于自注意力机制的架构，它需要保留与每个 Token 对应的键值（KV）状态，通常被称为KV缓存。自注意力机制的时间复杂度均匀为

，这意味着对于序列中的每个附加 Token ，计算工作量至少增加

。因此，在典型情况下，作者需要与

个缓存状态进行交互。换句话说，KV缓存所需的存储量随着 Token 的生成而线性增长。这构成了一个重大挑战。

KV缓存面临两大主要挑战：不断增长的体量和不友好的访问模式。巨大的体量要求越来越昂贵的硬件来高效存储和检索KV缓存，此外，由于每个推理请求都维护自己的专用KV缓存，使用批量处理加速系统是不可能的，导致芯片上的RAM带宽 Bottleneck 和计算资源的浪费（Williams等人，2009）。与此同时，不友好的内存访问模式是由于缓存大小频繁波动造成的。后者挑战可以通过更科学组织的缓存结构显著缓解，例如分页注意力（Kwon等人，2023）或异构推理系统如fastdecode（He和Zhai，2024）。然而，前者挑战仍然更为复杂。

针对KV缓存大小的问题，已提出多种方法。在注意力头 Level ，多 Query 注意力（MQA，Shazeer，2019）和分组 Query 注意力（GQA，Ainslie等人，2023）是有效且已被广泛验证的方法。在层 Level ，提出了跨层重用方法，例如仅缓存一次（You Only Cache Once，Sun等人，2024）和跨层注意力（Cross Layer Attention，Brandon等人，2024）。在token Level ，研究行人专注于驱逐和合并，方法包括重 hitter预言机（Heavy Hitter Oracle，Zhang等人，2023）、金字塔推理（PyramidInfer，Yang等人，2024b）、SirLLM（Yao等人，2024）以及Devoto等人提出的

范数方法。

尽管先前研究取得了显著进展，但直接减小Key和Value头的尺寸仍是一个较少探索的领域。在MHA（Multi-Head Attention）的框架下，每个注意力头是两个低秩变换的组合，第一个是K和Q^T的组合，第二个是V和O的组合。作者定义每个注意力头的维度为dh，头的数量为nh，模型的隐藏维度为d。作者观察到K和V代表了两个线性变换，它们将d维隐藏状态h下采样为两个dh维向量k和v。相应地，Q^T和O执行从dh到d维的上采样。KV Cache存储了由这两个低秩变换产生的潜向量。

通常情况下，作者假设

，但在考虑单个注意力头时，

和

不必遵循这一预设关系。MQA和GQA的工作以及其他近期研究（Yu等人，2024；DeepSeekAI等人，2024；Saxena等人，2024a）已经证明，保留的KV缓存无需完整的

维向量；低秩表示足以在token之间传递信息。然而，作者旨在进一步解除

的约束。llm-KV-Latent_2507涉及直接从现有模型中减小头大小，然后通过两阶段策略进行少量额外训练来恢复模型的性能，从而实现KV缓存减少的目标。由于作者将Key和Value映射到潜在空间，并通过Query转置和Output直接解码，因此作者将所llm-KV-Latent_2507命名为KV-Latent。

此外，作者观察到，即使在单个注意力头内部，

和

的低秩变换也不一定需要相同的维度。具体来说，作者可以将

分解为

和

，并且这两个维度不必相等。基于这一见解，作者探索了具有不同

和

值的各种降维策略，以研究它们对训练时间、推理效率以及最重要的方面——模型能力的影响。

最后，在作者的实验中，作者发现旋转位置嵌入（RoPE，Su等人）的稳定性在低维情况下会降低，从而影响其长距离能力。通过分析RoPE的采样机制，作者识别出当维度较低时，高频特征的噪声会主导。因此，作者通过以频率感知的方式修改RoPE的采样方法，优化了llm-KV-Latent_2507，即使在低维情况下也能保持稳定性。

作者的贡献包括：

1. 作者已证明，通过采用两阶段策略进行少量额外训练，可以将KV缓存拟合到潜在空间中，从而直接减少KV缓存的存储空间占用和带宽需求。
2. 通过使用

和

的不同组合，作者观察到模型的性能对

的敏感度比对

更高，这揭示了LLMs如何受到自注意力机制不同部分的影响，为优化LLMs的模型结构提供了启示。 3. 通过以频率感知的方式修改RoPE的采样机制，排除高频部分并放大低频部分，作者成功使RoPE在应用于低维Query和Key时更加稳定。

2 背景和相关工作

2.1 Transformer解码器

作者的主要关注点在于Transformer解码器（Vaswani等人，2017）的 Mask 自注意力机制。作者将

定义为第

层的输入在token Level 上的隐藏向量，而

则代表整个序列。作者的目标是计算

，它表示注意力模块的输出。该过程由公式1描述，其中

对应于第

个头的

、

、

和

变换的参数矩阵。而

、

则表示第

个头的KV缓存。在此上下文中，作者应用了右乘操作。

2.2 KV 缓存减少方法

2.2.1 层级头

MQA（Shazeer，2019）将所有Key和Value头合并为两个单一头，并使用不同的Query头多次 Query 单一Key头。在此基础上，GQA（Ainslie等人，2023）对所有注意力头进行预分组。在每个组内，多个Query头共享单个Key头，并相应地共享单个Value头。GQA引入了一个可调变量，即组数

以及每个组内的头数，找到了一种新的MQA与MHA（多头注意力）之间的权衡方法。这种方法在效率和性能之间提供了细粒度的平衡，提升了运算强度，并已广泛应用于LLaMA2（Touvron等人，2023）、LLaMA3（Dubey等人，2024）、Mistral（Jiang等人，2023，2024）和Qwen（Yang等人，2024a）等模型中。这些工作证明了KV缓存的低秩特性，从而保证了llm-KV-Latent_2507的有效性。

2.2.2 层级

跨层复用是另一个热门话题。You Only Cache Once（Sun等人，2024）和Cross Layer Attention（Brandon等人，2024）等方法通过跨不同解码器层复用相同的KV缓存状态，成功减小了KV缓存的大小。然而，由于复用内容随时间呈现非连续性，跨层复用无法有效优化计算强度，且数据交换带来的带宽限制依然存在，限制了推理速度的提升。

2.2.3 词元 Level

在token Level ，驱逐（Liu et al., 2023）和合并（Pang et al., 2024; Tang et al., 2025b）是最重要的方法，其核心思想是驱逐低注意力token，或将多个token的状态合并为一个。代表性工作包括Heavy Hitter Oracle（Zhang et al., 2023）、PyramidInfer（Yang et al., 2024b）、SirLLM（Yao et al., 2024）以及Devoto等人提出的

范数方法。token Level 压缩的潜在问题在于其依赖于注意力分数，因此无法与预填充加速方法（例如flash attention（Dao et al. 2022））结合使用。其他不依赖注意力分数的方法通常粒度不够精细，存在关键信息丢失的风险。实现大规模实际应用仍然是一个挑战。

2.3 旋转位置嵌入

旋转位置嵌入（RoPE），由Su等人于2021年提出，是一种增强解码器模型位置编码的方法。这种位置编码方式因其多种优良特性而得到广泛应用。首先，它遵循长程衰减特性：在序列中，两个相同向量之间的距离越远，它们的注意力连接就越弱。其次，RoPE是一种相对位置编码形式，意味着对于相同的相对距离，衰减保持一致。这一特性有助于提升泛化能力。最后，RoPE通过Sparse矩阵实现编码，从而提高计算效率。这些优良特性使其成为现代大语言模型几乎唯一的选择。然而，作者的实验发现，由于高频周期性噪声，RoPE在低维情况下表现出不稳定性。作者通过修改其频率采样方法缓解了这一问题。

3 方法

3.1 预备知识与符号说明

将RoPE应用于F1，作者得到F2。

指代单个token的隐藏状态，相应地，

表示整个序列的隐藏状态。自注意力机制的四个线性变换参数由

给出，这里的变换形式为右矩阵乘法。作者定义

为每个 Query 和键头的维度，

为值和输出头的维度。以

表示头的数量，作者可以得到

，

，

。在这种情况下，作者定义

为与第

个头（

）对应的参数，如公式3所示。

作者额外定义

为当前词元的位置，

为RoPE中定义的在第

个位置上的旋转矩阵，

。更具体地说，根据RoPE，

由公式4给出。

3.2 基于KV-Latent的两阶段训练

作者提出了KV-Latent范式，旨在通过直接修改预训练模型的

和

的形状来减少KV缓存。随后，作者通过使用更少的数据进行微调来恢复模型性能。该范式包括一个与RoPE兼容的注意力下采样策略和两阶段延续训练。

3.2.1 模型准备

在训练之前，作者需要在注意力头维度降低后初始化模型的副本（Ma等人，2025）。对于任何给定的注意力模型，仅随机采样就足以以适当平衡的方式保留注意力矩阵中的信息，因为每个注意力头内的通道表现出旋转对称性。这意味着，对于

和

，只需保留相同的通道即可。

然而，RoPE的引入未能采用这种方法，因为RoPE涉及在不同频率下旋转通道对。具体的实现方式，包括Sparse矩阵乘法和GPT-NeoX（Black等人，2022年）中发现的现代通道分组方法，在附录C中详细说明，其中均匀下采样足以用于权重初始化。利用这种方法，公式5描述了一个将

缩小一半并将

缩小四分之三的示例。

近期研究也应用奇异值分解（SVD）进行下采样（Saxena等人，2024b；Zhang等人，2024），然而这些方法面临主要困难，因为矩阵乘法不满足交换律，因此不适用于此处。

在降采样步骤之后，作者还希望训练 FFN （FFNs）以更好地使模型适应其修改后的注意力机制，但又不完全忘记其在预训练中学到的知识，因此作者应用低秩适配（LoRA，Hu等人，2022）到FFN的变换中，包括在大语言模型（LLM）中通常适配门控线性单元（GLU）作为FFN的Up、Down和Gate。图1描述了作者的降采样和模型构建过程。

3.2.2 第一阶段 - 层内蒸馏

在训练的第一阶段，作者的目标是保持两个解码器层中隐藏状态

之间的最大一致性。这种方法确保了作者最大限度地保留了模型的初始能力。为了实现这一点，作者采用了一种层内蒸馏方法，如图2所示。

作者定义

为初始模型中第

个Transformer解码器块的运算，以及

为其修改版本，该版本具有减小了

头大小的版本，并利用KV-Latent。作者首先使用原始的

进行推理，并在每一层之间保留中间隐藏状态

。对于第

层，作者定义三个形状相同的隐藏状态：

、

、

、

，这些状态由公式6获得。

作为输入、目标和预测的隐藏状态。作者希望最大化 H(l)a 和

之间的相似度。为此，作者使用均方误差（MSE）损失。作者将

定义为

的可训练权重。作者的优化目标在公式 7 中描述。

3.2.3 第二阶段 - 端到端训练/蒸馏

尽管进行了层内蒸馏，但在现代大语言模型上应用KV-Latent仍然面临挑战，这主要源于大语言模型的深度。即使是微小的扰动也会逐层放大，可能损害其模型的语言能力。为了解决这个问题，作者需要进行端到端的模型训练。在这个阶段，作者有两个选择：Next-Token-Prediction（NTP）训练和蒸馏。作者首先定义原始模型 $ \mathrmLM } ( \cdot )

和

作

者

的

模

型

\mathrm { L } \tilde { \mathrm { M } } ( \cdot )

，

并

将

\mathcal { C } = { c _ { 1 } , c _ { 2 } , \ldots , c _ {\mathcal { C| } }

定

义

为

用

于

训

练

的

语

料

库

，

其

中

c _i } = { t _ { 1 } , t _ { 2 } , \ldots , t _ {c _ { i| } } $ 。NTP训练是模型预训练的一部分，并采用交叉熵损失，如公式8所述。它需要较少的资源，然而，交叉熵损失提供的信息有限。

基于预测概率分布的蒸馏方法常用于模型恢复，该方法依赖于公式9中描述的KL散度损失。蒸馏帮助模型在相同的数据量下学习更多。然而，蒸馏需要额外的正向传递来计算概率分布，并需要维护一组额外的参数。

3.3 频率感知RoPE用于可变维度

3.3.1 动机

RoPE将位置信息引入注意力头的

和

分量中。在作者的初步实验中，作者观察到RoPE在应用于低维向量时表现出显著的数值不稳定性，如图3所示。具体来说，当维度

小于32时，振荡范围与预期衰减相当，表明位置编码能力丧失。根据Su等人所述，RoPE编码的向量即使在较大距离上也应保持一定程度上的自身相似性。作者可以通过使用特殊向量

来测量这一点。作者在公式10中定义

为两个相同向量在不同距离

下的相似性表示，其值理想情况下应始终为正，以表明其比两个随机向量更为相似。

3.4 有效性分析

为解释llm-KV-Latent_2507为何有效，作者给出以下推导过程。从图3中的理想曲线可以看出，随着

的增加，RoPE接近一条平滑的衰减曲线。该曲线的计算公式由公式12给出，详细内容见附录D.1。

3.3.2 模式发现

作者研究了不同

值对

的影响，如图3所示。作者观察到较小的

值会导致更大的噪声，并伴随负值出现的频率增加。作者将向量按通道分解。基于256维的情况，图4和图5分别展示了保留低编号通道和高编号通道的场景，而图6描绘了保留64个连续通道集（总量的四分之一）的RoPE函数。作者的研究结果表明，RoPE函数的低编号通道是噪声的主要来源，而高编号通道尽管衰减较慢，但相对稳定。这与一些先前的工作（bloc97, 2023; Peng et al., 2024）一致。

3.3.3 频率感知修改

作者实现了一种频率感知的修改策略，该策略包括对低频旋转进行密集采样并避免高频旋转采样，如公式11所述，因为编号较小的通道对应高频旋转，而编号较大的通道对应低频旋转。图7、8、9和10所示的结果表明，llm-KV-Latent_2507在提高稳定性的同时，也减少了负值的发生。

此时，作者将RoPE的稳定性问题转化为一个积分数值逼近问题。具体而言，对于函数

在长度为1的区间上的积分（如图11所示），作者通过执行

个采样来近似求解。当

取较大值时，该函数表现出剧烈的振荡，随着

的增加，积分窗口向右滑动，不可避免地进入这些剧烈振荡的区域。因此，为了精确求解积分，

必须足够大以适应

的增加。如果

太小，采样间隔可能短于振荡周期，导致数值逼近失去有效性并引入大量噪声。作者在附录E中提供了一个代码块来模拟这一现象。

此外，通过舍弃右侧的某些采样点，作者的改进在提高整体采样密度的同时，延缓了积分窗口进入剧烈振荡区域的时间，增强了RoPE的稳定性，从而降低了噪声幅度。此外，额外采样点的值通常接近1，而舍弃点的值在1和-1之间振荡。因此，感知频率的RoPE值几乎总是大于原始RoPE值，具体细节见附录D.2。

4 实验

4.1 训练

作者使用了FineWeb-edu（Lozhkov等人，2024年），它源自于FineWeb（Penedo等人，2024年），一个基于开放获取网页的网页数据集，包含150万亿个token。作者从FineWeb-edu中使用了1万亿个token的子集，这也是其他知名数据集（如minipile，Kaddour，2023年；Gao等人，2020年）常用的规模。作者的训练超参数在附录A中详细说明。作者在一个包含8块NVIDIA A100 80G SXM4 GPU的节点上进行了训练。

图4：仅低维

图5：仅高维

在模型层面，作者在LLaMA-3-8B（L3-8B）上训练了两个版本的KVLatent，其参数配置为

(64, 64)和(16, 16)，作为GQA的示例；另一个版本在LLaMA-2-7B（L2-7B）上训练，其参数配置为

(64, 64)，作为MHA的示例。

4.2 评估

作者从性能和效率两个角度对KV-Latent模型进行了测试。在性能方面，作者使用了O-shot MMLU（Hendrycks等人，2021年）、OBQA（Mihaylov等人，2018年）和AI2ARC（Clark等人，2018年）作为基准测试。此外，作者还进行了“大海捞针”（NIH）测试以评估信息检索能力。作者将一个短句随机放入一个包含3,840个token的上下文中，检查模型是否能够复述该句子，重复50次并计算通过率。在效率方面，作者在NIH实验中测量了KV缓存大小

（MB）以及第一个token的延迟

（ms）。作者还计算了改进率

和

。结果如表1所示，其中包含几个关键观察结果。首先，KV-Latent允许模型在减少KV缓存大小的情况下实现令人满意的性能。其次，尽管蒸馏传递了更多信息，但有限的训练量使其无法完全恢复模型的性能。第三，当

时，模型的性能未能恢复，这表明KV缓存大小的下限。最后，不使用GQA的LLaMA2在训练token较少时相对优于LLaMA3，这表明对于已经使用GQA训练的模型，采用KV-Latent将带来额外的挑战。

4.3 参数选择的影响

作者研究了不同的

、

和 LoRA 排序对模型性能的影响。作者使用 LLaMA-3-8B 基础模型进行了实验，并训练了多个配置不同的 KV-Latent 版本。默认情况下，作者设置

，LoRA 排序为 256。对于与效率相关的测试，作者基于 2048 的上下文长度生成了 256 个 Token ，重复该过程 15 次并取平均值。

4.3.1 QK和VO头大小的组合

作者测试了不同的组合

和

对模型性能和效率的影响。作者涵盖了三个方面：对数困惑度（log PPL），反映模型的语言建模能力；训练速度

，衡量训练效率；推理速度，包括生成第一个token的时间

，以及每个新token的毫秒数

。在空间方面，对于长度为4,000的token序列，在bf16浮点数下，KV Cache的大小

。为了更直观地表示，作者计算了在60GB内存下支持的最大KV Cache大小

，以80GB计算卡的情景为例，不包括大约15GB的模型参数和5GB的缓冲区。结果如表2所示。

作者发现，首先，与KV缓存相关的效率与其大小一致：总体容量越小，预填充和生成速度越快。然而，在比较减少

与

的结果时，在附录B的表6中，作者注意到将更多资源分配给

可以持续获得更好的效率和效果，这表明在KV缓存中，键（Keys）携带的信息不如值（Values）重要，因此更适合压缩。

4.3.2 LoRA Rank

LoRA的秩可能影响KV-Latent的性能和效率。作者专注于评估训练效率和对数困惑度（log PPL），因为LoRA在推理时没有额外成本。如表3所示，增加秩对应于增加训练时间。然而，作者注意到对数困惑度的变化并不显著。需要注意的是，LoRA的秩在更大规模的训练场景中可能具有更显著的影响。

4.3.3 跨方法可行性

在与其他方法的兼容性方面，KVLatent与Head-Level配合良好，这一点在LLaMA-3的测试中得到了验证。它也兼容Layer-Level，尽管这需要更高的训练资源要求。最后，llm-KV-Latent_2507也兼容Token-Level。表4展示了当与PyramidInfer（一种流行的token级压缩方法，压缩率为50%）结合使用时的结果，证明了作者的观点。KV-Latent与所有主流方法都兼容。

5 结论

作者提出了KV-Latent范式，通过两步训练过程直接降低模型的注意力头维度，从而减小KV缓存大小。该方法在仅使用少量额外训练token的情况下，实现了缓存减少和推理速度提升。作者证明了将

和

的关系解耦是可行的。值得注意的是，作者发现

对模型性能的影响更大，揭示了KV缓存中值与键之间信息的不平衡。

最后，通过修改频率采样方法，作者在保持RoPE衰减特性的同时增强了其稳定性。作者的工作可能为优化模型结构的研究做出贡献。

参考

[1]. KV-Latent: Dimensional-level KV Cache Reduction with Frequency-aware Rotary Positional Embedding