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

arXiv cs.AI Papers

Summary

KV-PRM introduces a process reward model that leverages KV-cache transfer to avoid re-encoding, achieving up to 5000x FLOP reduction while maintaining or improving performance on reasoning benchmarks.

arXiv:2607.09153v1 Announce Type: new Abstract: Process Reward Models (PRMs) have been proven to be highly effective in guiding test-time scaling (TTS) methods, which significantly boost the capabilities of LLM-based multi-agent systems. However, existing PRMs are text-based: they re-encode the entire trajectory text from scratch. In long multi-agent rollouts, the scoring cost, growing quadratically with respect to sequence length L, creates a severe computational bottleneck, severely limiting PRMs' application in long-context scenarios. To resolve this, we introduce KV-PRM, a highly efficient process reward model that eliminates the heavy text re-encoding by directly reading the KV cache produced naturally during the LLM's generation phase. By processing a single "verify token" against the pre-existing KV cache, KV-PRM reduces the scoring cost from O(L^2) to O(L). We formally prove that the KV cache contains strictly greater information capacity than text, and is more efficient for downstream reward modeling. Empirically, across the MATH, GSM8K, and AIME benchmarks, KV-PRM matches or strictly outperforms text-PRMs under various TTS methods such as Beam Search, MCTS, and Weighted Voting, with up to a 5,000x reduction in scoring FLOPs, a 37x reduction in latency, and a 34x reduction in per-sequence memory footprint compared to text-based PRMs.
Original Article
View Cached Full Text

Cached at: 07/13/26, 07:52 AM

