LaSeR: 最終トークン自己報酬付き強化学習

要旨

検証可能な報酬を伴う強化学習（Reinforcement Learning with Verifiable Rewards, RLVR）は、大規模言語モデル（Large Language Models, LLMs）の推論能力を向上させるための核心的なパラダイムとして近年注目を集めている。テスト時に検証信号が不足する問題に対処するため、従来の研究では、モデルの自己検証能力の訓練を標準的なRLVRプロセスに組み込み、単一のLLM内で推論能力と検証能力を統合するアプローチが取られてきた。しかし、これまでの手法では、LLMが2つの異なるプロンプトテンプレートを使用して解決策と自己検証を順次生成する必要があり、効率が大幅に低下していた。本研究では、自己検証のRL目的関数の閉形式解が驚くほど単純な形式に還元できることを理論的に明らかにする。具体的には、解決策の真の推論報酬は、その最後のトークンにおける自己報酬スコアに等しく、このスコアは、解決策の最後のトークンにおいて事前に指定された任意のトークンに割り当てられたポリシーモデルの次トークン対数確率と事前に計算された定数との差をKL係数でスケーリングしたものとして計算される。この知見に基づき、我々はLaSeR（Reinforcement Learning with Last-Token Self-Rewarding）を提案する。このアルゴリズムは、元のRLVR損失に、最後のトークン自己報酬スコアと検証器ベースの推論報酬を整合させるMSE損失を単純に追加し、LLMの推論能力と自己報酬能力を共同で最適化する。最適化された自己報酬スコアは、訓練とテストの両方でモデルの性能を向上させるために利用可能である。特に、我々のアルゴリズムは、生成直後の最後のトークンの次トークン確率分布からこれらのスコアを導出するため、追加のトークン推論の最小限のコストしかかからない。実験結果は、我々の手法がモデルの推論性能を向上させるだけでなく、顕著な自己報酬能力を備えさせることで、推論時のスケーリング性能を向上させることを示している。

English

Reinforcement Learning with Verifiable Rewards (RLVR) has recently emerged as a core paradigm for enhancing the reasoning capabilities of Large Language Models (LLMs). To address the lack of verification signals at test time, prior studies incorporate the training of model's self-verification capability into the standard RLVR process, thereby unifying reasoning and verification capabilities within a single LLM. However, previous practice requires the LLM to sequentially generate solutions and self-verifications using two separate prompt templates, which significantly reduces efficiency. In this work, we theoretically reveal that the closed-form solution to the RL objective of self-verification can be reduced to a remarkably simple form: the true reasoning reward of a solution is equal to its last-token self-rewarding score, which is computed as the difference between the policy model's next-token log-probability assigned to any pre-specified token at the solution's last token and a pre-calculated constant, scaled by the KL coefficient. Based on this insight, we propose LaSeR (Reinforcement Learning with Last-Token Self-Rewarding), an algorithm that simply augments the original RLVR loss with a MSE loss that aligns the last-token self-rewarding scores with verifier-based reasoning rewards, jointly optimizing the reasoning and self-rewarding capabilities of LLMs. The optimized self-rewarding scores can be utilized in both training and testing to enhance model performance. Notably, our algorithm derives these scores from the predicted next-token probability distribution of the last token immediately after generation, incurring only the minimal extra cost of one additional token inference. Experiments show that our method not only improves the model's reasoning performance but also equips it with remarkable self-rewarding capability, thereby boosting its inference-time scaling performance.

LaSeR: 最終トークン自己報酬付き強化学習

LaSeR: Reinforcement Learning with Last-Token Self-Rewarding

要旨

Support