MACRO: Markov Chain Routing of Transformer Layers

技術要約：MACRO – Transformerレイヤーのマルコフ連鎖ルーティング

問題提起
標準的な大規模言語モデル（LLM）は、最初のレイヤーから最後のレイヤーまで、固定された計算スケジュールに従ってトランスフォーマーブロックを逐次的に実行します。この静的なアプローチは、モデルの基礎となる重みを変更することなく、複雑なタスクに対する推論精度を向上させる可能性のある、動的な計算経路を無視しています。先行研究では、主に推論の加速（例：早期脱出、構造化されたレイヤー・ドロップ）を目的とした適応的深度が探索されてきましたが、凍結されたモデルの実行ルートを能動的に再構成して性能を高めることは、依然として困難な領域です。既存の動的ルーティング手法には大きなボトルネックが存在します。例えば、データ生成と学習のために高コストなオフラインのモンテカルロ木探索（MCTS）を必要とするもの（Dr.LLMなど）、推論時に正解ラベルを必要とするもの、あるいは計算コストが非常に高いサンプルごとの探索ループに依存するものがあります。

手法
著者らは、MACRO（Markov Chain Routing of Transformer Layers）を提案します。これは、凍結されたLLMアーキテクチャ上でタスク固有の実行ルートを学習するフレームワークです。MACROは、レイヤーの適用、スキップ、反復、および隠れ状態の加算のシーケンスを決定する、文脈依存のマルコフ方策 $\pi_\theta(a_t | s_t)$ としてレイヤールーティングをモデル化します。

• 状態と行動空間： 状態 $s_t$ は、現在のレイヤーインデックス（$\ell$）、残りの計算予算フェーズ（初期/中期/後期）、入力レイヤーの変位（$\delta$）、およびオペレータのコンテキスト（現在の隠れ状態がプレーンなレイヤーによって生成されたか、あるいは加算によるものか）を含むコンパクトな要約です。行動空間には以下が含まれます：

  1. ローカルレイヤー移動： 近傍のレイヤー $\ell + \Delta$ を適用する（スキップ、反復、または逆方向への再訪問を許容）。
  2. 加算と適用（Add-and-Apply）： 現在の隠れ状態 $h_t$ に、以前の隠れ状態 $h_{t-b}$ を係数 $\gamma > 0$ で加算してから、ローカルレイヤーを適用する。
  3. 再結合（Rejoin）： 学習されたプレフィックスを終了し、残りの標準的なサフィックス（接尾辞）を実行する。
     実現可能性マスク（Feasibility masks）により、サンプリングされたすべてのプログラムがレイヤーの範囲内および予算制約内に収まることを保証します。

• 学習手順： MACROは、分布推定アプローチを採用しています。前方バイアスを持つ事前分布（標準的な逐次パスが最も確率が高い状態）から出発し、フレームワークは候補となるルーティングプログラムを反復的にサンプリングし、それらを訓練データ・サブセット上で評価し、重み付き最大尤度を用いてマルコフ遷移行列を更新します。高性能なルートには指数関数的な重みが付けられ、方策を洗練させます。リプレイバッファは、過去のイテレーションからの例のカバー率を保持することで、安定した収束を保証します。

• デコーディングと選択： 確率的なサンプリングやテスト時ごとの探索に依存する手法とは異なり、MACROは学習された方策から最も確率の高い有効なルーティングプログラムを決定論的にデコードするために、トップk Viterbiアルゴリズムを使用します。構造的な実現可能性マスクがViterbi再帰に組み込まれています。最終的なルートは、保持された検証セット上でトップkの候補を評価し、最も精度の高いものを選択することで決定されます。

主な貢献

1. パラメータ効率の高いフレームワーク： MACROは、凍結されたおよび量子化されたLLMにおける適応的なレイヤールーティングのための、軽量で文脈依存のマルコフ方策を導入します。これにはモデルの重みの更新は必要ありません。
2. 厳密なデコーディングスキーム： 本論文は、構造的な実現可能性マスクを備えた厳密なトップk Viterbiデコーディングスキームを提示しています。これにより、サンプルごとのテスト時探索、補助的なニューラルネットワークの学習、または正解推論オーラクルを排除できます。
3. 実証的な優位性： 本フレームワークは、多様な推論および知識ベンチマークにおいて、非ルーティングのベースラインおよび既存の動的ルーティング手法に対して一貫した改善を示しています。

結果
GSM8K、MATH500、MedQA、MMLU-Proなどのベンチマークにおいて、複数のオープンウェイトLLM（Qwen3バリアント、Mixtral、DeepSeek-R1-Distillを含む）を用いて評価を行いました：

• 精度向上： MACROは、非ルーティングのベースラインに対して平均 +5.0% の精度向上を達成しました。この利得は、より小さなモデル（例：Qwen3-1.7Bで+12.68%）において最も顕著です。
• SOTAとの比較： MACROは、主要な動的ルーティング手法であるDr.LLMを、平均精度で +7.2% 上回っています。
• 効率性： ルート探索プロセスは大幅に効率化されており、Dr.LLMと比較して探索時間を 9.4倍 短縮しました（1ベンチマークあたり14.8時間から1.6時間へ）。
• 堅牢性： 統計的有意性テストにより、MACROが性能を低下させることなく、テストされたすべてのモデルにおいて性能を向上させることが確認されました。
• メカニズムの洞察： Logit-lens分析は、ルーティングがモデルの内部的な推論能力を変更するのではなく、すでに内部的に表現されている回答をより効果的に露出させ、出力付近での抑制を防ぐことで性能を向上させていることを示唆しています。

意義と主張
本論文は、MACROが凍結されたLLMにおいて動的計算を実用化するための重要な一歩であることを主張しています。ルーティングをマルコフ過程として定式化することで、著者らは従来のメソッド（MCTS駆動の学習など）に伴う計算コストとデータ要件を克服しました。本研究は、ルーティングが、モデルの学習された表現を根本的に変えるのではなく、内部に存在する回答をより効果的に「引き出す（elicit）」ためのメカニズムとして機能すると論じています。著者らは、このアプローチをさらに発展させ、入力依存の動的ルーティングを可能にすることで、LLMの柔軟なテスト時スケーリングの新たな次元を提供できると考えています。この手法は、小規模モデルに対する性能向上能力と、量子化モデルとの互換性において特に注目すべきであり、リソース制約のある環境における幅広い適用可能性を示唆しています。