创新之举：不丢 Token，大语言模型效率提升 35% !

大语言模型（LLM），如GPT和LLaMA系列，随着上下文窗口大小的增加，在处理复杂长文本任务方面的能力得到了提升，但这也导致了推理效率的降低，特别是在内存和计算复杂性方面的成本。现有的方法，包括选择性地保留Token和基于窗口的注意力机制，虽然提高了效率，但也存在舍弃未来文本生成所需的重要Token的风险。

在本文中，作者提出了一种方法，通过减少不重要Token的内存和计算负担而非直接舍弃它们，来增强LLM的效率而不损失Token。作者解决了两个挑战：

1）研究上下文中文本中重要Token的分布情况，发现近来的Token比远期的Token更具有重要性；

2）通过跨层共享注意力分数优化远程Token的资源分配。实验结果显示，作者的方法在不牺牲性能的情况下节省了35%的KV缓存。

注意力模块随上下文窗口大小呈二次增长。另一方面，KV缓存（Pope等，2023）的大小，一种常用的技术，旨在防止冗余计算，与上下文窗口大小呈线性关系。因此，提升具有扩展上下文窗口的大语言模型的效率至关重要。

在这一背景下，众多研究者提出了通过舍弃上下文中某些Token来提升大语言模型推理效率的方法。尤其是窗口注意力方法（Beltagy等，2020）保留了一个固定大小的窗口覆盖最近Token的KV状态。

LM-Infinite（Han等，2024）和StreamingLLM（Xiao等，2024）识别出了“注意力陷阱”现象，并保留了初始Token和最近的Token（参见图1-(a)）。（Zhang等，2023）考虑到了上下文中Token的不同重要性，并根据注意力分数选择性地保留KV缓存中最关键的Token。

虽然这些方法可以提高大语言模型处理长上下文的效率，但它们也引入了一个主要缺点（Tang等，2024）：在过程早期可能会不可逆地丢弃后续文本生成所需的关键Token。如图1-(a)所示，在重要Token（例如示例中的证据）落在窗口之外时，预测会失败。此外，StreamingLLM和在两个真实世界基准上的性能下降进一步证实了这一点（参见图1-(b)）。

1. Introduction

近年来，大语言模型（LLMs）不断增加的上下文窗口大小使得它们能够胜任需要深入探索长篇文本的复杂任务（Bairi等，2024；Mazumder和Liu，2024）。然而，这也对LLMs的计算能力和内存占用提出了挑战。具体来说，一方面，大多数LLMs基于Transformer（Vaswani等，2017）架构，其计算复杂度随着上下文窗口的增加而显著提升。本文旨在在减少性能下降的同时提高LLMs的效率。作者的核心动机在于，对于不那么重要的Token，分配较少资源而不是完全丢弃。

这带来了两个挑战：

1）一个Token要关注哪些重要的Token，重要Token的分布在哪里？

2）如何优化不那么重要的Token的内存和计算。作者通过两个关键观察来尝试解决这两个挑战：

观察1：相邻Token（初始Token+近期Token）比远处Token更重要（参见图1-(c)）。作者进行了实验并调查了当每个Token只关注固定数量的相邻Token时（相对于关注所有Token），下一Token预测结果相同的案例比例。图1-(c)显示，在同一输入序列中，即使只关注256个相邻Token，模型在80%的情况下对下一个Token的预测与关注所有Token的模型相同。这一现象证实了观察1。

观察2：相邻层之间的注意力分数相似。事实上，这一现象在较小的模型中之前已经被观察到（Xiao等，2019；Bhojanapalli等，2021），在这里作者将其扩展到现代大语言模型。如图1-(d)所示，红色框内的各层之间的注意力分数表现出较强的相似性。

为了实现这一目标，作者提出了PoD（近端Token优化远程Token）方法，在解码阶段优化推理效率。具体而言，该方法仅在层间共享远程Token的独特互层注意力分数，而保留近端Token不变，基于上述两个观察结果。

该方法包括三个主要阶段：

1）离线层间注意力共享的探索：确定哪些层可以共享注意力分数；

2）轻量级训练适应：根据识别出的层间注意力共享模式对密集模型进行有限数据的后训练调整；

