技術的サマリー：Next-Latent Prediction Transformers はコンパクトな世界モデルを学習する

問題定義

標準的なトランスフォーマーアーキテクチャは、再帰性を自己注意機構とシーケンス長に比例してスケーリングするメモリ機構に置き換える。これにより並列処理と予測能力が飛躍的に向上する一方で、一貫した遷移規則を持つコンパクトな潜在状態へと履歴を圧縮するという本質的な帰納的バイアスが失われる。その結果、トランスフォーマーは訓練データにはよく適合するが汎化性能が低い複雑なタスク固有のショートカットを学習することが多く、長期的な計画や推論に必要な整合的な「世界モデル」や信念状態を形成することに失敗する。標準的な次のトークン予測目的関数は、モデルの内部隠れ状態が過去の履歴に対する十分統計量として進化することを保証するものではなく、またこれらの状態が将来の観測を再帰的に予測できることを保証するものでもない。

手法：Next-Latent Prediction (NextLat)

著者らは、標準的な次のトークン学習に潜在空間における自己教師あり予測を追加するフレームワークNext-Latent Prediction (NextLat) を提案する。中核的なアイデアは、次のトークンが与えられたときの次の潜在状態を予測する潜在表現（具体的には最終層の隠れ状態 $h_t$）をトランスフォーマーに学習させることである。

理論的基盤

本論文は、モデルがNext-Token Consistency（$h_t$ から $X_{t+1}$ を予測すること）とTransition Consistency（$h_t$ と $X_{t+1}$ から $h_{t+1}$ を予測すること）の両方を最適化する場合、隠れ状態 $h_t$ が理論的に信念状態へと収束することが保証されることを確立している。信念状態とは、未来 $X_{t+1:T}$ を予測するための履歴 $X_{1:t}$ に対する十分統計量として定義される。過去のトークンを圧縮せずにアドホックな参照を許容する標準的な自己回帰的学習とは異なり、NextLat は、潜在状態が次の状態への遷移と次のトークンの予測に必要なすべての情報を保持しなければならないという再帰的な帰納的バイアスを強制する。

実装

NextLat は、標準的なトランスフォーマーの学習ループに軽量な補助目的関数を追加することで機能する：

1. 潜在ダイナミクスモデル ($p_\psi$): 現在の隠れ状態 $h_t$ と次のトークン $X_{t+1}$ が与えられたとき、次の隠れ状態 $\hat{h}_{t+1}$ を予測するように訓練される単純な多層パーセプトロン（MLP）。
2. 損失関数: 総損失は以下の組み合わせからなる：
   • Next-Token Loss ($L_{next-token}$): トークン予測のための標準的な交差エントロピー。
   • Next-Hidden State Loss ($L_{next-h}$): 予測状態 $\hat{h}_{t+1}$ を真の次の状態 $h_{t+1}$ に対して回帰させる Smooth L1 損失（表現の崩壊を防ぐため、ターゲットに対して stop-gradient を適用）。
   • KL 発散損失 ($L_{KL}$): 予測された潜在状態から予測されるトークン分布が、真の潜在状態からの分布と一致することを保証する補完的な目的関数であり、ソフトな教師信号として機能する。
3. マルチステップ監督: モデルは $d$ ステップにわたって潜在ダイナミクスを再帰的に展開するように訓練でき、長距離依存関係に対してより豊かな勾配信号を提供する。

重要なのは、NextLat がトランスフォーマーのアーキテクチャ、推論手順、並列学習の効率性を一切変更しない点である。潜在ダイナミクスモデルは表現を形成するために学習時のみ使用され、推論時にはモデルは通常通り自己回帰的にデコードする。

主要な貢献

