🤖 AI

KV-PRM: Efficient Process Reward Modeling via KV-Cache Transfer for Multi-Agent Test-Time Scaling

KV-PRM 是一种高效的过程奖励模型，它通过直接利用预计算的 KV 缓存，将评分复杂度从 O(L²) 降低到 O(L)，从而消除了基于文本重编码的计算瓶颈，在实现速度和内存大幅提升的同时，在多个推理基准测试中达到或超越了现有方法。

原作者： Peng Kuang, Haibo Jin, Xiaoyu Han, Yanli Wang, Xiaopeng Yuan, Ye Yu, Kaidi Xu, Haohan Wang

发布于 2026-07-13

📖 1 分钟阅读☕ 轻松阅读

原作者： Peng Kuang, Haibo Jin, Xiaoyu Han, Yanli Wang, Xiaopeng Yuan, Ye Yu, Kaidi Xu, Haohan Wang

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在主持一场规模宏大、赌注极高的游戏秀，一支由 AI 侦探组成的团队（我们称之为“多智能体小队”）正试图破解一个超级难的数学谜题。他们不仅仅是猜测答案，而是将问题分解成微小的步骤，并在彼此之间传递线索。为了确保他们不会误入歧途，他们配备了一位“法官”（过程奖励模型，简称 PRM），负责检查他们写下的每一个线索。

旧方法：强迫症式的重读者
在过去，每当法官想要检查一个线索时，他们都必须进行极其耗费精力的工作。想象一下，侦探们写了一个 5,000 字的故事。为了给它评分，法官必须从第一个字到最后一个字，重新坐下来阅读这整个 5,000 字的故事。如果故事变长了，阅读时间并不仅仅是略微增加，而是会发生爆炸式增长。如果你把故事长度翻倍，阅读时间会变为原来的四倍。这就是论文中所说的 $O(L^2)$ 成本。这就像是为了找一根针而不得不每次都重新搭建一遍整个干草堆。论文指出，这是一种巨大的能量和时间的浪费，尤其是在侦探们正在编写长篇、复杂的复杂故事时。

新方法：KV-PRM（“神奇记忆”阅读器）
作者 Peng Kuang 及其团队意识到，侦探们其实已经在免费地完成这项艰苦的工作了！当侦探们书写故事时，他们的脑海（AI 的内部“KV 缓存”）自然地存储了一份关于他们思考过的每一个词的高清、超详细的记忆。这就像是故事灵魂的一段完美、连续的录像，而不仅仅是印刷出来的文字。

它是如何工作的： 法官只需取一个微小的“验证标记”（你可以把它想象成一个神奇的问号“?”），然后询问记忆：“基于目前存储的所有信息，这条路径好吗？”
结果： 因为法官不需要重新阅读文本，成本从剧烈的爆炸式增长降为了简单的线性行走。论文在数学上证明了，这份记忆所包含的信息量严格大于文本本身。这就像拥有一个故事的 3D 全息图，而不是一张平面的纸；全息图在更小的空间内承载了更多的细节。

数据：它到底有多快？
团队在一些最难的数学谜题（如 MATH, GSM8K 和 AIME）上测试了不同的 AI 规模（0.6B, 4B 和 8B 参数）。以下是他们的测量结果：

速度： 在实际运行时间中，KV-PRM 快了高达 37 倍。对于一个长故事，旧法官需要 172.0 毫秒 才能完成的一次检查，新法官仅需 4.6 毫秒。
能量： 每次检查使用的计算步骤（FLOPs）减少了高达 5,000 倍。
内存： 它完成任务所需的计算机内存减少了 34.2 倍。
准确率： 尽管速度快了这么多，它不仅“跟上了”步伐，甚至往往比旧的、缓慢的法官得分更高。

他们明确排除了什么
论文非常明确地指出了哪些做法无效或不是答案：

