A Little Goes a Long Way: Efficient Long Context Training and Inference with Partial Contexts

作者/机构: Suyu Ge (伊利诺伊大学厄巴纳-香槟分校), Xihui Lin (微软), Yunan Zhang (微软), Jiawei Han (伊利诺伊大学厄巴纳-香槟分校), Hao Peng (伊利诺伊大学厄巴纳-香槟分校)

A1 主要贡献

本文旨在解决训练和部署长上下文大型语言模型（LLM）时遇到的巨大开销问题。通常，解决此问题需要两个关键步骤：首先，对一个预训练好的LLM，通过在长上下文数据上进行训练，进行一个独立的上下文长度扩展阶段；其次，通过架构上的修改来减少服务（serving）期间的KV缓存开销。本文认为，将长度扩展与一个对GPU友好的KV缓存削减架构相结合，不仅能减少长度扩展阶段的训练开销，还能实现更优的长上下文性能。

基于此，本文提出了 LONGGEN 方法，该方法在进行上下文长度扩展的同时，将预训练的LLM微调成一个高效的架构。LONGGEN的核心设计基于三个关键洞见：
1. 稀疏注意力模式的有效性：诸如窗口注意力（关注最近的token）、注意力汇点（attention sink，关注初始的token）和块状稀疏注意力（blockwise sparse attention，关注有步长的token块）等稀疏注意力模式非常适合构建高效的长上下文模型。这主要是因为它们具有对GPU友好的内存访问模式，从而不仅在理论上，在实践中也能带来效率提升。
2. 保留全局访问的重要性：模型必须能够直接访问所有token。本文发现，一个包含1/3全注意力层和2/3高效（稀疏）注意力层的混合架构，能够在效率和长上下文性能之间实现理想的平衡。
3. 轻量级训练的充分性：仅在50亿（5B）长上下文数据上进行轻量级训练，就足以将混合架构模型的上下文长度从4K扩展到128K。

LONGGEN 的架构概览如图1所示。它采用混合架构，在模型的顶部和底部2/3的层中使用KV简化的注意力机制，而在中间1/3的层中保留全注意力机制。论文探索了两种KV简化的注意力变体：注意力汇点（AttnSink）和块状稀疏（BlockSparse）。


图 1: LONGGEN 概览。左图：它使用一种混合架构，在顶部和底部的2/3层应用KV简化的注意力，同时保持中间1/3的层为全注意力。右图：探索了两种KV简化的注意力变体。

通过在Llama-2 7B和Llama-2 70B模型上的实验，LONGGEN证明了其在不同规模模型上的有效性。
* 训练效率：在128K长上下文的训练中，与全注意力基线相比，LONGGEN实现了1.55倍的训练速度提升，并将总训练时间减少了36%。
* 推理效率：在推理阶段，LONGGEN将KV缓存内存减少了62%，实现了1.67倍的预填充（prefilling）速度提升和1.41倍的解码（decoding）速度提升。
* 性能表现：与那些将KV缓存削减技术应用于全注意力长上下文LLM的基线方法相比，LONGGEN不仅在“大海捞针”（Needle-in-a-Haystack）检索任务上表现出显著更强的性能，在更具挑战性的长上下文推理任务（如BABILong和RULER）上也同样表现出色。

图2展示了在不同稀疏度水平下的训练和推理效率。实验结果表明，随着全注意力层数量的减少（即稀疏度的增加），训练速度和推理速度显著提升，同时KV缓存内存也大幅降低。


图 2: 不同稀疏度水平下的训练和推理效率。左图：训练总耗时加速比。中图：KV内存减少量。右图：推理加速比。我们将训练总耗时与FlashAttention进行比较，并在vLLM上进行推理基准测试。所有结果均在总共32层的Llama2-7B上测量。“1/7”、“1/5”、“1/3 Full”和“All Full”分别表示使用5、7、12和32个全注意力层。

A3 背景知识

