LLM強化学習の学習・推論ミスマッチをCISで補正

導入

検証可能な報酬を用いる強化学習（RLVR）では、回答の生成とモデルの更新が別の計算系で行われることが多い。推論エンジンがロールアウトをサンプリングし、学習エンジンがその結果を使って勾配を計算する構成だ。理想的には両者は同じ方策を表すはずだが、同じモデル重みを使っていても、数値精度、ハードウェア、並列化、演算実装の違いによって、同じトークンに異なる確率を割り当てる場合がある。

論文「Rethinking Training-Inference Mismatch in LLM Reinforcement Learning」は、この学習・推論ミスマッチを重要度サンプリングの観点から検討し、Calibrated Importance Sampling（CIS）を提案した。固定の重要度比クリップを適用するのではなく、確率差がどのように生じるかを調べ、その大きさをトークンの信頼度に合わせて調整する点が特徴だ。

主なポイント

• 問題は計算基盤の分離から生じる。 ロールアウトは推論エンジンから得られる一方、勾配は学習エンジンで計算される。両者の確率がずれると、重要度比が極端になり、方策更新の誤差や不安定さにつながる。
• 対数オッズ上の変位で差を表す。 著者らの実証分析では、ミスマッチは対数オッズに加わる変位として近似できる。この変位はsoftmax前の各logitの摂動によって決まり、選択トークンの信頼度に対しておおむね不変な分布を持つ。
• 信頼度に応じて切り詰める。 CISは変位の空間で大きな正の変位を一定の閾値で切り詰める。それを重要度比に戻すと、高信頼度トークンほど厳しい上限になり、低信頼度トークンでは有効な更新情報を残しやすくなる。
• 分散とバイアスを調整する。 理論上、CISは正確な重要度サンプリングで問題となる、場合によっては無界な二次モーメントを、定数で抑えられる項に置き換える。その代償として、切り詰めた超過量に応じたバイアスが生じる。また、小さな重要度重みを上向きにクリップすると、保持データの精度が下がるという診断結果も報告されている。

実験と意義

CISは3つの混合専門家（MoE）モデルと5つの数学推論ベンチマークで評価された。5ベンチマークの平均では、3モデルすべてで比較対象の中で最高の結果を達成した。追加の分析では、通常の切り詰め重要度サンプリングと比べ、低信頼度トークンに対する切り詰めバイアスが小さいことも示された。

この研究の意義は、新しい重み付け手法を示したことだけではない。推論基盤と学習基盤の小さな数値差が重要度比によって増幅され、方策学習の方向や大きさに影響し得ることを明確にした点にある。RLVRの設計では、学習フレームワーク、推論ランタイム、数値形式、並列計算方式の互換性も考慮する必要がある。

一方、今回の結果は論文で評価されたモデルと数学タスクに基づく。ほかのタスク、エンジン構成、学習条件でも同じ効果が得られるかは、今後の検証を要する。

出典：Hugging Face Daily Papers