低確率トークンでRL探索を深化！

紹介論文

今回紹介する論文はLow-probability Tokens Sustain Exploration in Reinforcement Learning
with Verifiable Rewardという論文です。

この論文を一言でまとめると

RLVRの課題である探索不足を、低確率だが重要なトークンに着目して解決するLp-Reg手法を解説。数式や実験結果を交え、その有効性と応用可能性を探ります。

RLVRの壁：探索の枯渇

近年、大規模言語モデル（LLM）は、その驚異的な性能で様々な分野に革新をもたらしています。特に、複雑な推論を必要とするタスクにおいて、LLMは目覚ましい成果を上げていますが、その裏には、RLVR（Reinforcement Learning with Verifiable Rewards：検証可能な報酬による強化学習）という強力な技術が存在します。

RLVRとは？

RLVRは、LLMが生成した推論過程や最終的な解答に対し、ルールベースの検証システムを用いて報酬を与えることで、モデルの学習を促進する手法です。従来の強化学習では、人間が報酬を与える必要がありましたが、RLVRでは検証システムが自動的に報酬を付与するため、学習の効率化と大規模化が可能になります。

RLVRのLLMへの応用

RLVRは、OpenAIのモデルやDeepSeek-R1など、最先端のLLMに採用され、その性能向上に大きく貢献しています。特に、数学オリンピックの問題や、複雑な論理パズルなど、高度な推論能力を必要とするタスクにおいて、RLVRはLLMの潜在能力を最大限に引き出すことに成功しています。

しかし、RLVRにも課題が…

しかし、RLVRを用いたLLMの学習は、決して平坦な道のりではありません。学習が進むにつれて、モデルの性能が頭打ちになり、最終的には崩壊してしまうという、探索不足という深刻な問題に直面することがあります。これは、モデルが特定の解法に偏り、新しい解法を探すことを怠ってしまうために起こります。

探索不足の原因：エントロピーの低下

探索不足の根本的な原因は、学習が進むにつれて、モデルのエントロピーが急速に低下することにあります。エントロピーとは、モデルの予測のランダム性を示す指標であり、エントロピーが低い状態は、モデルが特定の解法に過度に自信を持ち、多様な解法を試みなくなることを意味します。

従来手法の限界

従来の研究では、エントロピーを高く維持することで探索を促進しようとする試みがなされてきました。しかし、単にエントロピーを高めるだけでは、ノイズを増幅させ、かえって学習を不安定にしてしまうという問題があります。つまり、効果的な探索のためには、単にランダム性を追求するのではなく、より洗練されたアプローチが必要なのです。

本論文の着眼点：低確率トークンの重要性

本論文では、従来の探索手法とは異なる視点から、RLVRにおける探索不足の問題にアプローチします。著者らは、低確率トークン、つまり、モデルがあまり選択しないトークンの中に、重要な情報が隠されているという仮説を立て、その検証に取り組んでいます。次項では、この低確率トークンの重要性について、詳しく解説していきます。

重要トークンは低確率に眠る

大規模言語モデル（LLM）の性能を飛躍的に向上させる強化学習ですが、その過程で「探索の枯渇」という深刻な問題に直面することがあります。これは、モデルが多様な行動を試みなくなり、結果として最適な解にたどり着けなくなる現象です。従来の手法では、モデルの出力する確率分布のエントロピーを高く保つことで、この問題に対処しようとしてきました。しかし、エントロピーをただ高くするだけでは、本当に意味のある探索には繋がらないことが分かってきました。

従来のエントロピー最大化の問題点

従来の手法は、方策のエントロピーを高く保つことが探索に繋がると考えられていましたが、これは間接的で不正確なアプローチです。なぜなら、エントロピーが高い状態は、必ずしも有益な探索行動を意味しないからです。無差別にランダム性を最大化すると、ノイズが増幅され、学習が不安定になるリスクさえあります。具体的には、文法的に不自然な単語や、文脈にそぐわない単語の出現頻度が高まり、かえって学習を妨げてしまうのです。

低確率トークンという視点

