💬 NLP

Efficient and Trainable Language Model Test-Time Scaling via Local Branch Routing

本論文は、ローカル・ルックアヘッド・ツリーを拡張し、軽量なルーターを用いて最適な分岐を選択することで、言語モデルの推論を効率的に強化するトークンレベルのテスト時スケーリング・フレームワークであるLocal Branch Routing（LBR）を導入するものであり、これによりエンドツーエンドの強化学習を可能にし、数学的推論タスクにおいて既存の離散的な思考の連鎖（Chain-of-Thought）およびソフトトークンのベースラインを凌駕する。

原著者： Yutong Yin, Mingyu Jin, Jin Pan, Changyi Yang, Zijie Xia, Dhruv Pai, Shuming Hu, Zhen Zhang, Chenyang Zhao, Jinman Zhao, Wujiang Xu, Raymond Li, Xin Eric Wang, Julian McAuley, Zhaoran Wang

公開日 2026-06-25

📖 1 分で読めます☕ さくっと読める

CC BY 4.0

原著者： Yutong Yin, Mingyu Jin, Jin Pan, Changyi Yang, Zijie Xia, Dhruv Pai, Shuming Hu, Zhen Zhang, Chenyang Zhao, Jinman Zhao, Wujiang Xu, Raymond Li, Xin Eric Wang, Julian McAuley, Zhaoran Wang

原論文は CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) でライセンスされています。 ✨ これは以下の論文のAI生成解説です。著者が執筆または承認したものではありません。技術的な正確性については原論文を参照してください。免責事項の全文を読む

大きな問題：思考が遅すぎる、あるいは視野が狭すぎる

あなたが非常に難しい数学の問題を解こうとしているか、複雑な旅行の計画を立てていると想像してください。そこには、賢いアシスタント（AI）があなたを助けてくれています。

現在、AIアシスタントは通常、次の2つのどちらかの方法で動作しています：

「一本道」の歩行者： 彼らはステップバイステップで考え、最初に浮かんだアイデアに即座にコミットします。もし早い段階で道を間違えると、他の選択肢を見ようとしなかったために、行き詰まってしまう可能性があります。
「全マップ」の探索者： 彼らはあらゆる可能な解決策の経路を一度にすべて書き出し、それらすべてをチェックして、最善のものを選ぼうとします。これは非常に正確ですが、図書館にあるすべての本を読んで一つの文章を見つけ出そうとするようなもので、膨大な時間と計算能力を必要とします。

この論文の著者たちは、「ゴールドリックス（適度な）」解決策、つまり、すべてをチェックすることに手間取りすぎることなく、より良い選択をするためにいくつかの異なる可能性を見る方法を見つけたいと考えました。

解決策：ローカル・ブランチ・ルーティング (LBR)

著者らは、ローカル・ブランチ・ルーティング (Local Branch Routing) と呼ばれる新しい手法を提案しています。これは、**「先読みしてから決める」**戦略だと考えてください。

その仕組みを、ハイカーがトレイル（登山道）を選ぶという例えを使って、ステップごとに説明します。

1. 「先読み」 (ツリーの成長)

次の単語（またはトレイルの目印）を即座に選ぶ代わりに、AIは一旦停止します。まるでそれが実在するかのように、次の数単語（または数個の目印）を想像します。

論文の用語： 小さなローカル・ルックアヘッド・ツリーを展開する (Expands a small local lookahead tree)
例え： あなたが道の分岐点に立っていると想像してください。ただ道を選ぶのではなく、ルートAに3歩、ルートBに3歩、ルートCに3歩、と素早く進んでみます。まだどれにもコミット（決定）はしていません。ただ、地形がどのようなものかを見るために、頭の中で「歩いて」みるのです。

2. 「ルーター」 (意思決定者)

AIはこれらの短い経路を「歩いた」後、その結果を見ます。そして問いかけます。「これらの短い経路のうち、どれが最も有望に見えるだろうか？」

論文の用語： 軽量なルーターを使用して、深さ1のサブツリーを選択する (Uses a lightweight router to select the depth-1 subtree)
例え： スマートなガイド（ルーター）が、あなたが探索した地形を見ます。「ルートAは崖に続いており、ルートBは沼地だが、ルートCは美しい草原につながっている」といった具合です。ガイドはルートCを指差して、「よし、正式にこれに決まりだ」と言います。

