Next-Latent Prediction Transformers Learn Compact World Models

01 动机 · Motivation

Transformer 的自注意力机制允许模型在任意时刻直接"查询"过去的 token，从而不必把历史信息压缩成紧凑的内部状态。这带来了一个根本性的问题：

"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 的出发点。

不同预测机制对比 — **图 1：**四种预测机制对比。GPT 仅做 next-token 预测；MTP（Multi-Token Prediction）在 token 空间并行预测多步；JTP（Joint Token Prediction）联合预测；**NextLat** 在隐藏状态空间中通过轻量 dynamics model 迭代滚动预测，从而引入循环归纳偏置，且推理时完全不需要 dynamics model。

3.32×推理加速（Books 集，speculative decoding）

52.7NextLat 有效潜在秩（vs. GPT 的 160.1，压缩 3×）

98.7%世界建模任务 OOD 有效轨迹率（vs. GPT 97.0%）

~100%G₇,₇ 规划任务准确率（MTP/JTP 大幅失败）

02 方法 · Method

NextLat 在标准 Transformer 之外附加一个轻量 latent dynamics model（实验中为简单 MLP），通过三个损失函数联合训练，推理时 dynamics model 完全被丢弃，不增加任何推理开销。

训练目标

联合损失函数（公式 5）为：

ℒ_NextLat = ℒ_next-token + λ_next-h·ℒ_next-h + λ_KL·ℒ_KL

ℒ_next-token：标准 cross-entropy，与 GPT 完全相同。
ℒ_next-h（Smooth L1）：Dynamics model 在隐藏状态空间中做 d 步滚动预测，用 stop-gradient 防止表征坍缩。这一目标提供"比 token 级监督更密集的梯度信号"。
ℒ_KL：知识蒸馏约束——要求预测的隐藏状态与真实 token 预测分布保持一致，确保潜在预测在语义上有意义。

理论保证：Belief State 收敛（定理 3.2）

论文在 POMDP 框架下严格证明：若 NextLat 同时最优化 next-token 一致性和 transition 一致性，则隐藏状态必然成为 belief states——即"预测未来所需历史信息的充分统计量"（sufficient statistics of history for predicting future tokens）。这是首个证明潜在状态预测能诱导 Transformer 学习 belief state 的理论结果。

Self-Speculative Decoding — **图 2：**NextLat 支持的 variable-length self-speculative decoding。Dynamics model 在推理时充当 draft model，从当前隐藏状态出发滚动生成多个候选 token，再由主 Transformer 验证接受。与 MTP 固定 d 步草稿不同，NextLat 可动态超过训练 horizon，接受率更高。

与 MTP/JTP 的关键区别

MTP / JTP（token 空间预测）

在 token 输出空间并行预测多步
梯度信号稀疏，依赖 token 离散分布
Speculative decoding 固定在训练的 d 步
训练计算开销随 d 线性增长

NextLat（隐藏状态空间预测）

在连续隐藏状态空间中迭代滚动预测
梯度信号更密集，每步都携带未来多步信息
Speculative decoding 可超越训练 horizon
推理零开销（dynamics model 被丢弃）

03 实验 · Experiments

实验覆盖五类任务：世界建模（Manhattan 出租车）、推理（Countdown / Game of 24）、规划（Path-Star 图）、语言建模（FineWeb-Edu，1.3B 参数，100B tokens）、长程预测性（TinyStories）和状态追踪（A₅ 单词问题）。

世界建模：Manhattan 出租车轨迹

在有向图上随机游走数据集（OOD 测试）中，评估内部表征是否能重建地图结构：

模型	有效轨迹率 (OOD)	序列压缩比	有效潜在秩	绕路鲁棒性
GPT	97.0%	—	160.1	—
MTP (d=8)	98.1%	—	—	—
NextLat	98.7%	0.71（最高）	52.7（最低）	95.0%

