技術要約：言語モデルにおける潜在的な推論戦略の解明

問題設定

推論タスクに対して事前学習された言語モデル（LM）は、多くの場合、複数の異なる高レベルな戦略（例：背理法と帰納法、あるいは異なるソートアルゴリズムなど）を用いて問題を解決します。しかし、これらの戦略はモデルの応答分布 $p_\theta(y|x)$ の中に暗黙的に、かつ絡み合った状態で存在しています。生成中にどの戦略が使用されているかを特定したり制御したりするための、明示的な変数は存在しません。

本論文は、この応答分布を構造化された戦略条件付き表現へと分解する問題を定式化しています。目標は、ルーター・ジェネレーター分解 $(r_\phi(z|x), g_\phi(y|x, z))$ を学習することです。その要件は以下の通りです：

1. 周辺的な忠実度（Marginal Fidelity）： 統合された分布 $\int r_\phi(z|x)g_\phi(y|x, z) dz$ が、元のベースモデルの分布 $p_\theta(y|x)$ を近似すること。
2. 戦略のアライメント（Strategy Alignment）： 潜在変数 $z$ が情報量豊かで戦略にアライメントされていること。つまり、$z$ を変化させると高レベルな戦略が変化し、かつ同じ潜在値が異なる入力に対しても一貫した戦略的意味を持つこと。

核心となる課題： 標準的な潜在変数モデリングとは異なり、本研究の対象となるベースモデル $p_\theta$ は、潜在変数を伴わずにすでにターゲット分布を捉えています。そのため、潜在変数 $z$ を無視する分解（すなわち、すべての $z$ に対して $g_\phi(y|x, z) = p_\theta(y|x)$ となるもの）は、初期状態において利用可能であり、かつ周辺適合性のみを測定する目的関数に対してグローバルな最適解となります。結果として、標準的な変分推論（例：条件付きELBO）はこの設定において**後方分布の崩壊（posterior collapse）**に陥り、$z$ が戦略レベルの変動をエンコードすることを促すことができません。

手法

アーキテクチャ

著者らは、事前学習済みTransformerをルーター・ジェネレーター・フレームワークへと適応させています。

• ルーター ($r_\phi$): 入力 $x$ を読み取り、連続的な潜在戦略 $z$（対角ガウス分布としてモデル化）上の分布を出力します。
• ジェネレーター ($g_\phi$): $x$ と、サンプリングされた潜在変数の投影埋め込み $E(z)$ に条件付けられた自己回帰モデルです。
• 後方分布 ($q_\xi$): 学習時のみ使用され、完全な入力・応答ペア $(x, y)$ から $z$ を推論するために用いられます。
• パラメータ化: ルーターとジェネレーターはベースモデル $\theta$ から初期化されたパラメータを共有し、LoRA形式の低ランク適応を通じて更新されます。

モデル指向の再構成目的関数

標準的なELBOの過小指定問題を克服するため、著者らは、固定されたベースモデル $p_\theta$ を参照として用いて、潜在変数 $z$ がどの変動を説明すべきかを指示するモデル指向の再構成目的関数を提案しています。

目的関数 $J_\theta$ は、再構成項 $R_\theta$ とKLペナルティで構成されます：
$$J_\theta(x, y; \phi, \xi) = R_\theta(x, y; \phi, \xi) + \beta \text{KL}(q_\xi(\cdot|x, y) \parallel r_\phi(\cdot|x))$$

再構成項の主要な要素：

1. ベースモデルによる正規化: 再構成損失は、ベースモデルの期待応答損失（$c_\theta$）によって正規化され、「ベースモデルの参照単位」で表現されます。
2. 驚愕度指向の重み付け（Surprisal-Directed Weighting）: トークンへの重み付けを一様にするのではなく、ベースモデルのトークンレベルの驚愕度（$b_{\theta,t}$）によって重み付けを行います。
   • $b_{\theta,t}(x, y) = -\log p_\theta(y_t | x, y_{<t})$
   • 高い驚愕度を持つトークン（ベースモデルが不確実になる分岐点）には、より高い再構成圧力がかかります。
   • これにより、$z$ が、ベースモデルが複数の継続可能性を残している位置（通常、戦略の選択点に対応する）における変動をエンコードするように促されます。