LLM 上下文长度扩展。由于在长序列上训练Transformer的开销巨大，上下文扩展通常作为标准预训练之后的一个专门的后训练阶段。当前最先进的模型，如Gemini【【42，Gemini: a family of highly capable multimodal models，2023，arXiv】】、Llama系列【【43，Llama 2: Open foundation and fine-tuned chat models，2023，arXiv】；【8，The llama 3 herd of models，2024，arXiv】】和Qwen【【2，Qwen technical report，2023，arXiv】】，通常在大型短序列语料库上进行预训练，然后在相对较少的长序列数据上进行长度扩展。例如，Llama-3在15T个token（上下文长度为8K）上进行预训练，然后额外在800B个token上进行后训练，以将上下文长度扩展到128K【【8，The llama 3 herd of models，2024，arXiv】；【49，Effective long-context scaling of foundation models，2023，arXiv】】。一项相关的数据工程工作【【10，Data engineering for scaling language models to 128k context，2024，arXiv】】通过精心平衡数据源使其与预训练语料库相似，将上下文扩展所需的token数量显著减少到仅5B。通常，在长度扩展期间没有架构上的修改。因此，由于KV缓存带来的内存开销，这些长上下文模型的服务部署非常具有挑战性【【22，Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention，2024，arXiv】】。因此，需要技术来减少这些模型在推理时的内存开销。

推理时KV简化方法在长上下文场景的局限性。为了降低大型Transformer模型的推理成本，一个直接的想法是应用推理时KV简化方法。通过仅保存最受关注的KV缓存，这些方法可以减少推理内存消耗和前向计算的FLOPs。为了实现高效解码，上下文中的键（key）和值（value）向量被保存在GPU内存中，通常称为KV缓存。其内存开销已成为服务LLM的主要瓶颈。目前已开发出多种技术来减少KV缓存开销，通过只存储KV的一个子集并驱逐其余部分【【48，Efficient streaming language models with attention sinks，2024，ICLR 2024】；【27，Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time，2023，NeurIPS 2023】；【11，Model tells you what to discard: Adaptive KV cache compression for llms，2024，ICLR 2024】】。然而，如下文所示，这些方法在长上下文场景中通常表现不佳，而这恰恰是最需要KV缓存削减的场景。

现有KV简化方法的评估。我们在一个经过上下文扩展的Llama2-7B-128K模型上评估了两种成熟的KV缓存削减方法：Attention Sink【【48，Efficient streaming language models with attention sinks，2024，ICLR 2024】】和PyramidKV【【52，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】】。Attention Sink (AttnSink) 仅保留初始和局部的token，而PyramidKV则识别每层的独特注意力分布并分配动态的KV预算。与Heavy-Hitter (H2O)【【54，H2O: heavy-hitter oracle for efficient generative inference of large language models，2023，NeurIPS 2023】】类似，PyramidKV保留累积注意力分数最高的token。此外，我们还纳入了原始的Llama2-7B-128K基线和我们的LONGGEN，并在图3中报告了它们的性能。为了公平比较，我们为Attn Sink、PyramidKV和LONGGEN保持了相同的KV缓存预算，即原始全Transformer的60%。图3展示了在“大海捞针”（NIAH）检索任务上的结果。原始的Llama2-7B-128K基础模型和LONGGEN都能以100%的准确率通过针（needle）检索任务，我们将在实验部分的§4.2中对此进行分析。然而，两种KV驱逐方法在NIAH任务上表现不佳，尤其是在上下文长度超过32K时。


图 3: 推理时KV缓存削减方法在长上下文任务上失败。

现有方法失败原因分析。直观来看，它们表现不佳可归因于两个因素：
* 缺乏完整的上下文访问。由于在所有层都进行了KV削减，模型失去了对某些位置的直接访问，这使得模型执行精确检索变得更具挑战性。
* 缺乏对稀疏上下文的适应。在缓存削减期间，模型参数，特别是位置嵌入，保持不变。然而，已有观察指出，广泛使用的RoPE式位置嵌入【【35，Roformer: enhanced transformer with rotary position embedding，2021，arXiv】】在推广到未见过的位置范围时存在困难【【44，Resonance rope: Improving context length generalization of large language models，2024，arXiv】】。在KV削减后，位置范围变得不连续，与训练设置产生不匹配。这可能损害模型适应更长上下文的能力。

LONGGEN的设计动机。上述实验的教训激发了LONGGEN的关键设计选择：（1）在某些层保留完整的上下文，以及（2）训练模型以适应其他层的稀疏注意力。

A2 方法细节

