Skip to content

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

🕒 Published (v1): 2026-07-10 07:16 UTC · Source: Arxiv · link

Why this paper was selected

KV-cache transfer for process reward models cuts TTS cost in multi-agent systems; novel efficiency mechanism

Ask a follow-up

Open an assistant pre-loaded with this paper's context.

💬 Ask ChatGPT✦ Ask Claude

TL;DR

KV-PRM replaces text-based Process Reward Models (PRMs) with a KV-cache-transfer architecture that scores multi-agent reasoning trajectories by processing a single "verify token" against the generation KV cache rather than re-encoding trajectory text from scratch. This reduces per-call scoring cost from \(O(dL^2)\) to \(O(dL)\), yielding up to 5,000× fewer FLOPs, 37× lower latency, and 34× lower memory while matching or exceeding text-PRM accuracy on MATH, GSM8K, and AIME benchmarks.

Problem

Existing PRMs are text-based: they re-encode the entire trajectory from scratch for every scoring call, incurring \(O(dL^2)\) cost dominated by self-attention. In multi-agent systems (MAS) with beam width \(W\) and \(D\) agent steps, the PRM is invoked at least \(W \cdot D\) times over trajectories of \(L \gg d\) tokens, making PRM scoring cost comparable to generation itself and severely limiting long-context test-time scaling.

Method

KV-PRM reuses the KV cache \((\mathbf{K}, \mathbf{V}) \in \mathbb{R}^{N_L \times L \times d}\) already produced during autoregressive generation. A LoRA adapter \(\Delta\theta\) is appended to the base model; at scoring time, the adapter is activated and a single designated "verify token" \(v\) is processed against the pre-existing cache:

\[z = f_{\theta+\Delta\theta}(v \mid \mathbf{K}, \mathbf{V}), \quad s_j = P(+) = \text{softmax}([z_{[-]}, z_{[+]}])_1\]

This constitutes a depth-\(k=1\) readout in the paper's formal framework, costing \(O(dL)\) per call. The base model runs with the adapter disabled during generation, so no generation overhead is introduced. Training labels are MCTS Q-values from the MATH training split (R=64 rollouts, C=4 candidates/step), and only LoRA parameters are updated via MSE loss \(\mathcal{L} = \text{MSE}(P(+),\, (y+1)/2)\).

The theoretical grounding rests on two results: (1) under the Linear Representation Hypothesis, the KV cache is \(\Omega(d/\log|V|)\) times more information-dense per position than text tokens (Proposition 1 / Theorem 1); (2) marginal information gain from additional verify tokens decays exponentially as \(\Delta I_k \leq C_0 \cdot e^{-\alpha n_h N_L (k-1)}\), so \(k=1\) captures near-optimal reward signal (Theorem 2).

Key Contributions

  • Formal proof that KV cache contains strictly greater information capacity than decoded text, with an \(\Omega(d/\log|V|)\) capacity advantage per position.
  • Theorem showing the approximation gap of a depth-\(k\) readout decays exponentially in \(k\), formally justifying \(k=1\) near-optimality.
  • KV-PRM architecture: single verify token + LoRA adapter over the generation KV cache, reducing scoring from \(O(dL^2)\) to \(O(dL)\).
  • Drop-in compatibility with beam search, MCTS, and weighted majority voting in MAS pipelines.
  • Proof-of-concept "KV Steering": differentiability of KV-based scores enables gradient-based optimization of latent inter-agent messages, structurally impossible with text PRMs.

Results

  • FLOPs: up to 5,000× reduction vs. text-PRM (e.g., \(\sim\)4,900× at \(L \approx 4{,}900\) tokens).
  • Latency: 15–37× wall-clock speedup on NVIDIA GH200; at Qwen3-8B / \(L=4{,}096\), scoring takes 4.6 ms vs. 172.0 ms for text-PRM.
  • Memory: 34.2× reduction in average per-sequence memory during PRM inference.
  • MATH (sequential MAS): KV-PRM matches or exceeds text-PRM across Qwen3-0.6B/4B/8B under SBS, MCTS, weighted voting (e.g., MCTS Qwen3-8B: 69.57% KV-PRM vs. 68.92% text-PRM).
  • AIME 2024 (sequential): KV-PRM outperforms text-PRM by up to +6.7 pp in some configurations; gains more pronounced on harder problems.
  • GSM8K: near-parity with text-PRM, marginal improvements (e.g., 95.07% vs. 95.00% for Qwen3-8B beam search).
  • Efficiency vs. model-size ablation: KV-PRM-8B at ~3,000× lower cost than text-PRM-8B achieves 68.0% accuracy on MATH SBS, surpassing text-PRM at all sizes (0.6B, 4B, 8B).
  • Verify-token ablation: accuracy gain over \(k=1\) saturates rapidly; fitted curves confirm \(a(1-e^{-c(k-1)})\) form consistent with Theorem 2.
  • KV Steering (no search): +1.61 pp on MATH Qwen3-4B, +0.61 pp on GSM8K Qwen3-4B vs. no-steering baseline.

Limitations

  • KV-PRM requires the verifier and generator to share the same base model architecture (the LoRA adapter must be compatible with the generator's KV cache format), constraining verifier/generator decoupling.
  • Evaluated only on mathematical reasoning benchmarks (MATH, GSM8K, AIME); generalization to other reasoning domains (code, science, tool use) is undemonstrated.
  • KV cache must remain in memory and be transferred at agent handoffs; in distributed MAS deployments with heterogeneous hardware, KV transfer overhead may offset scoring savings.
  • Training data (MCTS Q-values) is generated on MATH training split only; out-of-distribution robustness of the LoRA adapter is not systematically studied.
  • KV Steering is presented only as a proof-of-concept with no convergence guarantees or study of optimization stability.

Relevance to Agentic AI / LLM Agents

KV-PRM directly addresses the computational scalability of process reward models in multi-agent test-time search, which is a central bottleneck for deploying capable reasoning agents under real-world resource constraints. By eliminating the \(O(L^2)\) re-encoding overhead, it makes PRM-guided beam search and MCTS tractable for the long-context trajectories that arise in multi-agent pipelines (planner → solver → verifier topologies). The KV Steering extension opens a qualitatively new direction: gradient-based optimization of latent inter-agent messages, enabling agents to coordinate in continuous representation space rather than through discrete text handoffs. This connects to broader work on latent-space communication, differentiable multi-agent systems, and test-time compute scaling as a complement to training-time scaling.