3. 「プルーニングとシフト」 (前進)

AIは、ルートCの最初のステップを正式に書き込みます。そして、ルートAやルートBのアイデアを捨て去ります（プルーニング）。次に、その最初のステップの終点に開始点を移動させ、プロセスを繰り返します。再び先読みを行い、次の最善のステップを選び、進んでいきます。

論文の用語： プルーン・シフト・グロウ・デコーディング・プロセス (Prune–shift–grow decoding process)
例え： あなたは草原へと続く道の最初の一歩を踏み出します。今、あなたは新しい場所にいます。再び先読みを行い、次の最善のステップを選び、歩き続けます。

なぜ他の手法よりも優れているのか？

論文では、これをAIの思考における他の2つの一般的な方法と比較しています。

対「離散的思考の連鎖 (Discrete Chain-of-Thought)」 (一本道の歩行者)：
- 問題点： 「一本道の歩行者」は、経路がどのようなものかを見る前に、どちらに進むかを決めなければなりません。それは、ドアを開ける前にドアを選ぶようなものです。
- LBRの利点： LBRは、決定を下す「前」にドアを開けます（経路を歩きます）。論文では、数歩歩いた後の「隠れ状態（パスの精神的なイメージ）」には、より良い決定を下すための貴重な手がかりが含まれていることが示されています。
対「ソフト・トークン・ブランチング (Soft-Token Branching)」 (ぼやけた混合物)：
- 問題点： 一部の手法は、すべての経路を「ぼやけた」平均として混ぜ合わせることで、一度にすべての経路を見ようとします。それは、3つの経路が重なり合った写真を見ているようなものです。どの単一の経路の詳細も明確に見ることができません。
- LBRの利点： LBRは、経路を**離散的（別々で明確）**に保ちます。LBRはルートAを歩み、次にルートBを歩み、次にルートCを歩み、それらを明確に比較します。論文では、経路を分離して保持することで、AIが「ぼやけた」混合物の中で失われてしまうような具体的な詳細（崖や草原など）を見ることができるようになることが判明しました。

結果：彼らは何を見出したのか？

著者らは、2種類のタスクでテストを行いました。

合成プランニング (作られたゲーム)： グラフ内をナビゲートしなければならないパズルを作成しました。その結果、LBRは、短い経路を歩くことで得られる「手がかり」を利用して正しい方向転換ができるため、はるかに優れた性能を示しました。
数学的推論 (実際の数学問題)： 難易度の高い数学ベンチマーク（競技会などで使われるもの）でLBRをテストしました。
- 結果： LBRは、標準的な「一本道の思考」法や「ぼやけた混合物」法よりも多くの問題を正解しました。
- 効率性： これを実現しながら、宇宙にあるすべての可能な解決策をチェックする必要はありませんでした。単にいくつかのローカルな選択肢をチェックし、賢い選択をし、そして進んだのです。

結論

ローカル・ブランチ・ルーティングは、AIに、決定を下す前に数歩先を覗き見ることができる「懐中電灯」を与えるようなものです。未来のすべてを見ようとする（それはコストがかかりすぎる）わけでも、盲目的に推測するわけでもありません。いくつかの短い可能性を調べ、それらを明確に比較し、最善のものを選ぶことで、AIは困難な推論問題を解くためのより賢く、より正確な存在になります。

技術要約：Local Branch Routingによる効率的かつ学習可能な言語モデルのテスト時スケーリング

問題提起

言語モデル（LLM）の推論における最近の進展は、推論時に追加の計算量を割り当てること（テスト時スケーリング）が性能を大幅に向上させることを示しています。しかし、既存のアプローチには決定的なトレードオフが存在します。

Chain-of-Thought (CoT) サンプリング: 効果的ではあるものの、標準的なCoTはシングルスレッドであり、代替案の結果を観察することなく、一つのトークンへのコミットを決定してしまいます。
探索ベースの手法: Tree of ThoughtsやMonte Carlo Tree Search (MCTS) のような手法は、中間的な思考のツリーを探索することで幅を導入します。しかし、これらは多くの場合、粒度が粗く、計算コストが高く、単一の学習可能なデコーディング・ポリシーとしてエンドツーエンドで最適化することが困難です。
ソフトトークン手法: Soft ThinkingやMultiplex Thinkingのような手法は、候補となる埋め込みを連続的な混合体へとマージすることで、トークンレベルの幅を提供します。これらは学習可能ですが、候補となるブランチの離散的なアイデンティティを曖昧にし、概念固有の情報を不明瞭にする可能性があります。

