DuoAttention 是什么？区分检索头和流式头来加速长文本推理

DuoAttention 是一个提升大模型长文本推理效率的框架，由麻省理工学院（MIT）韩松团队在 2024 年提出，论文标题是《DuoAttention：基于检索头和流式头的高效长上下文大模型推理》。

长文本的内存难题

大模型处理长文本时，需要为每一个 token 在每一层、每一个注意力头里保存它的键和值，也就是 KV 缓存。文本越长，KV 缓存越大，处理几十万 token 时，KV 缓存可能占用几十甚至上百 GB 的显存，成为长文本推理最主要的瓶颈。一种常见的压缩办法是只保留最开头的几个 token 和最近的一段 token，丢掉中间的部分，这种“流式”方法很省内存，但模型会忘掉中间的重要信息，在需要从长文中精确查找某个细节的任务上表现很差。

关键发现：两种注意力头

研究者发现，并不是所有的注意力头都需要看完整的上下文。注意力头可以分成两类：

检索头（Retrieval Heads）：只占一小部分，它们负责在整个长上下文中查找相关的信息，比如回答问题时去找文中提到的某个具体数字，因此需要完整地访问所有 token 的 KV 缓存。

流式头（Streaming Heads）：占大多数，它们主要关注最近的 token 和最开头的几个“注意力汇聚”token，并不需要看中间的大量内容。对这些头来说，只保留开头和最近的一小段 KV 缓存就足够了，不会影响效果。

DuoAttention 怎么做

基于这个发现，DuoAttention 只给检索头保留完整的 KV 缓存，而给流式头使用固定长度的小缓存，从而大幅减少内存占用。关键是如何准确地判断哪些头是检索头。DuoAttention 用一种轻量的、基于优化的方法来识别：用合成的数据，自动学出每个注意力头对完整上下文的依赖程度，依赖程度高的就被认定为检索头。这个识别过程只需要很少的计算。

效果

论文报告，DuoAttention 能把 KV 缓存的内存占用减少到原来的一半到四分之一左右，并相应地加快解码和预填充的速度，同时在长文本任务上的准确率几乎没有损失。结合量化技术，它能让一个 8B 的模型在单张 A100 显卡上处理高达 330 万 token 的上下文。