LONGGEN提升了长上下文的训练和推理效率。它在上下文长度扩展期间，将一个预训练的LLM微调成一个高效的架构。LONGGEN使用一种混合架构，在底层和顶层使用KV简化的注意力模式，同时在中间层保留完整的上下文（如图1所示）。我们将探索初始化稀疏注意力的各种策略，并通过对FLOPs和内存节省的全面分析来展示LONGGEN的理论效率优势。

3.1 使用高效架构进行上下文扩展

混合架构。先前的工作发现注意力在中间层相对专业化。一些中间层的注意力头被证明与检索和推理密切相关，表明在这些层中进行了精细的信息聚合【【47，Retrieval head mechanistically explains long-context factuality，2024，arXiv】；【9，How do language models put attention weights over long context，2024，Yao Fu’s Notion】；【53，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】】。受这些发现的启发，我们设计LONGGEN为沙漏形状——两端使用稀疏注意力以有效降低计算成本，中间使用全注意力以最大化长上下文性能。我们将在§4.4中展示为中间层保留完整上下文的必要性。我们使用与预训练相同的损失函数，在更长的序列上对这种混合架构进行后训练。

不同的KV简化注意力策略。LONGGEN中的高效注意力层可以使用各种KV缓存削减策略进行初始化。与那些为每个头自适应分配内存的注意力模式不同，LONGGEN更倾向于采用结构统一的注意力模式。具体来说，我们为稀疏注意力模式引入了两个标准：（1）位置的静态访问。对于每个头，要保留的位置应该是固定的，与输入token无关。这使得在训练中能够实现统一的GPU内存访问模式，从而允许优化的稀疏训练核函数和内存IO。在推理过程中，也可以强制执行固定的缓存管理策略，从而避免额外的计算和张量替换。从这两个方面来看，静态注意力模式都能带来更高的效率。（2）块级上下文处理。在GPU中，线程以块的粒度分配给流式多处理器（streaming multiprocessors），其中每个块处理连续的内存空间，这通常远大于每个token所占用的内存。因此，token级的KV设计会产生空闲线程并降低SRAM的利用率。因此，我们按块加载上下文，并将块大小设置为与每个线程块大小相似。在这两个原则的指导下，我们为LONGGEN配备了两种现有的KV缓存削减策略，并在图1中进行了说明。我们的第一个变体是Attention Sink，它将初始token与局部上下文窗口的块相结合【【48，Efficient streaming language models with attention sinks，2024，ICLR 2024】】。该技术建立在滑动窗口注意力的已被证明在长上下文训练中有效性的基础上，如Mistral所示【【21，Mistral 7b，2023，arXiv】】。另一种初始化策略是块稀疏注意力（block sparse attention），它将上下文划分为多个块，并根据预定的步长选择性地关注特定的块索引【【51，Big bird: Transformers for longer sequences，2020，NeurIPS 2020】；【31，Blockwise self-attention for long document understanding，2019，arXiv】】。

表 1: 传统全注意力与LONGGEN中稀疏注意力的效率比较。训练FLOPs计算了注意力前向和后向操作。$L_{full}$ 和 $L_{partial}$ 分别是全上下文和部分上下文层的数量。$N$ 和 $D$ 分别代表序列长度和头维度。在LONGGEN的attention sink变体中，$S$ 代表一个固定的窗口块大小。在块稀疏变体中，$S$ 代表一个固定的总块大小。在这两种变体中，$S$ 都不随上下文长度增长，且 $S \ll N$。

3.2 核函数级别的效率优化

基于FlashAttention的定制化核函数。由于我们的实现仅在稀疏注意力计算上与传统Transformer不同，LONGGEN可以受益于为提高Transformer效率而开发的工具，例如flash attention。实际上，我们是在flash attention的基础上构建我们的定制核函数。与FlashAttention【【7，Flashattention: Fast and memory-efficient exact attention with io-awareness，2022，NeurIPS 2022】】类似，LONGGEN接受一个注意力掩码矩阵作为输入。该注意力掩码以压缩稀疏行（compressed sparse row）格式表示，其内存消耗为$O(N)$，而不是原始的$O(N^2)$。需要注意的是，该稀疏掩码对所有头都是静态的，因此在构建动态掩码时没有开销。在前向和后向步骤中，如果某个位置的掩码值为零，我们就通过避免将其加载到HBM中来跳过对该位置的计算。这种实现节省了FLOPs，从而显著加速了训练和推理。

