AI-Papers
- 因果的自己注意を勾配の通り道として使い、報酬で重み付けした勾配を連続潜在状態へ直接割り当てるテスト時推論手法です
- 5つのLLMと3つの推論ベンチマークで平均64.5%を達成し、CoTを6.6ポイント、既存最強のLatentSeekを2.4ポイント上回りました
- 学習率への頑健性が標準偏差1.53から0.82へ改善し、潜在状態の影響が推論の接続語トークンに集中する解釈も示しました
研究の背景
大規模言語モデル（Large Language Model、LLM）の推論力を引き出す代表的な方法が、思考の連鎖（Chain-of-Thought、CoT）です。CoTは中間的な考えをテキストのトークンとして順に書き出すことで、複雑な問題を段階的に解かせます。ただしこのやり方には見過ごせない弱点があります。
中間推論を一度トークンとしてデコードすると、その値はとびとびの離散的なものになり、後から勾配で最適化しにくくなります。つまり「どの考えが正解にどれだけ貢献したか」を数値的にさかのぼって調整する経路が、途中で途切れてしまうのです。テスト時にサンプリング回数を増やす戦略（Sample More, Reflect Less）など多くの工夫が登場してきましたが、追加学習なしで推論の内部表現そのものを磨き込む方法には課題が残っていました。
先行手法のLatentSeekは、トークンではなく連続的な潜在状態を最適化するという着想を持ち込みました。しかし勾配の割り当て方が単純だったため、学習率などの設定にやや敏感で、性能が安定しにくいという弱点を抱えていたのです。GradCuitはこの割り当ての仕組みを根本から見直しています。
GradCuitの仕組み
GradCuitの核心は、Transformerの因果的自己注意（各トークンが直前までの情報だけを参照する仕組み）を「勾配を運ぶ回路」として捉え直した点にあります。手順としてはまず、選んだTransformer層に最適化可能な連続潜在状態を差し込みます。挿入されたこの潜在状態は、プロンプトの表現と生成される続きの間に置かれる、追加の作業メモリのような存在です。
自己注意を通すことで、生成されるすべての続きトークンの対数確率が、それ以前のすべての潜在状態へと微分可能な経路でつながります。この経路があるおかげで、生成結果の良し悪しを表す報酬で重み付けした勾配を、続き全体から潜在状態へさかのぼって直接届けることができます。
更新は、現在の潜在状態に報酬勾配を少しずつ足し込んでいくシンプルな方策勾配の形で行われ、推論の途中で潜在状態を繰り返し磨き込みます。重要なのは、この過程でモデル本体の重みを一切書き換えないことです。そのため追加学習が不要で、既存の学習済みモデルにそのまま適用できます。
従来のCoTが「言葉に翻訳してから考え直す」のに対し、GradCuitは「言葉になる前の内部表現のまま考えを研ぎ澄ます」と言い換えられます。離散化による情報の目減りを避けられる点が、精度と安定性の両面で効いてきます。
実験結果
著者らは5つの指示チューニング済みLLMバックボーンと3つの推論ベンチマーク、2種類の解答形式で評価を行いました。その結果、GradCuitは平均64.5%の正答率を達成しています。これは標準的なCoTプロンプトを6.6ポイント、当時最強の潜在推論手法であったLatentSeekを2.4ポイント上回る数字です。
さらに評価されているのが、設定の変化に対する頑健さです。7種類の学習率で試したところ、GradCuitは正答率の標準偏差を1.53から0.82へと下げ、LatentSeekよりも安定した挙動を見せました。更新方向をランダムに選ぶ簡易版でさえLatentSeekと肩を並べる場面があり、勾配クレジット割当という枠組み自体が効いていることがうかがえます。
解釈可能性と最適な挿入層
GradCuitは精度だけでなく、内部で何が起きているのかを読み解く手がかりも提供します。トークン単位で勾配の寄与を調べると、潜在状態の影響は「したがって」「なぜなら」といった推論の接続語トークンに集中していました。これは、潜在状態が論理の流れをつなぐ節目を重点的に調整していることを示唆します。
もう一つの重要な知見が、潜在状態を挿入する層の選び方です。層の深さを変えて調べたところ、序盤から中盤のTransformer層が最適な最適化空間として浮かび上がりました。深すぎる層では後段のブロックが少なく、潜在状態の調整がそのまま効果を伝播させて出力へ波及しにくくなるためだと考えられます。
まとめと今後の展望
GradCuitは、自己注意を勾配の通り道と見なすことで、報酬を潜在状態へきれいに割り当てるテスト時推論手法です。追加学習なしで幅広いモデルに適用でき、CoTや既存の潜在推論手法を上回る精度と、設定への高い頑健性を両立しました。解釈可能性の分析まで踏み込んでいる点も、実務での信頼につながります。
一方で課題も残ります。報酬をどう定義するかで性能が左右される点や、推論のたびに反復最適化を行うため追加の計算コストが生じる点は、実運用では無視できません。また評価は主に数学や科学系の推論ベンチマークに寄っており、より多様なタスクでの一般化は今後の検証課題でしょう。
とはいえ、離散トークンを介さずにモデルの内部表現を直接磨くという発想は、テスト時スケーリングの新しい方向性を示しています。勾配クレジット割当の考え方は、他の潜在推論手法や強化学習ベースの手法と組み合わせる余地も大きく、今後の発展が見込まれます。
