Microsoft Research · Preprint 2025 · arXiv:2511.05963
Next-Latent Prediction Transformers Learn Compact World Models
让 Transformer 在潜在空间里"预见未来"，从而自发学会紧凑世界模型Letting Transformers "foresee the future" in latent space so that a compact world model emerges on its own
Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Tim Pearce, Pratyusha Sharma, Akshay Krishnamurthy, Riashat Islam, Alex Lamb, John Langford · Microsoft Research
标准 next-token 预测缺乏压缩历史信息的内在动机，导致 Transformer 学到依赖 ad-hoc attention 查找的"捷径"，泛化能力有限。本文提出 NextLat——在不改变模型架构和推理流程的前提下，额外训练一个轻量 latent dynamics model，令每一步隐藏状态能预测下一步隐藏状态；理论上证明这使隐藏状态收敛为 belief states，并在世界建模、推理、规划和语言建模上全面提升性能，同时支持最高 3.3× 的 self-speculative decoding 加速。Standard next-token prediction lacks any intrinsic incentive to compress historical information, which leads Transformers to learn "shortcuts" that rely on ad-hoc attention lookups and generalize poorly. This paper proposes NextLat: without changing the model architecture or the inference pipeline, it additionally trains a lightweight latent dynamics model so that the hidden state at every step predicts the hidden state at the next step. It is proved theoretically that this makes the hidden states converge to belief states, and performance improves across world modeling, reasoning, planning and language modeling, while supporting self-speculative decoding speedups of up to 3.3×.
world modelnext-latent predictionbelief statePOMDPspeculative decodingtransformer潜在状态预测latent state prediction世界模型world model压缩表征compressed representation
01动机 · MotivationMotivation
Transformer 的自注意力机制允许模型在任意时刻直接"查询"过去的 token，从而不必把历史信息压缩成紧凑的内部状态。这带来了一个根本性的问题：The self-attention mechanism of a Transformer lets the model directly "query" past tokens at any moment, so it never has to compress historical information into a compact internal state. This gives rise to a fundamental problem:
"Transformers lack an inherent incentive to compress history into compact latent states with consistent transition rules, which often leads to learning solutions that generalize poorly."
相比之下，循环神经网络（RNN）因为每步只能访问一个固定大小的隐藏状态，被迫学习紧凑的状态表征。如果能把这种"循环归纳偏置"（recurrent inductive bias）注入 Transformer，同时保留其并行训练优势，就能得到一个既高效又泛化能力强的模型。这正是 NextLat 的出发点。By contrast, a recurrent neural network (RNN) can access only a fixed-size hidden state at each step and is therefore forced to learn compact state representations. If this "recurrent inductive bias" could be injected into a Transformer while keeping its parallel-training advantage, the result would be a model that is both efficient and strong at generalization. This is exactly the starting point of NextLat.
图 1：四种预测机制对比。GPT 仅做 next-token 预测；MTP（Multi-Token Prediction）在 token 空间并行预测多步；JTP（Joint Token Prediction）联合预测；NextLat 在隐藏状态空间中通过轻量 dynamics model 迭代滚动预测，从而引入循环归纳偏置，且推理时完全不需要 dynamics model。Figure 1: Comparison of four prediction mechanisms. GPT performs only next-token prediction; MTP (Multi-Token Prediction) predicts several steps in parallel in token space; JTP (Joint Token Prediction) predicts jointly; NextLat rolls out predictions iteratively in hidden-state space through a lightweight dynamics model, thereby introducing a recurrent inductive bias, and needs no dynamics model at all at inference time.
98.7%世界建模任务 OOD 有效轨迹率（vs. GPT 97.0%）Valid-trajectory rate on the OOD world-modeling task (vs. GPT 97.0%)
~100%G₇,₇ 规划任务准确率（MTP/JTP 大幅失败）Accuracy on the G₇,₇ planning task (MTP/JTP fail badly)
02方法 · MethodMethod
NextLat 在标准 Transformer 之外附加一个轻量 latent dynamics model（实验中为简单 MLP），通过三个损失函数联合训练，推理时 dynamics model 完全被丢弃，不增加任何推理开销。On top of a standard Transformer, NextLat attaches one lightweight latent dynamics model (a simple MLP in the experiments) and trains it jointly through three loss terms; at inference the dynamics model is discarded entirely, adding no inference overhead.
训练目标Training objective
联合损失函数（公式 5）为：The joint loss function (Equation 5) is:
ℒnext-token：标准 cross-entropy，与 GPT 完全相同。ℒnext-token: the standard cross-entropy, exactly the same as in GPT.
ℒnext-h（Smooth L1）：Dynamics model 在隐藏状态空间中做 d 步滚动预测，用 stop-gradient 防止表征坍缩。这一目标提供"比 token 级监督更密集的梯度信号"。ℒnext-h (Smooth L1): the dynamics model performs a d-step rollout prediction in hidden-state space, using stop-gradient to prevent representation collapse. This objective provides "a denser gradient signal than token-level supervision".
ℒKL：知识蒸馏约束——要求预测的隐藏状态与真实 token 预测分布保持一致，确保潜在预测在语义上有意义。ℒKL: a knowledge-distillation constraint requiring the predicted hidden state to stay consistent with the true token prediction distribution, ensuring that the latent prediction is semantically meaningful.
理论保证：Belief State 收敛（定理 3.2）Theoretical guarantee: convergence to belief states (Theorem 3.2)
论文在 POMDP 框架下严格证明：若 NextLat 同时最优化 next-token 一致性和 transition 一致性，则隐藏状态必然成为 belief states——即"预测未来所需历史信息的充分统计量"（sufficient statistics of history for predicting future tokens）。这是首个证明潜在状态预测能诱导 Transformer 学习 belief state 的理论结果。Within the POMDP framework the paper proves rigorously that if NextLat optimizes both next-token consistency and transition consistency, the hidden states necessarily become belief states, that is, "sufficient statistics of history for predicting future tokens". This is the first theoretical result showing that next-latent prediction can induce a Transformer to learn belief states.
图 2：NextLat 支持的 variable-length self-speculative decoding。Dynamics model 在推理时充当 draft model，从当前隐藏状态出发滚动生成多个候选 token，再由主 Transformer 验证接受。与 MTP 固定 d 步草稿不同，NextLat 可动态超过训练 horizon，接受率更高。Figure 2: The variable-length self-speculative decoding that NextLat supports. At inference the dynamics model acts as the draft model, rolling out several candidate tokens from the current hidden state, which the main Transformer then verifies and accepts. Unlike MTP’s fixed d-step drafts, NextLat can dynamically go beyond the training horizon and reaches a higher acceptance rate.
与 MTP/JTP 的关键区别Key differences from MTP/JTP
MTP / JTP（token 空间预测）MTP / JTP (prediction in token space)
在 token 输出空间并行预测多步Predicts several steps in parallel in the token output space
梯度信号稀疏，依赖 token 离散分布Sparse gradient signal, dependent on the discrete token distribution
Speculative decoding 固定在训练的 d 步Speculative decoding is fixed to the d steps used in training
训练计算开销随 d 线性增长Training compute cost grows linearly with d
NextLat（隐藏状态空间预测）NextLat (prediction in hidden-state space)
在连续隐藏状态空间中迭代滚动预测Rolls out predictions iteratively in a continuous hidden-state space
梯度信号更密集，每步都携带未来多步信息Denser gradient signal; every step carries information about multiple future steps
Speculative decoding 可 超越 训练 horizonSpeculative decoding can go beyond the training horizon
推理零开销（dynamics model 被丢弃）Zero inference overhead (the dynamics model is discarded)
03实验 · ExperimentsExperiments
实验覆盖五类任务：世界建模（Manhattan 出租车）、推理（Countdown / Game of 24）、规划（Path-Star 图）、语言建模（FineWeb-Edu，1.3B 参数，100B tokens）、长程预测性（TinyStories）和状态追踪（A₅ 单词问题）。The experiments cover five families of tasks: world modeling (Manhattan taxi), reasoning (Countdown / Game of 24), planning (Path-Star graphs), language modeling (FineWeb-Edu, 1.3B parameters, 100B tokens), long-range predictivity (TinyStories) and state tracking (A₅ word problems).
在有向图上随机游走数据集（OOD 测试）中，评估内部表征是否能重建地图结构：On a dataset of random walks over a directed graph (OOD test), the internal representation is evaluated for whether it can reconstruct the structure of the map:
模型Model
有效轨迹率 (OOD)Valid trajectory rate (OOD)
序列压缩比Sequence compression ratio
有效潜在秩Effective latent rank
绕路鲁棒性Detour robustness
GPT
97.0%
—
160.1
—
MTP (d=8)
98.1%
—
—
—
NextLat
98.7%
0.71（最高）0.71 (highest)
52.7（最低）52.7 (lowest)
95.0%
NextLat 的有效潜在秩（52.7）约为 GPT（160.1）的 1/3，说明其表征更紧凑，确实在内部维持了类似"地图"的世界模型。NextLat’s effective latent rank (52.7) is about 1/3 of GPT’s (160.1), showing that its representation is more compact and that it does maintain internally a world model resembling a "map".
图 3：对 NextLat 隐藏状态做线性探测后重建的内部"地图"。与 GPT/MTP/JTP 相比，NextLat 重建出的街道布局最接近真实曼哈顿网格，说明其隐藏状态中确实编码了紧凑的环境结构（belief state）。Figure 3: The internal "map" reconstructed by linear probing of NextLat’s hidden states. Compared with GPT/MTP/JTP, the street layout NextLat reconstructs is the closest to the real Manhattan grid, showing that its hidden states really do encode a compact structure of the environment (a belief state).
推理与规划Reasoning and planning
任务Task
GPT
MTP (d=8)
NextLat (d=1)
Countdown（Game of 24）准确率Countdown (Game of 24) accuracy
在 G₇,₇ 规划任务上，NextLat 接近满分，而 MTP/JTP 大幅失败，说明 NextLat 的潜在空间监督有效避免了 token 空间方法的"捷径学习"（shortcut learning）问题。On the G₇,₇ planning task NextLat is close to a perfect score whereas MTP/JTP fail severely, showing that NextLat’s latent-space supervision effectively avoids the "shortcut learning" problem of token-space methods.
NextLat 在保持 next-token 性能的同时（与 GPT 相近），在 LAMBADA 长程依赖基准上显著优于 MTP/JTP。While preserving next-token performance (on par with GPT), NextLat is markedly better than MTP/JTP on the LAMBADA long-range dependency benchmark.
图 4：Speculative decoding 加速比对比。NextLat（蓝）在所有 draft 长度下均大幅领先 MTP（橙）和 JTP（绿），且可超越训练 horizon 继续采样，因此接受 token 数（4.86）远超 MTP（1.83）。Figure 4: Comparison of speculative decoding speedups. NextLat (blue) leads MTP (orange) and JTP (green) by a wide margin at every draft length, and it can keep sampling beyond the training horizon, so its number of accepted tokens (4.86) far exceeds that of MTP (1.83).
图 5：在冻结表征上用线性探测预测 offset=1–20 步后的 token。NextLat（d=1 和 d=8）在所有偏移量上均最强，而 MTP/JTP 在超过 5–10 步后性能急剧下降，说明 NextLat 的表征确实编码了更长程的预测信息。Figure 5: Linear probing on frozen representations to predict the token at offset=1–20 steps ahead. NextLat (d=1 and d=8) is the strongest at every offset, while MTP/JTP degrade sharply beyond 5–10 steps, showing that NextLat’s representation really does encode longer-range predictive information.
训练效率对比Training efficiency comparison
模型Model
训练速度（steps/sec）Training speed (steps/sec)
GPT
3.09
NextLat (d=1)
3.09（与 GPT 相同）3.09 (same as GPT)
MTP (d=1)
2.80
NextLat (d=8)
1.73
BST
0.89（慢 3.5×）0.89 (3.5× slower)
d=1 时 NextLat 与 GPT 训练速度完全相同，梯度计算复杂度为 O(Td)，远优于 BST 的 O(T²)。At d=1 NextLat trains at exactly the same speed as GPT, and the gradient computation complexity is O(Td), far better than BST’s O(T²).
04局限性 · LimitationsLimitations
Note：以下局限性均由论文作者在正文中明确陈述（stated）。Note: all limitations below are explicitly stated by the authors in the body of the paper.
Dynamics model 设计简单，未探索更强的架构The dynamics model design is simple; stronger architectures are unexplored
所有实验均使用简单 MLP 作为 latent dynamics model，"more expressive architectures unexplored"。更复杂的 dynamics model（如 Transformer）是否能带来进一步提升尚不清楚。All experiments use a simple MLP as the latent dynamics model, and "more expressive architectures unexplored". Whether a more complex dynamics model (a Transformer, say) could bring further gains remains unclear.
Stop-gradient、KL loss 等关键设计选择"guided by small-scale ablations, not principled principles"，且没有系统性研究 d>1 和 KL 监督在大规模场景下是否必要。Key design choices such as stop-gradient and the KL loss are "guided by small-scale ablations, not principled principles", and there is no systematic study of whether d>1 and KL supervision are necessary at large scale.
与最新 MTP 变体（如 DeepSeek-v3）的比较有限Limited comparison with the latest MTP variants (e.g. DeepSeek-v3)
论文未系统比较 DeepSeek-v3 等近期更强的 MTP 变体，使得 NextLat 在最新工业实践中的相对优势尚不明确。The paper does not systematically compare against recent, stronger MTP variants such as DeepSeek-v3, which leaves NextLat’s relative advantage in the latest industrial practice unclear.
"Speculative decoding uses fixed draft lengths per prompt, not adaptive"。每个 prompt 的最优草稿长度应自适应调整，但当前实现使用固定长度，未能充分挖掘 NextLat 可变长度的潜力。"Speculative decoding uses fixed draft lengths per prompt, not adaptive". The optimal draft length should be adapted per prompt, but the current implementation uses a fixed length and does not fully exploit NextLat’s variable-length potential.
对优化动态敏感，表征语义分析不足Sensitive to optimization dynamics; the semantics of the representation are under-analyzed
损失轨迹因优化器（AdamW vs. Muon）而异；NextLat 学到的表征语义结构"not thoroughly analyzed"，缺乏对内部 belief state 的更深入可解释性研究。Loss trajectories differ across optimizers (AdamW vs. Muon); the semantic structure of the representations NextLat learns is "not thoroughly analyzed", and deeper interpretability studies of the internal belief state are lacking.
