Next-Latent Prediction Transformers Learn Compact World Models¶

会议: NeurIPS 2026
arXiv: 2511.05963
代码: https://github.com/JaydenTeoh/NextLat
领域: 模型压缩
关键词: 下一隐状态预测、世界模型、信念状态、自推测解码、递归表征压缩

一句话总结¶

提出 NextLat（Next-Latent Prediction）自监督学习框架，在自回归 Transformer 中引入下一隐状态预测与自蒸馏辅助目标，促使注意力表征收敛为紧凑的信念状态世界模型，在零架构改动下实现高达 3.3× 的变长自推测解码推理加速。

研究背景与动机¶

托勒密的地心说能够精准拟合地球视角的行星运行观测，但其本轮模型结构繁复且缺乏全局一致性；而哥白尼日心说以极度紧凑简洁的模型解释了相同规律并获得了超越局部视角的泛化能力。经典学习理论同样表明，对训练观测越简洁紧凑的解释越不易过拟合，其泛化界越强。然而，现代自回归 Transformer 架构摒弃了循环神经网络（RNN）的固定状态瓶颈，其上下文记忆随序列长度线性增长，并通过自注意力机制随时回溯过去所有 token。这种灵活性虽然带来了极强的并行训练吞吐和拟合容量，却彻底消除了将历史信息压缩为紧凑隐状态的归纳偏置，致使 Transformer 极易学到依赖表面 n-gram 或局部捷径的虚假关联，在长程规划与非分布（OOD）迁移中表现出严重的不一致性。

现有尝试让 Transformer 学习紧凑世界模型或长程规划能力的方法面临计算成本或理论保障的双重瓶颈。例如，信念状态 Transformer（BST）虽然通过双向跨度预测理论上诱导了信念状态，但需要联合训练双编码器并积累 $O(T^2)$ 复杂度的前缀-后缀梯度，训练吞吐降低 3 倍以上且内存开销巨大；联合多 token 预测（JTP）与多 token 预测（MTP）依然停留在离散 token 空间计算交叉熵损失，仅在系统满足严格的 $k$-可观测性且预测视界大于 $k$ 时才能近似局部信念状态，且离散 token 监督本质上是稀疏的独热标签，无法向隐空间提供丰富且连续的动态流形梯度。

本文的切入角度是：能否在完全保留标准 Transformer 纯自回归架构、并行训练效率与单步解码流程的前提下，仅凭轻量辅助目标向隐表征注入循环动力学的归纳偏置？核心 idea：通过引入轻量 MLP 隐状态转移模型，显式训练 Transformer 预测“给定下一 token 时的下一层前隐状态”，在理论上证明该目标促使中间隐状态必然收敛至充分统计量性质的信念状态，并在推理阶段直接利用隐状态递归展开实现免独立草稿模型的变长自推测解码。

方法详解¶

整体框架¶

NextLat 的核心机制是在标准的自回归 Transformer 顶层隐藏状态（Pre-logit hidden states）上附加自监督动力学预测任务。给定输入序列 $X_{1:T}$，主干 Transformer 在各时间步输出隐藏表示 $\mathbf{h}_t$ 并通过输出分类头 $p_\theta$ 计算标准 Next-Token 交叉熵损失。与此同时，NextLat 引入一个极小型的多层感知机隐状态动力学模型 $p_\psi$。该动力学模型以当前隐藏状态 $\mathbf{h}_t$ 和真实下一步输入 token $X_{t+1}$ 的嵌入拼接作为输入，预测下一个时间步的隐藏状态 $\hat{\mathbf{h}}_{t+1}$。为了支持更长视界的监督，动力学模型可以在隐空间自递归展开 $d$ 步。训练时，模型在隐状态空间计算 Smooth L1 回归损失，并在预测状态的输出概率分布上施加停止梯度的 KL 散度自蒸馏损失。在推理时，若进行标准生成，辅助动力学模型被完全丢弃；若进行投机推理，该动力学模型可在隐空间自回归推演生成草稿 token，并由主干 Transformer 一次性并行验证。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 tokens"] --> B["自回归 Transformer 主干<br/>并行前向编码输出隐藏状态"]
    B --> C["输出分类头<br/>下一 token 交叉熵损失"]
    B --> D["隐动力学递归预测<br/>当前状态与下一步动作拼接过 MLP"]
    D --> E["双重隐空间与分布对齐<br/>Smooth L1 回归与停止梯度 KL 自蒸馏"]
    D -.->|推理阶段递归推演| F["变长自推测解码<br/>无额外草稿头生成多 token 草稿"]

