KV-PRM: 通过KV缓存传输实现多智能体测试时间扩展的高效过程奖励建模

# KV-PRM：通过KV缓存传输实现多智能体测试时扩展的高效过程奖励建模
来源：https://arxiv.org/html/2607.09153
彭旷¹，晋海波¹，韩晓宇¹，王艳丽²，袁晓鹏¹，余烨¹，徐凯迪³，王浩涵¹  
¹伊利诺伊大学厄巴纳-香槟分校，²帝国理工学院，³香港城市大学  
[email protected], [email protected], [email protected]

###### 摘要

过程奖励模型（PRMs）已被证明在引导测试时扩展（TTS）方法方面极为有效，能够显著提升基于LLM的多智能体系统的能力。然而，现有的PRM本质上是*基于文本*的：它们需从头重新编码整个轨迹文本。在长多智能体交互过程中，评分成本随序列长度$L$呈二次增长，形成严重的计算瓶颈，极大限制了PRM在长上下文场景中的应用。为解决此问题，我们提出KV-PRM，一种高效的过程奖励模型，通过直接读取LLM生成阶段自然产生的KV缓存，消除了繁重的文本重新编码。通过针对预先存在的KV缓存处理单个“验证令牌”，KV-PRM将评分成本从$O(L^2)$降低到$O(L)$。我们形式化证明了KV缓存包含严格多于文本的信息容量，且对于下游奖励建模更为高效。在MATH、GSM8K和AIME基准测试上的实验表明，KV-PRM在各种TTS方法（如束搜索、蒙特卡洛树搜索和加权投票）下匹配或严格优于基于文本的PRM，评分FLOPs降低高达$5000\times$，延迟降低$37\times$，每条序列的内存占用减少$34\times$。

## 1 引言

