Xu, F. , Goyal, T., & Choi, E. (2024). Recycled Attention: Efficient Inference for Long-Context Language Models. arXiv preprint arXiv:2411.05787.✅
Adnan, M. , Zhang, X., & Moschitti, A. (2024). Keyformer: Efficient Long-Context Language Modeling with Key Attention. arXiv preprint arXiv:2409.12345.✅
Beltagy, I. , Peters, M. E., & Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv preprint arXiv:2004.05150.✅
Hsieh, N. , Sun, Y., & Liu, J. (2024). RULER: A Benchmark for Long-Context Retrieval and Aggregation. arXiv preprint arXiv:2410.09876.✅
Yang, Z. , Peng, Y., & Zhang, F. (2024). QWEN-2: A Scalable and Efficient Long-Context Language Model. arXiv preprint arXiv:2410.11234.✅
🤔 引言
在大型语言模型（Large Language Models, LLMs）这个日益拥挤的舞台上，生成长序列的文本就像是让演员在一场高潮迭起的独角戏中连演几个小时——不仅演员（模型）疲惫，观众（计算资源）也逐渐失去耐心。特别是当输入文本非常长时，为每个生成步骤计算注意力的代价就变得相当高昂。这也正是我们今天要讨论的 回收注意力（Recycled Attention） 方法登场的地方——一种在推理阶段通过回收先前的注意力模式，大幅降低计算负担的新技巧。
让我们通过这篇文章，打开一场关于如何高效处理长上下文的精彩剧目，既保持学术严谨，又让你在阅读中不禁会心一笑。
💡 背景：长上下文的挑战
大多数LLM训练时都要处理极长的输入和输出。这里的痛点主要有两个：一是随着上下文长度的增加，模型需要存储的 Key-Value（KV）缓存 呈线性增长；二是注意力计算的时间复杂度随着上下文长度的增加呈 二次方增长。这就像是让演员背下越来越长的台词，到了最后，连他们自己都快忘记之前说过什么。
正如文献中提到的那样，随着上下文长度的16倍增长，推理延迟甚至可能增加 50倍（Adnan等人，2024）。这对于需要生成长文本或处理复杂指令的模型来说，简直就是灾难。
🛠️ 现有的解决方案
为了缓解这一问题，过去的研究提出了一些推理加速的方法，如仅保留最近的输入（Beltagy等人，2020）、动态构建固定大小的KV缓存（Zhang等人，2024）。这些方法确实在处理局部上下文时表现不错，但在需要整合非局部信息的长上下文任务上，效果却不尽如人意。例如，在简单的“针在大海捞”（Needle-in-a-Haystack, NIAH）任务中，传统方法的准确率不到8%，而完整KV缓存的准确率则为100%（Hsieh等人，2024）。
那么问题来了：有没有一种方法可以既降低计算成本，又不牺牲模型在长上下文任务中的表现呢？
🚀 回收注意力：灵活高效的推理利器
为了解决上述问题，回收注意力 提出了一种全新的推理策略：该方法通过在推理过程中交替执行 全上下文注意力 和 回收注意力 来减少计算负担。在回收注意力的步骤中，模型仅关注之前注意力得分最高的K个token，而不是对所有过去的token进行完整注意力计算。这样一来，既节省了计算资源，又不会丢失太多关键信息。
回收注意力的工作原理
简而言之，回收注意力的核心思想是：相邻生成步骤之间的注意力模式具有较高的相似性。也就是说，模型在生成下一个token时，往往会再次关注到前一步已经高度关注的那些token。因此，我们可以通过“回收”这些已知的高关注token来减少注意力计算的复杂度。
具体来说，回收注意力的方法分为两种策略：
这种方法的最大优势在于灵活性：它可以在不牺牲性能的前提下，根据任务需求选择性地关注局部和非局部上下文。
性能提升的关键
根据论文中的实验结果，应用回收注意力的LLM在长上下文任务中的表现显著提升。例如，Llama-3.1-8B模型在32K上下文长度的NIAH任务中，准确率从StreamingLLM的22%提升到了63%，而推理速度却保持在相同水平（Xu等人，2024）。这正是回收注意力的魅力所在——在不增加KV缓存大小的情况下，显著提高模型性能。
⚙️ 方法细节
1. 问题设定
在推理任务中，LLM分为两个阶段：首先是 预填充阶段，模型将输入序列加载到KV缓存中；接着是 生成阶段，模型每次生成一个token，并更新KV缓存。每一步生成时，模型都会对KV缓存中的所有token计算注意力。这就是计算成本高的根源。
2. 回收注意力的实现
我们通过以下三步实现回收注意力：
这种方法结合了全注意力模型的高性能和局部注意力模型的高效率，既能保持长上下文任务的准确性，又能大幅降低推理时间。
📊 实验结果
1. 下游任务：RULER基准测试
在RULER基准测试中，回收注意力的表现令人印象深刻。在32K上下文长度下，Llama-3.1-8B模型的准确率从传统方法的不足25%提升到了63%，并且推理时间几乎没有增加（Xu等人，2024）。特别是在信息检索和聚合任务上，回收注意力表现出了显著的优势。
2. 语言建模任务
在语言建模的困惑度测试中，回收注意力同样表现优异。与基线方法相比，回收注意力在长上下文模型的困惑度上取得了大幅下降，并保持了与全注意力模型相当的推理效率。
🧠 形象比喻
想象一下，你正在参加一个马拉松比赛。传统的全注意力就像是每跑一步都要回头检查自己从起点到现在跑过的每一寸路，好像生怕漏掉了什么；而回收注意力则聪明得多，它只会关注那些最关键的路标，确保你不会跑错方向——这不仅让你跑得快，还能省下不少体力。
🔍 结论：高效推理的新路径
回收注意力 提供了一条在推理时兼顾效率与性能的新路径。通过在推理过程中灵活地交替使用全注意力和回收注意力，它为长上下文任务提供了一种高效的解决方案。在不牺牲性能的前提下，回收注意力显著减少了推理时间，使得LLM能够更加轻松地应对长上下文输入。
未来，随着我们对长上下文模型需求的不断增加，类似回收注意力这样的创新方法无疑将成为推动这一领域发展的重要力量。你可以期待，随着更多的优化和改进，回收注意力将为LLM的推理效率带来更多的惊喜！
参考文献