本論文は、言語モデルが、完全な解レベルの探索に伴うコストやソフトな混合による曖昧さを回避しつつ、ローカルな将来の候補の隠れ状態に基づいて意思決定を行うことができる、軽量なトークンレベルのテスト時スケーリングメカニズムから恩恵を受ける可能性があると主張しています。

手法：Local Branch Routing (LBR)

著者らは、Local Branch Routing (LBR) という、学習可能なトークンレベルの分岐およびルーティング（branch-and-route）デコーディング・フレームワークを提案しています。LBRは、標準的な自己回帰生成に、ローリング式のローカル・ルックアヘッド・ツリーを拡張します。

デコーディング・パイプライン

各デコーディングステップ $t$ において、コミットされた接頭辞 $x_{<t}$ が与えられたとき、LBRは4段階の「prune–shift–grow（枝刈り・シフト・成長）」プロセスを実行します。

Grow (成長): 単一の次トークンをサンプリングする代わりに、LBRは $x_{<t}$ を根とする幅 $K$ 、深さ $L$ のローカルツリー $\mathcal{A}_t$ を展開します。サンプリングされたすべてのノードは、その隠れ状態 $h_t(v)$ を計算するために言語モデルへと**フォワード（前方伝播）**されます。これにより、単なる候補文字列の集合ではなく、ローカルな隠れ状態のツリーが作成されます。
Route (経路選択): 軽量なルーター $\rho_\phi$ がフォワードされたツリーを観察し、コミットすべき深さ1のサブツリーを一つ選択します。ルーターは、セット・アテンション・メカニズムを使用して、各候補サブツリー（例えば、 $L=1$ の場合は単一トークン、 $L>1$ の場合はトークンとその継続）をベクトルへとエンコードします。これらのベクトルはクロス・サブツリー・アテンションを介して比較され、相対的なスコアを生成することで、ルーターが兄弟関係にある将来の選択肢間で対照的な決定を下せるようにします。
Prune (枝刈り): ルーターは最良のサブツリー（インデックス $k^*_t$ ）を選択します。他のすべてのサブツリーは破棄されます。選択されたサブツリーのルート・トークンが、コミットされた次のトークン $x_t$ となります。
Shift and Regrow (シフトと再成長): 選択されたサブツリーは前方へとシフトされます（そのルートが除去され、その子孫が次のステップの新しい接頭辞となります）。そして、ツリーの深さ $L$ を回復するために新しいレイヤーが成長させられ、サイクルが繰り返されます。

学習目的関数

LBRは、検証可能な報酬（RLVR）を用いたエンドツーエンドの強化学習を可能にする、扱いやすいツリー・トラジェトリー（木構造の軌跡）尤度を定義します。

確率性: 2つの操作のみが確率的です。新たに成長したノード $G_t$ をベースモデルからサンプリングすることと、ルーターの決定 $k^*_t$ をサンプリングすることです。枝刈り、シフト、およびノードの再利用は決定論的です。
尤度: 軌跡 $F$ の対数確率は、新しく成長したトークンの対数確率と、ルーターの選択確率の総和です。
$\log p_{\theta,\phi}(F | q) = \sum_{t=1}^T \sum_{v \in G_t} \log \tilde{\pi}_\theta(\text{tok}(v) | \text{ctx}(v)) + \sum_{t=1}^T \log \rho_\phi(k^*_t | x_{<t}, \mathcal{A}_t; \theta)$
最適化: このフレームワークは、検証器（verifier）がベースモデルのトークン生成とルーターの選択決定の両方にアドバンテージ $A(F)$ を与える尤度比目的関数（例：GRPO）を用いて学習されます。

主な貢献