1. 理論的保証: 本論文は、次のトークンと次の潜在状態の一貫性を最適化することが、トランスフォーマーの表現を信念状態へと確実に変化させることを証明しており、これは次のトークン予測のみでは保証されない特性である。
2. 実用的フレームワーク: NextLat は、トランスフォーマーの並列学習または推論構造を変更することなく、再帰的な帰納的バイアスを注入する計算効率の高い補助目的関数を導入する。
3. 可変長自己仮説的デコーディング: 整合的な潜在ダイナミクスを学習することで、NextLat は潜在空間における再帰的予測を可能にする。これにより、モデルが固定された訓練ホライズンを超えてトークンをドラフトできる可変長自己仮説的デコーディングが可能となり、推論が大幅に加速される。

実験結果

著者らは、世界モデル化、推論、計画、言語モデル化の 4 つのドメインで NextLat を評価した。

• 世界モデル化 (Manhattan Taxi Rides): NextLat は、標準的な GPT、MTP、JTP ベースラインと比較して、真の世界モデル（街路トポロジー）と著しく整合性の高い内部マップを学習する。シーケンス圧縮率が高く（GPT の 0.65 対 0.71）、有効な潜在ランクが大幅に低い（GPT の 160.1 対 52.7）ことを示しており、よりコンパクトで効率的な表現であることを示唆している。
• 推論 (Countdown): 数学的推論タスクにおいて、NextLat はベースラインを大幅に上回る（$d=1$ で 54.8% 対 GPT の 33.1%）。先読み計画に優れており、モデルが目標に一致させるために無効な最終式を強制する「後悔を伴う妥協」を減少させている。
• 計画 (Path-Star Graphs): NextLat は、さまざまなグラフトポロジーにわたってほぼ 100% の解決率を維持するのに対し、BST や JTP などのベースラインはより大きなグラフ（$G_{7,7}$）で失敗する。これは、NextLat がトークン空間予測に共通する「クレバー・ハンス」的なショートカット学習を回避していることを示唆している。
• 言語モデル化 (FineWeb-Edu):
  • ダウンストリーム精度: NextLat は、MTP や JTP よりも次のトークンのパープレキシティをよりよく維持しつつ、ダウンストリームのゼロショット精度で modest な向上を達成している（GPT の 58.82% 対平均 59.21%）。
  • 推論速度: 可変長自己仮説的デコーディングを活用することで、NextLat は標準的な自己回帰サンプリングと比較して最大3.3 倍の高速な推論を達成する。この速度向上は、訓練ホライズンを超えてドラフトする場合でも持続し、固定されたトークン空間ホライズンを持つ MTP や JTP には欠けている機能である。

意義と主張

本論文は、NextLat をトランスフォーマーにおいてコンパクトで予測性があり、汎化可能な表現を学習するためのシンプルかつ効果的なパラダイムとして位置づけている。その主な意義は、コンパクトな世界モデルの再帰的帰納的バイアスとトランスフォーマーのスケーラビリティおよび並列性を統合する点にある。

著者らは、NextLat が次のトークン予測の「近視眼的」な性質に対する解決策を提供し、モデルが浅いトークン相関に依存するのではなく、構造化された予測的内部状態を形成することを促すと主張している。アーキテクチャの変更なしに可変長の仮説的デコーディングを可能にすることで、NextLat はより高速な推論とより優れた長期的推論への実用的な道筋を提供する。本論文は、NextLat が複雑なシステムのより単純でコンパクト、したがってより汎化可能な表現を自己回帰モデルに付与する学習目的への一歩であると結論づけている。

著者らは将来の方向性について控えめに述べており、複雑な潜在ダイナミクスアーキテクチャ（単純な MLP のみを使用）を探索しなかったこと、および目的関数の最適設計（例えば、大規模スケールにおける KL 損失の必要性）にはさらなる体系的な研究が必要であることを指摘している。また、NextLat が表現を改善する一方で、固定深度トランスフォーマーの回路複雑性は TC0 クラスに制限されたままだが、共同訓練された潜在ダイナミクスモデル（RNN）はこれらの限界を超えて一般化できる能力を持っていることを認めている。