Next-Latent Prediction Transformers Learn Compact World Models
Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Tim Pearce, Pratyusha Sharma, Akshay Krishnamurthy, Riashat Islam, Alex Lamb, John Langford
cs.LG
2025-11-08
NextLat在next-token损失上加隐状态自预测，逼Transformer把历史压成信念状态。曼哈顿有效秩从160.1降到52.7，1.3B模型自投机解码最高加速3.3倍。
Transformer 拿随长度增长的 KV 记忆和随时回头翻历史的 self-attention，换掉了循环网络那套「每一步必须把过去压进固定大小状态」的约束。训练只盯下一个 token，模型可以靠临时查找走捷径：训练集上对得漂亮，换分布就露馅。Vafa 等人用曼哈顿出租车轨迹做过演示，next-token 合法转弯可以到 100%，内部地图却会出现不可能的道路朝向，甚至叠在别的路上的立交桥。
Belief State Transformer 已经证明，只优化 next-token 一致性，隐状态不必是信念状态，也就是对未来预测充分的历史压缩。NextLat 要补的就是这块压力：架构、并行训练、推理流程都不动，只在训练里加一条隐空间自预测，把「压缩历史、转移一致」的循环偏置灌回去。
Transformer 照常产出每步最后一层隐状态 ht。再挂一个潜在动力学模型 pψ，论文里就是两三层 MLP。它吃当前隐状态和下一个 token（当成动作），预测下一步隐状态。
总损失三项：
Theorem 3.2 给出保证：next-token 一致性和转移一致性都到位时，ht 必须是信念状态。d=1 已经够这条定理，多步监督只是加更密的梯度。推理时主模型可以单独解码。pψ 主要用来塑形表示；因为它能在隐空间递归展开，也可以拿来做可变长 self-speculative decoding，草稿长度不必锁死在训练时的 d。
信号比 one-hot token 稠。预测的是整段下一步预测分布所在的隐状态，而且递归一致性会把更远的未来间接写进当前表示。
曼哈顿随机游走有 91M 条、4.7B token，训练 6 个 epoch，多步 horizon d=8。所有方法的 next-token 测试都是 100%。真正分开的是世界模型质量：
| 方法 | 合法轨迹 | 序列压缩 | 有效隐秩↓ | 绕路鲁棒 |
| GPT | 97.0% | 0.65 | 160.1 | 85.0% |
| MTP | 98.1% | 0.64 | 57.7 | 95.0% |
| JTP | 97.1% | 0.32 | 215.8 | 87.0% |
| NextLat | 98.7% | 0.71 | 52.7 | 95.0% |
有效隐秩比 GPT 小三倍以上，重建地图的错误边更少、更局部。曼哈顿真图只有 4580 个路口，压得扁是合理的。
Countdown 上 d=1 就到 54.8%，同 horizon 的 MTP 是 39.2%、JTP 是 39.0%；拉到 d=8 是 58.7%，略高于 MTP 的 57.3%。最后一步等式有效率 54.8%，次优基线 42.3%，少了很多「算到最后硬凑答案」的 regretful compromise。Path-Star 用原论文更难的设置（200k 固定样本，节点从 1 到 100 抽样），G2,10 / G5,5 / G7,7 三种拓扑 NextLat 都接近 100%；BST 在更大的 G7,7 上开始掉。TinyStories 线性探针里，MTP、JTP、BST 会伤 next-token，远程预测掉得快；NextLat 跟 GPT 持平 next-token，并在最多 20 个 token 的前瞻上最强。
1.3B 模型、100B FineWeb-Edu token，下游平均准确率 GPT 58.82、NextLat（d=2）59.21，只高 0.39，而且任务间不稳定。FineWeb-Edu 困惑度 GPT 10.52，NextLat 10.83 到 10.88，略差于 GPT，但好于 MTP 和 JTP。拉开差距的是自投机解码：NextLat（d=2）在 Wikipedia / Books / Code / Math 分别加速 3.21× / 3.32× / 2.38× / 2.87×，单步接受的草稿 token 远超训练 horizon。d=1 时单卡 B200 训练速度与 GPT 同为 3.09 iter/s，BST 只有 0.89。
A5 字问题里，2 层 Transformer 自己推不出训练长度 12 以外；同训出来的 2.62M 参数 RNN，用 Transformer 初始化第一步之后独立跑，36 token 准确率超过 95%。直接在 36 token 上训的 GPT 也解不了这题。
这是训练目标上的补丁，不是新骨架。现有 decoder-only 都能挂一个 MLP。世界模型、组合搜索、需要躲开 n-gram 捷径的规划任务上，增益是实的。语言建模选择题几乎没动，但同一套动力学能当可变长草稿用：训练 horizon 可以很浅，推理 draft 可以拉到 10。这比靠加层堆 MTP 头更划算。
1.3B 规模上，它更接近「表示更干净、推理能加速」，谈不上 benchmark 全面涨点。已经在训 Transformer、愿意加一条辅助损失的人可以试；指望它单独把选择题拉开，目前证据不够。
作者自己列了几条。动力学模型只用简单 MLP，没扫隐层宽度这个容量瓶颈；Path-Star 和 Countdown 上反而更小的潜在维度更好。stop-gradient、KL、Smooth L1 来自小规模消融，大规模时 d>1 和 KL 还要不要，没有答案。没有对照更新的 MTP 变体。投机解码的 draft 长度是静态选最优，没有做自适应。预训练里 Smooth L1 还会往上走，对优化器敏感。
另外，曼哈顿合法轨迹只比 GPT 高 1.7 个百分点，主证据是压缩和隐秩，不是规划成功率的大跳。Countdown 和 Path-Star 都是合成任务。1.3B 下游只高 0.39，论文自己也写可能要更大模型才看得出。A5 那条 RNN 外推很漂亮，但停在 2 层短序列，不能直接说生产模型逃出了 TC0。