关键设计¶

1. 信念状态收敛理论与隐动力学建模：将连续状态转移作为第一类监督信号 针对标准自回归 Transformer 仅拟合表面离散 token 转移而缺乏紧凑环境模型的痛点，NextLat 在隐空间构建显式转移一致性约束。在理论层面，若 Transformer 隐藏状态 $\mathbf{h}_t$ 同时满足下一 token 预测一致性 $p_\theta(X_{t+1} \mid \mathbf{h}_t) = \mathbb{P}(X_{t+1} \mid X_{1:t})$ 以及转移一致性 $p_\psi(\mathbf{h}_{t+1} \mid \mathbf{h}_t, X_{t+1}) = \mathbb{P}(\mathbf{h}_{t+1} \mid X_{1:t+1})$，通过逆向数学归纳法可严格证明 $\mathbf{h}_t$ 必定为历史序列 $X_{1:t}$ 预测未来任意时间步 $X_{t+1:T}$ 的充分统计量，即严格的信念状态（Belief State）。隐状态动力学模型 $p_\psi$ 由 3 层 GELU 激活的 MLP 构成，接收 LayerNorm 后的隐藏状态与 token 嵌入拼接向量，以残差形式预测状态增量： $$\hat{\mathbf{h}}_{t+1} = p_\psi(\mathbf{h}_t, X_{t+1}) = f_\psi(\mathbf{h}_t, X_{t+1}) + \mathbf{h}_t$$ 相比 token 空间稀疏离散的交叉熵监督，隐藏状态本身参数化了完整的未来后验概率分布，因此向主干网络反向传播了更加稠密、平滑的流形梯度，大幅提升样本效率。

2. 停止梯度的多步回归与 KL 散度自蒸馏：防止表征坍缩并对齐语义流形 自预测表征学习常面临所有表征退化为常数向量的“表征坍缩”风险。NextLat 在多步递归滚动预测时设计了严密的停止梯度机制。对于多步视界 $d$，模型使用教师强制的 token 序列 $X_{t+1:t+d}$ 在隐空间展开得到预测状态序列 $\hat{\mathbf{h}}_{t+i}$。在计算隐状态 Smooth L1 损失时，对目标真实隐状态施加停止梯度算子 $\mathbf{sg}[\mathbf{h}_{t+i}]$： $$\mathcal{L}_{\text{next-h}}(\theta,\psi; d) = \mathbb{E}_t \left[ \frac{1}{d} \sum_{i=1}^d \text{SmoothL1Loss}\big(\mathbf{sg}[\mathbf{h}_{t+i}], \hat{\mathbf{h}}_{t+i}\big) \right]$$ 同时，为了确保预测出的隐藏状态在语义分类空间严格忠实于真实状态，NextLat 引入互补的 KL 散度自蒸馏损失。将输出分类头参数冻结为 $p_\theta^{\mathbf{sg}}$，迫使动力学生成的隐状态所解码出的 logits 分布逼近真实隐藏状态解码出的软概率标签： $$\mathcal{L}_{\text{KL}}(\theta,\psi; d) = \mathbb{E}_t \left[ \frac{1}{d} \sum_{i=1}^d D_{\text{KL}}\!\left(p_\theta^{\mathbf{sg}}(\cdot \mid \mathbf{sg}[\mathbf{h}_{t+i}]) \;\|\; p_\theta^{\mathbf{sg}}(\cdot \mid \hat{\mathbf{h}}_{t+i})\right) \right]$$ 消融实验证实，停止梯度彻底消除了额外反向传播通过庞大词表输出头的计算开销，而 Smooth L1 与 KL 散度联合约束构成了长程稳定预测的关键支柱。