NextLat 的有效潜在秩（52.7）约为 GPT（160.1）的 1/3，说明其表征更紧凑，确实在内部维持了类似"地图"的世界模型。

NextLat 重建的曼哈顿地图 — **图 3：**对 NextLat 隐藏状态做线性探测后重建的内部"地图"。与 GPT/MTP/JTP 相比，NextLat 重建出的街道布局最接近真实曼哈顿网格，说明其隐藏状态中确实编码了紧凑的环境结构（belief state）。

推理与规划

任务	GPT	MTP (d=8)	NextLat (d=1)
Countdown（Game of 24）准确率	33.1%	57.3%	54.8%
G₇,₇ 规划（Path-Star 图）	部分失败	大幅失败	~100%

在 G₇,₇ 规划任务上，NextLat 接近满分，而 MTP/JTP 大幅失败，说明 NextLat 的潜在空间监督有效避免了 token 空间方法的"捷径学习"（shortcut learning）问题。

语言建模（FineWeb-Edu，1.3B 参数）

模型	FW-Edu PPL ↓	Wiki PPL ↓	LAMBADA PPL ↓
GPT	10.52	17.93	20.26
MTP (d=2)	11.00	18.61	18.34
NextLat (d=2)	10.88	18.44	17.83

NextLat 在保持 next-token 性能的同时（与 GPT 相近），在 LAMBADA 长程依赖基准上显著优于 MTP/JTP。

Self-Speculative Decoding 加速

模型 (d=2)	Books 加速比	Books 平均接受 token 数	Wiki 加速比
MTP	1.72×	1.83	—
JTP	1.90×	1.89	—
NextLat	3.32×	4.86	3.21×

推理加速对比 — **图 4：**Speculative decoding 加速比对比。NextLat（蓝）在所有 draft 长度下均大幅领先 MTP（橙）和 JTP（绿），且可超越训练 horizon 继续采样，因此接受 token 数（4.86）远超 MTP（1.83）。

长程预测性（TinyStories）

TinyStories 线性探测结果 — **图 5：**在冻结表征上用线性探测预测 offset=1–20 步后的 token。NextLat（d=1 和 d=8）在所有偏移量上均最强，而 MTP/JTP 在超过 5–10 步后性能急剧下降，说明 NextLat 的表征确实编码了更长程的预测信息。

训练效率对比

模型	训练速度（steps/sec）
GPT	3.09
NextLat (d=1)	3.09（与 GPT 相同）
MTP (d=1)	2.80
NextLat (d=8)	1.73
BST	0.89（慢 3.5×）

d=1 时 NextLat 与 GPT 训练速度完全相同，梯度计算复杂度为 O(Td)，远优于 BST 的 O(T²)。

04 局限性 · Limitations

Note：以下局限性均由论文作者在正文中明确陈述（stated）。

Dynamics model 设计简单，未探索更强的架构

所有实验均使用简单 MLP 作为 latent dynamics model，"more expressive architectures unexplored"。更复杂的 dynamics model（如 Transformer）是否能带来进一步提升尚不清楚。

超参设计缺乏原则性依据

Stop-gradient、KL loss 等关键设计选择"guided by small-scale ablations, not principled principles"，且没有系统性研究 d>1 和 KL 监督在大规模场景下是否必要。

与最新 MTP 变体（如 DeepSeek-v3）的比较有限

论文未系统比较 DeepSeek-v3 等近期更强的 MTP 变体，使得 NextLat 在最新工业实践中的相对优势尚不明确。

Speculative decoding 使用固定草稿长度

"Speculative decoding uses fixed draft lengths per prompt, not adaptive"。每个 prompt 的最优草稿长度应自适应调整，但当前实现使用固定长度，未能充分挖掘 NextLat 可变长度的潜力。

对优化动态敏感，表征语义分析不足

损失轨迹因优化器（AdamW vs. Muon）而异；NextLat 学到的表征语义结构"not thoroughly analyzed"，缺乏对内部 belief state 的更深入可解释性研究。