微软研究院路保同：用向量检索重塑模型注意力｜Attention

Retrieval Attention 的核心观点认为，每个 Query 实际上只需要和一小部分 Key 进行强交互即可，剩下的注意力都是冗余的；而注意力本身，其实也是天然稀疏的。

因此，研究团队的核心做法是：把大部分 KV 向量从 GPU 下放到 CPU，每次模型要做推理的时候，就用近似最近邻检索（ANN Search）来找出对当前 Query 最相关的少量 Key，比如只找前 1%  再与 GPU 上少量 “可预测” 的 KV 一起并行计算注意力并合并结果。

实测在 RTX4090（24GB）上，8B 级模型可在 128K 上下文下稳定生成，每 token~0.188s，且与全注意力精度几乎一致（∞-Bench、RULER 等基准）。

后续工作 RetroInfer 继续扩大了性能优势，在 A100 GPU 上相比于全注意力（Full attention）实现了 4.5 倍的解码吞吐，并在 1M token 的上下文时相比于其它 GPU-CPU 稀疏注意力系统实现了 10.5 倍的吞吐。

Retrieval Attention 的核心创新在于，它将向量检索机制引入注意力计算路径，并在计算层面实现了真正意义上的动态稀疏化。

具体而言，作者构建了一个注意力感知的向量索引（attention-aware retrieval index）：通过在预填阶段学习 Query 与 Key 的精确映射，并将这种关系投影到 Key–Key 图结构中，使检索过程能够以极低的扫描比例（约 1–3%）实现高召回率。这一设计让每个 Query 能在极小的计算代价下动态定位最关键的 Key，从而摆脱传统 ANN 检索在应用到注意力后因为 Query-key 分布不同而造成的低效问题。

在系统架构上，Retrieval Attention 进一步提出 CPU–GPU 协同的双路注意力机制：GPU 负责保留少量“可预测”的局部 KV 缓存，而 CPU 以检索方式动态调用大规模 KV 存储，两路计算独立并行，最终通过数值稳定的重标定公式融合结果，实现了显存占用与推理延迟的双重下降。

值得注意的是，整个机制无需对模型进行重新训练。Retrieval Attention 以可插拔模块的形式接入现有 Transformer，仅修改注意力层的前向逻辑，即可在不牺牲精度的前提下显著加速长上下文推理。这种 “免训练的动态稀疏注意力” 思路，为长上下文语言模型的可扩展性提供了一条切实可行的工程路径。

以下是绿洲与路博士的部分对话内容：

绿洲：请先介绍一下您的研究背景，特别是您在数据库管理系统（Database Management System）方面的研究，以及当时是基于怎样的思考，形成了这一设计思路？

路博士：我主要的研究方向是数据库管理系统，也就是如何更高效地存储、检索和管理海量数据。从学科上看，数据库系统和机器学习好像是两个相对独立的领域，一个偏向系统与结构，另一个偏向模型与优化。但它们在底层其实有一些共通的问题，比如在资源有限的情况下，怎样让信息组织得更高效。

在博士阶段，我主要做的是研究新的硬件架构，比如持久内存和分离式内存，对数据库设计会带来哪些新的挑战。具体包括如何针对这些硬件设计新的数据结构、数据库索引，以及存储系统在查询性能上的优化。可以说，这个方向的核心是把底层硬件的变化和数据管理的性能问题结合起来看。

绿洲：您刚刚提到，自己的研究核心是让系统在不同硬件架构下实现更高效的数据管理。那在这个过程中，是如何从传统的数据库优化问题，逐步延伸到今天与模型推理相关的研究方向的？

路博士：传统数据库系统的演进，其实就是一系列技术不断被完善，用来支撑更高效的数据检索与存取。核心目标很简单，如何让查询更快、更稳定。比如，索引结构的优化、缓存机制的提升，都是为了让系统在面对持续增长的数据量时，依然能维持查询性能。

但随着硬件架构的变化，比如持久内存、分离式内存、或者分层存储系统，这些新型硬件与传统模式之间其实有很多不兼容的地方。于是我们需要重新思考：在这种新的架构下，数据库系统应该怎样设计，才能最大化地发挥硬件的潜能？

我在博士期间的研究主要就是围绕这一点展开的。

