arXiv cs.AI· Peng Kuang, Haibo Jin, Xiaoyu Han, Yanli Wang, Xiaopeng Yuan, Ye Yu, Kaidi Xu, Haohan Wang·· 11 小时前AI 评分38
KV-PRM：通过 KV-Cache 传输实现多智能体测试时扩展的高效过程奖励建模
KV-PRM: Efficient Process Reward Modeling via KV-Cache Transfer for Multi-Agent Test-Time Scaling
AI 导读
KV-PRM 通过直接读取 LLM 生成阶段产生的 KV cache 完成过程奖励打分，将评分成本从 O(L^2) 降至 O(L)。在 MATH、GSM8K、AIME 上，它在 Beam Search、MCTS、Weighted Voting 等 TTS 方法下持平或优于文本 PRM，评分 FLOPs 最多降低 5,000 倍、延迟降低 37 倍、单序列内存占用降低 34 倍。
正文 · 原文
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.
| Comments: | Accepted at NeurIPS 2026. 27 pages, including appendices |
| Subjects: | Artificial Intelligence (cs.AI) |
| Cite as: | arXiv:2607.09153 [cs.AI] |
| (or arXiv:2607.09153v2 [cs.AI] for this version) | |
| https://doi.org/10.48550/arXiv.2607.09153 arXiv-issued DOI via DataCite |
Submission history
From: Peng Kuang [view email]
[v1]
Fri, 10 Jul 2026 07:16:43 UTC (212 KB)
[v2]
Wed, 4 Nov 2026 05:26:47 UTC (235 KB)
来源：arXiv cs.AI · arxiv.org