推理优化。在推理过程中，瓶颈在于尽可能快地加载KV缓存。我们没有加载整个头维度，而是在不同的线程块之间拆分头维度的加载。我们凭经验发现，这比原始的FlashAttention更有效，可能是因为每次加载周期减少了SRAM的使用量。对于attention sink初始化，我们遵循Mistral的做法，为局部滑动窗口token使用滚动缓冲区缓存（rolling buffer cache）【【21，Mistral 7b，2023，arXiv】】。

3.3 理论训练和推理优势

理论训练优势。由于每个token都需要被存储以计算损失函数，LONGGEN在训练过程中不提供任何内存节省。主要的效率增益来自于训练FLOPs的大幅减少。表1表明，当上下文长度$N$显著超过块大小$S$时（$S \ll N$），稀疏注意力层所需的FLOPs变得可以忽略不计。因此，与原始Transformer模型相比，LONGGEN有效地将总FLOPs减少到$L_{full}/(L_{full} + L_{sparse})$的比例。实证研究表明，$L_{full} : L_{sparse} = 1 : 2$的比例提供了一个最佳平衡，既实现了显著的效率增益，又没有损害模型的长上下文性能。

理论推理优势。在推理阶段，LONGGEN在预填充时间和KV缓存内存使用方面都表现出显著的改进，如表1的最后两列所示。在全注意力层模型中，KV缓存通常随序列长度线性增长，而预填充时间则呈二次方增长。然而，LONGGEN有效地缓解了这些因素。与训练阶段的分析类似，当$S \ll N$时，由稀疏注意力层引起的成本变得可以忽略不计，将这两项指标都降低到原始值的$L_{full}/(L_{full} + L_{sparse})$。

定制核函数的实际效果。通过我们定制的triton核函数，LONGGEN可以将理论上的效率增益转化为实际的总耗时加速和GPU内存节省，我们将在§ 4.3中进一步说明。

A4 实验

4.1 实验环境

  • 模型架构: Llama-2-7B-base（常规多头注意力）和Llama-2-70B-base（分组查询注意力）。
  • 数据集与训练: 在SlimPajama数据集【【34，Slimpajama-dc: Understanding data combinations for llm training，2023，arXiv】】上进行持续预训练，该数据集是Llama-2预训练语料库的开源复现。通过拼接Slimpajama中的短文本，并用<bos> token标记文档边界，来构建128K长度的长序列训练数据。模型在50亿个token上进行训练，这仅占其2.4T预训练语料库的0.2%。
  • 硬件配置: 使用32块NVIDIA A100-80G GPU进行后训练。7B和70B全注意力模型的后训练分别消耗74和387个GPU小时。

  • 软件与超参数:

    • 学习率：2e-5
    • 全局批量大小：32
    • RoPE base：5M
    • 自定义核函数块大小：每个核函数块处理64个token。
    • LONGGEN配置：对于Attention Sink变体，保留第一个块用于sink token，最近的32个块用于局部上下文。对于Block Sparse变体，步长设为64个块。这确保了在稀疏层数量相同时，两种变体具有相似的KV缓存预算（对于128K序列，约保留2K个token）。
    • 推理框架：在vLLM基准测试上运行推理并比较效率。

  • 评测基准:

    • 长上下文:

      1. Needle-in-a-Haystack (NIAH): 检索任务。
      2. BABILong【【24，Babilong: Testing the limits of llms with long context reasoning-in-a-haystack，2024，arXiv】】: 要求模型从长上下文中检索多条证据并进行推理。使用3个不同的随机种子为每个任务生成2250个测试样本，并对前五个任务的性能进行平均。
      3. RULER【【20，Ruler: What’s the real context size of your long-context language models?，2024，arXiv】】: 仅在其两个QA子任务上进行评估，因为其他合成任务的性能方差巨大，而多跳QA任务更接近实际应用。

    • 短上下文:

      1. MMLU【【18，Measuring massive multitask language understanding，2020，arXiv】】
      2. Math【【19，Measuring mathematical problem solving with the math dataset，2021，arXiv】】
      3. BigBenchHard (BBH)【【38，Challenging big-bench tasks and whether chain-of-thought can solve them，2022，arXiv】】

