推論LLMの「どの思考が効いたか」に報酬を配り直す ― 連鎖思考の報酬再配分「RREDCoT」

本研究(RREDCoT)は、推論する大規模言語モデル(LLM)を強化学習(RL)で鍛える際の「どの思考ステップが効いたか」を見極めて報酬を配り直す手法を提案する。

近年の推論モデルの進歩は、RL微調整によって駆動されてきた。多くはGRPO(Group Relative Policy Optimization)アルゴリズムやその改良を用いて、連鎖思考(Chain-of-Thought, CoT)のトレースを生成するようモデルを誘導する。問題は、最終的な答えはCoTトレースが完了して初めて検証でき、報酬もそのとき割り当てられる点だ——これは「遅延報酬(delayed reward)」問題である。GRPOやその改良は、標準的なRLでいうモンテカルロ法に相当し、高分散に悩まされることが知られている。

この問題への解決策の一つが、クレジット割当(credit assignment)による報酬の再配分である。すなわち、望ましい解にたどり着くうえで重要なCoTトレースのセグメントを、より高い報酬を割り当てて強調する。モンテカルロ標本化は中間状態価値の不偏推定を与えうるが、その計算負荷の大きさから、長い文脈で高い粒度の学習時クレジット割当には不向きだ。

著者らが導入するRREDCoT(Reward REDistribution for Chain of Thoughts)は、追加の生成を行うことなく、モデル自身を用いて最適な報酬再配分を近似する。著者らは、MC標本化やいくつかの帰属(attribution)法と比較してその利点を検証し、さらに再配分の構成に関わる側面——CoTトレースのセグメント分割や状態価値推定など——も分析している。なお著者にはLSTMの考案者として知られるSepp Hochreiterら(JKU Linz)が名を連ねる。