仅仅缩小法官的规模： 团队尝试使用较小的 AI（0.6B 或 4B）作为旧式的文本阅读型法官以节省成本。他们发现，虽然速度变快了，但并不够聪明。即使是一个微小的 8B KV-PRM 也大幅度超越了庞大的 8B 文本阅读型法官。论文认为，解决办法不仅仅是缩小模型，而是改变其“阅读方式”。
阅读更多的标记： 团队想过：“如果我们用不止一个问号来检查故事会怎样？”他们的数学推导（定理 2）和实验表明，第一个问号已经捕捉到了几乎所有的有用信息。增加第二个或第三个问号几乎不会带来额外的收益，但却会增加成本。因此，坚持只使用一个验证标记才是最完美的平衡点。

一个额外的技巧：“KV 转向”
因为新的法官观察的是“神奇记忆”（一个平滑、连续的信号）而非仅仅是文本（一个破碎、离散的信号），团队发现了一个酷炫的副作用。他们可以通过数学梯度，在侦探们思考的过程中对其进行“微调”，从而将对话引导向更好的答案。他们称之为 KV 转向 (KV Steering)。论文显示，这种方法在 AIME 谜题上的概念验证中表现出色，在没有运行搜索的情况下，将准确率提升了高达 3.33 个百分点。论文指出，由于你无法像引导记忆信号那样去“引导”一张纸，因此这在旧的基于文本的法官身上在结构上是不可能实现的。

底线
这篇论文并不仅仅是暗示这可能奏效，他们通过多个数据集和模型规模进行了测量，并在数学上证明了这一点。他们发现，通过复用 AI 自身的“神奇记忆”而非重新阅读文本，我们可以更快、更便宜、通常也更准确地解决复杂问题。这是一种从“重读整本书”到“查阅作者完美笔记”的转变。

技术摘要：KV-PRM

问题陈述

本文解决了大语言模型（LLM）在测试时扩展（Test-Time Scaling, TTS），特别是多智能体系统（Multi-Agent Systems, MAS）和推理流水线中面临的关键计算瓶颈。虽然过程奖励模型（Process Reward Models, PRMs）对于通过评估中间推理步骤来引导搜索算法（如束搜索、MCTS）至关重要，但现有的 PRM 是基于文本的。它们需要为每一次评分调用重新对整个轨迹文本进行编码。

由于自注意力机制的复杂度随序列长度（ $L$ ）呈二次方增长，单次 PRM 评分调用的计算成本为 $O(L^2)$ 。在长上下文 MAS 场景中，当轨迹跨越数千个 token 且搜索涉及数十个候选对象时，这种冗余的重新编码会产生严重的瓶颈，其计算量往往会达到甚至超过生成本身所需的 FLOPs。这限制了 PRM 在长上下文推理中的实际应用。

方法论：KV-PRM

作者提出了 KV-PRM，一种高效的过程奖励模型，它通过直接利用 LLM 自回归生成阶段自然产生的 键值（Key-Value, KV）缓存，消除了文本重新编码的问题。

核心机制

KV 缓存传输： KV-PRM 不再处理文本 token，而是在序列末尾附加一个单一的**“验证 token”**（例如“?”）。然后，它通过一个轻量级的 LoRA 适配器 处理该 token，该适配器以基础模型的预存 KV 缓存为条件进行计算。
复杂度降低： 通过读取预计算的缓存，评分成本从 $O(L^2)$ （文本重新编码）降至 $O(L)$ （针对 $L$ 个缓存位置处理单个 token）。
架构： 基础模型（ $f_\theta$ ）生成轨迹和 KV 缓存。奖励模型（ $f_{\theta+\Delta\theta}$ ）使用相同的基础权重，并配合一个激活的 LoRA 适配器来解释缓存，并通过验证 token 的 logits 输出分数。在推理过程中，没有梯度流经基础模型；只有 LoRA 参数在训练期间进行更新。

理论基础

