技術要約: MUX - 多重化されたトークンによる連続的推論

1. 問題提起

現代の自己回帰型言語モデルは、複雑な問題（数学、コーディング、常識）を解決するために、自然言語による中間的な推論ステップ（Chain-of-Thought）を明示することで解決を図っています。しかし、この離散的なアプローチは計算効率が低いです。各推論ステップは単一のサブワードしか伝達できず、厳格な帯域幅のボトルネックを生み出します。さらに、多くの推論トークンは、最適化された計算というよりも、人間のコミュニケーションパターンを模倣した冗長なものとなっています。

連続的な潜在表現による推論（Continuous latent reasoning）は、モデルがサブワードの代わりに連続的なベクトルを予測する代替案を提供します。これは、重ね合わせ（superposition）を通じて複数のサブワードを同時にエンコードすることを可能にし、理論的にはより高い帯域幅を実現できます。しかし、既存の潜在推論手法には大きな課題があります。

• グローバルな手法（最終回答や軌跡の終点のみを監督する手法）は、「ショートカット行動」を招きやすく、潜在トークンが情報の乏しいプレースホルダーになってしまう傾向があります。
• ローカルな手法（各潜在トークンを離散的な推論のスパンに対して監督する手法）は、複雑な補助アーキテクチャ（例：自己回帰デコーダ）を必要とするか、あるいは探索タスクに必要な多様な仮説を維持できないことがよくあります。

核心となる課題は、離散的な推論に対して忠実であり、潜在的な崩壊を防ぎ、かつ過度なアーキテクチャの複雑化なしに並列的な探索をサポートできる、連続的な潜在推論のための監督対象（supervision target）を定義することです。

2. 手法: MUX

著者らは、多重化されたトークン（multiplexed tokens）を用いて、言語モデルに連続的な潜在推論を行わせるためのローカル蒸留手法であるMUXを提案しています。

コア概念

MUXは、ローカル蒸着タスクを、可変長のカテゴリ信号（離散的な推論ステップのスパン）をエンコードする固定次元の連続信号を学習することとして扱います。通信システムにおける符号分割多重化（code-division multiplexing）から着想を得て、MUXは離散的なサブワードのスパンを、語彙空間におけるone-hotエンコーディングの加重線形重ね合わせとして表現します。

技術的実装

1. 多重化されたターゲット: 離散的な推論スパン $r_i = (r^1_i, \dots, r^{S_i}_i)$ に対して、学習ターゲットは以下のように構築されます：
   $$\text{mux}(r_i) := \sum_{j=1}^{S_i} \alpha^{(i)}_j \cdot \text{onehot}(r^j_i)$$
   ここで、$\alpha^{(i)}_j$ は正規化された正の位数的重みです。
2. 位数的重み付け: 重ね合わせが**ロスレス（lossless）**であること（すなわち、元のスパンを一意に復元できること）を保証するために、著者らは特定の重み付けスキームを提案しています：
   • 幾何学的（Geometric）: $w_j = \rho^{j-1}$ （指数関数的減衰）。
   • 正弦波的（Sinusoidal）: スパンの終点付近でピークを持つ正弦関数に基づく。
   • 回転的（Rotary）: 回転位置エンコーディングに類似した、平均余弦成分に基づく。
   • 理論的保証: これらの重み付けは「部分和分離条件（subset-sum separation condition）」を満たし、異なる位置の集合が任意のサブワードに対して異なる総質量を生み出すことを保証します。これにより、サブワードが複数回出現する場合でも曖昧さを防ぎます。
3. 学習目的関数: モデルは、線形ソフトマックスヘッドとKLダイバージェンス損失（$L_{local}$）を用いて、これらの多重化されたターゲットに適合するように訓練されます。これは、標準的な回答予測損失（$L_{answer}$）および、隠れ特徴量を離散的な推論モデルに整合させるためのオプションの軌跡レベルの損失（$L_{global}$）と組み合わされます。
4. 推論: 推論時、モデルは隠れ空間において潜在トークンを自己回帰的に生成します。語彙投影は、訓練中の監督のためにのみ使用され、最終的な出力トークンの生成には使用されません。

3. 主な貢献

(1) 多重化されたトークンによる潜在推論

