推論は何の目的を解決しているのか?
What objective is reasoning solving? (alexwang.ai)
要約
この記事では、AIモデルにおける推論を「トークン空間における経路探索」として捉え直します。強化学習(RL)は、正しい答えにたどり着く経路を見つける確率を最大化することを目指しており、個々の経路が必ずしも最適解に収束するわけではないと説明しています。この視点は、推論モデルにおける「思考を長くする」ことの効果や、報酬ハッキングといった現象を理解するのに役立ちます。
全文翻訳
推論は何の目的を解決しているのか?
トークン空間における経路探索としての推論。
推論は、開始コンテキスト $x$ から正しい答え $y(z)$ への経路を見つけようとする各試み $z$ として捉えられます。$\Scal_x$ は正しい答えの集合です。
推論について直感的で単純化された見方があります。それは、モデルに長く考えさせれば、より正しくなる可能性が高いというものです。これは人間の直感ともよく合致します。より多くの時間と労力を思考に費やせば、正しい答えにたどり着く可能性が高まるでしょう。
しかし、この見方だけでは全てを説明できないようです。RLトレーニング中、モデルはより長い思考連鎖(chain of thought)を書くことを学習し、その精度も向上します。しかし、訓練済みのモデルを取り出し、任意の思考連鎖を強制的に長く実行させても、しばらくは役立ちますが、その後は効果がなくなります。1 場合によっては、まるで考えすぎているかのように、答えはさらに悪化することさえあります。2 もし思考連鎖が本当に何らかの最適値に収束しているのであれば、ステップ数を増やしても悪化するはずがありません。では、推論時に役立つのは本当に各連鎖の長さなのでしょうか?それとも、もっと微妙な何かがあるのでしょうか?
この記事は、推論をトークン空間を通る経路探索として見ることで、この混乱(主に私自身の混乱)を明確にしようとする試みです。RLは $x$ から $\Scal_x$ への経路を見つける確率を最大化しますが、個々の経路が正しい答えに収束することは保証されません。さらに、この視点は、思考連鎖の不忠実さや報酬ハッキングといった、推論モデルに関するいくつかの奇妙な観察結果を明らかにします。
経路探索としての推論#
同じ数学の問題を推論モデルに10回尋ねると、10通りの異なる思考連鎖が得られます。推論を答えに向かう単一の神秘的なプロセスと考えるのではなく、トークン空間における経路探索として考えてみましょう。問題から正しい答えへの経路は数多く存在し、モデルは各サンプリングでそのうちの1つをたどります。
問題「$1 + 2 + 3 + \cdots + 10$ は?」を考えてみましょう。1つの解法は、数字を1つずつ足していくことです。1、3、6、10、...、55となります。別の方法は、外側からペアにしていくことです。$1 + 10$、$2 + 9$、$3 + 8$、$4 + 7$、$5 + 6$ はそれぞれ11になり、11が5つで55になります。これらの経路はどちらも正しい答えである55に到達します。
推論をこの一般的なパターンに従うこととして考えることができます。私たちはプロンプト、つまり $x$ と呼ばれる開始点を与えられます。そして、 $x$ から始まり、正しい答え $y$ で終わる経路 $z$ を見つけたいのです。経路が正しい場所に到達したかどうかを判断するために、 $y$ が正しい場合は1、そうでない場合は0を返す検証者 $r(x, y)$ があると仮定します。
したがって、開始点 $x$ が与えられた場合、集合に到達するトークンの経路 $z$ を見つけたいのです。
$$\Scal_x = {y : r(x, y) = 1}.$$
この記事の冒頭にある図は、この絵を描いています。各灰色の線はサンプリングされた経路です。それは左側のプロンプトから始まり、$z$ をサンプリングしてトークン空間をさまよい、これらの経路の一部が $\Scal_x$ に到達します。
多くの経路が $\Scal_x$ に到達するため、成功確率はすべての経路にわたって加算されます。モデルがプロンプト $x$ を与えられたときに経路 $z$ を生成する確率を $\pi(z \mid x)$、 $z$ の終点での最終的な答えを $y(z)$ とすると、成功確率は、各経路を生成する確率に、その経路が $\Scal_x$ に到達する確率を掛けたものを、すべての経路について合計したものです。
$$p_x = \prob{\text{a path lands in } \Scal_x \mid x} = \sum_{z} \underbrace{\pi(z \mid x)}_{\text{samples path } z}; \underbrace{\prob{y(z) \in \Scal_x \mid z}}_{\text{path } z \text{ lands in }\Scal_x}.$$
2番目の因子は、経路が $\Scal_x$ に到達すれば1、そうでなければ0なので、 $\Scal_x$ に到達する経路の確率のみを合計します。
$$p_x = \sum_{z} \pi(z \mid x), \ind{y(z) \in \Scal_x} = \sum_{z ,:, y(z) \in \Scal_x} \pi(z \mid x).$$
これを 成功確率 と呼びます。この確率を最大化することは、モデルが $\Scal_x$ に到達する経路をサンプリングする可能性が高まることを意味します。
この確率を増加させるには2つの方法があることに注意してください。モデルが同じ答えに対して複数の経路を、それぞれ低い確率でサンプリングするか、モデルが単一の経路を非常に高い確率でサンプリングするかです。成功確率はどちらでも構いませんが、実際には特定の経路を優先したり、サンプルの多様性を確保したりすることを優先するかもしれません3。
成功確率をどう高めるか?#
ベストオブ-$k$ サンプリング#
経路探索の最も簡単な方法は、モデル $\pi$ に $\Scal_x$ に到達する経路が見つかるまで、経路 $z$ をサンプリングし続けることです。
検証者がいると仮定したので、ベストオブ-$k$ サンプリングを実行できます。これは $k$ 個の経路をサンプリングし、成功を探します。各経路が失敗する確率は $1 - p_x$ なので、$k$ 個すべてが失敗する確率は $(1-p_x)^k$ となり、成功確率は $p_x$ から以下のように上昇します。
$$\prob{\text{at least one of } k \text{ paths lands in } \Scal_x \mid x} = 1 - \underbrace{(1-p_x)^k}_{\text{all } k \text{ paths miss}}.$$
これは計算スケールにおける保証された軸です。 $p_x > 0$ である限り、より多くの経路をサンプリングするにつれて成功確率は単調に増加するため、より多くのサンプルに計算リソースを費やすことは常に役立ちます。これは、Brown et al. 2024 によって早期に示された通りです。しかし、ここでのスケーリングは、経路の長さ(長く考えること)ではなく、サンプルの数に関するものです。
より良いサンプルを生成するようにモデルを訓練する#
これまでのところモデルは固定されていたため、$p_x$ も固定されていました。モデルを変更して $p_x$ 自体を増やすことはできるでしょうか?
簡単な方法があります。答えを確認できる多くのプロンプトを収集します。それぞれについて、モデルからいくつかの経路をサンプリングします。$\Scal_x$ に到達した経路を保持し、それらを使ってモデルを訓練します。そして繰り返します。各ラウンドで、モデルがサンプリングした成功した経路により多くの重みを付けるように訓練します。このループは拒否サンプリングファインチューニング(rejection-sampling fine-tuning)と呼ばれ、STaR は推論におけるその初期バージョンでした。
検証可能な報酬を持つRLは、今日の推論モデルが訓練される方法であり、実質的に同じことを行います。モデルの重みを $\theta$ とするモデルを $\pi_\theta$ とすると、$r(x, y(z)) = 1$ は $y(z) \in \Scal_x$ と同値なので:
$$ \begin{aligned} p_x &= \sum_{z ,:, y(z) \in \Scal_x} \pi_\theta(z \mid x) \ &= \sum_{z} \pi_\theta(z \mid x), \ind{y(z) \in \Scal_x} \ &= \sum_{z} \pi_\theta(z \mid x), r(x, y(z)) \ &= \ev[z \sim \pi_\theta(\cdot \mid x) ]{r(x, y(z))}, \end{aligned} $$
これは、期待されるバイナリ報酬が成功確率 $p_x$ に等しいことを意味します。
対数微分トリック(log-derivative trick)を使用して成功確率の勾配を計算すると、次のようになります。
$$\grad_\theta, p_x = \ev[z \sim \pi_\theta(\cdot \mid x) ]{r(x, y(z)), \grad_\theta \log \pi_\theta(z \mid x)},$$
これはまさにポリシー勾配(policy gradient)であり、REINFORCE推定値とも呼ばれます。これを $N$ 個のサンプリングされた経路 $z_1, \dots, z_N$ でこの勾配を推定します。失敗した経路は $r = 0$ で消え、$\Scal_x$ に到達した経路は $r = 1$ となります。 $N$ 個の経路のうち $M$ 個が成功した場合、推定値は次のようになります。
$$\grad_\theta, p_x \approx \frac{1}{N} \sum_{i ,:, \text{landed}} \grad_\theta \log \pi_\theta(z_i \mid x) = \underbrace{\frac{M}{N}}_{\text{success rate}} \cdot \underbrace{\frac{1}{M} \sum_{i ,:, \text{landed}} \grad_\theta \log \pi_\theta(z_i \mid x)}_{\text{log likelihood of the successful paths}}.$$
この勾配は、経験的な成功率 $M/N$ と、成功した経路の平均対数尤度という2つの項を含みます。
言い換えれば、バイナリ報酬を持つREINFORCEの1ステップは、現在の成功率 $M/N$ でスケーリングされた、モデル自身のサンプリングされた経路に対する拒否サンプリングファインチューニングのステップです。
これはまた、導入部の最初の観察結果、つまりRLトレーニング中に思考連鎖が長くなる理由も説明します。各ステップはモデルを到達した経路の方に引き寄せます。もし成功した経路が失敗した経路よりも長い傾向があった場合(例えば、後退行動を含んでいたり、環境からより多くのフィードバックを集めたり、単に運が良かったりした場合)、モデルはより長い経路を書くように学習します。
償却と汎化#
RLトレーニング中、私たちは良い経路を探索するために多くの計算リソースを費やし、その後、モデルがそれらの経路を直接生成するように訓練します。訓練が完了すると、モデルはサンプリングに余分な計算リソースをすべて費やすことなく、答えを生成できるようになります。これは検索の償却(amortizing)と呼ばれます。私たちはトレーニング中に一度だけそのコストを支払うため、モデルに質問するたびに再度支払う必要がありません。4 多くのプロンプトに対して単一のモデルを訓練するため、モデルは新しいプロンプトに対して汎化(generalize)できる可能性があります。
下の図は、それぞれ独自の $\Scal_x$ を持つ4つのプロンプトに対するこの効果を示しており、トレーニング前とトレーニング後でモデルが切り替わります。トレーニング中、私たちは各プロンプトから経路をサンプリングし、成功したものを保持し、それらに対してSFTを行います。パーセンテージは、各プロンプトに対する経験的な成功率を示します。空のプロンプトは、モデルが以前に見たことのない新しいプロンプトですが、モデルはトレーニング後にこの新しいプロンプトに汎化します。新しいプロンプトがトレーニング分布から離れるほど、その精度は低下します(ドラッグしてみてください!)。
新しいプロンプトをドラッグしてみてください。
推論に関するいくつかの観察結果を再訪する#
推論モデルを、任意の $x$ に対して $\Scal_x$ に到達するように訓練された分布として見ると、これらのモデルに関するいくつかの奇妙な事実が明らかになります。
思考連鎖は必ずしも忠実ではない#
上記の検証者は、答えが正しいかどうかだけをチェックします。 $y(z) \in \Scal_x$ の場合にのみ報酬が1となり、それ以外の場合は0となります。したがって、中間トークン $z$ は、モデルが答えに到達するための単なる手段です。これらのトークンが真実であること、または後続の答えを説明することを要求するものは何もありませんが、もちろん、モデルの事前訓練のためにこれらはしばしば相関しています。
意味のある思考連鎖が必要な場合は、報酬またはトレーニング目的を調整する必要があります。
報酬ハッキングは検証者とサンプラーの両方から生じる#
検証者が決して間違えず、$r(x, y)=1$ が $y$ が実際に正しい答えであることと必要十分条件であると仮定しました。しかし実際には検証者は不完全である可能性があります。その場合、検証者が受け入れる答えの集合である $\Scal_x$ は、正しい答えの集合と同じではありません。それは間違った答えを含み、正しい答えを省略する可能性があります。この誤った仕様は、モデルが意図したものとは異なる $\Scal_x$ の領域に到達することを学習する可能性を開きます。
したがって、報酬ハッキングは、私たちが本当に望むものを捉え損ねる不完全な検証者と、これらのショートカット領域から離れられない不完全なサンプリング手順の両方から生じます。
下の図では、破線円は正しい答えの集合であり、赤い楕円は $\Scal_x$ です。楕円はプロンプトに向かって伸びているため、その間違った答えは到達するのが最も簡単であり、トレーニング後にはほとんどの経路がそこに到達します。
[図: Interactive figure: One prompt and a dashed circle of correct answers, shown before and after training. The verifier's region, a red ellipse, overlaps part of the circle and stretches back toward the prompt, so it accepts wrong answers on that side and rejects the correct answers it misses. With it, more walks land on wrong answers the verifier accepts than on correct ones; with an exact verifier, most walks land on correct answers.]
この視覚化は、RLをスケールアップすること(より多くのサンプルとより長いサンプルを生成すること)が、なぜ報酬ハッキングをより可能性が高くするのかも示しています。検証者をどれだけ広範囲に「ぼかす」かによって、誤って仕様化された $\Scal_x$ の領域に到達する経路を見つける可能性が高まります。モデルがこのショートカットを見つけると、トレーニングはそのモデルを同様の行動に向かわせます。
では、推論は何の目的を解決しているのか?#
推論は実際には何の目的も解決していません。経路が正しい答えに収束することは保証されません。成功確率が高くなるのは、経路の分布全体で集計した場合のみです。
RLトレーニングは、成功する経路を探索するために多くの計算リソースを費やし、その後モデルをそれらを再現するように訓練することによって、成功確率 $p_x$ を直接向上させます。「長く考える」ことによる進歩は、推論時ではなく、トレーニング中に費やされた探索によるものです。長い思考連鎖は、トレーニング中に成功と相関する傾向があり、それがモデルに訓練されます。サンプリングのコストは重みに償却されるため、推論時にはモデルは単一のサンプルで成功する可能性が高いですが、個々のサンプルが成功することは保証されません。
- Muennighoff et al. 2025 は、モデルが停止しようとするたびに「Wait」を追加することで、モデルに思考を続けさせました。精度は最初に上昇し、その後平坦化しました。↩︎
- Gema et al. 2025 が、長い推論が精度低下につながるタスクを見つけたという興味深い例があります。↩︎
- これは将来の記事で再検討するかもしれません。↩︎
- 高価な検索を使用してより良い答えを見つけ、その後ネットワークを訓練してその答えを安価に出力させるというこのアイデアは、言語モデル以外でも見られます。AlphaZero は自己対戦の各手でツリー検索を実行し、見つかった手を優先するようにポリシーネットワークを訓練します(expert iteration も同じアイデアです)。物理学者や天気予報士はシミュレーターを使用します。高価な手作りのシミュレーションを実行して答えを得てから、ニューラルネットワークを訓練してその答えを直接予測します。↩︎