# KV-PRM: Efficient Process Reward Modeling via KV-Cache Transfer for Multi-Agent Test-Time Scaling
Source: [https://arxiv.org/html/2607.09153](https://arxiv.org/html/2607.09153)
Peng Kuang1, Haibo Jin1, Xiaoyu Han1, Yanli Wang2, Xiaopeng Yuan1, Ye Yu1,Kaidi Xu3,Haohan Wang1 1University of Illinois Urbana\-Champaign,2Imperial College London, 3City University of Hong Kong pengk2@illinois\.edu,kaidixu@cityu\.edu\.hk, haohanw@illinois\.edu

###### Abstract

Process Reward Models \(PRMs\) have been proven to be highly effective in guiding test\-time scaling \(TTS\) methods, which significantly boost the capabilities of LLM\-based multi\-agent systems\. However, existing PRMs are*text\-based*: they re\-encode the entire trajectory text from scratch\. In long multi\-agent rollouts, the scoring cost, growing quadratically with respect to sequence lengthLL, creates a severe computational bottleneck, severely limiting PRMs’ application in long\-context scenarios\. To resolve this, we introduce KV\-PRM, a highly efficient process reward model that eliminates the heavy text re\-encoding by directly reading the KV cache produced naturally during the LLM’s generation phase\. By processing a single "verify token" against the pre\-existing KV cache, KV\-PRM reduces the scoring cost fromO​\(L2\)O\(L^\{2\}\)toO​\(L\)O\(L\)\. We formally prove that the KV cache contains strictly greater information capacity than text, and is more efficient for downstream reward modeling\. Empirically, across the MATH, GSM8K, and AIME benchmarks, KV\-PRM matches or strictly outperforms text\-PRMs under various TTS methods such as Beam Search, MCTS, and Weighted Voting, with up to a 5,000×\\timesreduction in scoring FLOPs, a 37×\\timesreduction in latency, and a 34×\\timesreduction in per\-sequence memory footprint compared to text\-based PRMs\.

## 1Introduction

Scaling test\-time compute has recently emerged as a highly effective direction for improving the capabilities of Large Language Models \(LLMs\) on complex reasoning tasks\[[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)\]\. Instead of relying on a single greedy decoding pass, advanced reasoning pipelines leverage multi\-agent systems \(MAS\)\[[7](https://arxiv.org/html/2607.09153#bib.bib6)\]and test\-time search \(TTS\) algorithms\[[15](https://arxiv.org/html/2607.09153#bib.bib34)\], such as beam search and Monte Carlo Tree Search \(MCTS\)\[[5](https://arxiv.org/html/2607.09153#bib.bib15)\], to explore diverse solution trajectories\. A critical component guiding this search is the Process Reward Model \(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)\]\. By evaluating the correctness of intermediate reasoning steps, PRMs guide the search algorithm to allocate compute toward more promising partial solutions\.

However, scaling test\-time search exposes a critical computational bottleneck: the cost of the PRM itself\. Existing PRMs are fundamentally*text\-based*\. For every candidate trajectory proposed by the generator, a text\-PRM must encode the entire sequence of text tokens from scratch to output a reward\. As self\-attention scales quadratically with sequence length, this introduces anO​\(L2\)O\(L^\{2\}\)computational cost per scoring call, whereLLis the sequence length\. In modern MAS workflows, reasoning trajectories routinely span thousands or tens of thousands of tokens\. When searching over dozens of candidates at multiple agent handoffs, the redundant text re\-encoding by the PRM becomes an even heavier burden, often matching the FLOPs required for the generation itself\. Such an architectural bottleneck limits the application of PRMs in long\-context reasoning scenarios\.

In this work, we identify a missed opportunity in current test\-time scaling systems: the generation process of agents inherently computes a rich, high\-dimensional representation of the trajectory, the Key\-Value \(KV\) cache\. As autoregressive generation requires caching historical states to predict the next token, a complete record of the model’s intermediate representations across all layers and positions is produced naturally during the LLM’s generation phase\. In Section[3](https://arxiv.org/html/2607.09153#S3), we further theoretically prove the representational advantage of KV cache and its efficient exploitation compared to text\. Specifically, we establish a theoretical framework demonstrating that the KV cache contains strictly greater information capacity than text, and requires much less compute for downstream reward modeling\. We formally show that the KV cache provides anΩ​\(d/log⁡\|𝒱\|\)\\Omega\(d/\\log\|\\mathcal\{V\}\|\)capacity advantage per position over text tokens, and prove that the approximation gap decays exponentially in readout depthkk, with efficientk=1k\\\!=\\\!1capturing most of extractable reward information for practical architectures \(Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)\)\. Re\-encoding the decoded text from scratch discards this high\-fidelity, continuous representation, forcing the verification model to rely solely on discrete text tokens that collapse the rich internal state and contain strictly less information\.

Building upon this insight, we propose KV\-PRM, a highly efficient architecture for process reward modeling\. Instead of processing text, KV\-PRM scores trajectories via*KV\-cache transfer*\. By appending a single “verify token” to the end of the sequence and processing it through a lightweight LoRA adapter\[[10](https://arxiv.org/html/2607.09153#bib.bib21)\]conditioned on the base model’s pre\-existing KV cache, KV\-PRM achieves orders\-of\-magnitude complexity reduction\. The per\-call scoring cost drops fromO​\(L2\)O\(L^\{2\}\)toO​\(L\)O\(L\), resulting in up to a5,000×5\{,\}000\\timesreduction in scoring FLOPs for typical MAS trajectories\.

Empirically, we evaluate KV\-PRM across three model scales \(Qwen3\-0\.6B, 4B, 8B\) on MATH, GSM8K, and AIME benchmarks in Section[5](https://arxiv.org/html/2607.09153#S5)\. We find that KV\-PRM not only recovers the performance of text\-based PRMs but frequently outperforms them, closing the theoretical approximation gap while accelerating wall\-clock scoring latency by1515–37×37\\times, and reducing the per\-sequence memory footprint by 34\.2×\\times\. Finally, as KV\-PRM evaluates the continuous KV cache rather than discrete text, its reward signal is fully differentiable\. We explore this property through a proof\-of\-concept technique called KV Steering \(Section[6](https://arxiv.org/html/2607.09153#S6)\), demonstrating that gradient\-based optimization of the latent messages between agents is structurally possible and presents a promising direction for test\-time scaling in the latent space\. In summary, our contributions are:

- •Theoretical framework:We formalize the representational advantage of the KV cache over text tokens for verification, and prove that the approximation gap decays exponentially in readout depth, formally justifying the efficiency of KV cache\-based reward modeling\.
- •Algorithm \(KV\-PRM\):We introduce a novel PRM architecture that reuses the generation KV cache, reducing scoring complexity fromO​\(L2\)O\(L^\{2\}\)toO​\(L\)O\(L\)by addressing the architectural bottleneck of text\-based PRMs\.
- •Extensive empirical validation:We demonstrate that KV\-PRM matches or exceeds the accuracy of text\-based PRMs across multiple model scales and search algorithms while delivering37×37\\timeswall\-clock speedups and 34\.2×\\timesreduction in per\-sequence memory footprint\.

## 2Preliminaries and Notations

Multi\-agent reasoning systems\.A multi\-agent system \(MAS\) decomposes a problem acrossDDspecialized LLM agentsa1,…,aDa\_\{1\},\\ldots,a\_\{D\}\[[7](https://arxiv.org/html/2607.09153#bib.bib6),[9](https://arxiv.org/html/2607.09153#bib.bib7)\]\. Each agentaja\_\{j\}generates output𝐨j\\mathbf\{o\}\_\{j\}conditioned on preceding agents’ outputs and the original question, using a decoder\-only transformerfθf\_\{\\theta\}withNLN\_\{L\}layers,nhn\_\{h\}attention heads per layer, and hidden dimensiondd\. LetLjL\_\{j\}denote the total sequence length after agentaja\_\{j\}completes, andL=LDL=L\_\{D\}the final trajectory length\.

Process reward models and test\-time search\.A process reward model \(PRM\) predicts a scoresj∈\[0,1\]s\_\{j\}\\in\[0,1\]for the trajectory up to agent stepjj, indicating the likelihood of reaching a correct answer\[[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 scores guide test\-time search \(TTS\), such as beam search, MCTS\[[13](https://arxiv.org/html/2607.09153#bib.bib16)\], or weighted majority voting\[[25](https://arxiv.org/html/2607.09153#bib.bib13)\], to select high\-quality trajectories\[[23](https://arxiv.org/html/2607.09153#bib.bib12)\]\. With beam widthWWandDDagent steps, the PRM is invoked at leastW⋅DW\\cdot Dtimes\. Existing PRMs are*text\-based*: they perform a full forward pass over the trajectory text𝐱j\\mathbf\{x\}\_\{j\}of lengthLjL\_\{j\}, incurring costFforward​\(Lj\)F\_\{\\text\{forward\}\}\(L\_\{j\}\)per call\.

KV cache in autoregressive generation\.During autoregressive generation, the transformer avoids redundant recomputation by caching the key and value projections from all previous tokens at every layer\. After generatingLLtokens, this*KV cache*\(𝐊,𝐕\)\(\\mathbf\{K\},\\mathbf\{V\}\)with𝐊,𝐕∈ℝNL×L×d\\mathbf\{K\},\\mathbf\{V\}\\in\\mathbb\{R\}^\{N\_\{L\}\\times L\\times d\}constitutes a complete record of the model’s intermediate representations across all layers and positions\. Each new token is generated by attending to this cache at costO​\(d⋅L\)O\(d\\cdot L\)per token\. Crucially, the KV cache is a*byproduct*of generation that exists whether or not it is used for any downstream purpose\.

The computational bottleneck\.The dominant cost in a forward pass over a sequence of lengthLLis:

Fforward​\(L\)=NL⋅\(cattn⋅d⋅L2⏟self\-attention\+cffn⋅d2⋅L⏟feed\-forward\)F\_\{\\text\{forward\}\}\(L\)=N\_\{L\}\\cdot\\big\(\\underbrace\{c\_\{\\text\{attn\}\}\\cdot d\\cdot L^\{2\}\}\_\{\\text\{self\-attention\}\}\+\\underbrace\{c\_\{\\text\{ffn\}\}\\cdot d^\{2\}\\cdot L\}\_\{\\text\{feed\-forward\}\}\\big\)\(1\)For the long trajectories in MAS \(L≫dL\\gg d\), self\-attention dominates, and each text\-PRM scoring call costsO​\(d⋅L2\)O\(d\\cdot L^\{2\}\)\. When the PRM and the LLM are at the same scale, PRM scoring roughly*doubles*the total system compute\. This motivates our central question:*can we reduce the per\-call scoring cost fromO​\(d⋅L2\)O\(d\\cdot L^\{2\}\)toO​\(d⋅L\)O\(d\\cdot L\), a reduction ofL×L\\times, without a loss in scoring quality?*We begin by analyzing the information in the KV cache to establish that this reduction is not only computationally feasible but theoretically well\-grounded in Section[3](https://arxiv.org/html/2607.09153#S3), before presenting our method in Section[4](https://arxiv.org/html/2607.09153#S4)\.

## 3Theoretical Understanding

We begin by reviewing the representational advantage of the KV cache over text tokens in Section[3\.1](https://arxiv.org/html/2607.09153#S3.SS1), which serves as the foundation for our main theoretical contributions: a verification error decomposition showing that the Bayes\-optimal floor favors KV\-cache scoring \(Theorem[1](https://arxiv.org/html/2607.09153#Thmtheorem1)\), and a diminishing\-returns analysis proving that a single\-token readout readily captures most of the extractable reward information \(Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)\)\. Full proof in Appendix[A\.3](https://arxiv.org/html/2607.09153#A1.SS3)\.

### 3\.1Representational Advantage of KV Cache

As a foundation for the results that follow, we first establish that the KV cache produced during generation carries strictly richer information than decoded text\.

###### Assumption 1\(Linear Representation Hypothesis\[[22](https://arxiv.org/html/2607.09153#bib.bib19),[35](https://arxiv.org/html/2607.09153#bib.bib35)\]\)\.

The hidden embeddingsh∈ℝdh\\in\\mathbb\{R\}^\{d\}produced by the transformer are linear combinationsh=∑i=1dci​sih=\\sum\_\{i=1\}^\{d\}c\_\{i\}s\_\{i\}of a linearly independent semantic basis\{s1,…,sd\}⊂ℝd\\\{s\_\{1\},\\ldots,s\_\{d\}\\\}\\subset\\mathbb\{R\}^\{d\}with ternary coefficientsci∈\{0,±1\}c\_\{i\}\\in\\\{0,\\pm 1\\\}, whereci=0c\_\{i\}=0indicates absence of semanticii, andci=±1c\_\{i\}=\\pm 1its positive or negative presence\.

###### Proposition 1\(Representational Advantage of KV Cache\)\.

Letfθf\_\{\\theta\}be an autoregressive transformer that, during generation, produces KV cache𝐇=\(𝐊,𝐕\)\\mathbf\{H\}=\(\\mathbf\{K\},\\mathbf\{V\}\)and text tokens𝐱=\(x1,…,xL\)\\mathbf\{x\}=\(x\_\{1\},\\ldots,x\_\{L\}\)via𝐱=decode​\(𝐇\)\\mathbf\{x\}=\\mathrm\{decode\}\(\\mathbf\{H\}\)\. LetYYdenote any target variable of interest \(e\.g\., trajectory correctness\)\. Then, under Assumption[1](https://arxiv.org/html/2607.09153#Thmassumption1), if the information in the KV cache of a length\-LLtrajectory is to be expressed losslessly through text, the required number of tokens is at least

m′=Ω​\(d⋅Llog⁡\|𝒱\|\),m^\{\\prime\}\\;=\\;\\Omega\\\!\\left\(\\frac\{d\\cdot L\}\{\\log\|\\mathcal\{V\}\|\}\\right\),\(2\)where\|𝒱\|\|\\mathcal\{V\}\|is the vocabulary size\. Equivalently, the KV cache isΩ​\(d/log⁡\|𝒱\|\)\\Omega\(d/\\log\|\\mathcal\{V\}\|\)times more information\-dense per position than text\.

### 3\.2Efficient Verification via KV\-Cache Readout

Theorem[1](https://arxiv.org/html/2607.09153#Thmproposition1)establishes that the KV cache is an information\-rich representation for verification\. We now formalize how this richness can be exploited efficiently\.

###### Definition 1\(Depth\-kkReadout\)\.

A*depth\-kkreadout*ℛk:\(𝐊,𝐕\)↦s∈\[0,1\]\\mathcal\{R\}\_\{k\}:\(\\mathbf\{K\},\\mathbf\{V\}\)\\mapsto s\\in\[0,1\]processeskkquery vectors against the pre\-existing KV cache of lengthLLthrough a parameterized attention functionfψf\_\{\\psi\}\. We denote byℱk\\mathcal\{F\}\_\{k\}the function class of all depth\-kkreadouts with bounded parameter capacity‖ψ‖≤B\\\|\\psi\\\|\\leq B\. Text re\-encoding corresponds tok=Lk=L; KV\-cache readout usesk≪Lk\\ll L\.

###### Definition 2\(Approximation Gap\)\.

For a function classℱk\\mathcal\{F\}\_\{k\}operating on KV cache𝐇\\mathbf\{H\}, the*approximation gap*is:

Gap​\(ℱk,𝐇\)=inff∈ℱk𝔼​\[ℓ​\(f​\(𝐇\),Y\)\]−inff𝔼​\[ℓ​\(f​\(𝐇\),Y\)\],\\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\)\],\(3\)where the second infimum is over all measurable functions\.Gap​\(ℱk,𝐇\)\\mathrm\{Gap\}\(\\mathcal\{F\}\_\{k\},\\mathbf\{H\}\)is monotonically non\-increasing inkk\.

###### Theorem 1\(Verification Error Decomposition\)\.

For any depth\-kkreadoutℛk∈ℱk\\mathcal\{R\}\_\{k\}\\in\\mathcal\{F\}\_\{k\}, the verification error decomposes as:

ϵ​\(ℛk\)=ϵBayes​\(𝐇\)\+Gap​\(ℱk,𝐇\),\\epsilon\(\\mathcal\{R\}\_\{k\}\)=\\epsilon\_\{\\mathrm\{Bayes\}\}\(\\mathbf\{H\}\)\+\\mathrm\{Gap\}\(\\mathcal\{F\}\_\{k\},\\,\\mathbf\{H\}\),\(4\)whereϵBayes​\(𝐇\)≤ϵBayes​\(𝐱\)\\epsilon\_\{\\mathrm\{Bayes\}\}\(\\mathbf\{H\}\)\\leq\\epsilon\_\{\\mathrm\{Bayes\}\}\(\\mathbf\{x\}\)by Proposition[1](https://arxiv.org/html/2607.09153#Thmproposition1)\. That is, the Bayes\-optimal predictor from the KV cache achieves equal or lower error than the Bayes\-optimal predictor from text\. The approximation gapGap​\(ℱk,𝐇\)\\mathrm\{Gap\}\(\\mathcal\{F\}\_\{k\},\\mathbf\{H\}\)is controlled by the readout depthkkand is monotonically non\-increasing inkk\.

###### Assumption 2\(Low\-Rank Reward Structure\)\.

The rewardYYdepends on the KV cache𝐇\\mathbf\{H\}through a linear projectionΦ​\(𝐇\)=Wϕ​pool​\(𝐇\)∈ℝr\\Phi\(\\mathbf\{H\}\)=W\_\{\\phi\}\\,\\mathrm\{pool\}\(\\mathbf\{H\}\)\\in\\mathbb\{R\}^\{r\}, whererris the effective reward dimensionality andWϕ∈ℝr×dW\_\{\\phi\}\\in\\mathbb\{R\}^\{r\\times d\}\. The singular valuesσ1≥σ2≥⋯≥σr\\sigma\_\{1\}\\geq\\sigma\_\{2\}\\geq\\cdots\\geq\\sigma\_\{r\}of the covariance ofΦ​\(𝐇\)\\Phi\(\\mathbf\{H\}\)projected onto the reward\-relevant subspace decay asσj=O​\(e−α​j\)\\sigma\_\{j\}=O\(e^\{\-\\alpha j\}\)for some spectral decay rateα\>0\\alpha\>0\[[6](https://arxiv.org/html/2607.09153#bib.bib33),[20](https://arxiv.org/html/2607.09153#bib.bib32)\]\.

###### Theorem 2\(Diminishing Returns of Readout Depth\)\.

Under Assumptions[1](https://arxiv.org/html/2607.09153#Thmassumption1)and[2](https://arxiv.org/html/2607.09153#Thmassumption2), letℛk∈ℱk\\mathcal\{R\}\_\{k\}\\in\\mathcal\{F\}\_\{k\}be a depth\-kkreadout over KV cache𝐇\\mathbf\{H\}, and define the marginal information gainΔ​Ik=I​\(ℛk​\(𝐇\);Y∣ℛk−1​\(𝐇\)\)\\Delta I\_\{k\}=I\(\\mathcal\{R\}\_\{k\}\(\\mathbf\{H\}\);\\,Y\\mid\\mathcal\{R\}\_\{k\-1\}\(\\mathbf\{H\}\)\)\. Then:

1. \(a\)Exponential decay\.The marginal gain decays exponentially: Δ​Ik≤C0⋅e−α⋅nh​NL⋅\(k−1\),\\Delta I\_\{k\}\\;\\leq\\;C\_\{0\}\\cdot e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\\cdot\(k\-1\)\},\(5\)whereC0=∑j=1r12​log⁡\(1\+SNR⋅σj2\)C\_\{0\}=\\sum\_\{j=1\}^\{r\}\\frac\{1\}\{2\}\\log\(1\+\\mathrm\{SNR\}\\cdot\\sigma\_\{j\}^\{2\}\)is the total extractable information, withSNR=Var​\(Y\)/σϵ2\\mathrm\{SNR\}=\\mathrm\{Var\}\(Y\)/\\sigma\_\{\\epsilon\}^\{2\}the signal\-to\-noise ratio\.
2. \(b\)k=1k\\\!=\\\!1near\-optimality\.The fraction of reward information captured atk=1k=1satisfies: I​\(ℛ1​\(𝐇\);Y\)I​\(𝐇;Y\)≥1−e−α⋅nh​NL1−e−α⋅nh​NL\.\\frac\{I\(\\mathcal\{R\}\_\{1\}\(\\mathbf\{H\}\);\\,Y\)\}\{I\(\\mathbf\{H\};\\,Y\)\}\\;\\geq\\;1\-\\frac\{e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\}\}\{1\-e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\}\}\.\(6\)
3. \(c\)Cost–information tradeoff\.The information gained per additional FLOP decays as: Δ​IkFreadout​\(k,L\)−Freadout​\(k−1,L\)=O​\(e−α⋅nh​NL⋅\(k−1\)d⋅L\)\.\\frac\{\\Delta I\_\{k\}\}\{F\_\{\\mathrm\{readout\}\}\(k,L\)\-F\_\{\\mathrm\{readout\}\}\(k\\\!\-\\\!1,L\)\}=O\\\!\\left\(\\frac\{e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\\cdot\(k\-1\)\}\}\{d\\cdot L\}\\right\)\.\(7\)

## 4Method

Motivated by the theoretical analysis in Section[3](https://arxiv.org/html/2607.09153#S3), we introduce KV\-PRM, a process reward model that instantiates the KV\-cache readout framework \(Definition[1](https://arxiv.org/html/2607.09153#Thmdefinition1)\) atk=1k=1, achieving the maximumΩ​\(L\)\\Omega\(L\)speedup by scoring trajectories through the KV cache already produced during generation\.

![Refer to caption](https://arxiv.org/html/2607.09153v1/x1.png)Figure 1:Text\-PRM re\-encodes the full trajectory for each scoring call atO​\(d​L2\)O\(dL^\{2\}\)FLOPs\. KV\-PRM reuses the KV cache from generation and scores with a single verify token atO​\(d​L\)O\(dL\)FLOPs, a reduction ofL×L\\times, which is three orders of magnitude for typical MAS trajectories\.### 4\.1KV\-PRM: Scoring via KV\-Cache Transfer

During generation, the LLM naturally produces a KV cache, encoding the full reasoning trajectory\. KV\-PRM reuses this pre\-existing cache to score trajectory quality, without re\-encoding\.

Architecture\.Letfθf\_\{\\theta\}denote the base language model with parametersθ\\theta, and letfθ\+Δ​θf\_\{\\theta\+\\Delta\\theta\}denote the same model with LoRA adapter parametersΔ​θ\\Delta\\theta\[[10](https://arxiv.org/html/2607.09153#bib.bib21)\]\. We designate a*verify token*vv\(the "?" token\), and two*judgment tokens*: "\+\+" \(positive\) and "−\-" \(negative\)\. We use text tokens to keep the output aligned with the base model’s pre\-trained vocabulary distribution\[[28](https://arxiv.org/html/2607.09153#bib.bib31)\]\.

Scoring procedure\.As shown in Figure[1](https://arxiv.org/html/2607.09153#S4.F1), after agentaja\_\{j\}completes generation, the accumulated KV cache\(𝐊,𝐕\)\(\\mathbf\{K\},\\mathbf\{V\}\)already exists as a byproduct of the generation process\. KV\-PRM scores the trajectory by processing a single verify token with the LoRA adapter activated:

𝐳=fθ\+Δ​θ​\(v∣𝐊,𝐕\),sj=P​\(\+\)=softmax​\(\[𝐳\[−\],𝐳\[\+\]\]\)1\\mathbf\{z\}=f\_\{\\theta\+\\Delta\\theta\}\(v\\mid\\mathbf\{K\},\\mathbf\{V\}\),\\qquad s\_\{j\}=P\(\+\)=\\text\{softmax\}\\big\(\[\\mathbf\{z\}\_\{\[\-\]\},\\;\\mathbf\{z\}\_\{\[\+\]\}\]\\big\)\_\{1\}\(8\)The LoRA adapter learns to interpret the base model’s KV cache, which encodes position information, attention patterns, and hidden state interactions of the full trajectory, and outputs a quality judgment through the verify token’s logits\. Crucially, the KV cache is produced by the base model with the adapter*disabled*during normal generation, so KV\-PRM introduces no overhead to the generation process itself\.

Cost analysis\.The verify\-token forward pass processes a single query token againstLLcached positions acrossNLN\_\{L\}layers:

Fscorekv=NL⋅\(cattn⋅d⋅L\+cffn⋅d2\)=O​\(d⋅L\)F\_\{\\text\{score\}\}^\{\\text\{kv\}\}=N\_\{L\}\\cdot\(c\_\{\\text\{attn\}\}\\cdot d\\cdot L\+c\_\{\\text\{ffn\}\}\\cdot d^\{2\}\)=O\(d\\cdot L\)\(9\)A text\-based PRM must re\-encode the full trajectory from scratch, costingFscoretext=O​\(d⋅L2\)F\_\{\\text\{score\}\}^\{\\text\{text\}\}=O\(d\\cdot L^\{2\}\)\(Equation[1](https://arxiv.org/html/2607.09153#S2.E1)\)\. The per\-call speedup is:

FscoretextFscorekv=O​\(d⋅L2\)O​\(d⋅L\)=O​\(L\)\\frac\{F\_\{\\text\{score\}\}^\{\\text\{text\}\}\}\{F\_\{\\text\{score\}\}^\{\\text\{kv\}\}\}=\\frac\{O\(d\\cdot L^\{2\}\)\}\{O\(d\\cdot L\)\}=O\(L\)\(10\)For typical MAS trajectories \(L≈5,000L\\approx 5\{,\}000tokens\), this yields∼5,000×\{\\sim\}5\{,\}000\\timesfewer FLOPs per scoring call\. The same factor\-of\-LLreduction extends to peak activation memory: Text\-PRM materializesO​\(L2\)O\(L^\{2\}\)attention maps per layer, while KV\-PRM requires onlyO​\(L\)O\(L\)\. In test\-time search, where the PRM is invokedO​\(W⋅D\)O\(W\\cdot D\)times, this per\-call reduction makes the total scoring cost negligible relative to generation, eliminating the computational bottleneck identified in Section[2](https://arxiv.org/html/2607.09153#S2)\. This instantiation corresponds tok=1k\\\!=\\\!1in the readout framework, achieving the maximum theoretical speedup\.

Inference\.KV\-PRM is a drop\-in replacement for text\-based PRMs in all search algorithms described in Section[2](https://arxiv.org/html/2607.09153#S2): beam search, MCTS, and weighted majority voting\. The scoring mechanism changes: KV\-PRM reads the pre\-existing KV cache instead of re\-encoding text, but the search algorithms themselves remain unchanged\.

### 4\.2Training

Training data generation\.We generate training labels using Monte Carlo Tree Search \(MCTS\)\[[5](https://arxiv.org/html/2607.09153#bib.bib15)\]over the multi\-agent reasoning process\. For each training question, we runRRMCTS rollouts withCCcandidates per agent step\. Terminal trajectories are evaluated by exact\-match against the ground truth, producing binary rewards\. Backpropagation through the search tree assigns Q\-valuesy∈\[−1,1\]y\\in\[\-1,1\]to each intermediate agent step\. Training data is generated separately for each \(topology, execution mode\) combination\. Each \(trajectory, agent step\) pair produces one training sample\(𝐱j,yj\)\(\\mathbf\{x\}\_\{j\},y\_\{j\}\)\.

KV\-PRM training\.For each sample\(𝐱,y\)\(\\mathbf\{x\},y\), the base model \(adapter disabled\) processes𝐱\\mathbf\{x\}to produce KV cache\(𝐊,𝐕\)\(\\mathbf\{K\},\\mathbf\{V\}\)without gradients\. The model \(adapter active\) processes the verify token conditioned on\(𝐊,𝐕\)\(\\mathbf\{K\},\\mathbf\{V\}\), producing scoreP​\(\+\)P\(\+\)via Equation[8](https://arxiv.org/html/2607.09153#S4.E8)\. The loss is:

ℒKV=MSE​\(P​\(\+\),y\+12\)\\mathcal\{L\}\_\{\\text\{KV\}\}=\\text\{MSE\}\\\!\\left\(P\(\+\),\\;\\frac\{y\+1\}\{2\}\\right\)\(11\)Only the LoRA parametersΔ​θ\\Delta\\thetaare updated\. Gradients flow only through the single verify\-token forward pass\. Algorithm[1](https://arxiv.org/html/2607.09153#alg1)summarizes the procedure\.

Algorithm 1KV\-PRM Training0:Training set

𝒟=\{\(𝐱i,yi\)\}\\mathcal\{D\}=\\\{\(\\mathbf\{x\}\_\{i\},y\_\{i\}\)\\\}with trajectories and MCTS Q\-values, base model

fθf\_\{\\theta\}, LoRA parameters

Δ​θ\\Delta\\theta, learning rate

η\\eta
1:foreach mini\-batch

ℬ⊂𝒟\\mathcal\{B\}\\subset\\mathcal\{D\}do

2:foreach

\(𝐱,y\)∈ℬ\(\\mathbf\{x\},y\)\\in\\mathcal\{B\}do

3:

\(𝐊,𝐕\)←fθ​\(𝐱\)\(\\mathbf\{K\},\\mathbf\{V\}\)\\leftarrow f\_\{\\theta\}\(\\mathbf\{x\}\)// Encode trajectory with adapter off \(no grad\)

4:

𝐳←fθ\+Δ​θ​\(v∣𝐊,𝐕\)\\mathbf\{z\}\\leftarrow f\_\{\\theta\+\\Delta\\theta\}\(v\\mid\\mathbf\{K\},\\mathbf\{V\}\)// Forward verify token with adapter on

5:

P​\(\+\)←softmax​\(\[𝐳\[−\],𝐳\[\+\]\]\)1P\(\+\)\\leftarrow\\text\{softmax\}\\big\(\[\\mathbf\{z\}\_\{\[\-\]\},\\;\\mathbf\{z\}\_\{\[\+\]\}\]\\big\)\_\{1\}
6:

ℒ←MSE​\(P​\(\+\),\(y\+1\)/2\)\\mathcal\{L\}\\leftarrow\\text\{MSE\}\\\!\\big\(P\(\+\),\\;\(y\+1\)/2\\big\)
7:endfor

8:

Δ​θ←Δ​θ−η​∇Δ​θ1\|ℬ\|​∑ℒ\\Delta\\theta\\leftarrow\\Delta\\theta\-\\eta\\,\\nabla\_\{\\Delta\\theta\}\\,\\frac\{1\}\{\|\\mathcal\{B\}\|\}\\sum\\mathcal\{L\}// Update only LoRA parameters

9:endfor

## 5Experiments

### 5\.1Setup

Datasets and metrics\.We evaluate on four mathematical reasoning benchmarks: MATH\[[8](https://arxiv.org/html/2607.09153#bib.bib22)\], GSM8K\[[1](https://arxiv.org/html/2607.09153#bib.bib23)\], AIME 2024, and AIME 2025, reporting exact\-match accuracy with numeric equivalence checking\. We evaluate three search algorithms: step\-level beam search \(SBS\), MCTS, and weighted majority voting, with per\-dataset configurations reported in Tables[1](https://arxiv.org/html/2607.09153#S5.T1)–[2](https://arxiv.org/html/2607.09153#S5.T2)\. Scoring cost is normalized to Text\-PRM=1×=1\\times\. Both KV\-PRM and Text\-PRM are trained with LoRA \(r=256r\\\!=\\\!256,α=32\\alpha\\\!=\\\!32\) on MCTS\-generated labels \(R=64R\\\!=\\\!64rollouts,C=4C\\\!=\\\!4candidates per step\) from the MATH training split\. Full hyperparameters and implementation details are in Appendix[B](https://arxiv.org/html/2607.09153#A2)\.

Models and MAS configurations\.We use three scales of the Qwen3 family\[[29](https://arxiv.org/html/2607.09153#bib.bib20)\]: 0\.6B, 4B, and 8B, with thinking mode disabled, under two representative multi\-agent topologies:Sequential\(Reader→\\toPlanner→\\toSolver→\\toVerifier\) andHierarchical\(Math/Science/Code Agents→\\toTask Summarizer\)\.

Baselines\.We compare KV\-PRM against three scoring strategies: \(1\)No PRM\(Random Sampling / Majority Voting\), where the MAS generates without process reward guidance; \(2\)Policy Log\-prob, which ranks trajectories by the generator’s own token\-level log\-probabilities at negligible additional cost; and \(3\)Text\-PRM, a LoRA adapter trained with the same architecture, data, and objective as KV\-PRM \(Equation[11](https://arxiv.org/html/2607.09153#S4.E11)\), but scoring by re\-encoding the full trajectory\[𝐱;v\]\[\\mathbf\{x\};v\]in a single forward pass with the adapter active on all tokens atO​\(d​L2\)O\(dL^\{2\}\)cost per call\. Text\-PRM training requires gradient checkpointing due to full\-sequence gradient computation\.

### 5\.2Main Results

KV\-PRM outperforms Text\-PRM across a broad range of settings while reducing scoring cost by 3 orders of magnitude\.Table[1](https://arxiv.org/html/2607.09153#S5.T1)and Table[2](https://arxiv.org/html/2607.09153#S5.T2)show that KV\-PRM achieves this accuracy\-efficiency trade\-off with dramatically lower scoring FLOPs per call, with relative cost reductions ranging from roughly9×1029\\times 10^\{2\}to4\.9×103×4\.9\\times 10^\{3\}\\timescompared to Text\-PRM, depending on the benchmark, model scale, and search configuration\. The accuracy results show that replacing full\-text re\-encoding with single\-token KV\-cache readout improves, or at least preserves, scorer quality while making verification far cheaper\.

KV\-PRM significantly reduces the inference latency by 1\-2 orders of magnitude\.We validate the theoretical complexity reduction with wall\-clock latency measurements on an NVIDIA GH200 GPU\. As shown in the left panel of Figure[3](https://arxiv.org/html/2607.09153#S5.F3)\. KV\-PRM consistently delivers substantial per\-sequence speedups over Text\-PRM, ranging from 15×\\timesto 37×\\timesdepending on model scale and sequence length\. The gap widens as sequences become longer and models become larger, in line with the expectedO​\(L\)O\(L\)vs\.O​\(L2\)O\(L^\{2\}\)scoring\-cost difference\. At 8B andL=4096L=4096, a single KV\-PRM scoring call takes only 4\.6 ms, compared with 172\.0 ms for Text\-PRM\. This confirms that KV\-cache transfer translates the theoretical FLOP savings into large practical reductions in end\-to\-end scorer latency\.

KV\-PRM significantly reduces the memory consumption during inference by 1\-2 orders of magnitude\.As shown in the right panel of Figure[3](https://arxiv.org/html/2607.09153#S5.F3), KV\-PRM reduces the average memory increase per sequence by up to34\.2×34\.2\\times, relative to Text\-PRM under matched batch sizes\. This gain comes from avoiding full\-sequence re\-encoding and the associated attention\-map materialization, allowing KV\-PRM to score trajectories with substantially lower activation overhead\. In practice, the smaller memory footprint enables larger verification batches on the same hardware, providing a complementary systems advantage beyond latency alone\.

Table 1:Test\-time search accuracy \(%\) and relative scoring cost onMATHandAIME 2024under sequential multi\-agent execution\. Scoring cost normalized to Text\-PRM=1×=1\\times; Policy Log\-prob cost is negligible \(—\); KV\-PRM cost=1/L=1/LwhereLLis the maximum mean per\-step trajectory length\.MATHAIME 2024Qwen3\-0\.6BQwen3\-4BQwen3\-8BQwen3\-4BQwen3\-8BSearchConfigScorerAccCostAccCostAccCostAccCostAccCostAvg\.Random Samplingn=1n\\\!=\\\!1—28\.0054\.1556\.7020\.0020\.0035\.77Majority Votingn=40n\\\!=\\\!40—34\.6559\.2060\.5536\.6733\.3344\.88Beam SearchW=1W\\\!=\\\!1Policy Log\-prob32\.50—58\.70—63\.70—23\.33—23\.33—40\.31Text\-PRM32\.751×1\\times62\.001×1\\times66\.621×1\\times26\.671×1\\times20\.001×1\\times41\.61KV\-PRM31\.9012409\\frac\{1\}\{2409\}65\.5512543\\frac\{1\}\{2543\}67\.0212782\\frac\{1\}\{2782\}20\.0012007\\frac\{1\}\{2007\}30\.0012367\\frac\{1\}\{2367\}42\.89MCTSn=50n\\\!=\\\!50Policy Log\-prob29\.40—63\.90—61\.90—20\.00—23\.33—39\.71Text\-PRM35\.551×1\\times65\.951×1\\times68\.921×1\\times30\.001×1\\times36\.671×1\\times47\.42KV\-PRM35\.1013598\\frac\{1\}\{3598\}67\.1513771\\frac\{1\}\{3771\}69\.5714367\\frac\{1\}\{4367\}33\.3313652\\frac\{1\}\{3652\}40\.0013944\\frac\{1\}\{3944\}49\.03Weighted Votingn=10n\\\!=\\\!10Policy Log\-prob33\.75—58\.45—59\.95—30\.00—33\.33—43\.10Text\-PRM34\.201×1\\times58\.751×1\\times60\.511×1\\times36\.671×1\\times23\.331×1\\times42\.69KV\-PRM34\.4012444\\frac\{1\}\{2444\}59\.3512749\\frac\{1\}\{2749\}60\.4613137\\frac\{1\}\{3137\}33\.3312360\\frac\{1\}\{2360\}33\.3312575\\frac\{1\}\{2575\}44\.17n=200n\\\!=\\\!200Policy Log\-prob35\.35—59\.40—60\.85—36\.67—33\.33—45\.12Text\-PRM36\.051×1\\times59\.951×1\\times61\.261×1\\times36\.671×1\\times33\.331×1\\times45\.45KV\-PRM36\.1512441\\frac\{1\}\{2441\}60\.1512721\\frac\{1\}\{2721\}61\.0613086\\frac\{1\}\{3086\}36\.6712298\\frac\{1\}\{2298\}36\.6712539\\frac\{1\}\{2539\}46\.14

Table 2:Test\-time search accuracy \(%\) and relative scoring cost onGSM8KandAIME 2025under sequential multi\-agent execution\. Same format as Table[1](https://arxiv.org/html/2607.09153#S5.T1)\.GSM8KAIME 2025Qwen3\-0\.6BQwen3\-4BQwen3\-8BQwen3\-4BQwen3\-8BSearchConfigScorerAccCostAccCostAccCostAccCostAccCostAvg\.Random Samplingn=1n\\\!=\\\!1—50\.8091\.2891\.4316\.6720\.0054\.04Majority Votingn=40n\\\!=\\\!40—62\.7793\.1093\.0320\.0016\.6757\.11Beam SearchW=10W\\\!=\\\!10Policy Log\-prob52\.99—90\.98—93\.40—13\.33—16\.67—53\.47Text\-PRM58\.681×1\\times93\.781×1\\times95\.001×1\\times26\.671×1\\times26\.671×1\\times60\.16KV\-PRM58\.0011341\\frac\{1\}\{1341\}92\.8011311\\frac\{1\}\{1311\}95\.0711295\\frac\{1\}\{1295\}26\.6714266\\frac\{1\}\{4266\}30\.0014392\\frac\{1\}\{4392\}60\.51MCTSn=30n\\\!=\\\!30Policy Log\-prob52\.92—89\.99—91\.13—13\.33—16\.67—52\.81Text\-PRM57\.771×1\\times93\.331×1\\times95\.151×1\\times23\.331×1\\times23\.331×1\\times58\.58KV\-PRM54\.7411357\\frac\{1\}\{1357\}92\.8711386\\frac\{1\}\{1386\}94\.6211402\\frac\{1\}\{1402\}23\.3314877\\frac\{1\}\{4877\}26\.6714840\\frac\{1\}\{4840\}58\.45Weighted Votingn=20n\\\!=\\\!20Policy Log\-prob61\.64—92\.72—92\.87—16\.67—16\.67—56\.11Text\-PRM62\.401×1\\times93\.181×1\\times93\.181×1\\times13\.331×1\\times16\.671×1\\times55\.75KV\-PRM63\.611908\\frac\{1\}\{908\}92\.8711006\\frac\{1\}\{1006\}93\.0311036\\frac\{1\}\{1036\}20\.0013402\\frac\{1\}\{3402\}20\.0013415\\frac\{1\}\{3415\}57\.90n=100n\\\!=\\\!100Policy Log\-prob63\.00—92\.80—92\.95—16\.67—16\.67—56\.42Text\-PRM63\.381×1\\times93\.181×1\\times93\.251×1\\times20\.001×1\\times16\.671×1\\times57\.30KV\-PRM63\.761934\\frac\{1\}\{934\}93\.181997\\frac\{1\}\{997\}93\.2511025\\frac\{1\}\{1025\}23\.3313428\\frac\{1\}\{3428\}16\.6713380\\frac\{1\}\{3380\}58\.04

### 5\.3Analysis

![Refer to caption](https://arxiv.org/html/2607.09153v1/x2.png)
![Refer to caption](https://arxiv.org/html/2607.09153v1/x3.png)

Figure 2:Analysis of KV\-PRM efficiency and readout depth\.*Left*: Scaling up the Text\-PRM from 0\.6B to 8B yields diminishing returns; KV\-PRM surpasses all sizes at∼103×\{\\sim\}10^\{3\}\\timeslower cost\.*Right*: Accuracy gain overk=1k\\\!=\\\!1as a function of the number of verify tokenskkaveraged across search hyperparameters\. Fitted curvesa​\(1−e−c​\(k−1\)\)a\(1\-e^\{\-c\(k\-1\)\}\)confirm exponentially diminishing returns\.Can smaller Text\-PRMs close the efficiency gap?A natural strategy for reducing Text\-PRM scoring cost is to use a smaller verifier model, which is a common practice in the literature\[[23](https://arxiv.org/html/2607.09153#bib.bib12)\]\. We evaluate this on the MATH dataset, using Qwen3\-8B as the base generator and scoring the same search trees with Text\-PRMs of sizes 0\.6B, 4B, and 8B alongside KV\-PRM of size 8B\. Figure[2](https://arxiv.org/html/2607.09153#S5.F2)\(Left\) plots accuracy against relative scoring cost per call for SBSW=2W\\\!=\\\!2\. Scaling up the Text\-PRM from 0\.6B to 4B to 8B improves accuracy, but with diminishing returns and increasing cost\. KV\-PRM achieves the*highest*accuracy \(68\.0%\), surpassing all Text\-PRM sizes, at∼3,000×\{\\sim\}3\{,\}000\\timeslower cost than Text\-PRM\-8B, demonstrating that KV\-cache transfer is a fundamentally more efficient strategy than reducing the size of the PRM\.

What happens with more verify tokens \(k\>1k\>1\)?Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)makes two concrete predictions: \(a\) the marginal information gainΔ​Ik\\Delta I\_\{k\}from each additional verify token decays exponentially inkkat rateα⋅nh​NL\\alpha\\cdot n\_\{h\}N\_\{L\}\(Equation[5](https://arxiv.org/html/2607.09153#S3.E5)\), and \(b\) a single token \(k=1k\\\!=\\\!1\) already captures most extractable reward information \(Equation[6](https://arxiv.org/html/2607.09153#S3.E6)\)\. Figure[2](https://arxiv.org/html/2607.09153#S5.F2)\(Right\) tests these predictions on Qwen3\-0\.6B \(nh⋅NL=448n\_\{h\}\\\!\\cdot\\\!N\_\{L\}\\\!=\\\!448\), plotting accuracy gain overk=1k\\\!=\\\!1averaged across search hyperparameters\. Both curves are well fit bya​\(1−e−c​\(k−1\)\)a\(1\\\!\-\\\!e^\{\-c\(k\-1\)\}\), consistent with the exponential decay forme−α⋅nh​NL⋅\(k−1\)e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\\cdot\(k\-1\)\}predicted by Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)\(a\)\. The rapid saturation validates thek=1k\\\!=\\\!1near\-optimality bound \(Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)b\): even at the smallest model scale, a single verify token captures most of the available reward signal\. By the cost, information tradeoff \(Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)c\), the information gained per additional FLOP decays asO​\(e−α⋅nh​NL⋅\(k−1\)/\(d​L\)\)O\(e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\\cdot\(k\-1\)\}/\(dL\)\), so even atk=4k\\\!=\\\!4the total scoring cost of4⋅O​\(d​L\)4\\cdot O\(dL\)remains∼103×\{\\sim\}10^\{3\}\\timescheaper than text re\-encoding atO​\(d​L2\)O\(dL^\{2\}\)\.

![Refer to caption](https://arxiv.org/html/2607.09153v1/x4.png)Figure 3:Wall\-clock efficiency on NVIDIA GH200 120GB GPU\.*Left*: Per\-sequence scoring latency \(log scale\)\. KV\-PRM \(solid\) is11–22orders of magnitude faster than Text\-PRM \(dashed\), reaching37×37\\timesat 8B \(L=4,096L\\\!=\\\!4\{,\}096\)\.*Right*: Per\-sequence average memory increase during PRM inference\. KV\-PRM reduces memory by up to34\.2×34\.2\\times, enabling substantially larger batches on the same hardware\.Does KV\-PRM work on more MAS topologies?We additionally evaluate KV\-PRM under a hierarchical multi\-agent topology\. Table[4](https://arxiv.org/html/2607.09153#S5.T4)reports results across all four datasets\. KV\-PRM matches or outperforms Text\-PRM in the majority of search configurations\. On MATH, the advantage is consistent but modest \(\+0\.5\+0\.5–0\.80\.8pp\), while on AIME benchmarks the gaps are more pronounced \(up to\+6\.7\+6\.7pp\), reflecting greater headroom for scorer quality on harder problems\. These results confirm that the KV\-cache scoring mechanism works effectively for non\-sequential MAS topologies\.

Table 3:Hierarchical MAS accuracy \(%\) across datasets and model scales\. KV\-PRM matches or exceeds Text\-PRM on all dataset–model combinations\.Δ\\Delta= KV\-PRM−\-Text\-PRM\.
Table 4:KV Steering proof\-of\-concept: accuracy with gradient\-based KV\-cache optimization at agent handoffs, without any search\. This is structurally impossible with text\-based PRMs\.

## 6Discussion on Broader Applications of KV\-PRM

As KV\-PRM scores are differentiable with respect to the KV cache, we explore whether gradient\-based optimization of latent inter\-agent messages\[[11](https://arxiv.org/html/2607.09153#bib.bib36)\]can improve generation quality, a capability structurally impossible with text\-based PRMs\. KV Steering performs gradient ascent on the accumulated KV cache at each agent handoff to maximize the KV\-PRM score before the next agent begins generation \(algorithm details in Appendix[C](https://arxiv.org/html/2607.09153#A3)\)\. Table[4](https://arxiv.org/html/2607.09153#S5.T4)presents proof\-of\-concept results\. KV Steering improves accuracy on MATH by\+1\.6\+1\.6pp \(Qwen3\-4B\) and\+1\.2\+1\.2pp \(Qwen3\-8B\) without any search\. On AIME benchmarks with Qwen3\-8B, steering yields\+3\.3\+3\.3pp gains on both AIME 2024 and AIME 2025\. While preliminary, these results demonstrate that the KV\-native architecture opens a qualitatively new direction for inference\-time optimization beyond search\.

## 7Related Work

Process reward models\.PRMs\[[12](https://arxiv.org/html/2607.09153#bib.bib27),[33](https://arxiv.org/html/2607.09153#bib.bib28),[34](https://arxiv.org/html/2607.09153#bib.bib29),[14](https://arxiv.org/html/2607.09153#bib.bib37)\]provide step\-level supervision for mathematical reasoning, pioneered byLightmanet al\.\[[19](https://arxiv.org/html/2607.09153#bib.bib1)\]who demonstrated that step\-level human feedback improves over outcome\-based reward models\. Math\-Shepherd\[[24](https://arxiv.org/html/2607.09153#bib.bib2)\]automates PRM label generation using MCTS\. OmegaPRM\[[21](https://arxiv.org/html/2607.09153#bib.bib3)\]further scales automated PRM training\. GenRM\[[32](https://arxiv.org/html/2607.09153#bib.bib4)\]frames verification as generation\. All existing PRMs operate on text representations; KV\-PRM is the first to score via KV\-cache transfer, reducing the computational cost by orders of magnitude\.

LLM\-based multi\-agent systems\.Multi\-agent systems have been explored in debate frameworks\[[4](https://arxiv.org/html/2607.09153#bib.bib8)\], collaborative problem\-solving\[[17](https://arxiv.org/html/2607.09153#bib.bib9)\], and structured pipelines\[[9](https://arxiv.org/html/2607.09153#bib.bib7)\]\. MASPRM\[[31](https://arxiv.org/html/2607.09153#bib.bib5)\]introduced PRMs for multi\-agent systems using text\-based scoring\. We address the computational bottleneck that MASPRM’s text\-based approach creates in test\-time search\.

KV\-cache methods\.Prior work has explored KV\-cache compression\[[18](https://arxiv.org/html/2607.09153#bib.bib10)\], eviction\[[27](https://arxiv.org/html/2607.09153#bib.bib11)\], and memory management via paging\[[16](https://arxiv.org/html/2607.09153#bib.bib30)\]for efficient inference\. These methods aim to reduce the*size*or improve the*allocation*of the KV cache; our work instead uses the KV cache as an*input representation*for reward modeling, a fundamentally different application\.

Test\-time compute scaling\.Scaling test\-time compute through search\[[23](https://arxiv.org/html/2607.09153#bib.bib12)\], self\-consistency\[[25](https://arxiv.org/html/2607.09153#bib.bib13)\], tree\-of\-thoughts\[[30](https://arxiv.org/html/2607.09153#bib.bib14)\], MCTS\[[5](https://arxiv.org/html/2607.09153#bib.bib15)\], and reinforcement learning for reasoning\[[3](https://arxiv.org/html/2607.09153#bib.bib25)\]has shown consistent gains\. However, the cost of scoring candidates, particularly ones with long trajectories, has received limited attention\. Our work directly addresses this bottleneck\.

## 8Limitations and Conclusion

We presented KV\-PRM, a process reward model that scores trajectories through a single verify\-token forward pass over the generator’s pre\-existing KV cache, reducing per\-call verifier cost fromO​\(d​L2\)O\(dL^\{2\}\)toO​\(d​L\)O\(dL\)\. Our theoretical framework proves that the KV cache is a strictly richer representation than decoded text for verification \(Theorem[1](https://arxiv.org/html/2607.09153#Thmproposition1)\), and that the marginal information gain from additional readout tokens decays exponentially, making the efficientk=1k\\\!=\\\!1design near\-optimal \(Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)\)\. Empirically, KV\-PRM matches or outperforms text\-based PRMs across beam search, MCTS, and weighted voting on four benchmarks and two MAS topologies, while delivering up to∼5,000×\{\\sim\}5\{,\}000\\timesfewer scoring FLOPs,37×37\\timeslower latency, and34×34\\timesless memory per sequence\. The differentiability of KV\-cache scoring further enables KV Steering, a proof\-of\-concept for gradient\-based inter\-agent message optimization that is structurally impossible with text\-based PRMs\. Limitations include the requirement that generator and verifier share the same architecture, reliance on the Linear Representation Hypothesis in the theoretical bounds, and the preliminary nature of KV Steering\. We hope this work encourages a shift from redundant text re\-encoding toward representation reuse in verifier design\.

## References

- \[1\]K\. Cobbe, V\. Kosaraju, M\. Bavarian, M\. Chen, H\. Jun, L\. Kaiser, M\. Plappert, J\. Tworek, J\. Hilton, R\. Nakano, C\. Hesse, and J\. Schulman\(2021\)Training verifiers to solve math word problems\.arXiv preprint arXiv:2110\.14168\.Cited by:[§5\.1](https://arxiv.org/html/2607.09153#S5.SS1.p1.5)\.
- \[2\]\(2006\)Elements of information theory\.2nd edition,Wiley\-Interscience\.Cited by:[§A\.1](https://arxiv.org/html/2607.09153#A1.SS1.p1.10)\.
- \[3\]DeepSeek\-AI, D\. Guo, D\. Yang, H\. Zhang, J\. Song, P\. Wang,et al\.\(2025\)DeepSeek\-R1: incentivizing reasoning capability in LLMs via reinforcement learning\.arXiv preprint arXiv:2501\.12948\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1),[§7](https://arxiv.org/html/2607.09153#S7.p4.1)\.
- \[4\]Y\. Du, S\. Li, A\. Torralba, J\. B\. Tenenbaum, and I\. Mordatch\(2023\)Improving factuality and reasoning in language models through multiagent debate\.arXiv preprint arXiv:2305\.14325\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p2.1)\.
- \[5\]X\. Feng, Z\. Wan, M\. Wen, S\. M\. McAleer, Y\. Wen, W\. Zhang, and J\. Wang\(2024\)Alphazero\-like tree\-search can guide large language model decoding and training\.Proceedings of the International Conference on Machine Learning \(ICML\)\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1),[§4\.2](https://arxiv.org/html/2607.09153#S4.SS2.p1.4),[§7](https://arxiv.org/html/2607.09153#S7.p4.1)\.
- \[6\]S\. Fu, Y\. Wang, S\. Hu, P\. Wang, and D\. Tao\(2026\)Theoretical guarantees for iterative alignment of self\-rewarding language models\.External Links:[Link](https://openreview.net/forum?id=Dhd7nHdAaf)Cited by:[Assumption 2](https://arxiv.org/html/2607.09153#Thmassumption2.p1.9.9)\.
- \[7\]T\. Guo, X\. Chen, Y\. Wang, R\. Chang, S\. Pei, N\. V\. Chawla, O\. Wiest, and X\. Zhang\(2024\)Large language model based multi\-agents: a survey of progress and challenges\.Proceedings of the International Joint Conference on Artificial Intelligence \(IJCAI\)\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1),[§2](https://arxiv.org/html/2607.09153#S2.p1.11)\.
- \[8\]D\. Hendrycks, C\. Burns, S\. Kadavath, A\. Arora, S\. Basart, E\. Tang, D\. Song, and J\. Steinhardt\(2021\)Measuring mathematical problem solving with the MATH dataset\.InAdvances in Neural Information Processing Systems \(NeurIPS\),Cited by:[§5\.1](https://arxiv.org/html/2607.09153#S5.SS1.p1.5)\.
- \[9\]S\. Hong, M\. Zhuge, J\. Chen, X\. Zheng, Y\. Cheng, C\. Zhang, J\. Wang, Z\. Wang, S\. K\. S\. Yau, Z\. Lin, L\. Zhou, C\. Ran, L\. Xiao, C\. Wu, and J\. Schmidhuber\(2024\)MetaGPT: meta programming for a multi\-agent collaborative framework\.arXiv preprint arXiv:2308\.00352\.Cited by:[§2](https://arxiv.org/html/2607.09153#S2.p1.11),[§7](https://arxiv.org/html/2607.09153#S7.p2.1)\.
- \[10\]E\. J\. Hu, Y\. Shen, P\. Wallis, Z\. Allen\-Zhu, Y\. Li, S\. Wang, L\. Wang, and W\. Chen\(2022\)LoRA: low\-rank adaptation of large language models\.arXiv preprint arXiv:2106\.09685\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p4.3),[§4\.1](https://arxiv.org/html/2607.09153#S4.SS1.p2.7)\.
- \[11\]H\. Jin, K\. Peng, Y\. Yu, X\. Yuan, and H\. Wang\(2026\)Agent primitives: reusable latent building blocks for multi\-agent systems\.External Links:2602\.03695,[Link](https://arxiv.org/abs/2602.03695)Cited by:[§6](https://arxiv.org/html/2607.09153#S6.p1.3)\.
- \[12\]M\. Khalifa, R\. Agarwal, L\. Logeswaran, J\. Kim, H\. Peng, M\. Lee, H\. Lee, and L\. Wang\(2025\)Process reward models that think\.arXiv preprint arXiv:2504\.16828\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[13\]L\. Kocsis and C\. Szepesvári\(2006\)Bandit based Monte\-Carlo planning\.InEuropean Conference on Machine Learning \(ECML\),pp\. 282–293\.Cited by:[§2](https://arxiv.org/html/2607.09153#S2.p2.8)\.
- \[14\]P\. Kuang, X\. Wang, W\. Liu, J\. Dong, and K\. Xu\(2025\)TIM\-prm: verifying multimodal reasoning with tool\-integrated prm\.External Links:2511\.22998,[Link](https://arxiv.org/abs/2511.22998)Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[15\]P\. Kuang, Y\. Wang, X\. Han, Y\. Liu, K\. Xu, and H\. Wang\(2026\)Optimal aggregation of LLM and PRM signals for efficient test\-time scaling\.InThe Fourteenth International Conference on Learning Representations,External Links:[Link](https://openreview.net/forum?id=x85kiYqL4y)Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1)\.
- \[16\]W\. Kwon, Z\. Li, S\. Zhuang, Y\. Sheng, L\. Zheng, C\. H\. Yu, J\. E\. Gonzalez, H\. Zhang, and I\. Stoica\(2023\)Efficient memory management for large language model serving with PagedAttention\.InProceedings of the ACM SIGOPS 29th Symposium on Operating Systems Principles \(SOSP\),Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p3.1)\.
- \[17\]G\. Li, H\. A\. A\. K\. Hammoud, H\. Itani, D\. Khizbullin, and B\. Ghanem\(2023\)CAMEL: communicative agents for “mind” exploration of large language model society\.Advances in Neural Information Processing Systems\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p2.1)\.
- \[18\]H\. Liet al\.\(2024\)A survey on large language model acceleration based on KV cache management\.arXiv preprint arXiv:2412\.19442\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p3.1)\.
- \[19\]H\. Lightman, V\. Kosaraju, Y\. Burda, H\. Edwards, B\. Baker, T\. Lee, J\. Leike, J\. Schulman, I\. Sutskever, and K\. Cobbe\(2023\)Let’s verify step by step\.arXiv preprint arXiv:2305\.20050\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1),[§2](https://arxiv.org/html/2607.09153#S2.p2.8),[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[20\]Y\. Liu\(2026\)The spectral geometry of thought: phase transitions, instruction reversal, token\-level dynamics, and perfect correctness prediction in how transformers reason\.External Links:2604\.15350,[Link](https://arxiv.org/abs/2604.15350)Cited by:[Assumption 2](https://arxiv.org/html/2607.09153#Thmassumption2.p1.9.9)\.
- \[21\]L\. Luo, Y\. Liu, R\. Liu, S\. Phatale, H\. Lara, Y\. Li, L\. Shu, Y\. Zhu, L\. Meng, J\. Sun, and A\. Rastogi\(2024\)Improve mathematical reasoning in language models by automated process supervision\.arXiv preprint arXiv:2406\.06592\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[22\]K\. Park, Y\. J\. Choe, and V\. Veitch\(2024\)The linear representation hypothesis and the geometry of large language models\.Proceedings of the International Conference on Machine Learning \(ICML\)\.Cited by:[Assumption 1](https://arxiv.org/html/2607.09153#Thmassumption1)\.
- \[23\]C\. Snell, J\. Lee, K\. Xu, and A\. Kumar\(2024\)Scaling LLM test\-time compute optimally can be more effective than scaling model parameters\.arXiv preprint arXiv:2408\.03314\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1),[§2](https://arxiv.org/html/2607.09153#S2.p2.8),[§5\.3](https://arxiv.org/html/2607.09153#S5.SS3.p1.2),[§7](https://arxiv.org/html/2607.09153#S7.p4.1)\.
- \[24\]P\. Wang, L\. Li, Z\. Shao, R\. Xu, D\. Dai, Y\. Li, D\. Chen, Y\. Wu, and Z\. Sui\(2024\)Math\-shepherd: verify and reinforce LLMs step\-by\-step without human annotations\.arXiv preprint arXiv:2312\.08935\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1),[§2](https://arxiv.org/html/2607.09153#S2.p2.8),[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[25\]X\. Wang, J\. Wei, D\. Schuurmans, Q\. Le, E\. Chi, S\. Narang, A\. Chowdhery, and D\. Zhou\(2023\)Self\-consistency improves chain of thought reasoning in language models\.arXiv preprint arXiv:2203\.11171\.Cited by:[§2](https://arxiv.org/html/2607.09153#S2.p2.8),[§7](https://arxiv.org/html/2607.09153#S7.p4.1)\.
- \[26\]J\. Wei, X\. Wang, D\. Schuurmans, M\. Bosma, B\. Ichter, F\. Xia, E\. Chi, Q\. Le, and D\. Zhou\(2022\)Chain\-of\-thought prompting elicits reasoning in large language models\.Advances in Neural Information Processing Systems \(NeurIPS\)\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1)\.
- \[27\]G\. Xiao, Y\. Tian, B\. Chen, S\. Han, and M\. Lewis\(2023\)Efficient streaming language models with attention sinks\.arXiv preprint arXiv:2309\.17453\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p3.1)\.
- \[28\]W\. Xiong, H\. Zhang, N\. Jiang, and T\. Zhang\(2024\)An Implementation of Generative PRM\.GitHub\.Note:Publication Title: GitHub repositoryExternal Links:[Link](https://github.com/RLHFlow/RLHF-Reward-Modeling)Cited by:[§4\.1](https://arxiv.org/html/2607.09153#S4.SS1.p2.7)\.
- \[29\]A\. Yanget al\.\(2025\)Qwen3 technical report\.arXiv preprint arXiv:2505\.09388\.Cited by:[§5\.1](https://arxiv.org/html/2607.09153#S5.SS1.p2.4)\.
- \[30\]S\. Yao, D\. Yu, J\. Zhao, I\. Shafran, T\. L\. Griffiths, Y\. Cao, and K\. Narasimhan\(2023\)Tree of thoughts: deliberate problem solving with large language models\.arXiv preprint arXiv:2305\.10601\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p4.1)\.
- \[31\]M\. Yazdani, M\. Mostajabdaveh, Z\. Zhou, and Y\. Xiong\(2025\)MASPRM: multi\-agent system process reward model\.arXiv preprint arXiv:2510\.24803\.Cited by:[§2](https://arxiv.org/html/2607.09153#S2.p2.8),[§7](https://arxiv.org/html/2607.09153#S7.p2.1)\.
- \[32\]L\. Zhang, A\. Hosseini, H\. Bansal, M\. Kazemi, A\. Kumar, and R\. Agarwal\(2024\)Generative verifiers: reward modeling as next\-token prediction\.arXiv preprint arXiv:2408\.15240\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[33\]J\. Zhao, R\. Liu, K\. Zhang, Z\. Zhou, J\. Gao, D\. Li, J\. Lyu, Z\. Qian, B\. Qi, X\. Li, and B\. Zhou\(2025\)GenPRM: scaling test\-time compute of process reward models via generative reasoning\.arXiv preprint arXiv:2504\.00891\.Cited by:[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[34\]C\. Zheng, J\. Zhu, Z\. Ou, Y\. Chen, K\. Zhang, R\. Shan, Z\. Zheng, M\. Yang, J\. Lin, Y\. Yu, and W\. Zhang\(2025\)A survey of process reward models: from outcome signals to process supervisions for large language models\.arXiv preprint arXiv:2510\.08049\.Cited by:[§1](https://arxiv.org/html/2607.09153#S1.p1.1),[§7](https://arxiv.org/html/2607.09153#S7.p1.1)\.
- \[35\]J\. Zou, X\. Yang, R\. Qiu, G\. Li, K\. Tieu, P\. Lu, K\. Shen, H\. Tong, Y\. Choi, J\. He, J\. Zou, M\. Wang, and L\. Yang\(2025\)Latent collaboration in multi\-agent systems\.External Links:2511\.20639,[Link](https://arxiv.org/abs/2511.20639)Cited by:[Assumption 1](https://arxiv.org/html/2607.09153#Thmassumption1)\.

## Appendix ATheoretical Proofs

### A\.1Proof of Theorem[1](https://arxiv.org/html/2607.09153#Thmproposition1)

Part \(a\): Information sufficiency\.The generation process defines the Markov chainY→𝐇→𝐱Y\\to\\mathbf\{H\}\\to\\mathbf\{x\}: trajectory qualityYYinfluences the generation dynamics that produce KV cache𝐇=\(𝐊,𝐕\)\\mathbf\{H\}=\(\\mathbf\{K\},\\mathbf\{V\}\), and text𝐱=decode​\(𝐇\)\\mathbf\{x\}=\\mathrm\{decode\}\(\\mathbf\{H\}\)is a deterministic function of𝐇\\mathbf\{H\}\(each tokenxt=argmaxv\[Woutht\]vx\_\{t\}=\\arg\\max\_\{v\}\[W\_\{\\mathrm\{out\}\}h\_\{t\}\]\_\{v\}or is sampled fromsoftmax​\(Wout​ht\)\\mathrm\{softmax\}\(W\_\{\\mathrm\{out\}\}h\_\{t\}\), wherehth\_\{t\}is determined by𝐇\\mathbf\{H\}\)\. By the Data Processing Inequality\[[2](https://arxiv.org/html/2607.09153#bib.bib17), Theorem 2\.8\.1\], for any Markov chainY→𝐇→𝐱Y\\to\\mathbf\{H\}\\to\\mathbf\{x\}:

I​\(𝐇;Y\)≥I​\(𝐱;Y\)\.I\(\\mathbf\{H\};\\,Y\)\\;\\geq\\;I\(\\mathbf\{x\};\\,Y\)\.\(12\)Equality holds if and only if𝐱\\mathbf\{x\}is a sufficient statistic of𝐇\\mathbf\{H\}forYY, i\.e\.,Y⟂𝐇∣𝐱Y\\perp\\mathbf\{H\}\\mid\\mathbf\{x\}\. Sincedecode​\(⋅\)\\mathrm\{decode\}\(\\cdot\)maps fromℝNL×L×d\\mathbb\{R\}^\{N\_\{L\}\\times L\\times d\}to\{1,…,\|𝒱\|\}L\\\{1,\\ldots,\|\\mathcal\{V\}\|\\\}^\{L\}—a many\-to\-one projection at each position—distinct hidden states that carry different information aboutYYare generically collapsed to the same token, so equality does not hold in general\.□\\square

Part \(b\): Capacity bound\.Under Assumption[1](https://arxiv.org/html/2607.09153#Thmassumption1), the set of hidden embeddings is

ℋ=\{∑i=1dci​si:c1,…,cd∈\{0,±1\}\},\\mathcal\{H\}=\\left\\\{\\sum\_\{i=1\}^\{d\}c\_\{i\}s\_\{i\}:c\_\{1\},\\ldots,c\_\{d\}\\in\\\{0,\\pm 1\\\}\\right\\\},\(13\)with\|ℋ\|=3d\|\\mathcal\{H\}\|=3^\{d\}\(since the semantic basis is linearly independent\)\. For a trajectory of lengthLL, the KV cache encodesLLsuch hidden states \(considering the final\-layer representation per position\), giving\|ℋL\|=3d​L\|\\mathcal\{H\}^\{L\}\|=3^\{dL\}distinct configurations\.

To represent these losslessly via text tokens from vocabulary𝒱\\mathcal\{V\}, a text sequence of lengthm′m^\{\\prime\}can encode at most\|𝒱\|m′\|\\mathcal\{V\}\|^\{m^\{\\prime\}\}distinct configurations\. Lossless representation requires:

\|𝒱\|m′≥3d​L\.\|\\mathcal\{V\}\|^\{m^\{\\prime\}\}\\;\\geq\\;3^\{dL\}\.\(14\)Taking logarithms:

m′≥d​L⋅log⁡3log⁡\|𝒱\|=Ω\(d⋅Llog⁡\|𝒱\|\)\.□m^\{\\prime\}\\;\\geq\\;\\frac\{dL\\cdot\\log 3\}\{\\log\|\\mathcal\{V\}\|\}\\;=\\;\\Omega\\\!\\left\(\\frac\{d\\cdot L\}\{\\log\|\\mathcal\{V\}\|\}\\right\)\.\\quad\\square\(15\)

### A\.2Proof of Theorem[1](https://arxiv.org/html/2607.09153#Thmtheorem1)

Verification error decomposition\.For any representationRRand function classℱ\\mathcal\{F\}, the expected loss of the best function inℱ\\mathcal\{F\}decomposes as:

inff∈ℱ𝔼​\[ℓ​\(f​\(R\),Y\)\]=inff𝔼​\[ℓ​\(f​\(R\),Y\)\]⏟ϵBayes​\(R\)\+inff∈ℱ𝔼​\[ℓ​\(f​\(R\),Y\)\]−inff𝔼​\[ℓ​\(f​\(R\),Y\)\]⏟Gap​\(ℱ,R\)\.\\inf\_\{f\\in\\mathcal\{F\}\}\\mathbb\{E\}\[\\ell\(f\(R\),Y\)\]=\\underbrace\{\\inf\_\{f\}\\,\\mathbb\{E\}\[\\ell\(f\(R\),Y\)\]\}\_\{\\epsilon\_\{\\mathrm\{Bayes\}\}\(R\)\}\+\\underbrace\{\\inf\_\{f\\in\\mathcal\{F\}\}\\mathbb\{E\}\[\\ell\(f\(R\),Y\)\]\-\\inf\_\{f\}\\,\\mathbb\{E\}\[\\ell\(f\(R\),Y\)\]\}\_\{\\mathrm\{Gap\}\(\\mathcal\{F\},R\)\}\.\(16\)This is a direct decomposition of the achievable error\. SettingR=𝐇R=\\mathbf\{H\}andℱ=ℱk\\mathcal\{F\}=\\mathcal\{F\}\_\{k\}yields Equation[4](https://arxiv.org/html/2607.09153#S3.E4)\. The inequalityϵBayes​\(𝐇\)≤ϵBayes​\(𝐱\)\\epsilon\_\{\\mathrm\{Bayes\}\}\(\\mathbf\{H\}\)\\leq\\epsilon\_\{\\mathrm\{Bayes\}\}\(\\mathbf\{x\}\)follows from Theorem[1](https://arxiv.org/html/2607.09153#Thmproposition1)\(a\): sinceI​\(𝐇;Y\)≥I​\(𝐱;Y\)I\(\\mathbf\{H\};Y\)\\geq I\(\\mathbf\{x\};Y\), a predictor with access to𝐇\\mathbf\{H\}can achieve at least as low a Bayes error as one with access to𝐱\\mathbf\{x\}only\.

Monotonicity ofGap​\(ℱk,𝐇\)\\mathrm\{Gap\}\(\\mathcal\{F\}\_\{k\},\\mathbf\{H\}\)inkk: sinceℱk⊆ℱk′\\mathcal\{F\}\_\{k\}\\subseteq\\mathcal\{F\}\_\{k^\{\\prime\}\}fork≤k′k\\leq k^\{\\prime\}\(a depth\-k′k^\{\\prime\}readout can ignorek′−kk^\{\\prime\}\-kof its query tokens\), we haveinff∈ℱk′≤inff∈ℱk\\inf\_\{f\\in\\mathcal\{F\}\_\{k^\{\\prime\}\}\}\\leq\\inf\_\{f\\in\\mathcal\{F\}\_\{k\}\}, so the gap is non\-increasing\.□\\square

### A\.3Proof of Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)

We prove each part of Theorem[2](https://arxiv.org/html/2607.09153#Thmtheorem2)under Assumptions[1](https://arxiv.org/html/2607.09153#Thmassumption1)and[2](https://arxiv.org/html/2607.09153#Thmassumption2)\.

Setup\.Under Assumption[2](https://arxiv.org/html/2607.09153#Thmassumption2), the rewardYYdepends on𝐇\\mathbf\{H\}throughrrfeaturesϕ1​\(𝐇\),…,ϕr​\(𝐇\)\\phi\_\{1\}\(\\mathbf\{H\}\),\\ldots,\\phi\_\{r\}\(\\mathbf\{H\}\)withϕj=\(Wϕ\)j​pool​\(𝐇\)\\phi\_\{j\}=\(W\_\{\\phi\}\)\_\{j\}\\,\\mathrm\{pool\}\(\\mathbf\{H\}\), where the covariance eigenvalues satisfyσj=O​\(e−α​j\)\\sigma\_\{j\}=O\(e^\{\-\\alpha j\}\)\. Under a Gaussian channel model, the mutual information contributed by thejj\-th feature is:

Ij=12​log⁡\(1\+SNR⋅σj2\),I\_\{j\}=\\tfrac\{1\}\{2\}\\log\(1\+\\mathrm\{SNR\}\\cdot\\sigma\_\{j\}^\{2\}\),\(17\)whereSNR=Var​\(Y\)/σϵ2\\mathrm\{SNR\}=\\mathrm\{Var\}\(Y\)/\\sigma\_\{\\epsilon\}^\{2\}\. The total extractable information isC0=∑j=1rIjC\_\{0\}=\\sum\_\{j=1\}^\{r\}I\_\{j\}\.

Feature extraction capacity of depth\-kkreadouts\.A depth\-kkreadout processeskkquery tokens throughNLN\_\{L\}transformer layers withnhn\_\{h\}attention heads each\. Each attention head computes a linear combination of the cached value vectors via its attention weights:headi=softmax​\(qi⊤​K/dk\)​V\\mathrm\{head\}\_\{i\}=\\mathrm\{softmax\}\(q\_\{i\}^\{\\top\}K/\\sqrt\{d\_\{k\}\}\)V, which is a linear function of𝐇\\mathbf\{H\}for fixed queryqiq\_\{i\}\. Withkkquery tokens, the readout has access tok⋅nh⋅NLk\\cdot n\_\{h\}\\cdot N\_\{L\}such linear projections of𝐇\\mathbf\{H\}\. Under Assumption[1](https://arxiv.org/html/2607.09153#Thmassumption1), these suffice to extractk⋅nh⋅NLk\\cdot n\_\{h\}\\cdot N\_\{L\}independent linear features from the reward\-relevant subspace\.

Greedy ordering\.An optimal readout extracts features in decreasing order of information contribution\. Afterkkdepth steps, the readout has extracted featuresϕ1,…,ϕk⋅nh​NL\\phi\_\{1\},\\ldots,\\phi\_\{k\\cdot n\_\{h\}N\_\{L\}\}\(or allrrfeatures ifk⋅nh​NL≥rk\\cdot n\_\{h\}N\_\{L\}\\geq r\)\. The residual \(unexplained\) information is:

Rk=∑j=k⋅nh​NL\+1rIj≤∑j=k⋅nh​NL\+1r12​log⁡\(1\+SNR⋅σj2\)\.R\_\{k\}=\\sum\_\{j=k\\cdot n\_\{h\}N\_\{L\}\+1\}^\{r\}I\_\{j\}\\leq\\sum\_\{j=k\\cdot n\_\{h\}N\_\{L\}\+1\}^\{r\}\\tfrac\{1\}\{2\}\\log\(1\+\\mathrm\{SNR\}\\cdot\\sigma\_\{j\}^\{2\}\)\.\(18\)
Part \(a\): Exponential decay\.The marginal gain at depthkkis bounded by the residual:

Δ​Ik≤Rk−1−Rk=∑j=\(k−1\)​nh​NL\+1k⋅nh​NLIj\.\\Delta I\_\{k\}\\leq R\_\{k\-1\}\-R\_\{k\}=\\sum\_\{j=\(k\-1\)n\_\{h\}N\_\{L\}\+1\}^\{k\\cdot n\_\{h\}N\_\{L\}\}I\_\{j\}\.\(19\)Sinceσj=O​\(e−α​j\)\\sigma\_\{j\}=O\(e^\{\-\\alpha j\}\), we haveIj=O​\(e−2​α​j\)I\_\{j\}=O\(e^\{\-2\\alpha j\}\)for largejj\(usinglog⁡\(1\+x\)≤x\\log\(1\+x\)\\leq xfor smallxx\)\. Thus:

Δ​Ik≤∑j=\(k−1\)​nh​NL\+1k⋅nh​NLO​\(e−2​α​j\)=O​\(e−2​α​\(k−1\)​nh​NL\)⋅∑i=1nh​NLO​\(e−2​α​i\)=O​\(e−2​α​\(k−1\)​nh​NL\)\.\\Delta I\_\{k\}\\leq\\sum\_\{j=\(k\-1\)n\_\{h\}N\_\{L\}\+1\}^\{k\\cdot n\_\{h\}N\_\{L\}\}O\(e^\{\-2\\alpha j\}\)=O\(e^\{\-2\\alpha\(k\-1\)n\_\{h\}N\_\{L\}\}\)\\cdot\\sum\_\{i=1\}^\{n\_\{h\}N\_\{L\}\}O\(e^\{\-2\\alpha i\}\)=O\(e^\{\-2\\alpha\(k\-1\)n\_\{h\}N\_\{L\}\}\)\.\(20\)Sincee−2​α≤e−αe^\{\-2\\alpha\}\\leq e^\{\-\\alpha\}, we obtain the stated boundΔ​Ik≤C0⋅e−α⋅nh​NL⋅\(k−1\)\\Delta I\_\{k\}\\leq C\_\{0\}\\cdot e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\\cdot\(k\-1\)\}, where the constantC0C\_\{0\}absorbs the geometric series prefactor\.□\\square

Part \(b\):k=1k\\\!=\\\!1near\-optimality\.The information captured atk=1k=1isI​\(ℛ1;Y\)≥C0−R1I\(\\mathcal\{R\}\_\{1\};Y\)\\geq C\_\{0\}\-R\_\{1\}\. The residual afterk=1k=1is:

R1=∑j=nh​NL\+1rIj≤C0⋅∑m=1∞e−α⋅nh​NL⋅m=C0⋅e−α⋅nh​NL1−e−α⋅nh​NL\.R\_\{1\}=\\sum\_\{j=n\_\{h\}N\_\{L\}\+1\}^\{r\}I\_\{j\}\\leq C\_\{0\}\\cdot\\sum\_\{m=1\}^\{\\infty\}e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\\cdot m\}=C\_\{0\}\\cdot\\frac\{e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\}\}\{1\-e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\}\}\.\(21\)Therefore:

I​\(ℛ1;Y\)I​\(𝐇;Y\)≥C0−R1C0≥1−e−α⋅nh​NL1−e−α⋅nh​NL\.□\\frac\{I\(\\mathcal\{R\}\_\{1\};Y\)\}\{I\(\\mathbf\{H\};Y\)\}\\geq\\frac\{C\_\{0\}\-R\_\{1\}\}\{C\_\{0\}\}\\geq 1\-\\frac\{e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\}\}\{1\-e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\}\}\.\\quad\\square\(22\)
Part \(c\): Cost–information tradeoff\.The incremental cost of increasing depth fromk−1k\-1tokkis:

Freadout​\(k,L\)−Freadout​\(k−1,L\)=NL⋅\(cattn⋅d⋅L\+cffn⋅d2\)=Θ​\(d⋅L\)\.F\_\{\\mathrm\{readout\}\}\(k,L\)\-F\_\{\\mathrm\{readout\}\}\(k\\\!\-\\\!1,L\)=N\_\{L\}\\cdot\(c\_\{\\mathrm\{attn\}\}\\cdot d\\cdot L\+c\_\{\\mathrm\{ffn\}\}\\cdot d^\{2\}\)=\\Theta\(d\\cdot L\)\.\(23\)Dividing part \(a\) by this constant\-per\-step cost:

Δ​IkΘ​\(d⋅L\)=O\(e−α⋅nh​NL⋅\(k−1\)d⋅L\)\.□\\frac\{\\Delta I\_\{k\}\}\{\\Theta\(d\\cdot L\)\}=O\\\!\\left\(\\frac\{e^\{\-\\alpha\\cdot n\_\{h\}N\_\{L\}\\cdot\(k\-1\)\}\}\{d\\cdot L\}\\right\)\.\\quad\\square\(24\)

## Appendix BImplementation Details

### B\.1Training Hyperparameters

Table 5:Training hyperparameters for KV\-PRM and Text\-PRM\.
### B\.2Verify Token Details

The verify token is “?” \(token ID determined by the tokenizer\)\. The judgment tokens are “\+” and “\-”\. During scoring, we extract logits at these two positions and apply softmax to obtainP​\(\+\)P\(\+\)\. Labels are mapped fromy∈\[−1,1\]y\\in\[\-1,1\]toy^=\(y\+1\)/2∈\[0,1\]\\hat\{y\}=\(y\+1\)/2\\in\[0,1\]\.

## Appendix CKV Steering Algorithm

KV Steering exploits the differentiability of KV\-PRM’s KV\-cache scoring to optimize inter\-agent messages via gradient ascent on the process reward\. At each agent handoff, the accumulated KV cache is refined before the next agent begins generation\.

Algorithm 2KV Steering at agent handoffj→j\+1j\\to j\+10:KV cache

kvj\\text\{kv\}\_\{j\}after agent

aja\_\{j\}, step size

η\\eta, number of steps

TT
1:

kv\(0\)←kvj\\text\{kv\}^\{\(0\)\}\\leftarrow\\text\{kv\}\_\{j\}
2:for

t=0t=0to

T−1T\-1do

3:

kv~\(t\)←clone​\(kv\(t\)\)\\widetilde\{\\text\{kv\}\}^\{\(t\)\}\\leftarrow\\text\{clone\}\(\\text\{kv\}^\{\(t\)\}\), set

requires\_grad←True\\text\{requires\\\_grad\}\\leftarrow\\text\{True\}
4:

ℓ←−log⁡P​\(\+\)​\(kv~\(t\)\)\\ell\\leftarrow\-\\log P\(\+\)\(\\widetilde\{\\text\{kv\}\}^\{\(t\)\}\)// Negative log\-prob of positive judgment

5:

kv\(t\+1\)←kv~\(t\)−η⋅∇kv~ℓ\\text\{kv\}^\{\(t\+1\)\}\\leftarrow\\widetilde\{\\text\{kv\}\}^\{\(t\)\}\-\\eta\\cdot\\nabla\_\{\\widetilde\{\\text\{kv\}\}\}\\ell// Gradient ascent on reward

6:endfor

7:return

kv\(T\)\\text\{kv\}^\{\(T\)\}

Similar Articles

CacheRL:Multi-Turn Tool-Calling Agents via Cached Rollouts and Hybrid Reward

arXiv cs.CL

CacheRL trains small agent foundation models for multi-step tool-calling tasks, achieving 92% process accuracy (approaching GPT-5's 94%) with 100x less compute using cached rollouts and hybrid reward shaping, with innovations in knowledge transfer, cache-aware rewards, and iterative SFT/GRPO training.

SCI-PRM: A Tool Aware Process Reward Model for Scientific Reasoning Verification

arXiv cs.AI

SCI-PRM introduces a tool-aware Process Reward Model for scientific reasoning, trained on the SCIPRM70K dataset featuring 'Chain-of-Tool' trajectories that interleave reasoning with scientific tool execution. It enables effective test-time scaling and serves as a dense reward signal in reinforcement learning, outperforming proprietary models like GPT-5-Mini on tool-calling steps across scientific benchmarks.