3）高效推理：层间共享远程Token的注意力分数，使得作者可以将关键状态保存在一个层的KV缓存中。此外，作者还可以预先识别出仅需使用近端Token来预测下一个Token的情况，从而消除对远程Token的注意力计算。

作者在Needle in a Haystack和两个实际场景中的长上下文基准上评估了PoD的表现，分析了其效率，并检查了关键超参数的影响。还进行了案例研究。大量实验表明，PoD可以在不牺牲模型性能的情况下节省35%的KV缓存空间。

总结起来，作者的贡献包括：

1）作者提出了一种基于Token在上下文中的位置来评估其重要性的想法，并通过减少对较不重要Token资源的分配来提高推理效率。

2）作者引入了PoD，这是一种新模型，能够适应远距离Token间的跨层共享注意力分布。

3）作者进行了广泛实验以证明PoD的有效性，并计划在未来开源作者的代码和模型。

2. Methodology

作者的方法包含三个关键步骤（图2）：

1）分析给定长上下文大语言模型中各层注意力分数的相似性，并将连续相似的层分组为块；

2）在每个块内应用注意力共享并对大语言模型进行后训练；

3）通过使用后训练的大语言模型进行高效的推理。

2.1. Offline Inter-Layer Attention Sharing Exploration

2.2. Lightweight Training Adaptation

为了增强模型在同一块内部的各层之间共享注意力的能力，作者首先介绍了所需的后训练过程，以使模型适应这一机制。

2.3.Efficient Inference

3. Experiments

在本节中，作者探讨了两个关键问题：1) PoD 是否出现性能下降；2) PoD 是否提高了长上下文推理的效率。

实施细节 在数据方面，作者从Dolma（Soldaini等，2024）中采样了总共5亿个Token进行后训练，确保每个长度区间内的总Token数量保持一致（GLM等，2024）。在模型方面，首先使用LLaMA3-8B进行了初始化，并对5亿个Token进行了后训练，最大序列长度为32K，得到了LLaMA3-8B-32K。随后，作者从LLaMA3-8B-32K进行了初始化，并对相同的5亿个Token进行了后训练，最大序列长度仍为32K，最终获得了PoD模型，其参数设置为和。层相似性阈值设置为0.5，对应保留35%的KV缓存状态。在训练过程中，批量大小设为4百万个Token，学习率为1e-5，并采用余弦学习率调度器。此外，在RoPE（旋转位置嵌入）（Su等，2023）中，基底增加到，与(Xiong等，2024)中的方法一致。在代码实现方面，作者利用了HuggingFace（Wolf等，2020）以及DeepSpeed（Rasley等，2020），结合了ZeRO-3（Rajbhandari等，2020）和Ulysses（Jacobs等，2023）序列并行技术，并采用了来自PyTorch（Paszke等，2019）的高效FlexAttention。

baseline 作者主要考虑以下三种 Baseline ：

1. 基于Token选择的方法：1) SnapKV（Li等人，2024）：重要Token基于注意分数进行选择，并且仅在预填充阶段缓存这些重要Token。2) PyramidKV（Zhang等人，2024b）：这项工作扩展了SnapKV，在不同层中缓存的Token数量有所不同。3) Quest（Tang等人，2024）：这项工作不减少KV缓存的大小，而是通过高效的Token选择方法减少了从KV缓存中读取参与注意计算的Token的数量。
2. 基于Token驱逐的方法：1) 窗口注意力（WA）（Beltagy等，2020）：每个Token只会关注有限数量的相邻Token。2) 窗口注意力+CPT（连续预训练）：继续在相同的5亿个Token上对LLaMA3-8B-32K进行窗口注意力的预训练。3) Stream 大语言模型（Xiao等，2024）：除了相邻的Token外，每个Token还会关注初始的几个Token。4) LM-Infinite（Han等，2024）：每个Token关注的Token与 Stream 大语言模型相同，但位置嵌入有所不同。5) ：每个Token不仅专注于相邻的Token，还会根据解码过程中的注意分数动态添加重要Token并移除较不重要的Token。
3. 基于层共享的方法：CLA（Brandon等人，2024）通过在相邻层之间共享 Key和Value 状态来减少KV缓存。

3.1.Performance Evaluation

为了评估PoD的表现，作者在两个领域进行了实验：1）针尖 haystack 和 2）实用的长上下文基准。

