Next-Latent Prediction Transformers Learn Compact World Models

TL;DR

Next-Latent Prediction Transformers introduce a self-supervised latent space prediction task that encourages learning compact, belief-state-like representations, improving generalization and performance in sequence modeling tasks.

Contribution

The paper proposes NextLat, a novel auxiliary training objective for transformers that induces belief-state-like latent representations, enhancing their ability to model and generalize in sequential tasks.

Findings

Significant improvements in downstream accuracy across benchmarks.

Enhanced representation compression and interpretability.

Better performance in reasoning, planning, and language modeling tasks.

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 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 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 also 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 targeting core sequence modeling competencies -- world modeling, reasoning, planning, and language modeling -- NextLat demonstrates significant gains over standard next-token training in downstream accuracy, representation compression, and lookahead planning. NextLat stands as a simple and efficient paradigm for shaping transformer representations toward stronger generalization.