3. 变长自推测解码：突破固定视界约束的轻量推理加速 传统投机采样（Speculative Decoding）通常需要训练专用的草稿小模型或维护复杂的投机树，而基于多 token 预测（MTP）的自投机方案其投机长度被严格绑定在训练阶段设定的固定视界 $d$（例如 $d=2$ 或 $4$），且每个额外视界都需要新增额外的 Transformer 解码层。NextLat 彻底解除了这一限制：由于隐动力学模型学到了连续环境的内生转移规律，即使模型在预训练时仅采用极浅的监督视界（$d=1$ 或 $d=2$），在推理时也能以完全自回归的方式在隐空间无限递归生成多步隐藏状态 $\hat{\mathbf{h}}_{t+1}, \hat{\mathbf{h}}_{t+2}, \dots$，并借由共享输出头生成候选草稿 token。在验证阶段，主模型仅需单次前向传播即可并行校验整段草稿。这种动态调整草稿长度的能力显著降低了投机验证的循环次数，在各类语言建模任务中实现了超出训练视界的显著加速。

一个完整示例¶

以曼哈顿网格出租车轨迹规划任务为例，说明 NextLat 如何实现状态压缩与推演。 1. 输入与主干编码：给定一段随机漫步轨迹的前缀坐标序列 $X_{1:t} = [\text{Street } 42, \text{Avenue } 5, \dots]$，Transformer 主干将其映射为顶层隐藏向量 $\mathbf{h}_t \in \mathbb{R}^{384}$。 2. 潜在转移预测：当下一步动作 token 为“向北转”（$X_{t+1}$）时，隐状态动力学 MLP 将 $\mathbf{h}_t$ 与动作嵌入拼接，直接计算残差增量输出预测状态 $\hat{\mathbf{h}}_{t+1} \in \mathbb{R}^{384}$。 3. 状态充分性表现：若存在另一条完全不同的绕路轨迹 $X'_{1:j}$，只要其物理终点同样到达“第 42 街与第 5 大道交叉口”，经过 NextLat 训练后两者的隐状态在空间中高度重合。模型对其有效隐维度（Effective Latent Rank）由基线 GPT 的 160.1 骤降至 52.7，序列压缩一致性达到 0.71（GPT 仅 0.65）。 4. 自推测展开：在推理时，动力学模型直接以 $(\hat{\mathbf{h}}_{t+1}, \text{East}) \to \hat{\mathbf{h}}_{t+2}$ 展开数步，连续提出后续 5 个路口转弯动作，主干网络单次前向完成批量判定，验证接受率远超传统 MTP。

损失函数 / 训练策略¶

NextLat 的总体训练目标为三项损失的加权和： $$\mathcal{L}_{\text{NextLat}}(\theta, \psi; d, \lambda_{\text{next-h}}, \lambda_{\text{KL}}) = \mathcal{L}_{\text{next-token}}(\theta) + \lambda_{\text{next-h}} \mathcal{L}_{\text{next-h}}(\theta, \psi; d) + \lambda_{\text{KL}} \mathcal{L}_{\text{KL}}(\theta, \psi; d)$$ 在标准实验中，超参数设置极为鲁棒：通常取 $\lambda_{\text{next-h}} = 1.0$（部分回归值尺度极小的任务取 $2.0$），$\lambda_{\text{KL}} = 1.0$（曼哈顿任务取 $0.1$）。在语言模型预训练（如 FineWeb-Edu 100B tokens）中，优化器采用 AdamW，最大学习率 $4 \times 10^{-4}$，并配合 Warmup-Stable-Decay（WSD）学习率调度。对于多步训练，仅需在隐状态上展开极低开销的轻量 MLP，在 $d=1$ 时的训练吞吐与标准 GPT 几乎完全一致（3.09 it/s vs. 3.09 it/s）。

实验关键数据¶

主实验¶

论文在曼哈顿出租车世界模型评测基准与 1.3B 规模语言模型多领域推测解码加速上开展了深入对比。