MUXは、潜在トークンが加重線形重ね合わせを通じて離散的な推論のスパンを表現する、シンプルなローカル蒸留フレームワークを導入します。従来のローカル手法とは異なり、補助的な自己回帰デコーダや圧縮キャッシュターゲットを必要とせず、語彙空間への投影に依存しています。

(2) ロスレスな多重化

本論文は、特定の位数的重み付け（幾何学的、正弦波的、回転的）がロスレスな多重化を保証するという理論的証明を提供しています。これは、離散的な推論スパンが連続的な表現の中に完全に保持されることを意味します。著者らは、このロスレス性が、グローバルな潜在推論手法でよく見られる失敗モードである潜在的な崩壊（latent collapse）（トークンが意味的に均質化してしまう現象）を防ぐことを証明しています。

(3) 重ね合わせによる並列探索

理論的に、著者らは多重化されたトークンが自然に**並列幅優先探索（BFS）**を実装できることを示しています。トークンが重ね合わされているため、単一の潜在トークンは、複数の仮説（探索のフロントライン）全体を同時に表現し、更新することができます。これにより、モデルは直列的な監督フレームワークの中で、複数の問題解決パスを並列に探索することが可能になります。

(4) 実証的な優位性

32の評価設定（4つの言語モデル、2つの訓練コーパス、4つのテストセットにわたる）において、MUXは離散的および連続的な推論の両方で強力なベースラインを上回りました。特筆すべきは、MUXが蒸留元である離散的なChain-of-Thought（SFT-CoT）のベースラインを、大幅に少ないトークン数（約2.4倍から5.9倍少ない）でありながら、しばしば凌駕している点です。

4. 実験結果

• 数学的推論: GSM8K-AUG, GSM8K-AUG-NL, SVAMP, GSM-Hard, および MultiArith において、MUXは GPT-2, LLaMA 3.2 1B, 3B, 3.1 8B の各モデルにおいて最高の潜在推論性能を達成しました。MUXは頻繁に SFT-CoT を上回り、コンパクトな重ね合わせ推論が元の離散的なトレースよりも優れた汎化性能を発揮できることを示しました。
• 探索ベンチマーク: MNNS（Minimum Nonnegative Sum）および Game of 24 タスクにおいて、MUXはベースライン（Coconut や CoT2 を含む）を大幅に上回り、多重化されたトークンが並列探索能力を可能にするという理論的主張を裏付けました。
• アブレーション研究:
  • ローカル損失: グローバル損失（$\gamma=0$）を取り除き、多重化されたローカル損失のみに依存した場合でも、SIM-CoTと比較して優れた結果が得られ、ローカルターゲットの有効性が確認されました。
  • 重み付け: ロスレスな重み付け（幾何学的、正弦波的、回転的）は、一貫してロスのある一様重み付けよりも優れた性能を示しましたが、高度に構造化されたデータセットではその差は緩やかでした。
  • チャンキング: 離散的なスパンを潜在トークンへとランダムにチャンキングする方法が最も優れた性能を示し、これが構造化されたデータ拡張として機能しました。

5. 意義と主張

本論文は、ロスレスな重ね合わせをローカルな学習ターゲットとして用いることが、強力で効率的、かつ解釈可能な連続的な潜在推論を実現するための十分条件であると主張しています。

• 効率性: MUXは、離散的なCoTよりも少ない推論ステップで高い性能を達成し、計算的な帯域幅のボトルネックに対処しています。
• 簡潔さ: 本手法は、補助デコーダのような従来のローカル蒸留アプローチのアーキテクチャ上の複雑さを回避しながら、性能を維持または向上させています。
• 解釈可能性: プロービング分析により、学習された潜在トークンが忠実で解釈可能な推論内容をエンコードしていることが確認されました。これは、潜在推論が単なる解釈不能なプレースホルダーであるという懸念に反するものです。
• 並列性: 結果は、多重化されたターゲットを使用することで、直列的な監督から並列的な探索ダイナミクスが自然に創発できることを示唆しており、言語モデルにおけるプランニングと探索を向上させるための新しい経路を提示しています。

著者らは、シンプルで解釈可能なローカルターゲットを用いることで、潜在推論をより強力かつ訓練しやすくできると結論付けており、連続空間の効率性と離散的な推論の忠実性の間の溝を埋めています。