本論文では、これまで見過ごされてきた「低確率トークン」に着目することで、この問題を解決しようと試みています。低確率トークンとは、モデルが出力する可能性が低いと判断した単語や記号のことです。しかし、これらの低確率トークンの中には、モデルが新しい推論の道筋を発見するための「推論の火花（Reasoning Sparks）」となるものが含まれているのです。

「推論の火花」とは何か？

「推論の火花」の具体例としては、「待って」「しかし」「たぶん」などの言葉が挙げられます。これらの言葉は、一見すると重要ではないように思えますが、多様な推論経路を開始する論理的な接続詞や不確実性の表現として機能します。例えば、数学の問題を解く際に「しかし、別の解法もあるかもしれない」と考えることは、新しい探索の第一歩となるでしょう。また、プログラミングにおいて「たぶん、この関数は別の引数を受け取るべきだ」と考えることは、より良いコードを生み出すきっかけになるかもしれません。

低確率トークンが排除されるメカニズム

それでは、なぜこのような重要な「推論の火花」が、学習の過程で排除されてしまうのでしょうか？
論文では、標準的なGRPO（Group Relative Policy Optimization）トレーニングにおいて、重要な探索的トークンの低確率サンプリングが抑制されることを指摘しています。GRPOは、モデルの出力と正解とのずれを小さくするように学習を進めますが、その過程で、正解に繋がりそうもない低確率トークンを過度に抑制してしまうのです。さらに、無差別にランダム性を高めると、数学的推論タスクの文脈から外れた無関係なトークン（”cost”、”fine”など）の低確率サンプリングが増幅されます。これらの無関係なトークンの増幅はノイズとなり、ベースラインよりも速い性能崩壊につながるのです。

このように、従来の手法では、探索の重要性を理解しているにも関わらず、その具体的なメカニズムを捉えきれていませんでした。その結果、重要な低確率トークンを排除してしまい、モデルの探索能力を損なっていたのです。次のセクションでは、本論文で提案されている、この問題を解決するための新しい手法、Lp-Regについて解説します。

Lp-Reg：低確率トークン活用法

前のセクションでは、従来の探索手法では見過ごされがちだった「低確率トークン」こそが、大規模言語モデル（LLM）の推論能力を飛躍的に向上させる鍵となる可能性について解説しました。このセクションでは、その可能性を最大限に引き出すために提案された手法、Lp-Reg（Low-probability Regularization）の仕組みを、数式を交えながら分かりやすく解説していきます。

Lp-Regは、一言で言うと、「価値のある低確率トークンを保護する」ための技術です。従来の探索手法とは異なり、単にエントロピーを最大化するのではなく、LLM自身が持つ予測分布を活用することで、ノイズとなるトークンを除去し、本当に重要なトークンに焦点を当てます。これにより、探索の効率と安定性を両立させることが可能になります。

Lp-Regの概要：ノイズ除去と選択的保護

Lp-Regの核心は、無差別に低確率トークンを扱うのではなく、その中でも「推論の火花（Reasoning Sparks）」となりうるトークンを選択的に保護することにあります。具体的な手順は以下の通りです。

1. ノイズトークンのフィルタリング：確率が閾値τを下回るトークンを、ノイズとして除去します。この閾値τは、固定値またはピーク確率に対する相対値として設定可能です。
2. 確率の再正規化：除去されたトークンの確率を、残りの候補トークンに再分配します。これにより、価値のある低確率トークンの相対的な重要度が増します。
3. 正則化：元のポリシーと、ノイズが除去されたプロキシ分布との乖離にペナルティを課します。これにより、重要なトークンが不当に排除されるのを防ぎます。

Lp-Regを支える数式

Lp-Regの動作をより深く理解するために、論文に記載されている目的関数を見てみましょう。

J(θ) = E_{B~D,(q,a)~B,{o_i}~π_θold(·|q)} [ 1/G Σ_{i=1}^{G} Σ_{t=1}^{|o_i|} clip(r_{i,t}(θ), ε, U) · A_{i,t}
    - β · I[π_θ(o_{i,t}|q,o_{i, 0 ∧ A_{i,t} < 0]
    · D_{KL}(π_proxy(·|q,o_{i,