4.2 整体性能

基线方法比较。我们将LONGGEN与三类基线方法进行了比较：(1) 全注意力: 作为性能上限参考；(2) KV缓存驱逐方法: 包括Attention Sink【【48，Efficient streaming language models with attention sinks，2024，ICLR 2024】】、Heavy Hitter (H2O)【【54，H2O: heavy-hitter oracle for efficient generative inference of large language models，2023，NeurIPS 2023】】、RazorAttention【【39，Razorattention: Efficient kv cache compression through retrieval heads，2024，arXiv】】和PyramidKV【【52，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】】；(3) 长上下文训练方法: 包括滑动窗口注意力（如Mistral使用）和YOCO【【37，You only cache once: Decoder-decoder architectures for language models，2024，arXiv】】。为了公平比较，所有方法的KV缓存预算都保持一致。

长上下文性能评估。

  • NIAH检索任务 (表2): 分析结果得出两个关键发现：(1) 长上下文训练方法（如滑动窗口、YOCO、LONGGEN）的性能一致优于推理时KV驱逐技术（如Attn Sink、H2O等）。这验证了§2中的观点，即确保对上下文窗口的完全访问和训练模型以适应更长上下文至关重要。(2) 在长上下文训练方法中，只有LONGGEN实现了完美的100%“捞针”成功率。与滑动窗口相比，LONGGEN保留了整个序列；与YOCO相比，LONGGEN的架构更接近预训练时的原始Transformer架构，并利用了中间层的信息聚合能力。

表 2: 不同方法在“大海捞针”检索任务上的平均通过率。所有方法均基于Llama2-7B模型。评估上下文长度范围从512到128K。

  • BABILong推理任务 (表3): 为了评估更高级的检索和推理能力，我们在BABILong上进行了测试。虽然趋势与NIAH相似，但BABILong能更好地分辨不同方法的性能差异。在所有高效模型中，LONGGEN表现最佳，平均准确率达到0.27。值得注意的是，尽管在NIAH上取得了100%的准确率，LONGGEN在更具挑战性的任务上仍略逊于全注意力模型。这凸显了在具有挑战性的推理任务上评估模型的重要性。

表 3: 在BABILong前五个任务上的评估结果。列对应序列长度，行表示模型。每个数字表示模型在给定序列长度下5个任务的平均准确率，基于1000个样本计算。

与全注意力模型的全面比较 (表4)。我们将LONGGEN-7B和LONGGEN-70B与其全注意力对应模型（性能上限）进行了比较。由于LONGGEN w/ AttnSink和LONGGEN w/ BlockSparse的性能差异不大，我们以AttnSink变体的结果为例。
* 性能表现: 尽管LONGGEN的性能略低于全注意力模型，但它在各种上下文规模上都取得了可比的结果，同时提供了巨大的效率提升。
* 关键发现:
1. 性能差距在长上下文场景中比短上下文场景中更为显著。这可能是因为大多数短上下文可以被LONGGEN中的局部块完全捕获，而来自更远距离的信息仅被稀疏地表示。
2. 随着模型规模的增加，LONGGEN与性能上限之间的差距在缩小。这可能是因为更大的模型表现出更高程度的信息稀疏性，从而突显了将我们的方法应用于更大参数量模型的潜力。

表 4: LONGGEN与全注意力基线的整体长上下文性能比较。评估上下文长度从512到128K。NIAH衡量模型在“大海捞针”检索任务上的平均正确回答率。BABILong指在BABILong基准测试前五个任务上的平均准确率得分。RULER报告在RULER的两个多跳问答子任务上的平均准确率。

4.3 效率提升

训练效率。我们使用总耗时训练加速比来评估训练效率，测量了在256个A100-80G GPU上训练Llama2-7B模型（128K序列长度）的延迟。如图2左侧所示，通过增加稀疏层，LONGGEN显著提高了训练吞吐量。我们选择的“1/3 Full”设置（即1/3的层是全注意力层）在总训练时间上实现了1.55倍的加速。

