KV-PRM：别再让奖励模型重读一遍长推理轨迹

一个多智能体推理系统最尴尬的时刻，不一定发生在模型“想不出来”的时候，而可能发生在模型已经想了很久、写了几千个 token 之后：系统还要把这整段轨迹再喂给另一个奖励模型，让它从头读一遍，判断这一步走得好不好。

如果只验证一次，这听起来还能接受。但 test-time scaling 的核心恰恰不是“一次生成，一次验证”。它要展开很多候选路径，要做 beam search、MCTS、weighted voting，要在多个 agent 交接处反复打分。于是一个本来用来节省错误探索的 PRM，自己变成了新的计算黑洞。

KV-PRM 这篇论文抓住的就是这个工程缝隙：生成模型在写出长轨迹时，其实已经顺手算出了一份高维“过程记录”——KV cache。既然这份记录已经在显存里，为什么奖励模型还要从离散文本重新编码一遍？

先把几个概念放在桌面上

大模型生成文字，是一个 token 接一个 token 地往前走。为了不在每一步都重算前面所有内容，Transformer 会把历史 token 的 key 和 value 存下来。这个存下来的东西就是 KV cache。它不像最终文本那样是“写给人看的句子”，更像模型一路推理时留下的中间神经活动记录。

PRM，也就是 Process Reward Model，可以理解成推理过程里的裁判。它不只看最后答案对不对，而是在中间步骤给分：这条路看起来是不是通向正确答案？在 test-time scaling 里，系统会生成多条候选轨迹，再用 PRM 把算力分给更有希望的路径。

Multi-agent test-time scaling 则把这个问题放大了。一个问题可能先由 Reader 读题，再由 Planner 拆解，再由 Solver 求解，最后由 Verifier 检查。每个 agent 的输出都会接到后面，轨迹越来越长；搜索算法又会复制出很多候选分支。于是 PRM 不再是偶尔调用的“小裁判”，而是每一步都要出场的基础设施。

KV-PRM 的核心想法非常朴素：不要让裁判重新读文字，让裁判直接看生成时留下的 KV cache。论文把这个动作做成一个单独的 verify token：在轨迹结束处追加一个“?”，只让这个 token 带着 LoRA adapter 去 attention 已经存在的 KV cache，然后从“+”和“-”两个 judgment token 的 logits 里得到正向概率 P(+)。

直觉上，为什么 Text-PRM 会这么贵？

我们先想一个最普通的 Text-PRM。它拿到一段长度为 L 的推理轨迹，把整段文本从第一个 token 到最后一个 token 重新做一遍 transformer forward，然后输出一个分数。

这里的关键不是“多跑了一次模型”这么简单，而是自注意力的成本。论文把 full forward 的主项写成两部分：self-attention 是 O(d·L²)，feed-forward 是 O(d²·L)。在长多智能体轨迹里，L 可能是几千甚至上万，注意力项会变成真正的主导成本。论文明确说，在 PRM 和生成 LLM 同规模时，PRM scoring 可能大致把整个系统 compute 翻倍。

更糟糕的是，搜索会乘上调用次数。如果 beam width 是 W，有 D 个 agent step，PRM 至少被调用 W·D 次。MCTS 和 weighted voting 也会反复对候选轨迹评分。也就是说，Text-PRM 的问题不是“有点慢”，而是它把每条候选长轨迹都当成第一次见到，从头理解一遍。

这时我们可能会问：能不能把 PRM 做小？论文确实把这个可能性拿来做了对照。在 Qwen3-8B 作为生成器、MATH 数据集、SBS W=2 的设置下，更小的 Text-PRM 从 0.6B 到 4B 到 8B 准确率会上升，但成本也上升，而且收益递减。KV-PRM-8B 的准确率达到 68.0%，超过所有 Text-PRM 尺寸，同时比 Text-PRM-8B 便宜约 3000×。这说明瓶颈不只是“PRM 太大”，而是“重读文本”这个架构动作本身太浪费。

KV cache 不是副产品，而是更接近现场的证据

KV-PRM 最有意思的地方，是它没有把 KV cache 只当成加速推理的工程缓存。论文把它重新解释成一种 reward modeling 的输入表示。

最终文本是离散的。模型内部可能在某个位置同时激活了多个语义方向，可能对某一步推理有不确定性，可能在 attention pattern 里留下了“这里依赖了前面哪一步”的痕迹。但解码成 token 之后，许多连续状态会被压成同一个词。就像一个学生做题时草稿纸上有擦改、犹豫、依赖关系，最后交上来的答案只剩一行整洁文字。

论文用一个理论框架表达这个直觉。它假设隐藏表示可以看成若干线性语义基的组合；在这个假设下，如果要把长度为 L 的 KV cache 信息无损表达成文本，至少需要 Ω(d·L / log|V|) 个 token。换句话说，每个位置上，KV cache 相对文本有 Ω(d / log|V|) 的信息密度优势。