表 1：曼哈顿出租车网格世界建模性能评测（多步视界 $d=8$） 评测指标包括：Next-Token 准确率、分布外（OOD）有效轨迹比例（Valid Trajectories）、相同终点序列压缩一致性（Sequence Compression）、隐状态有效秩（Effective Latent Rank，越低说明表征越紧凑）、以及 75% 随机绕路下的鲁棒性（Detour Robustness）。

模型	Next-Token 准确率 (↑)	有效轨迹比例 (↑)	序列压缩度 (↑)	有效隐维度 (↓)	绕路鲁棒性 (↑)
真实世界模型	100%	100%	1.00	—	100%
GPT (标准自回归)	100%	97.0%	0.65	160.1	85.0%
MTP ($d=8$)	100%	98.1%	0.64	57.7	95.0%
JTP ($d=8$)	100%	97.1%	0.32	215.8	87.0%
NextLat (Ours, $d=8$)	100%	98.7%	0.71	52.7	95.0%

表 2：1.3B 参数语言模型在多领域上的自推测解码加速比与单步平均接受 Token 数（8× NVIDIA B200 GPUs） 所有模型在 100B FineWeb-Edu tokens 上预训练。其中“接受 Token 数”不包含必被接受的当前首个预测。

模型	Wikipedia 加速比	Wiki 接受 Token	Books 加速比	Books 接受 Token	Code 加速比	Code 接受 Token	Math 加速比	Math 接受 Token
JTP ($d=1$)	1.46×	0.96	1.47×	0.97	1.47×	0.98	1.46×	0.97
JTP ($d=2$)	1.88×	1.84	1.90×	1.89	1.88×	1.85	1.89×	1.86
MTP ($d=1$)	1.38×	0.91	1.39×	0.95	1.40×	0.97	1.39×	0.95
MTP ($d=2$)	1.68×	1.72	1.72×	1.83	1.75×	1.91	1.72×	1.84
NextLat ($d=1$)	2.68×	3.52	2.72×	3.64	2.29×	2.66	2.30×	2.72
NextLat ($d=2$)	3.21×	4.59	3.32×	4.86	2.38×	2.83	2.87×	3.94

消融实验¶

在 TinyStories 合成故事建模基准上，作者固定模型结构并训练单层线性探针（Linear Probe），评估在冻结的隐藏状态下向未来跨步预测偏移量（1 到 20 个 tokens）的交叉熵损失差值（相对基线 GPT，负值越低表示预测能力越强），并分析不同损失与梯度截断设计。

实验配置	偏移 1 Token 相对交叉熵	偏移 10 Tokens 相对交叉熵	偏移 20 Tokens 相对交叉熵	机制分析说明
基线 GPT (标准自回归)	0.00	0.00	0.00	基线参考，隐状态缺乏长程预测约束
MTP ($d=8$)	+0.08	+0.15	+0.22	额外离散 token 目标损害了即时预测能力，长程严重衰减
JTP ($d=8$)	+0.10	+0.18	+0.25	离散联合预测同样带来局部负迁移
NextLat (仅 Smooth L1, $d=1$)	-0.01	-0.12	-0.18	纯状态回归已能强力诱导长程信念表征
NextLat (CE 替换 KL, $d=1$)	+0.02	-0.06	-0.10	离散硬交叉熵监督不如平滑概率分布对齐
NextLat (无 Target 截断梯度)	-0.00	-0.04	-0.05	缺失停止梯度使表征漂移，探针性能明显受损
NextLat (完整版, $d=1$)	-0.02	-0.14	-0.21	完美保留即时预测精度且显著增强 20 步长程预测
NextLat (完整版, $d=8$)	-0.02	-0.16	-0.24	扩大视界进一步强化长程相关性表征

关键发现¶