推理效率。为了验证LONGGEN在服务期间能带来实际的系统级改进，我们将其推理核函数与vLLM兼容。如图2中间和右侧所示，在单个128K序列上进行测试：

  • KV缓存内存: 我们选择的“1/3 Full”设置将内存消耗从69.2 GB减少到26.5 GB，减少了62%。
  • 推理延迟: 同时，它节省了40%的预填充（prefilling）时间（1.67倍加速）和29%的解码（decoding）时间（1.41倍加速）。
    这些结果证明了LONGGEN在高效长上下文服务方面的潜力。我们没有选择更稀疏的架构是为了保持最优的长上下文性能。

4.4 探究长上下文的关键因素

本节研究了全注意力层的数量和位置对长上下文性能的影响。

全注意力层的位置 (表5)。我们固定全注意力层的数量为12个，并改变它们的位置，探索了四种配置：(1) 堆叠在顶层，(2) 堆叠在中间层，(3) 堆叠在底层，(4) 在中间层交错排列。BABILong数据集上的结果表明，全注意力层的位置对长上下文性能有显著影响。值得注意的是，将全注意力层置于网络中间能产生最佳结果。这与先前研究表明中间层的注意力头在信息聚合中起关键作用的发现相符。

表 5: 全注意力层位置的消融实验。我们在总共32层的Llama2-7B基础上，在128K上下文长度上进行训练，并报告在BABILong上的结果。我们将全注意力层的预算固定为总层数的三分之一（总共12个全注意力层），并测试了全注意力层的不同位置。顶部、中间和底部设置将所有全注意力层堆叠在一起，而交错设置则在每两个稀疏层之间插入一个全注意力层。

全注意力层的数量 (表6)。我们进一步研究了全注意力层数量与模型性能之间的关系。目标是找到计算效率和准确性之间的最佳平衡点。表6展示了不同数量全注意力层的性能。分析表明，将全注意力层的数量维持在总网络深度的1/5（6层）到1/3（12层）之间，足以保持令人满意的长上下文准确性。

这些实证观察启发了LONGGEN的最终架构设计，即在网络中间设置1/3的全注意力层。

表 6: LONGGEN w/ AttnSink中全注意力层数量的消融实验。基础模型是总共32层的Llama2-7B-128K。为获得最佳性能，我们将所有全注意力层置于模型中间。All, 1/3, 1/5, 和 1/16 Full 分别表示使用32、12、6和2个全注意力层。

A7 补充细节

5.1 高效的长上下文训练架构

替代架构探索。为了克服Transformer的二次方复杂度，研究人员引入了几种模型架构替代方案。随着近期为提升训练效率所做的关键修改【【13，Mamba: Linear-time sequence modeling with selective state spaces，2023，arXiv】】，状态空间模型（SSM）因其仅需线性计算复杂度而崭露头角。Gu等人【【15，Hippo: Recurrent memory with optimal polynomial projections，2020，NeurIPS 2020】；【16，Combining recurrent, convolutional, and continuous-time models with linear state space layers，2021，NeurIPS 2021】】首先通过hippo矩阵和可并行的操作数缓解了序列模型的训练可扩展性问题。Mamba【【13，Mamba: Linear-time sequence modeling with selective state spaces，2023，arXiv】；【6，Transformers are ssms: Generalized models and efficient algorithms through structured state space duality，2024，ICML 2024】】向SSM中引入了输入依赖的状态参数化，并优化了硬件感知的实现。一些工作【【30，Can mamba learn how to learn? A comparative study on in-context learning tasks，2024，ICML 2024】；【56，Efficient long sequence modeling via state space augmented transformer，2022，arXiv】；【28，Mega: Moving average equipped gated attention，2023，ICLR 2023】；【32，Sparse modular activation for efficient sequence modeling，2023，NeurIPS 2023】；【12，Zamba: A compact 7b SSM hybrid model，2024，arXiv】；【33，Samba: Simple hybrid state space models for efficient unlimited context language modeling，2024，arXiv】】提出了通过将不同变体的Transformer层集成到SSM中来构建混合架构。然而，正如最近的研究表明，SSM在长上下文能力方面与Transformer相比竞争力较弱【【45，Rnns are not transformers (yet): The key bottleneck on in-context retrieval，2024，arXiv】；【1，Simple linear attention language models balance the recall-throughput tradeoff，2024，ICML 2024】】。