这个结论不能被读成“KV cache 一定包含所有正确性信息”，也不能被读成“任何 KV-based scorer 都一定更准”。更准确的说法是：如果我们把文本看成 KV cache 的一种离散投影，那么文本丢失信息是结构性的；让验证器直接看 KV cache，理论上有更低的 Bayes-optimal error floor。

接下来论文又问：既然 KV cache 很丰富，那是不是要用很多 token、很多层去读？它给出的 Theorem 2 是一个“收益递减”论证。在低秩 reward structure 假设下，readout depth k 增加时，边际信息增益按指数衰减；单个 query token，也就是 k=1，已经能捕获大部分可提取 reward 信息。实验里的 k>1 消融也支持这个趋势：多加 verify token 会有增益，但曲线快速饱和。

方法其实很克制：一个问号，一个 LoRA，一个二分类分数

KV-PRM 的实现没有引入复杂的新搜索算法。它要求生成器先正常生成轨迹，此时 adapter 关闭，KV cache 按常规方式产生。到了需要评分的时刻，打开 LoRA adapter，把 verify token “?” 放到当前序列末尾，让它对已有 KV cache 做一次 forward。

输出层只看两个 judgment token：“+” 和 “-”。论文用这两个 token 的 logits 做 softmax，得到 P(+)，也就是轨迹质量分数。训练时，标签来自 MCTS 产生的中间 Q-value，范围是 [-1, 1]，再映射到 [0, 1]；损失是 P(+) 和映射后标签之间的 MSE。训练只更新 LoRA 参数，base model 不动。

这个设计的工程含义很强。Text-PRM 评分一条长度 L 的轨迹，要重新编码 [x; v]，成本是 O(d·L²)。KV-PRM 的 verify token 只作为一个 query 去读 L 个缓存位置，成本是 O(d·L)。如果典型 MAS 轨迹长度约 5000 token，论文估算单次 scoring FLOPs 可以减少约 5000×。

这里还有一个容易忽略的限制：KV-PRM 需要 verifier 和 generator 共享架构，至少要能理解同一种 KV cache。它不像传统 Text-PRM 那样，理论上可以拿任意文本轨迹给另一个独立模型评分。这是 KV-PRM 的代价，也是它获得效率的原因：它放弃了“文本作为通用接口”，换来“表示复用”。

数字真正说明了什么？

论文的实验覆盖 Qwen3-0.6B、4B、8B，benchmark 包括 MATH、GSM8K、AIME 2024、AIME 2025，搜索方法包括 step-level beam search、MCTS 和 weighted voting。Text-PRM 与 KV-PRM 使用相同 LoRA rank r=256、α=32，并都用 MCTS 生成的标签训练。

先看准确率。顺序多智能体设置下，Table 1 的 MATH/AIME 2024 平均值里，Beam Search 的 Text-PRM 是 41.61，KV-PRM 是 42.89；MCTS 的 Text-PRM 是 47.42，KV-PRM 是 49.03；Weighted Voting n=10 是 42.69 对 44.17，n=200 是 45.45 对 46.14。不是每一个单元格 KV-PRM 都赢，比如 MATH Qwen3-0.6B 的 Beam Search 从 32.75 到 31.90 略降，但总体上它不是“用准确率换速度”。

Table 2 的 GSM8K/AIME 2025 也类似。Beam Search 平均值从 Text-PRM 60.16 到 KV-PRM 60.51；MCTS 是 58.58 对 58.45，几乎持平；Weighted Voting n=20 是 55.75 对 57.90；n=100 是 57.30 对 58.04。尤其在 AIME 这种样本少、难度高的设置里，单项数字会有波动，所以更合理的读法不是“KV-PRM 每格必胜”，而是：在大范围搜索配置中，KV-cache readout 可以恢复 Text-PRM 的评分质量，并经常略好。

效率数字更直接。论文说，按 scoring cost 归一到 Text-PRM=1×，KV-PRM 的相对成本大约从 1/908 到 1/4877，对应约 9×10² 到 4.9×10³ 的 FLOPs 降低。墙钟延迟在 NVIDIA GH200 120GB GPU 上测量，KV-PRM 比 Text-PRM 快 15× 到 37×；在 8B、L=4096 时，单次 KV-PRM scoring 是 4.6 ms，Text-PRM 是 172.0 ms。内存方面，平均每序列增量最高减少 34.2×。

这组数字的价值在于，它同时打掉了两个怀疑。第一，单 token readout 会不会太弱？实验显示在多数设置下足够。第二，FLOPs 理论降很多，真实系统会不会被别的开销吃掉？延迟和内存实验说明至少在 GH200 上，收益能落到实际 wall-clock。

