Next-Latent Prediction Transformers Learn Compact World Models

Transformer模型以随序列长度增长的记忆机制和可对过往标记进行即席查询的自注意力机制取代了循环结构。因此，这类模型缺乏一种内在动力，去将历史信息压缩为具有稳定转移规则的紧凑潜在状态，往往导致学习出泛化能力较差的解决方案。我们提出了“下一潜在状态预测”（NextLat）方法，在标准的下一个标记预测训练基础上，引入了在潜在空间中的自监督预测任务。具体而言，NextLat训练Transformer学习这样的潜在表征：在给定下一个输出标记的条件下，该表征能够预测其下一个潜在状态。理论上，我们证明这些潜在状态能够收敛到“信念状态”——即压缩后的历史信息，恰好包含预测未来所需的信息。这一简单的辅助目标同时向Transformer注入了一种循环结构的归纳偏置，而无需改变其原有架构、并行训练方式或推理过程。NextLat有效地促使Transformer构建起包含自身信念状态和状态转移动态的紧凑内部世界模型，而这一关键特性在标准的下一个标记预测Transformer中是缺失的。实验结果表明，在面向核心序列建模能力（包括世界建模、推理、规划和语言建模）的多个基准任务上，NextLat在下游任务准确率、表征压缩效率以及前瞻性规划能力方面均显著优于标准的下一个标记预测训练方法。NextLat作为一种简单高效的范式，能够引导Transformer的表征学习朝着更强的泛化能力发展。

标准Transformer在序列建模中依赖自注意力机制和可变长度的记忆，缺乏将历史信息压缩为紧凑且具有一致转移规律的隐状态的内在动机，导致学习到的表示泛化能力差。论文试图验证：通过引入对隐状态的预测任务，能否促使Transformer形成更紧凑、更具预测性的内部世界模型，从而提升其在序列建模中的泛化能力。这是一个重要但尚未被充分解决的问题，尤其在需要长期依赖和推理的任务中。

提出Next-Latent Prediction（NextLat），在标准的下一词元预测目标之外，增加一个自监督的隐空间预测任务：训练Transformer使其当前隐状态能预测下一个隐状态，条件是已知下一个输出词元。该方法理论上可使隐状态收敛到贝尔曼最优信念状态（即压缩后的历史信息），并注入循环系统的归纳偏置，而无需改变Transformer的结构、并行训练或推理效率。

理论证明NextLat的隐状态会收敛到信息论意义上的最小充分统计量（即信念状态）；实验涵盖世界建模、推理、规划和语言建模等多个基准，在下游准确率、表示压缩性和前瞻规划能力上均显著优于标准训练方式；方法简单高效，仅需额外轻量级头模块，不增加推理成本；代码已开源，便于复现与扩展；值得深入研究的方向包括将其应用于强化学习中的状态表示学习、长序列建模以及多模态序列预测。

1. A Structured State Space for Sequence Modeling (2023) 2. RetNet: The Return of the Implicit Attention (2023) 3. Hyena Hierarchy: Towards Larger Convolutional Language Models (2023) 4. State Space Models as Attractors in Transformer Dynamics (2024) 5. Efficient Transformers for Long-Sequence Modeling: A Survey (2023)