近期，扩展测试时计算已成为提升大语言模型（LLMs）在复杂推理任务上能力的极具前景的方向\[26 (https://arxiv.org/html/2607.09153#bib.bib24),23 (https://arxiv.org/html/2607.09153#bib.bib12),3 (https://arxiv.org/html/2607.09153#bib.bib25)\]。不再依赖单次贪婪解码，高级推理流程利用多智能体系统（MAS）\[7 (https://arxiv.org/html/2607.09153#bib.bib6)\]和测试时搜索（TTS）算法\[15 (https://arxiv.org/html/2607.09153#bib.bib34)\]（如束搜索和蒙特卡洛树搜索（MCTS）\[5 (https://arxiv.org/html/2607.09153#bib.bib15)\]）来探索多样化解轨迹。引导这种搜索的关键组件是过程奖励模型（PRM）\[19 (https://arxiv.org/html/2607.09153#bib.bib1),24 (https://arxiv.org/html/2607.09153#bib.bib2),34 (https://arxiv.org/html/2607.09153#bib.bib29)\]。通过评估中间推理步骤的正确性，PRM引导搜索算法将计算资源分配给更有前景的部分解。

然而，扩展测试时搜索暴露了一个关键的计算瓶颈：PRM自身的成本。现有PRM本质上是*基于文本*的。对于生成器提出的每个候选轨迹，文本PRM必须从头编码整个文本令牌序列以输出奖励。由于自注意力机制随序列长度呈二次增长，每次评分调用引入$O(L^2)$的计算成本，其中$L$为序列长度。在现代MAS流程中，推理轨迹通常包含数千甚至数万个令牌。当在多智能体交接时搜索数十个候选时，PRM冗余的文本重新编码成为更重的负担，往往匹配生成本身所需的FLOPs。这种架构瓶颈限制了PRM在长上下文推理场景中的应用。

在这项工作中，我们识别出现有测试时扩展系统中的一个被忽视的机会：智能体的生成过程本身计算了轨迹的丰富、高维表示——键值（KV）缓存。由于自回归生成需要缓存历史状态以预测下一个令牌，模型在LLM生成阶段自然产生了跨所有层和位置的完整中间表示记录。在第3节 (https://arxiv.org/html/2607.09153#S3)中，我们进一步从理论上证明了KV缓存的表示优势及其相比文本的高效利用。具体而言，我们建立了一个理论框架，证明KV缓存包含严格多于文本的信息容量，并且下游奖励建模所需的计算量更少。我们形式化证明KV缓存每个位置的容量优势为$\Omega(d/\log|\mathcal{V}|)$，并证明近似间隙随着读出深度$k$呈指数衰减，对于实际架构，$k=1$即可捕获大部分可提取的奖励信息（定理2 (https://arxiv.org/html/2607.09153#Thmtheorem2)）。从头重新编码解码后的文本会丢弃这种高保真、连续的表示，迫使验证模型仅依赖离散文本令牌，这坍缩了丰富的内部状态且包含严格更少的信息。

基于这一洞察，我们提出KV-PRM，一种用于过程奖励建模的高效架构。KV-PRM不处理文本，而是通过*KV缓存传输*对轨迹进行评分。通过在序列末尾附加单个“验证令牌”，并利用基于基础模型预先存在的KV缓存的轻量级LoRA适配器\[10 (https://arxiv.org/html/2607.09153#bib.bib21)\]进行处理，KV-PRM实现了数量级的复杂度降低。每次评分调用的成本从$O(L^2)$降至$O(L)$，对于典型的MAS轨迹，评分FLOPs减少高达$5000\times$。

在实验方面，我们在第5节 (https://arxiv.org/html/2607.09153#S5)中评估了KV-PRM在三个模型规模（Qwen3-0.6B、4B、8B）下的MATH、GSM8K和AIME基准测试表现。我们发现KV-PRM不仅恢复了基于文本PRM的性能，且经常超越它们，在加速壁钟评分延迟$15$–$37\times$的同时缩小了理论近似间隙，并将每条序列的内存占用减少$34.2\times$。最后，由于KV-PRM评估连续的KV缓存而非离散文本，其奖励信号完全可微。我们通过一项名为KV Steering的概念验证技术（第6节 (https://arxiv.org/html/2607.09153#S6)）探索了这一特性，表明对智能体间潜在消息进行基于梯度的优化在结构上是可行的，并代表了测试时扩展在潜在空间中的一个有前景的方向。总结而言，我们的贡献如下：

- •**理论框架**：我们形式化了KV缓存相比文本令牌在验证方面的表示优势，并证明近似间隙随读出深度指数衰减，为基于KV缓存的奖励建模的效率提供了理论依据。
- •**算法（KV-PRM）**：我们提出了一种新颖的PRM架构，重用生成KV缓存，将评分复杂度从$O(L^2)$降至$O(L)$，解决了基于文本PRM的架构瓶颈。
- •**广泛的经验验证**：我们证明KV-PRM在多个模型规模和搜索算法下匹配或超越基于文本PRM的准确性，同时实现$37\times$的壁钟加速和每条序列$34.2\times$的内存占用减少。

## 2 预备知识与符号

**多智能体推理系统**。多智能体系统（MAS）将问题分解给$D$个专用LLM智能体$a_1,\ldots,a_D$\[7 (https://arxiv.org/html/2607.09153#bib.bib6),9 (https://arxiv.org/html/2607.09153#bib.bib7)\]。每个智能体$a_j$基于前序智能体的输出和原始问题，使用具有$N_L$层、每层$n_h$个注意力头、隐藏维度$d$的仅解码器transformer $f_\theta$生成输出$\mathbf{o}_j$。令$L_j$表示智能体$a_j$完成后的总序列长度，$L=L_D$为最终轨迹长度。

**过程奖励模型与测试时搜索**。过程奖励模型（PRM）预测截至智能体步骤$j$的轨迹分数$s_j \in [0,1]$，指示达到正确答案的可能性\[19 (https://arxiv.org/html/2607.09153#bib.bib1),24 (https://arxiv.org/html/2607.09153#bib.bib2),31 (https://arxiv.org/html/2607.09153#bib.bib5)\]。PRM分数引导测试时搜索（TTS），如束搜索、MCTS\[13 (https://arxiv.org/html/2607.09153#bib.bib16)\]或加权多数投票\[25 (https://arxiv.org/html/2607.09153#bib.bib13)\]，以选择高质量轨迹\[23 (https://arxiv.org/html/2607.09153#bib.bib12)\]。使用束宽度$W$和$D$个智能体步骤，PRM至少被调用$W \cdot D$次。现有PRM是*基于文本*的：它们对长度为$L_j$的轨迹文本$\mathbf{x}_j$执行完整前向传播，每次调用产生成本$F_{\text{forward}}(L_j)$。

**自回归生成中的KV缓存**。在自回归生成过程中，transformer通过缓存所有先前令牌在每一层的键和值投影来避免冗余重计算。生成$L$个令牌后，这个*KV缓存* $(\mathbf{K},\mathbf{V})$，其中$\mathbf{K},\mathbf{V} \in \mathbb{R}^{N_L \times L \times d}$，构成了模型跨所有层和位置的完整中间表示记录。每个新令牌通过以$O(d \cdot L)$的成本关注此缓存来生成。关键的是，KV缓存是生成的*副产品*，无论是否用于任何下游目的都存在。

**计算瓶颈**。对长度为$L$的序列进行前向传播的主要成本为：

$$F_{\text{forward}}(L)=N_L\cdot\big(\underbrace{c_{\text{attn}}\cdot d\cdot L^2}_{\text{自注意力}}+\underbrace{c_{\text{ffn}}\cdot d^2\cdot L}_{\text{前馈}}\big)$$

对于MAS中的长轨迹（$L \gg d$），自注意力占主导，每次文本PRM评分调用成本为$O(d \cdot L^2)$。当PRM和LLM规模相同时，PRM评分大约*翻倍*了系统总计算量。这引出了我们的核心问题：*我们能否在不损失评分质量的前提下，将每次评分调用的成本从$O(d \cdot L^2)$降至$O(d \cdot L)$，即降低$L \times$，同时保持评分质量？* 我们首先在第3节 (https://arxiv.org/html/2607.09153#S3)中分析KV缓存中的信息，以确立这种降低不仅在计算上可行，而且在理论上有坚实基础，随后在第4节 (https://arxiv.org/html/2607.09153#S4)中提出我们的方法。

## 3 理论理解

我们首先在第3.1节 (https://arxiv.org/html/2607.09153#S3.SS1)中回顾KV缓存相比文本令牌的表示优势，这构成了我们主要理论贡献的基础：验证误差分解表明贝叶斯最优下限偏向KV缓存评分（定理1 (https://arxiv.org/html/2607.09153#Thmtheorem1)），以及收益递减分析证明单令牌读出即可捕获大部分可提取的奖励信息（定理2 (https://arxiv.org/html/2607.09153#Thmtheorem2)）。完整证明见附录A.3 (https://arxiv.org/html/2607.09153#A1.SS3)。

### 3.1 KV缓存的表示优势

作为后续结果的基础，我们首先建立生成过程中产生的KV缓存携带严格丰富于解码文本的信息。

###### 假设1（线性表示假说\[22 (https://arxiv.org/html/2607.09153#bib.bib19),35 (https://arxiv.org/html/2607.09153#bib.bib35)\]）。

transformer产生的隐藏嵌入$h \in \mathbb{R}^d$是线性无关语义基$\{s_1,\ldots,s_d\} \subset \mathbb{R}^d$的线性组合$h = \sum_{i=1}^d c_i s_i$，系数$c_i \in \{0,\pm1\}$为三元，其中$c_i=0$表示语义$i$缺失，$c_i=\pm1$表示其正或负存在。

###### 命题1（KV缓存的表示优势）。

令$f_\theta$为一个自回归transformer，在生成过程中产生KV缓存$\mathbf{H}=(\mathbf{K},\mathbf{V})$和文本令牌$\mathbf{x}=(x_1,\ldots,x_L)$，其中$\mathbf{x}=\mathrm{decode}(\mathbf{H})$。令$Y$表示任何感兴趣的目标变量（例如，轨迹正确性）。那么在假设1 (https://arxiv.org/html/2607.09153#Thmassumption1)下，如果长度为$L$的轨迹的KV缓存中的信息要无损地通过文本表达，所需令牌数至少为

$$m' = \Omega\left(\frac{d \cdot L}{\log|\mathcal{V}|}\right),$$

其中$|\mathcal{V}|$是词汇表大小。等价地，KV缓存每个位置的信息密度是文本的$\Omega(d / \log|\mathcal{V}|)$倍。

### 3.2 通过KV缓存读出的高效验证

定理1 (https://arxiv.org/html/2607.09153#Thmproposition1)确立KV缓存是用于验证的高信息量表示。我们现在形式化如何高效利用这种丰富性。

###### 定义1（深度$k$读出）。

*深度$k$读出* $\mathcal{R}_k: (\mathbf{K},\mathbf{V}) \mapsto s \in [0,1]$通过参数化注意力函数$f_\psi$处理$k$个查询向量，针对长度为$L$的预先存在的KV缓存。我们记$\mathcal{F}_k$为所有参数容量有界$\|\psi\| \leq B$的深度$k$读出的函数类。文本重新编码对应$k=L$；KV缓存读出使用$k \ll L$。

###### 定义2（近似间隙）。

对于作用于KV缓存$\mathbf{H}$的函数类$\mathcal{F}_k$，*近似间隙*为：

$$\mathrm{Gap}(\mathcal{F}_k,\mathbf{H}) = \inf_{f \in \mathcal{F}_k} \mathbb{E}[\ell(f(\mathbf{H}),Y)] - \inf_{f} \mathbb{E}[\ell(f(\mathbf{H}),Y)],$$

其中第二个下确界针对所有可测函数。$\mathrm{Gap}(\mathcal{F}_k,\mathbf{H})$随$k$单调非增。

###### 定理1（验证误差分解）。

对于任何深度$k$读出$\mathcal{R}_k \in \mathcal{F}_k$，验证误差分解为：

$$\epsilon(\mathcal{R}_k) = \epsilon_{\mathrm{Bayes}}(\mathbf{H}) + \mathrm{Gap}(\mathcal{F}_k,\mathbf{H}),$$

其中由命题1 (https://arxiv.org/html/2607.09153#Thmproposition1)得$\epsilon_{\mathrm{Bayes}}(\mathbf{H}) \leq \epsilon_{\mathrm{Bayes}}(\mathbf{x})$。即，来自KV缓存的贝叶斯最优预测器实现等于或低于来自文本的贝叶斯最优预测器的误差。近似间隙$\mathrm{Gap}(\mathcal{F}_k,\mathbf{H})$由读出深度$k$控制，且随$k$单调非增。

###### 假设2（低秩奖励结构）。

奖励$Y$通过一个线性投影$\Phi(\mathbf{H}) = W_{\varphi} \text{pool}(\mathbf{H})$依赖于KV缓存$\mathbf{H}$...