一方面，我关注系统在新硬件下的性能优化，比如如何减少访问延迟、提升索引效率；另一方面，也会研究如何让系统在处理大规模任务时更加智能化。这些经验后来也影响了我现在的研究方向。

现在我们团队更关注的是，如何把传统数据库系统中的 “检索” 逻辑迁移到模型层面。比如，当模型面对长上下文推理时，其实遇到的瓶颈和数据库很像，都是如何在庞大的存储中快速定位关键的信息。这让我们想到，也许可以把数据库的检索思想，与模型的注意力机制结合起来。

我们的做法，是尝试把数据库中成熟的向量检索方法，移植到语言模型的推理过程中，让模型在生成时只访问 “最相关” 的信息。换句话说，我们不是去重新训练模型，而是希望通过系统层的设计，让模型更高效地利用已有的记忆。这其实就是我们现在这篇论文的起点。

绿洲：那能不能具体讲讲，和硬件结合的这个研究思路是怎么形成的？包括当时在硬件或系统层面上，您是怎么去思考和结合的？

路博士：这个想法最早其实是从一个很具体的问题出发的。

我们发现注意力机制本质上是一种向量检索的过程，当模型在生成时，需要不断地从已有的上下文中 “查找” 相关的信息。而在这个过程中，Memory（也就是显存）成了最大的限制。

像现在的 GPU，它的 Memory 是非常有限的。我们发现，模型在长上下文推理时，会不断堆积大量的 KV 缓存，而这些缓存的管理方式，其实和传统数据库系统里 “索引+缓存” 的逻辑非常相似。所以我们就在想，既然数据库能够通过分层存储和检索机制高效管理海量数据，那模型的注意力机制是否也可以借鉴类似的思路？

于是，我们开始尝试把部分缓存从 GPU 卸载到更大的存储层，比如 CPU 内存上，再通过检索去动态调用需要的部分。这种 “分层+检索” 的结构让系统在长上下文下依然能保持高效，不需要付出额外的训练成本。

绿洲：其实您是把数据库系统里 “存储—索引—检索” 的那套思维，迁移到了模型的注意力计算上？

路博士：对，思路上很相似。数据库的目标是让数据访问变得更快，而我们关注的，是让模型在生成过程中 “找到信息” 的过程更快。不同的是，我们面对的是非结构化、动态变化的上下文信息，所以检索和缓存策略都要重新设计。

可以说，我们是把传统数据库系统在性能优化上的积累，重新应用到模型的注意力机制里，让它在计算上更接近真实世界的 “存取” 过程。

这也是 Retrieval Attention 最初的出发点。

绿洲：我们是不是可以从更基础的部分聊起，帮助读者理解陌生的概念？在数据库管理系统（Database Management System）中，“数据管理机制” 的核心逻辑是什么？它是否可以被视作一个 “大数据库” 的核心？

路博士：我可以先简单解释一下。在机器学习或者认知科学的语境下，“注意力机制（Attention Mechanism）” 最初的出发点，其实是对人类注意力的一种模拟。也就是在处理一连串信息时，模型如何判断哪些内容更重要、哪些可以被忽略。

举个例子，我们给模型一个问题，它不会平均地看待输入的所有信息，而是会自动分配 “注意力权重”，集中在与这个问题更相关的部分上。这就是注意力机制的核心思想：在一个序列中，根据任务目标，把计算重点放在那一小部分真正重要的元素上。

早期的序列模型，比如 RNN 或 Seq2Seq，其实已经具备这种 “选择性聚焦” 的特性，只不过那时它更依赖顺序结构。

而 Attention 的出现，使模型能够在任意位置之间建立联系：它不再受时间或位置的限制，而是可以直接捕捉全局的依赖关系。这也是为什么 Transformer 会成为一种里程碑式的架构。

从实现的角度来看，注意力机制可以理解为一种 “查询—匹配—加权” 的过程。模型会把输入分解为三类向量：Query、Key 和 Value。

• Query 代表当前需要处理的信息，也就是 “我现在在找什么”；

• Key 代表所有可能被匹配的线索；

• Value 则是这些线索对应的内容。

模型通过计算 Query 与各个 Key 的相似度，得到一组权重，再用这些权重去加权求和所有的 Value。结果就是模型认为 “最值得关注” 的部分，也就是最终的注意力输出。