万针之中寻一针 该任务在长上下文中随机插入一个陈述，并要求模型检索该陈述。图3展示了不同方法的搜索结果。作者发现，在“万针”之外的区域，StreamingLLM和无法正确检索到目标陈述。相比之下，作者的方法避免了token丢失，检索性能与密集模型相当，并且几乎可以定位到所有目标陈述。

长文本基准测试 为了确保PoD能够处理现实世界任务，作者在两个知名的长文本基准测试上进行了评估：LongBench（Bai等，2024）和LEval（An等，2024）。作者对LongBench中的14个数据集进行了测试，这些数据集涉及单文档问答、多文档问答、总结、少样本学习以及代码补全任务。LEval包含20个子任务，并分为两组：封闭域和开放域。封闭域组主要评估在较长上下文中进行推理和理解的能力，而开放域组则侧重于总结和问答等需要从长文档中聚合信息的任务。

表1列出了所有实验结果。为了确保公平性，所有 Baseline 注意力机制具有相同的窗口大小。对于PoD，作者还确保每个Token关注的近邻Token数量与该窗口大小一致。

作者得出了以下结论：

1）PoD优于基于Token移除的方法，这表明作者不丢失Token的做法确实是有效的。

2）即使在少量后训练数据的支持下，PoD也超越了经典的基于层共享的方法CLA，这证明了作者的模型在适应现有大语言模型方面具有优势。

3）PoD和基于Token选择的方法都能达到与标准密集模型相当的性能。此外，PoD与基于Token选择的方法是正交的。

3.2.Efficiency Evaluation

内存占用 在理论和实际层面都可以分析内存消耗的节约情况。理论上，作者可以根据离线分析获得的层共享结果来计算潜在的KV缓存大小减少量。实际操作中，作者可以通过端到端的评估来评估实际的节约效果。借鉴FlexGen（Sheng等人，2023）和LCKV（Wu & Tu, 2024）的方法，在给定长度为的 Prompt 下，让模型生成个Token。使用该GPU上能够达到的最大批量大小来评估模型的内存需求。较大的表明模型更具内存效率。表2展示了内存消耗的结果。作者发现PoD在不同输入文本长度的情况下，最大批量大小增加了超过30%，这与作者理论预测的KV缓存节省率35%非常接近，表明PoD有效地减少了内存使用。关于性能、KV缓存节省以及优化阶段的比较，详见附录A.2。

远距离Token的计算正如第2.3节所述，某一层的注意力输出是由最近和远距离Token加权输出得出的。权重系数可以在应用注意力机制之前，通过下一层共享的注意力分数提前获得。因此，当最近Token的权重系数大于时，作者在此层当前解码时间点认为远距离Token对预测下一个Token无关紧要，从而可以跳过它们的计算。

图4-(a)展示了在LEval上的计算节省比例与性能损失与值之间的关系。作者观察到，随着的减小，忽略远距离Token的计算变得更容易，从而导致计算节省增加，但同时也伴随着一些性能损失。然而，当时，性能下降的趋势减缓，而计算节省则更为显著。当时，计算成本降低了25%，同时只带来了5%的性能下降。

将 LLaMA3-8B-32K 的同一超参数配置应用于训练 LLaMA3.1-8B，并使用序列长度为 128K 进行后训练。在不同上下文大小下，InfiniteBench (Zhang et al., 2024a) 的 5 个子任务上的评估结果见表3。

与表1中的结论一致，作者的方法在性能下降方面比基于Token驱逐的方法要小。然而，一个显著的区别是，基于Token选择的方法在处理较长上下文情景时似乎难以维持模型性能。这一局限性也在结合了作者方法和基于Token选择方法的综合模型（PoD+SnapKV）中有所体现，该模型显示了性能下降的趋势。这在一定程度上表明，作者的方法对上下文长度更具有鲁棒性。

3.3.Additional Analysis

扩展到更长的上下文和其他大语言模型 为了探索作者方法的一般性，作者在LLaMA3.1-8B（Dubey等，2024）上进行了实验，该模型能够处理长达128K的上下文。作者从ProLong-data-512K（Ga0等，2024）数据集中采样了5B Tokens进行实验。

