KV-PRM: KVキャッシュ転送によるマルチエージェントテスト時スケーリングのための効率的なプロセス報酬モデリング

近年、プロセス報酬モデル（PRM）はテスト時スケーリング（TTS）手法のガイドとして非常に有効であることが証明されており、LLMベースのマルチエージェントシステムの能力を大幅に向上させています。PRMは各ステップで細かい報酬信号を提供することで、モデルが推論中により効果的な探索と意思決定を行うのを支援します。しかし、既存のPRMはほとんどがテキストベースであり、軌跡全体をゼロから再エンコードする必要があります。長いマルチエージェントロールアウトでは、シーケンス長Lが非常に長くなることがあり、スコアリングコストがシーケンス長に対して二次的に増大する（O(L²)の計算複雑性）ため、深刻な計算ボトルネックが生じます。これにより、長コンテキストシナリオ（例えば、マルチターンダイアログや複雑なタスク計画）へのPRMの適用が大幅に制限されています。この効率問題を解決するために、著者らはKV-PRMを提案します。KV-PRMの核心的な革新は、LLMの生成フェーズで自然に生成されるKVキャッシュを直接読み取ることで、重いテキスト再エンコードを排除することです。具体的には、事前に存在するKVキャッシュに対して単一の「検証トークン」を処理するだけでスコアリングが完了し、計算複雑性をO(L²)からO(L)に削減します。また、KVキャッシュはテキストよりも厳密に大きな情報容量を持ち、下流の報酬モデリングにおいてより効率的であることを正式に証明しています。この理論的発見は、KV-PRMの効率性に確固たる数学的基盤を提供します。実証評価では、MATH、GSM8K、AIMEベンチマークにおいて、ビームサーチ、MCTS、加重投票など様々なTTS手法のもとでKV-PRMをテストしました。結果は、KV-PRMがテキストPRMと同等以上の性能を示すことを示しています。同時に、テキストベースPRMと比較して、スコアリングFLOPsで最大5000倍、レイテンシで37倍、シーケンスあたりのメモリフットプリントで34倍の削減を達成しました。これらの顕著な効率向上により、KV-PRMは計算リソースが限られている環境やレイテンシに敏感な実システムへの展開に特に適しています。全体として、KV-PRMは長コンテキストマルチエージェントシステムの報酬モデリングのための効率的かつ実用的なソリューションを提供し、PRMのより広範な応用シナリオへの展開を促進する可能性があります。