技術要約：SuperThoughts: 重ね合わせ状態における推論トークン

問題提起

大規模言語モデル（LLM）は、明示的な思考の連鎖（Chain-of-Thought: CoT）シーケンスを生成することで複雑な問題を解決する。ここでは、各トークンが計算の単位（1回のフォワードパス）を表す。OpenAI o1やDeepSeek-R1のような最新のモデルは、テスト時計算量（test-time compute）を増やすことで性能が大幅に向上することを示しているが、この逐次的なトークン生成は計算コストが高い。

連続的な潜在空間で推論を行おうとする既存のアプローチ（COCONUTや潜在推論など）は、離散的なトークン生成を回避しようとする一方で、決定的な課題に直面している。それは「中間的な教師あり学習（supervision）の欠如」である。標準的なCoTの学習は、すべてのステップにおいてトークンレベルのクロスエントロピー損失に依存しており、密な勾配信号を提供する。しかし、推論が制約のない潜在空間で行われる場合、この教師あり信号が消失するため、学習の不安定化、表現のドリフト、および長期的なタスクにおける性能低下を招く。逆に、単にコンテキストを圧縮したり、ファインチューニングによって離散トークン数を削減したりする方法は、必要なフォワードパスの回数自体を減らすことはできず、推論のスループット向上という観点では限界がある。

本研究が取り組む中心的な問いは、**「モデルは、離散的なトークンの教師あり信号を維持したまま、より豊かな重ね合わせ状態（superposed）の潜在空間で推論するように学習できるか？」**である。

手法

著者らは、連続する一対のCoTトークンを単一の潜在表現へと圧縮し、モデルが1ステップで2つのトークンをデコードできるようにするフレームスワーク、SuperThoughtsを提案する。このアプローチは、長さ $L$ の推論チェーンに必要なフォワードパスの回数を $L/2$ ステップへと半減させると同時に、離散的なトークンによる教師あり学習を維持する。

アーキテクチャ

モデルは主に3つのコンポーネントで構成される：

1. Compressor（圧縮器）: トークン埋め込みのペア $(c_{2i-1}, c_{2i})$ を単一の潜在ベクトル $x_i$ にエンコードする。著者らは2つの実装を検討している：
   • 線形投影（Linear Projection）: 埋め込みを結合し、学習可能な行列を通じて投影する。
   • Transformerブロック: ペアを小さなTransformer層で処理し、2番目のトークンの隠れ状態を抽出する。
2. Main Module（メインモジュール）: 主要な推論バックボーン（事前学習済みLLMから初期化）。これは圧縮されたベクトル $x_i$ を処理して隠れ状態 $h_i$ を生成し、次の奇数インデックスのトークン ($c_{2i+1}$) を予測する。
3. Multi-Token Prediction (MTP) Module（マルチトークン予測モジュール）: 次の偶数インデックスのトークン ($c_{2i+2}$) を予測する軽量なモジュール。これは、直前の偶数トークン、新たに予測された奇数トークン、およびMainモジュールの隠れ状態を入力として受け取り、それらをベクトル $x'_i$ へ投影して次の予測を行う。両モジュールは同じ出力言語モデルヘッドを共有する。

学習戦略

潜在的な推論に伴う不安定性に対処するため、2段階の学習プロトコルを採用している：

1. 潜在蒸留（Stage 1: Latent Distillation）: Compressorを、圧縮された潜在空間を事前学習済みの言語多様体（language manifold）に適合させるように学習させる。「教師」モデルは離散的なトークンシーケンスを処理し、「生徒」モデルは圧縮されたペアを処理する。Compressorは、生徒の隠れ状態（$x_i$ を処理した後）と、対応するトークンペアを処理した後の教師の隠れ状態との間のSmooth-L1距離を最小化するように最適化される。これにより、圧縮された表現が元の離散シーケンスの推論状態を保持することが保証される。
2. 結合エンドツーエンド学習（Stage 2: Joint End-to-End Training）: Mainモジュール、MTPモジュール、およびCompressorを、予測されるすべてのトークン（MainおよびMTPの両方の出力）に対する標準的なクロスエントロピー損失を用いて共同で学習させる。