表征压缩与世界模型保真度：在曼哈顿出租车任务中，尽管所有模型均能达到 100% 的单步 Next-Token 准确率，但基线 GPT 重建出的内部地图充斥着荒谬的“高架桥”和断头路；而 NextLat 的有效隐状态秩仅为 52.7（相比 GPT 的 160.1 压缩了 3 倍以上），有效序列压缩度提升至 0.71，其内部拓扑几乎完全忠实于真实道路连通图。
长程规划杜绝“遗憾妥协”：在 Countdown 算术推理评测中，基线模型普遍在最终等式出现高达 50% 以上的计算错误（明知无法凑出目标数而强行输出错误等式）；NextLat 即使仅设置视界 $d=1$，准确率也从 GPT 的 33.1% 跃升至 54.8%，$d=8$ 时达到 58.7%，有效缓解了局部贪婪短视带来的规划崩溃。
超越训练视界的推测泛化：虽然 NextLat 仅以 $d=1$ 或 $d=2$ 训练，但在推测解码时单步平均接受 token 却高达 4.59～4.86 个，在 Books 领域实现 3.32× 的端到端吞吐提升，大幅超越训练了更多参数头的 MTP 与 JTP。

亮点与洞察¶

理论与实践的高度自洽：证明了隐藏状态满足单步下一 token 一致性与状态转移一致性即可收敛为充分统计量性质的信念状态，无需像传统强化学习世界模型那样构建复杂的对抗或重构框架。
以无循环之身行循环之实：利用 Transformer 全局自注意力在时间维度的高效并行性生成隐状态目标，同时用前向轻量 MLP 施加局部状态转移惩罚，兼得 Transformer 的训练吞吐与 RNN 的泛化归纳偏置。
突破性长度泛化发现：在非恒定深度不可表达的 $A_5$ 群交错置换代数问题中，一个仅在 12 步序列上联合训练的 2 层 Transformer 配合 NextLat 动力学 MLP，使得其动力学模块作为 RNN 独立推演时能够在 36 步超长序列上维持超过 95% 的极高状态追踪准确率，成功突破了常数深度 Transformer 的电路复杂性极限。

局限与展望¶

隐动力学架构设计探索不足：论文全篇仅采用朴素的 3 层 MLP 参数化隐动力学模型，未探讨更强表达能力的自适应瓶颈维度或多层跨步层次化隐结构对复杂逻辑推理的增益。
预训练后期 Smooth L1 损失翘尾现象：在 FineWeb-Edu 大规模预训练退火阶段，AdamW 优化器下出现隐状态回归损失不降反升的反常优化现象（虽然推测接受率持续上升），其几何流形演化成因仍有待深入探索。
自适应投机策略留白：目前推理时的投机采样长度仍为固定预设（静态 2 到 10 个 token），尚未实现在生成过程中根据隐状态不确定性动态调节草稿步长的弹性策略。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 巧妙将自预测强化学习思想与自回归序列建模融合，建立了严谨的隐状态信念理论。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖曼哈顿世界模型、Countdown 推理、Path-Star 图规划、TinyStories 探针以及 100B Token 级 LLM 多任务与推测推理。
写作质量: ⭐⭐⭐⭐⭐ 理论证明清晰，直觉比喻生动贴切，图表设计极其详实完整。
价值: ⭐⭐⭐⭐⭐ 不改推理架构即可显著提升紧凑世界模型表征质量并直接变现为 3.3× 的无损推理加速，具备广泛工业与学术应用前景。

模型	Wikipedia 加速比	Wiki 接受 Token	Books 加速比	Books 接受 Token	Code 加速比	Code 接受 Token	Math 加速比	Math 接受 Token
JTP (\(d=1\))	1.46×	0.96	1.47×	0.97	1.47×	0.98	1.46×	0.97
JTP (\(d=2\))	1.88×	1.84	1.90×	1.89	1.88×	1.85	1.89×	1.86
MTP (\(d=1\))	1.38×	0.91	1.39×	0.95	1.40×	0.97	1.39×	0.95
MTP (\(d=2\))	1.68×	1.72	1.72×	1.83	1.75×	1.91	1.72×	1.84
NextLat (\(d=1\))	2.68×	3.52	2.72×	3.64	2.29×	2.66	2.30×	2.72
NextLat (\(d=2\))	3.21×	4.59	3.32×	4.86	2.38×	2.83	2.87×	3.94