如果我们把注意力机制放到数据库系统的框架里来看，其实可以把它理解为一个动态的信息检索系统。模型在生成每一个新 token 的时候，都需要在已有的语义空间里 “查询” 最相关的信息，这和数据库在执行查询请求的过程非常相似。

在这个类比中，注意力权重就代表一次查询的 “匹配程度”。权重越高，说明当前生成的 token 越应该 “关注” 那一部分输入，也就是从对应的 Value  中提取更多信息。换句话说，模型每生成一个词，其实都在做一次 “查找并汇总” 操作。

如果再映射到数据库的语义上，Query 就是用户发出的查询语句，Key 相当于索引表中的字段，Value 则是存储的具体内容。模型通过计算 Query 与 Key 的相似度，找到最匹配的 Value，然后加权求和得到输出。每一步生成，都像是在数据库中进行一次动态的、多字段的模糊查询。

从系统角度看，这种机制非常像一种 “隐式的检索器”，它在模型内部完成信息的索引与匹配，但没有显式的结构化过程。我们的研究正是从这里出发，希望把这种隐性的、不可见的检索逻辑显性化，用系统的方式去优化它。

简单来说，我们想让模型的注意力机制变得更像一个 “可控的数据库”：

每次生成时，模型不再被动地遍历全部上下文，而是能主动去查询、筛选、调用它真正需要的信息。

绿洲：Retrieval Attention，这也是你们论文的核心命题。当时在立项的时候，你们最想解决的具体问题是什么？研究过程中有没有遇到或发现一些意料之外的收获？

路博士：我们想解决的，其实是长上下文推理中无法避免的性能瓶颈。

传统的注意力机制在上下文变长时，计算复杂度随序列长度平方增长，同时还要存储大量的 Key–Value（KV）向量，显存消耗极高。以 Llama-3 8B 为例，若上下文达到 100 万个 token，仅 KV 缓存就需要超过 125GB 显存，几乎超出单张 GPU 的上限。

为了降低显存占用，很多系统会把部分 KV 缓存转移到 CPU，但新的问题随之出现：模型在每一步生成时，应该从成千上万个缓存中取哪些最关键的信息？这些 “重要 token” 是动态变化的，无法提前预测。

Retrieval Attention 正是在这个问题上提出了解法。我们把全部 KV 存在更大的 CPU 内存中，而 GPU 只保留一小部分 “活跃缓存”。每当生成新 token 时，系统会通过向量检索在 CPU 内存中找出最相关的 KV，再动态加载回 GPU 计算。这样显存占用降至原来的约 1/10，同时几乎不损失精度。

真正的挑战在于如何判断 “哪些 KV 是最有价值的”。为此，我们设计了基于近似最近邻（ANN）的动态检索机制，并构建了 CPU–GPU 协同调度系统，让检索与计算能并行完成。最终，模型在生成过程中不再被动遍历全部上下文，而是能主动 “查询” 它所需的信息，实现了在长上下文下的高效推理。

绿洲：这也是这篇工作的主要贡献之一？

路博士：是的，我觉得这篇工作的主要贡献可以分成两个层面。

第一个层面是理论上的。我们提出了一个新的视角：注意力机制本质上可以被看作是一种检索系统。 传统向量数据库中的查询通常假设查询和数据处于相同分布，而在注意力机制中，查询（query）和数据 (key) 天然在不同的分布。因此，我们重新设计了注意力的“检索结构”，让模型能够更精确地找到那部分真正重要的上下文信息。我们的后续工作 RetroInfer 中进一步观察到不同模型层以及注意力头的稀疏程度（重要数据的数量）也不同，我们后续的向量索引设计能高效的覆盖变化的稀疏程度，实现了和全注意力近乎相同的精度。

第二个层面是系统与工程实现。Retrieval Attention 不仅是算法层面的改进，更是一种 CPU–GPU 资源调度的系统优化。我们将传统的 “线性缓存” 重构为动态分配结构，使模型能在不同存储层之间高效切换，从而在大规模推理场景下显著降低显存和计算开销。

除此之外，我们的后续工作 RetroInfer 还针对线上服务的可用性进行了改进。传统系统在构建索引时往往开销很大，而我们的工作通过轻量化的检索结构，实现了在线动态索引与查询，大幅减少了构建和维护成本。

总体而言，这项工作的价值在于：它不仅提出了一个新的理论框架，也让注意力机制在系统层面上具备了可操作性和可扩展性。