適応的推論（Adaptive Inference）

すべての推論ステップが等しく困難ではないことを考慮し、著者らは信頼度ベースの適応メカニズムを導入している。各ステップにおいて、MTPモジュールの予測信頼度（最大確率）が閾値 $\tau$ を下回る場合、モデルはMTPの予測を拒否する。その後、そのステップに対して標準的なデコーディングへとフォールバックし、単一のトークンを直接Mainモジュールに供給する。これにより、モデルは「難しい」トークンにはより多くの計算資源を割り当て、「容易な」トークンでは効率性を維持することができる。

主な貢献

1. SuperThoughtsアーキテクチャ: トークンペアを重ね合わせ状態の埋め込みへと融合し、MainモジュールとMTPモジュールを用いて2つのトークンを1ステップでデコードすることで、必要なフォワードパスを実質的に半減させる新しいフレームワーク。
2. 2段階学習プロトコル: 離散的な教師あり学習の前に蒸留を通じて潜在空間を整合させることで、潜在的な推論を安定化させる手法。
3. 適応的推論メカニズム: 不確実な推論ステップにおいて、スループットと精度のトレードオフを行う動的なフォールバックシステムであり、複雑なタスクにおける性能低下を防ぐ。
4. 実証的評価: 数学的および科学的推論ベンチマーク（MATH500, AMC23, OlympiadBench, GPQA-Diamond）における包括的な評価により、精度を最小限の低下に抑えつつ、大幅なCoTの長さ削減を実現した。

結果

著者らはQwen2.5-Mathモデル（1.5B, 7B, 14Bの各バリアント）を用いて複数のベンチマークでSuperThoughtsを評価した：

• 長さの削減: 適応メカニズムを用いることで、SuperThoughtsはCoTの長さを20〜35%削減する。
• 精度: ほとんどのタスクにおいて、標準的な離散CoTのベースラインから1〜2ポイント以内の精度を維持している。例えば、1.5Bモデルにおいて $\tau=0.999$ の場合、MATH500の精度は73.0%に達した（ベースラインの72.4%に対し）、CoTの長さを36%削減しながら達成された。
• スケール効果: 大きなモデル（7B）は、小さなモデル（1.5B）よりも積極的な圧縮に耐えることができる。適応的推論がない場合、1.5Bモデルは14〜21ポイントの大幅な精度低下を招いたが、7Bモデルの低下はわずか5〜12ポイントであった。
• HAMburgerとの比較: SuperThoughtsはあらゆる圧縮レベルにおいてHAMburger（Liu & Zhang, 2025）のベースラインを上回り、同じCoTの長さに対してより高い精度を実現するか、あるいは同じ精度に対してより短いCoT長を実現した。
• 推論の高速化: 理論的なスピードアップは長さの削減に比例するが、実際のウォーククロックタイムの高速化は、CompressorおよびMTPモジュールのオーバーヘッドにより、それよりわずかに低くなる。しかし、このオーバーヘッドはモデルサイズが大きくなるにつれて減少する。14Bモデルでは、約33%の長さ削減が約28%の生成時間の削減につながった。

意義

本論文は、SuperThoughtsが、離散的なトークンの教師あり学習による安定性を損なうことなく、モデルが重ね合わせ状態の潜在空間で推論できることを実証したと主張している。トークンペアを圧縮し、適応的推論を利用することで、このフレームワークは長期間の推論タスクにおける計算コストを削減するための実用的な道筋を提示している。この研究は、より大きなモデルは重ね合わせ状態のトークンを確実に扱うための十分なステップあたりの容量を備えていることを示唆しており、将来的なスケーリングによってさらなる圧縮（例：3つ以上のトークンの重ね合わせ）が可能になることを示唆している。さらに、著者らは、ネイティブなMTPモジュールが将来のLLM（例：Qwen3-Next, MiMo）の標準となるにつれ、SuperThoughtsの適用がより簡単かつ効果的になるだろうと述べている。