不过也要诚实：论文没有证明 KV-PRM 在所有类型任务上都更强。它的主实验集中在数学推理 benchmark 和两个多智能体拓扑上；训练标签来自 MCTS 与 exact-match；Text-PRM baseline 也是 LoRA 版本。代码生成、开放工具使用、长程网页任务、含噪人类偏好标签下的表现，仍需要更多证据。

最有想象力的部分：分数可以对 KV cache 求梯度

传统 Text-PRM 的评分对象是文本。文本是离散 token，PRM 可以告诉你“这段轨迹分数高不高”，但很难直接告诉你“把中间状态往哪个连续方向推一点”。KV-PRM 不一样。因为它读的是 KV cache，P(+) 对 cache 是可微的。

论文在第 6 节做了一个 proof-of-concept：KV Steering。每次 agent 交接时，不马上让下一个 agent 继续写，而是把当前累计 KV cache clone 出来，打开 requires_grad，对负的 log P(+) 做梯度下降，也就是沿着提高正向奖励的方向轻推 cache，然后再交给下一个 agent。

这不是传统意义上的搜索，因为它没有展开更多文本候选；也不是训练，因为模型参数不变。它更像在潜空间里给多智能体消息做一次微调。结果不算巨大，但方向清楚：MATH 上 Qwen3-4B 从 60.23 到 61.84，+1.61；Qwen3-8B 从 60.79 到 62.00，+1.21；AIME 2024 上两个规模都是 +3.33；AIME 2025 上 Qwen3-8B 也是 +3.33，Qwen3-4B 没涨。

这部分应该被当成“结构性可能性”，而不是成熟技术。论文自己也说 KV Steering 是 preliminary。它证明的是：一旦奖励模型原生工作在连续表示上，test-time scaling 不一定只能靠采样更多文本、搜索更多分支，也可以探索对 latent message 本身做优化。

这篇论文证明了什么，又没有证明什么？

它证明了三件相当具体的事。

第一，在当前设定下，重复文本编码确实是 PRM 用于多智能体 test-time scaling 的关键瓶颈。论文从复杂度推导、FLOPs 归一化、GH200 延迟和内存测量四个层面都给了证据。

第二，读取生成时已有 KV cache 的单 verify token 方案，在数学推理任务上可以达到接近或超过 Text-PRM 的评分效果。这个结论由 MATH、GSM8K、AIME 2024/2025，以及 Beam Search、MCTS、Weighted Voting 多组配置支撑。

第三，KV-native scorer 带来一个 Text-PRM 没有的性质：分数对 cache 可微。这使 KV Steering 这种潜空间 test-time optimization 至少在结构上可行，并在小规模实验中出现正收益。

它没有证明的东西同样重要。

第一，它没有证明 KV cache scoring 对所有模型架构、所有任务、所有 PRM 训练范式都优于文本评分。论文实验使用 Qwen3 系列，且主要是数学推理 benchmark。

第二，它没有消除理论假设。代表性优势依赖 Linear Representation Hypothesis，k=1 近似最优依赖低秩 reward structure 和谱衰减假设。这些假设帮助我们理解为什么方法可能有效，但不是无条件定理。

第三，它没有解决跨模型验证器的通用性问题。Text-PRM 的接口是文本，天然可移植；KV-PRM 的接口是某个模型的内部 cache，更高效，但更绑定。

所以这篇论文最准确的定位不是“PRM 的终局形态”，而是一个很强的系统设计信号：当 LLM 系统越来越依赖长上下文、多候选、多 agent、反复验证时，把中间表示丢掉再从文本重建，可能会成为越来越不合理的默认选择。

Big Picture：推理优化正在从“多生成”走向“少浪费”

过去两年 test-time scaling 的主旋律是：多想一点，多采样一点，多搜索一点。这个方向确实有效，但它有一个隐含前提：每一份额外计算都应该尽量花在探索新可能性上，而不是重复理解已经生成过的东西。

KV-PRM 的启发就在这里。它没有让模型更会数学，也没有发明新的搜索树策略。它只是指出：生成过程已经支付过一次表示成本，验证过程应该复用它。这个想法一旦成立，影响就不只在 PRM。未来的 verifier、critic、router、memory scorer，甚至 agent 间通信模块，都可能从“读文本”转向“读表示”。

当然，文本接口不会消失。文本可解释、可保存、可跨模型、可审计。KV cache 则高效、连续、贴近模型现场，但更难标准化。真正的问题不是二选一，而是什么时候应该用文本作为公共合同，什么时候应该在同一系统内部复用表示。

KV-PRM 给出的答案很清楚：如果验证器和生成器本来就在同一个推理系统里，如果轨迹很长、评分很频繁、成本已经压过生成本身，那就不要再让裁判重读一遍作文了。让它看草稿纸。让它看模型刚刚留下的现场。