再構成項は、ベースモデルの損失の、重み付けされた正規化された分数として定義されます：
$$R_\theta = \frac{1}{c_\theta} \mathbb{E}_{z \sim q_\xi} \left[ \sum_{t} a_{\theta,t}(x, y) [-\log g_\phi(y_t | x, z, y_{<t})] \right]$$
ここで、重み $a_{\theta,t}$ は、一様な圧力と驚愕度指向の圧力の間を補間します。

実験設定とベンチマーク

著者らは、戦略のアライメントを測定可能にするため、マルチ戦略アルゴリズムタスクの制御されたベンチマークを導入しました。

• タスク: リストの総和、ソートアルゴリズム、グリッド経路探索、一次方程式の解法、基数変換、および多桁加算。
• 設計: 各タスクは、同一の問題インスタンスに対して複数の有効な解のトレース（戦略）を許容します。正解の戦略ラベルは生成されますが、学習中は隠匿され、評価の際にのみ使用されます。
• 評価指標:
  1. 分布的忠実度（Distributional Fidelity）: 生成されたトレースのうち、タスクの有効な解として残っている割合。
  2. 戦略のアライメント（類推的一貫性 / Analogical Consistency）: ソース入力 $X$ から潜在変数 $z$ をサンプリングし、それを関連するターゲット入力 $X'$ に適用した際、得られる解のトレース $Y'$ が、$X$ から生成された $Y$ と同じ戦略を使用する確率。

主な結果

1. 暗黙的な情報の存在: ベースモデルの隠れ状態に対する線形プローブは、戦略ラベルをデコード可能であり、分解を行う前からターゲットとなる構造が暗黙的に存在していることを確認しました。
2. 標準的なELBOの失敗: 標準的な変分目的関数は、分布的忠実度は維持しますが、**潜在的な崩壊（latent collapse）**を引き起こし、$z$ は不活性なままとなり、戦略とのアライメントに失敗します。
3. モデル指向の再構成の成功: 提案された目的関数は、以下の特性を持つ潜在コードの回収に成功しました：
   • 高い分布的忠実度を維持する（ベースモデルの有効な出力を保持する）。
   • 高い類推的一貫性を達成する（$z$ が異なる入力間でも再利用可能な戦略的意味を運んでいることを示す）。
   • 学習が進むにつれ、潜在空間内でタスク固有および戦略固有の領域へと分離される。
4. アブレーション研究:
   • グローバルスケール（正規化）またはトークン重み付け（驚愕度の方向）を除去すると、戦略のアライメントが低下します。
   • 逆驚愕度による重み付け（低驚愕度のトークンを強調する）は、アライメントを弱め、戦略の変動を回収するためには、不確実な／分岐するトークンに圧力をかける必要があることを裏付けています。

理論的分析

本論文は、回復問題に関する理論的分析を開始しています。目的関数を、相互情報量 $I(Y; Z|X)$ から近似複雑度ペナルティを差し引いたプロファイルスコアを最大化するものとして定式化しています。著者らは、最適解と劣解の間の「回復ディスクレパンシー（recovery discrepancy）」が正である場合に限り、複雑度ペナルティ付きの変分目的関数が、ターゲット戦略構造 $S^*$ にアライメントされた潜在分解を選択するという十分条件を提示しています。

意義と主張

本論文は、ゼロからの学習ではなく、既存のモデルを適応させる際に特有の、病理的な後方分布の崩壊に対処していると主張しています。モデル指向の再構成目的関数を導入することで、著者らは以下の手法を実現しました：

• 暗黙的な推論戦略を、明示的で制御可能な潜在変数へと**分離（Disentangle）**する。
• ベースモデルの応答分布を保持しながら、潜在変数を情報豊かなものにする。
• 介入の実現: 回復された潜在構造により、ベースモデルを再学習することなく、特定の高レベルな戦略（例：特定の証明手法の強制）に向けて生成を誘導することが可能になります。

著者らは、本手法が正解の戦略を持つ合成アルゴリズムタスクによって検証されていることを述べています。また、理論的なオラクル選択基準と実際の学習ダイナミクスとの関連性を完全に特徴付けることは、今後の課題として、オープンエンドな推論タスクへの拡張を含めて明示的に残されています。言及されている潜在的な応用分野には、モデルの解釈可能性、制御可能性、および強化学習のポストトレーニングにおける探索の方向付けなどが含まれます。