技術的サマリー：メモリ内推論（RiM）

問題定義

大規模言語モデル（LLM）の推論能力を向上させる現在の手法は、Chain-of-Thought（CoT）推論のように、最終的な答えを生成する前に中間トークンを生成することで、テスト時の計算量をスケーリングすることに大きく依存しています。しかし、このアプローチは内部計算を自動回帰的なテキスト生成に直接結びつけています。その結果、モデルは「声に出して考える」ことを強要され、純粋な内部推論ではなく、文法的に流暢な自然言語を生成することに計算リソースの大部分を割かざるを得ません。

近年の「潜在推論」アプローチは、連続表現を用いることで自然言語の制約を回避しようとしていますが、それらは依然としてステップごとの自動回帰的生成パラダイムを大部分維持しています。これらの手法では、次のステップがそれらに条件付けられる前に、中間計算が依然として逐次的に外部化されます。本論文は、人間の認知が優れた設計原理である「作業記憶」を提供すると主張します。人間は、すべての中間思考を外部化することなく、タスクに関連する情報を内部に保持・操作することができます。著者らは、LLM が同様の内部ワークスペースを利用するように訓練され、推論をトークン生成の逐次的なボトルネックから切り離すことができる posit しています。

手法：メモリ内推論（RiM）

著者は、推論ステップの自動回帰的生成を固定されたメモリブロックに置き換える潜在推論手法、メモリ内推論（RiM） を提案します。

中核メカニズム

推論トークンのシーケンスを生成する代わりに、モデルには $K$ 個の固定メモリブロック $m_{1:K}$ が提供されます。各ブロックは、事前定義された特殊トークンのシーケンスで構成されます：
$$m_k = [<b>, <m>, \dots, <m>, </b>]$$
ここで、$<b>$ と $</b>$ はブロックの境界を区切り、$<m>$ トークンは潜在ワークスペースを形成します。

• 固定入力: これらのブロックは、固定トークンとして入力質問に付加されます。
• 単一フォワードパス: ブロックは生成された出力ではなく固定入力であるため、拡張されたシーケンス（質問＋メモリブロック）全体が単一のフォワードパスで処理されます。これにより、自動回帰的生成に伴う逐次的デコードのボトルネックが排除されます。
• 学習された表現: トークンの同一性と位置は固定されていますが、その文脈表現はタスク依存性を持つように学習され、中間推論ステップを符号化します。

2 段階のカリキュラム

モデルにこれらの固定ブロックを作業記憶として使用させるために、RiM は 2 段階のトレーニングカリキュラムを採用します：

ステージ 1：推論ステップの監督

• 目的: 推論ステップに計算役割を割り当てることで、メモリブロックを接地（グラウンディング）します。
• メカニズム: $T$ ステップの推論トレースに対して、モデルは $T$ 個のメモリブロックで訓練されます。各ブロック $m_t$ の後、モデルは次の明示的な推論ステップ $r_{t+1}$ を予測する必要があります。
• 制約: カスタム注意マスクにより、ステップ $t+1$ の読み出しが質問と以前のメモリブロックに注意を払うことはできますが、以前の明示的な推論ステップには注意を払えないようにします。これにより、モデルは必要な中間計算をメモリブロックの潜在ワークスペース内に符号化することを強制されます。
• 監督: 監督強度を徐々に低下させるアニーリングスケジュールを用いた、重み付き負の対数尤度目的関数を使用します。

ステージ 2：最終答えの洗練

• 目的: 接地されたメモリブロックを、最終答えの予測のための潜在ワークスペースに変換します。
• メカニズム: 明示的な推論トレースは破棄されます。モデルは $K$ 個の固定メモリブロックを使用します。各ブロックの後、モデルは最終答えを直接予測するように訓練されます。
• 目的: これにより、中間ステップの予測から、より多くのメモリブロック（潜在計算）が利用可能になるにつれて最終答えを漸進的に洗練する監督へとシフトします。
• 推論: テスト時には、モデルは質問とメモリブロックを 1 パスで処理し、最終ブロックの読み出しが答えを提供します。

注意マスク

