Next-Latent Prediction Transformers Learn Compact World Models
Jayden Teoh ⋅ Manan Tomar ⋅ Kwangjun Ahn ⋅ Edward Hu ⋅ Tim Pearce ⋅ Pratyusha Sharma ⋅ Akshay Krishnamurthy ⋅ Riashat Islam ⋅ Alex Lamb ⋅ John Langford
Abstract
Transformers replace recurrence with a memory that grows with sequence length and self-attention that enables ad-hoc look ups over past tokens. Consequently, they lack an 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 output token. Theoretically, we show that these latents provably converge to *belief states*, compressed information of the history necessary to predict the future. This simple auxiliary objective injects a recurrent inductive bias into transformers, while leaving their architecture, parallel training, and inference unchanged. NextLat effectively encourages the transformer to form compact internal world models with its own belief states and transition dynamics—a crucial property absent in standard next-token prediction transformers. Empirically, across benchmarks in world modeling, reasoning, planning, and language modeling, NextLat demonstrates significant gains over standard next-token training in downstream accuracy, representation compression, and lookahead planning. Furthermore, NextLat enables *variable-length self-speculative decoding*, accelerating inference by up to $3.3\times$ in the language domain. NextLat stands as a simple and efficient paradigm for shaping transformer representations toward stronger generalization.
Chat is not available.
Successful Page Load