PoD模型性能与两个关键超参数的关系：最近邻Tokens的数量和KV缓存保存率之间的关系 从LLaMA3-8B-32K初始化开始，作者继续使用2B数据进行训练以开展实验。如图4-(b)所示，随着最近邻Tokens数量的增加，PoD的性能稳步提升。当Token计数达到4K时，使用2B数据进行训练的性能可以接受，与LLaMA3-8B-32K阶段相当。一些工作试图通过仅缓存重要的输入Tokens来减少生成的KV缓存大小，如SnapKV、PyramidKV和LazyLLM（Fu等，2024）。同时，MInference（Jiang等，2024a）和RetrievalAttention 利用Transformer固有的稀疏注意力机制来最小化预填充延迟。此外，还有一些方法通过直接压缩输入Prompt的长度来提高效率。本工作与这些方法相互独立，并主要集中在优化解码阶段。

使用了50亿条数据进行训练。在平衡性能和效率之后，作者最终选择了使用4096个近邻Token。

关于模型性能与KV缓存节省之间的关系，图4-c显示，随着节省率的增加，性能会下降。为了在性能和效率之间找到平衡，作者将KV缓存压缩至35%，并在相同的训练数据下，保持性能与LLaMA3-8B-32K相当。

案例研究 在图5中，作者提供了四种不同的代表性案例，用于比较StreamingLLM、和作者的PoD。对于案例(a)，正确答案位于相邻Token的窗口内，使所有方法在解码过程中都能注意到它，因此所有三种方法都能做出准确的预测。对于案例(b)，答案位于开头位置，且StreamingLLM和PoD可以在解码过程中注意到它，从而导致正确的预测。

然而，对于来说，紧跟在答案后的冗余长序列使其误判了最初的答案，导致错误的预测。对于案例(c)，答案位于中间，StreamingLLM无法在解码过程中注意到它，因此导致错误的预测。不过，答案之后与问题相关联的文字使得能够将其包含在其注意力窗口内，从而做出正确的预测。由于PoD可以注意到所有Token，它也做出了正确的预测。对于案例(d)，只有PoD能够回答“针叶林中的刺”的示例，因为其他两种方法忽略了答案Token。

4. Related Work

长上下文大语言模型在推理过程中由于参数量庞大和需要处理的长序列导致内存消耗和延迟问题，现面临重大挑战。从优化的角度来看，作者大致将这些方法分为以下几类： 近年来的研究从两个方向优化了上下文计算（Fu, 2024），即预填充和解码。在预填充阶段，（Zhang等人，2023）根据注意力得分丢弃KV缓存中的不重要Token；而在解码阶段，LM-Infite（Han等人，2024）和StreamingLLM（Xiao等人，2024）则在每次解码步骤中仅保留最近的几个Token和若干初始Token。

尽管所有这些方法提升了效率，但也存在舍弃对未来文本生成至关重要的Token的风险。本工作旨在提高效率的同时确保不遗漏任何重要Token。

减小隐藏状态维度并进行量化。作者还提出了另一种方法线，包括隐藏尺寸缩减和量化，尽管这些技术与作者工作的重点并不一致。一方面，MQA（Shazeer, 2019）和GQA（Ainslie等，2023）通过将多个键值对状态的头组合成一对来减少隐藏状态的维度。MLA（Liu等，2024a）则将一对键值状态压缩为一个低秩潜在向量。另一方面，AwQ（Lin等，2024）和QLLM（Liu等，2024d）将模型权重和激活值转换为低位宽格式，从而减少内存使用和计算开销。

另一条与作者研究工作密切相关的方法旨在通过减少Transformer各层之间的冗余性来提高效率。LCKV（Wu & Tu，2024）、CLA 和MiniCache 通过在层之间共享键值状态，利用了键值之间的相似性。

与这些方法相比，作者的方法有两个主要的区别特征：

1. 作者利用了 Transformer 各层之间注意力得分的相似性（Xiao等人，2019；Bhojanapalli等人，2021），并将这一现象扩展到大语言模型中；
2. 而不像他们的共享策略涉及相邻层或最终层，作者的共享策略是按头进行的，并来源于一个搜索过程。

参考

[0]. Compressing KV Cache for Long-Context LLM Inference with Inter-Layer Attention Similarity .