カスタム因果注意マスク（論文の図 2）は、この手法にとって不可欠です。以下の条件を満たすことで情報の漏洩を防ぎます：

1. メモリブロックは、質問と以前のメモリブロックに注意を払う。
2. 対象となる読み出し（ステージ 1 では推論ステップ、ステージ 2 では最終答え）は、質問と利用可能なメモリブロックに注意を払うが、他の対象読み出しには注意を払わない。
   これにより、モデルが将来の正解を読み取ることで「不正」を働くことなく、すべての対象を単一のフォワードパスで予測することが可能になります。

主要な貢献

1. メモリ内推論（RiM）: 中間計算のための専用かつ固定された潜在ワークスペースを提供することで、LLM の作業記憶容量を解放し、自動回帰的生成から切り離す新しい潜在推論手法。
2. 2 段階カリキュラム: 明示的なステップ監督を通じてメモリブロックを接地し、その後、最終答えの予測のためにそれらを洗練する構造化されたトレーニングアプローチ。これにより、複雑な補助目的なしに潜在空間を効果的に利用可能にする。
3. 効率性と性能: 既存の明示的および潜在推論ベースラインと同等かそれ以上の性能を示し、かつ直接答えモデルの推論速度（逐次的生成のオーバーヘッドなし）を維持することを示す実証的証拠。

実験結果

著者は、GSM8K（分布内）およびGSM-Hard（分布外）の数学的推論ベンチマークにおいて、GPT-2およびLlama-3.2モデル（1B および 3B パラメータ）を使用して RiM を評価しました。

• 性能: RiM は、最も強力な潜在推論ベースライン（Coconut）および直接答えの教師あり微調整（SFT）を一貫して上回りました。
  • GSM8K（Llama-3.2-1B）において、RiM は貪欲法精度で42.1%を達成し、Coconut の36.9%、CoT なし SFT の**23.9%**を上回りました。
  • GSM-Hard において、RiM は**10.5%**を達成し、Coconut（8.5%）および CoT なし SFT（5.3%）を凌駕しました。
  • RiM はまた、優れたPass@8スコアを示し、正しい解が潜在軌道全体に高い確率で分布していることを示唆しました。
• レイテンシ: RiM は、CoT なし SFT と同等の推論速度を達成しました（Llama-3.2-1B の場合、最初のトークンまでの時間約 16ms）。一方、Coconut は約 7 倍遅く、CoT 付き SFT は自動回帰的生成により約 27 倍遅くなりました。
• 表現分析: メモリブロック表現の PCA 投影は、モデルがサンプル依存の明確な潜在ワークスペースを作成することを学習していることを示しました。表現はトレーニング中に体系的に進化し、ブロックが単なる固定のプレースホルダーではなく、能動的な計算状態であることを確認しました。
• ロバスト性: 推論時のメモリ予算（ブロック数）が変化しても性能は安定しており、この手法が展開の制約に対してロバストであることを示唆しています。

意義と主張

本論文は、LLM は、推論プロセスをテキストとして外部化するのではなく、内部化することによって、作業記憶を潜在推論の有効なメカニズムとして使用するように訓練できると主張しています。

主な意義は、計算とコミュニケーションの分離にあります。中間思考の生成を固定メモリブロックに置き換えることで、RiM はモデルが自然言語や逐次的連続表現の生成というオーバーヘッドなしに、高密度で並列化可能な内部計算を実行することを可能にします。これにより、以下の特性を持つ手法が実現されます：

1. 計算効率性: 自動回帰的生成の逐次的ボトルネックを回避し、単一パス推論を可能にします。
2. 有効性: 数学的ベンチマークにおいて、潜在推論手法の中で最先端の性能を達成します。
3. スケーラビリティ: 複雑な再帰的アーキテクチャやマルチパストレーニング目的を必要とすることなく、異なるモデルファミリーやサイズで機能します。

著者らは、このアプローチが、推論プロセスを外部化してテキストとして表現するのではなく、内部化することによって人間の認知的効率を模倣する、より効率的な推論システムへの実現可能な道筋を示していると結論付けています。