论文提供了一个理论框架来证明该方法的合理性：

信息容量（命题 1）： KV 缓存包含的有效信息容量严格大于解码后的文本。若要通过文本无损地表达长度为 $L$ 的 KV 缓存中的信息，需要 $\Omega(d \cdot L / \log |V|)$ 个 token，这意味着缓存的信息密度显著更高。
误差分解（定理 1）： 验证误差被分解为贝叶斯最优下界和近似差距。由于 KV 缓存具有更丰富的表示，基于 KV 缓存的贝叶斯最优预测器的误差等于或低于基于文本的预测器。
收益递减（定理 2）： 增加更多验证 token（ $k$ ）带来的边际信息增益呈指数级衰减。单个验证 token（ $k=1$ ）即可捕捉到大部分可提取的奖励信息，这使得 $O(L)$ 的设计接近最优。

核心贡献

理论框架： 形式化了 KV 缓存相对于文本在验证方面的表示优势，并证明了近似差距随读取深度呈指数级衰减，从而证明了单 token 读取的效率。
算法 (KV-PRM)： 引入了一种新型 PRM 架构，能够复用生成过程中的 KV 缓存，将评分复杂度从 $O(L^2)$ 降低到 $O(L)$ ，消除了基于文本的 PRM 的架构瓶颈。
实证验证： 证明了在多个模型规模（Qwen3-0.6B, 4B, 8B）和搜索算法（束搜索、MCTS、加权投票）下，KV-PRM 的表现能够匹配或超越基于文本的 PRM，并在 MATH、GSM8K 和 AIME 基准测试中得到了验证。
KV 引导 (KV Steering)： 探索了 KV 缓存评分的可微性，以实现对潜在智能体间消息的基于梯度的优化，这是基于离散文本的 PRM 在结构上无法实现的。

实验结果

在 MATH、GSM8K、AIME 2024 和 AIME 2025 上进行的评估显示：

准确率： KV-PRM 匹配或严格优于基于文本的 PRM。例如，在 MATH 任务中使用 MCTS 和 Qwen3-8B 时，KV-PRM 实现了 69.57% 的准确率，而 Text-PRM 为 68.92%。
效率 (FLOPs)： 实现了单次调用评分 FLOPs 高达 5,000 倍 的降低（例如，对于 MATH 上的 8B 模型，其成本仅为 Text-PRM 的 $1/2782$ ）。
延迟： 将实际运行中的评分延迟降低了 15 到 37 倍。在 8B 参数且序列长度为 4,096 的情况下，KV-PRM 的评分时间为 4.6 ms，而 Text-PRM 为 172.0 ms。
内存： 通过避免构建 $O(L^2)$ 的注意力图，将单序列内存占用减少了 34.2 倍。
读取深度： 实验确认，增加验证 token 数量（ $k > 1$ ）会产生收益递减，验证了 $k=1$ 已足够充分的理论主张。

重要性与主张

论文声称，KV-PRM 代表了验证器设计的根本转变，即从冗余的文本重新编码转向表示复用。

可扩展性： 通过将评分成本降低到生成成本的极小部分，KV-PRM 使得在长上下文、多智能体推理场景中实际应用 PRM 成为可能，而此前这些场景被认为在计算上是难以实现的。
新能力： KV 缓存评分信号的可微性为测试时优化（KV Steering）开辟了新方向，允许通过基于梯度的方法对智能体交互进行精炼，而无需进行搜索，这在离散文本输出的情况下是无法实现的。
局限性： 作者指出，该方法要求生成器和验证器共享相同的架构，其理论边界依赖于线性表示假设（Linear Representation Hypothesis），且 KV Steering 目前仍处于概念验证阶段。

总之，KV-PRM 通过利用生成过程中已经计算出的高保真潜在表示，解决了测试时扩展的计算瓶颈，在不牺牲（甚至提升）推理性能的前提下，实现了数量级的效率提升。