让Transformer学会压缩历史,生成更通用的预测模型。
Next-Latent Prediction Transformers Learn Compact World Models
- 在自注意力基础上增加隐状态预测任务,引导模型学习紧凑的内部世界模型。
- 在多个任务中提升准确率,实现3.3倍的推理加速和更强的规划能力。
- 适合需要高效推理与泛化能力的场景,如语言建模与决策规划。
Transformer通过自注意力机制替代循环结构,但缺乏将历史信息压缩为一致状态转移的内在激励,常导致泛化能力差。本文提出下一隐状态预测(NextLat),在标准下一个词预测基础上,引入隐空间中的自监督预测任务:训练Transformer根据下一个词预测其下一隐状态。理论上,这些隐状态收敛至信念状态,即预测未来所需的最小历史信息。该简单辅助目标注入了循环归纳偏置,不改变Transformer架构、并行训练效率和推理方式。实验表明,NextLat在世界建模、推理、规划和语言建模等多个基准上,显著优于标准下一个词预测及其他基线,在下游准确率、表示压缩和前瞻规划方面均有提升。此外,支持可变长度自推测解码,使语言建模推理速度最高加快3.3倍。代码已开源。
原文摘要 · Abstract (English)
Transformers replace recurrence with a memory that grows with sequence length and self-attention that enables ad-hoc lookups over past tokens. Consequently, they lack an inherent incentive to compress history into compact latent states with consistent transition rules. This often leads to learning solutions that generalize poorly. We introduce Next-Latent Prediction (NextLat), which extends standard next-token training with self-supervised predictions in the latent space. Specifically, NextLat trains a transformer to learn latent representations that are predictive of its next latent state given the next token. Theoretically, we show that these latents provably converge towards belief states, compressed information about the history necessary to predict the future. This simple auxiliary objective injects a recurrent inductive bias into transformers while leaving their architecture, parallel training efficiency, and inference unchanged. NextLat effectively encourages transformers to form compact internal world models with coherent belief states and transition dynamics -- crucial properties not guaranteed by standard next-token prediction alone. Empirically, across benchmarks in world modeling, reasoning, planning, and language modeling, NextLat demonstrates significant gains over standard next-token prediction and other baselines in downstream accuracy, representation compression, and lookahead planning. Furthermore, NextLat enables variable-length self-speculative decoding, accelerating inference by up to 3.3x in language modeling. NextLat offers a simple yet effective paradigm for learning compact, predictive representations in transformers that generalize better. Our code is available at https://github.com/JaydenTeoh/NextLat.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。