arXiv:2607.09153cs.AI2026-07被引 2

用缓存重用技术让奖励模型提速5000倍,大幅降低长对话计算开销。

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

  • 直接读取生成过程中的键值缓存,避免重复编码整个对话文本。
  • 在多个数学推理任务上表现优于或持平传统方法,计算量降低5000倍。
  • 适合需要长序列推理的多智能体系统,尤其关注效率与资源受限场景。

过程奖励模型(PRMs)已被证明能有效引导测试时扩展(TTS)方法,显著提升基于大语言模型的多智能体系统能力。然而,现有PRMs均为基于文本:需从头重新编码完整轨迹文本。在长多智能体回放中,评分成本随序列长度L呈平方增长,造成严重计算瓶颈,极大限制了其在长上下文场景的应用。为此,我们提出KV-PRM,一种高效的过程奖励模型,通过直接读取大语言模型生成阶段自然产生的键值缓存(KV cache),消除冗余文本重编码。仅需处理一个“验证标记”,即可利用预存的KV缓存,将评分成本从O(L²)降至O(L)。我们形式化证明了键值缓存的信息容量严格高于文本,且更适用于下游奖励建模。实验表明,在MATH、GSM8K和AIME基准上,KV-PRM在各类TTS方法(如束搜索、蒙特卡洛树搜索、加权投票)下表现匹配或超越文本型PRMs,评分浮点运算量减少高达5000倍,延迟降低37倍,每序列内存占用减少34倍。

原文摘要 · Abstract (English)

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.

奖励模型长序列效率优化多智能体

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。