離散的アイデンティティを持つトークンレベルの分岐: 候補を連続的な混合体へとマージするソフト・トークン手法とは異なり、LBRは候補を離散的なフォワード・ブランチとして保持します。これにより、ルーターは候補生成後の隠れ状態——すなわち、候補トークンが生成され、モデルによって処理された後に初めて明らかになる情報——を利用することが可能になります。
扱いやすいエンドツーエンド学習: ツリーの軌跡に対する特定の尤度を定義することで、LBRは、離散的なトークン推論システムで使用されるものと同じRLVRの原理の下で、ベースモデルとルーターを共同で最適化することを可能にします。
効率的なテスト時スケーリング: 本フレームワークは、浅いルックアヘッド（ $L$ は小さく、例えば1または2）を伴うトークンレベルの幅を導入しており、完全な深さの探索に伴う計算オーバーヘッドを回避しつつ、ルートの次トークン分布を超えた証拠を活用できます。
セット・アテンション・ルーター: 著者らは、クロス・サブツリー・アテンションを使用して兄弟関係にある候補を比較するルーター・アーキテクチャを導入し、ルーティングが独立した再ランキングではなく、対照的な決定であることを示しました。

実験結果

合成的な階層的プランニング

ラジックス変換されたグラフ到達可能性に関する合成タスクにおいて：

性能: LBRは最高の目標精度（0.967）を達成し、Discrete CoT（0.914）およびSoft Thinking（0.845）を上回りました。
メカニズム分析: プロービング実験により、LBRにおける候補生成後の隠れ状態は、分岐前の状態よりも最終的なターゲットを大幅に予測しやすいことが明らかになりました。これは、候補をフォワードすることで、標準的な自己回帰デコーディングでは利用不可能なダウンストリームの情報にアクセスできることを裏付けています。
表現: 連続的な混合によって概念のアイデンティティを曖昧にするSoft Thinkingとは異なり、LBRは離散的なブランチのアイデンティティを保持し、正しいパスの復元性を高めていました。

数学的推論ベンチマーク

DeepSeek-R1-Distill-Qwenモデル（1.5Bおよび7B）を用い、6つのベンチマーク（Minerva, AIME, MATH500など）で評価を行いました。

性能向上: LBRは、以下に対してPass@1およびPass@32の両方で性能を向上させました。
- Discrete Chain-of-Thought
- バニラの離散トークンRLVR
- RL互換のソフトトークン分岐ベースライン（例：Multiplex Thinking）
深さのアブレーション: ルックアヘッドの深さを $L=1$ から $L=2$ に増やすことで、特に7Bバックボーンにおいて性能がさらに向上しました。これは、より深いローカルな継続がより多くのルーティング証拠を提供することを示唆しています。
ルーターのアブレーション: クロス・サブツリー・アテンションを除去した場合（独立したルーターを使用した場合）、性能が低下しました。これは、兄弟関係にあるブランチの対照的な比較が極めて重要であることを検証しています。
学習ダイナミクス: RLVR学習中、ルーターがより決定的になるにつれてルーターのエントロピーは減少し、同時にルーティングされた位置におけるベースモデルのエントロピーも減少しました。これは、モデルがより識別可能なローカル・ブランチを生成するように学習するという共適応（co-adaptation）を示しています。

意義と主張

本論文は、Local Branch Routing が、効率的で学習可能、かつ離散的な形式の言語モデルのテスト時スケーリングを提供すると主張しています。その主な意義は、以下の点にあります。

候補生成後の証拠は価値がある: ルーティングの決定は、単なる分岐前の分布だけでなく、候補トークンによって誘発される隠れ状態から恩恵を受ける。
離散的分岐はソフトな混合よりも優れている: 離散的なブランチのアイデンティティを保持することで、概念のアイデンティティを不明瞭にする可能性があるソフト・トークンの混合と比較して、より明確な意思決定が可能になる。
スケーラビリティ: 本フレームワークは、シングルスレッドのCoTと高コストなフルツリー探索の中間領域を提供し、軽量なローカル分岐を通じて効果的なスケーリングを可能にする。

著者らは、**局所性（locality）**という制限についても認めています。LBRは浅いルックアヘッド・ウィンドウに制約されており、完全な解に対するグローバルな探索は行いません。今後の課題として、適応的な深さや選択的なグローバル探索を組み合わせることが示唆されています。