早期的稀疏注意力工作。更早的时候，一些工作尝试向Transformer中注入稀疏性以提高在长上下文设置中的效率【【41，Efficient transformers: A survey，2023，ACM Comput. Surv.】；【4，Generating long sequences with sparse transformers，2019，arXiv】；【3，Longformer: The long-document transformer，2020，arXiv】；【50，Big bird: Transformers for longer sequences，2020，NeurIPS 2020】】。然而，正如Dao等人【【7，Flashattention: Fast and memory-efficient exact attention with io-awareness，2022，NeurIPS 2022】】所指出的，由于与加速器的兼容性差，这些算法很难带来实际的运行时间加速。与此同时，Sun等人【【37，You only cache once: Decoder-decoder architectures for language models，2024，arXiv】】复兴了编码器-解码器架构来节省KV缓存，其中所有解码器层通过交叉注意力消耗来自编码器侧的相同KV向量。然而，尚不清楚如何将该架构用于上下文扩展，因为它与默认的仅解码器模型有显著差异。

5.2 推理时KV缓存削减

现有方法概述。最近，许多尝试在推理时驱逐KV缓存，以实现更快的长上下文服务场景【【29，Towards efficient generative large language model serving: A survey from algorithms to systems，2023，arXiv】】。Zhang等人【【54，H2O: heavy-hitter oracle for efficient generative inference of large language models，2023，NeurIPS 2023】】和Liu等人【【27，Scissorhands: Exploiting the persistence of importance hypothesis for LLM KV cache compression at test time，2023，NeurIPS 2023】】提出基于累积的注意力分数丢弃不太重要的KV向量。Han等人【【17，Lm-infinite: Simple on-the-fly length generalization for large language models，2023，arXiv】】和Xiao等人【【48，Efficient streaming language models with attention sinks，2024，ICLR 2024】】发现LLM倾向于将重要的全局信息存储到初始的token中，这被称为注意力汇点（attention sink）。通过仅在缓存中保留注意力汇点和最近的token，LLM可以保持相当不错的性能。Ge等人【【11，Model tells you what to discard: Adaptive KV cache compression for llms，2024，ICLR 2024】】和Li等人【【26，Snapkv: LLM knows what you are looking for before generation，2024，arXiv】】设计了算法以混合驱逐策略自适应地维护KV缓存。PyramidKV【【53，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】】发现不同层的注意力稀疏度不同，并提出为每层动态分配KV缓存预算。

现有方法的局限性。这一系列工作表明，LLM在推理时可以使用部分上下文工作【【26，Snapkv: LLM knows what you are looking for before generation，2024，arXiv】；【53，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】】。然而，Jiang等人【【22，Minference 1.0: Accelerating pre-filling for long-context llms via dynamic sparse attention，2024，arXiv】】和Zhang等人【【53，Pyramidkv: Dynamic kv cache compression based on pyramidal information funneling，2024，arXiv】】的研究表明，许多方法在长上下文任务中会出现显著的性能下降。此外，复杂的驱逐方法是否与现实的服务系统兼容也值得怀疑，因为在PagedAttention【【25，Efficient memory management for large language model serving with pagedattention，2023，SOSP 2023】】和Prefix Caching【【55，Efficiently programming large language models using sglang，2023，arXiv】】中实现计算累积注意力分数和释放任意内存是具有挑战性的。

A5 结论

本文提出了LONGGEN方法，在上下文长度扩展期间将一个预训练的LLM微调成一个高效的架构。我们集成了对GPU友好的稀疏注意力机制，以实现实际的效率增益。此外，我们发现采用一个包含1/3全注意力层和2/3高效层的混合架构，可以在效率和长上下文性能之间实现良好的平衡。通过轻量级训练，我们将Llama2-7B和70B的上下文长度从4K扩展到了128K。在多个长上下文基准上的评估表明，LONGGEN在长上下文检索和推理方面达到了与全注意力模型相当的性能。与此同时，它将总训练时间减少了36%，KV缓存减少了62%，在预填充阶段实现了1.67倍的加速，在解码阶段实现了1.41倍的加速。我们的研究为未来高效Transformer架构和低成本长上下文扩展方法的研究提供了动力。