Next-Latent Prediction Transformers Learn Compact World Models
Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Tim Pearce, Pratyusha Sharma, Akshay Krishnamurthy, Riashat Islam, Alex Lamb, John Langford
cs.LG
2025-11-08
NextLat adds next-hidden-state prediction so transformers form belief states. Manhattan latent rank falls from 160.1 (GPT) to 52.7; 1.3B models get up to 3.3x speculative speedup.
Transformers swapped recurrence for a KV cache that grows with length and attention that can look up any past token on demand. That removes the pressure to crush history into a fixed-size state with a consistent update rule. Next-token training only asks for the next symbol to be right, so models can fit the training set with ad-hoc lookups and still fall apart off distribution. Vafa et al. showed this on Manhattan taxi traces: legal-turn accuracy can hit 100% while the model's internal map invents impossible street orientations and flyovers.
Belief State Transformers already showed that next-token consistency alone does not force hidden states to be belief states, i.e. sufficient statistics of the past for predicting the future. NextLat puts that pressure back without touching architecture, parallel training, or the inference path. It adds a self-supervised next-hidden-state objective so the model has to keep compact, recurrent-like dynamics in latent space.
The transformer still emits a last-layer hidden state ht at each step. A small latent dynamics model pψ, an MLP in all experiments, takes (ht, next token) and predicts h{t+1}. The token is treated as an action.
The training loss has three terms:
Theorem 3.2 says that if both next-token consistency and transition consistency hold, ht must be a belief state. The guarantee already holds at d=1; longer horizons only densify the signal. At inference the transformer can decode on its own. pψ is there to shape representations, and because it can be rolled in latent space it also supports variable-length self-speculative decoding, with draft lengths that are not capped at the training horizon.
The learning signal is denser than a one-hot next token. The target is a hidden state that parameterizes a full next-token distribution, and recursive consistency leaks information about later states into earlier ones.
On Manhattan random walks (91M sequences, 4.7B tokens, 6 epochs, d=8), every method scores 100% on the next-token test. World-model quality is where they split:
| Method | Valid trajectories | Sequence compression | Effective latent rank ↓ | Detour robustness |
| GPT | 97.0% | 0.65 | 160.1 | 85.0% |
| MTP | 98.1% | 0.64 | 57.7 | 95.0% |
| JTP | 97.1% | 0.32 | 215.8 | 87.0% |
| NextLat | 98.7% | 0.71 | 52.7 | 95.0% |
Effective rank is more than 3x smaller than GPT. Reconstructed maps have fewer, mostly local, illegal edges. The true graph has 4,580 intersections, so a flat latent space is the right shape.
On Countdown, NextLat at d=1 already reaches 54.8%, versus 39.2% MTP and 39.0% JTP at the same horizon; at d=8 it is 58.7%, just above MTP's 57.3%. Final-equation validity is 54.8% against 42.3% for the next-best baseline, fewer last-step regretful compromises. On Path-Star, using the harder original setup (200k fixed samples, nodes drawn from 1 to 100), NextLat stays near 100% on G2,10, G5,5, and G7,7; BST starts failing on G7,7. TinyStories linear probes show MTP, JTP, and BST hurting next-token prediction and dropping off at longer offsets; NextLat matches GPT at offset 1 and is strongest out to 20 tokens.
A 1.3B model trained on 100B FineWeb-Edu tokens gets 59.21 average zero-shot accuracy for NextLat (d=2) versus 58.82 for GPT: a 0.39 gain that does not hold on every task. FineWeb-Edu perplexity is 10.52 for GPT and 10.83-10.88 for NextLat, a bit worse than GPT but better than MTP and JTP. The real gap is self-speculative decoding: NextLat (d=2) speeds up Wikipedia / Books / Code / Math by 3.21x / 3.32x / 2.38x / 2.87x, with accepted draft tokens far beyond the training horizon. At d=1, training throughput on one B200 matches GPT at 3.09 iter/s; BST sits at 0.89.
On the A5 word problem, a 2-layer transformer cannot length-generalize past the 12-token training window. The co-trained 2.62M-parameter RNN, after one transformer-initialized step, exceeds 95% at 36 tokens. A GPT trained directly on 36-token sequences still fails.
This is an auxiliary loss, not a new backbone. Any decoder-only stack can attach an MLP. Gains are real on world modeling, combinatorial search, and planning tasks that punish n-gram shortcuts. Multiple-choice LM scores barely move, but the same dynamics model doubles as a variable-length drafter: train at a shallow horizon, draft up to 10 tokens at inference. That is cheaper than stacking extra MTP layers.
At 1.3B this looks like cleaner representations plus faster decoding, not a broad benchmark lift. It is worth trying if you already train transformers and can afford one extra loss. It is not yet a reason to expect MMLU-style jumps.
The authors stay conservative. Dynamics are always simple MLPs; they never sweep hidden width, even though that width is a capacity bottleneck, and smaller latent sizes helped on Path-Star and Countdown. Stop-gradient, KL, and Smooth L1 come from small ablations; whether d>1 or the KL term still matter at scale is open. They skipped newer MTP variants. Speculative draft length was a static sweep from 2 to 10, not an adaptive policy. Smooth L1 can rise during pretraining, so the objective is optimizer-sensitive.
Manhattan valid-trajectory gains are only 1.7 points over GPT; compression and rank carry the world-model claim. Countdown and Path-Star are synthetic. The 1.3B downstream bump is 0.39 and the paper itself says larger models may be needed. The A5 RNN result is striking but sits on a 2-layer, short-sequence setup; it does not show that production transformers have left TC0.