【DL輪読会】Next-Latent Prediction Transformers Learn Compact World Models

DEEP LEARNING JP [DL Papers] Next-Latent Prediction Transformers World Models Learn Compact Kohsei Matsutani, Matsuo Lab http://deeplearning.jp/ 1

• http://deeplearning.jp/

書籍情報 • Title: Next-Latent Prediction Transformers Learn Compact World Models • Authors: Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Pratyusha Sharma, Riashat Islam, Alex Lamb, John Langford (Microsoft Research) • Preprint • ArXiv: https://arxiv.org/abs/2511.05963 TL;DR • Next tokenではなく, 次のlatent stateの予測を自己教師あり学習で, Transformerに sequentialな信念状態 (compact world model) を学習させるNext-Latent Prediction (NextLat) を提案. • 追加の推論コスト無しで, world modeling, reasoning, planning, language modelingで性 能が向上した. 2

• https://arxiv.org/abs/2511.05963

Background • そもそも観測したデータをシンプルに表現する方が汎化する – 従来のTransformerは過去のhistoryをすべて見るため, それを圧縮することは学 習していない – その結果, 複雑で, タスク特化のshortcutを学ぶ • Reasoning, planningでのNext Token Prediction (NTP) の限界 – Teacher forcingによるError accumulation (Bachmann et al., 2024) – 直前のprefixから次のtokenを予測するのは, “local coherence”を保つのには有効 だが、複数ステップ先を見越した構造設計 (lookahead)は弱い (Nagarajan et al., 2025.) • Self-Supervised Learning (SSL) は, 強化学習 (RL) などでは使われているが, 言語モデル ではあまり使われていない 3

Methodology (1/3) • NextLatは, next-hidden state predictionを行う • NextLatは信念状態を学習できる – は, history の十分統計量 4

Methodology (2/3) • NextLatが, Next-Token ConsistencyとTransition Consistencyが満たされるには, は の信念状態に収束しなければならない – Next-Token Consistencyを最適化しただけでは, は信念状態に収束しない (普通の Transformer) 5

Methodology (3/3) • Transformerのアーキテクチャを変えず, 最終層のhidden state なし そのsequential方向の遷移構造を学習 • 損失関数 をlatent stateとみ (NTPを交差エントロピー最小化で学習) (Latent stateの時間変化を回帰問題とし て扱う) (d-step horizon) (KL-divでtokenの予測分布を近づける) 6

Experiments (1/4) • World Modeling – マンハッタンの道を予測するタスク Latent Stateの圧縮度合 OOD汎化 Vafa et al., 2024.に 従って内部状態から 地図を再構成すると, GPTは地図が破綻 7

Experiments (2/4) • Reasoning – Countdown (与えられた数字から目標の数字を四則演算で作るタスク) Validityが高い → NextLatは,. Planningがうま くいっている 8

Experiments (3/4) • Planning – Path-Star Graph (グラフのエッジリスト/開始ノード/終了ノードから正しい道順 を生成するタスク) 9

Experiments (4/4) • Language Modeling – TinyStories – Token offset (先のtoken数) をhidden stateからtokenをlinear probingし, そのGPTと の交差エントロピー損失の差分を比較 – 低いほど, 未来を予測できている 10

Discussion • メリット – NextLatは, BSTのT (sequence長) の二乗を回避 – dに依存せず信念状態を学習できる (dは1つの – Token-levelの近視性を改善 – JTPなどより強い からいくつ先を予測するか) • RNNとの比較 – RNNはbackpropでO(T)のsequential dependances – NextLatは, O(d) d<<T • SSMとの比較 – NextLatは非線形に設計できるので, SSMより表現力は高い (ただ, Transformerは超えない) 11

Discussion (本論文に記載なし) • • RLでLLMのreasoningをスケールさせる上での課題 – LLMはaction spaceが大きすぎる • Yue et al., 2025. NIPS Oral のDiscussionから引用 • action space in language models is exponentially larger than that of Go or Atari games (Ramamurthy et al., 2023). RL algorithms were not originally designed to handle such a vast action space 関連研究 – Jia et al., Controlling Large Language Model with Latent Actions . ICML 2025 - Language World Model: history (state列) + latent actionから次のstate (token) を予測 - Policy Model: historyから次のlatent actionを予測 12

Summary • • • • Next-Token Prediction (NTP) は, 長期構造が学習されにくい NextLatはコンパクトな潜在表現で複数ステップ先 (長いhorizon) も予測する 長いhorizonがコンパクトな信念状態の獲得を誘導 World modeling, reasoning, planning, language modelingでGPTなどを上回る • あくまでtoyタスクで言語タスクなどでスケールでできるかはまだ未検証 – そもそもLLMがworld modelを獲得しているか議論の余地 • Horizonを長くすると勾配伝播が難しく, 学習が困難 13

References • • • • • Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Pratyusha Sharma, Riashat Islam, Alex Lamb, John Langford. Next-Latent Prediction Transformers Learn Compact World Models. arXiv preprint, 2025. Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. ICML 2024. Vaishnavh Nagarajan, Chen Henry Wu, Charles Ding, and Aditi Raghunathan. Roll the dice & look before you leap: Going beyond the creative limits of next-token prediction. arXiv preprint arXiv:2504.15266, 2025. Yue, Yang; Chen, Zhiqi; Lu, Rui; Zhao, Andrew; Wang, Zhaokai; Yue, Yang; Song, Shiji; Huang, Gao. Does Reinforcement Learning Really Incentivize Reasoning Capacity in LLMs Beyond the Base Model. NeurIPS 2025 (Oral). Jia et al., Controlling Large Language Model with Latent Actions . ICML 2025. 14