🤖 machine learning

Next-Latent Prediction Transformers Learn Compact World Models

The paper introduces Next-Latent Prediction (NextLat), a simple auxiliary training objective that compels transformers to learn compact, predictive latent world models by enforcing self-supervised predictions in the latent space, thereby improving generalization, representation compression, and inference speed without altering the model architecture.

Original authors: Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Tim Pearce, Pratyusha Sharma, Akshay Krishnamurthy, Riashat Islam, Alex Lamb, John Langford

Published 2026-05-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Tim Pearce, Pratyusha Sharma, Akshay Krishnamurthy, Riashat Islam, Alex Lamb, John Langford

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Library" vs. The "Diary"

Imagine you are trying to learn a new language. You have two ways to study:

The Library Approach (Standard Transformers): You keep a massive library of every single book you've ever read. When you need to write the next sentence, you run a search engine through your entire library to find the perfect word based on the last few words you wrote. It's powerful, but your library grows infinitely big as you read more. You never really "summarize" what you've learned; you just keep looking things up.
The Diary Approach (Recurrent Models): You keep a small diary. Every day, you read the news, update your diary with a short summary of what happened, and then write the next day's entry based only on that diary. Your diary stays the same size, no matter how many years you live.

The Issue: Modern AI (Transformers) uses the "Library Approach." It's incredibly fast and smart, but because it doesn't have to summarize everything into a small diary, it often learns "shortcuts." It might memorize specific patterns in the training data without actually understanding the underlying rules of the world. This makes it bad at planning ahead or handling new situations it hasn't seen before.

The Solution: Next-Latent Prediction (NextLat)

The researchers introduced a new training method called Next-Latent Prediction (NextLat). Think of this as forcing the "Library" AI to start keeping a "Diary" without changing how it reads books.

Here is how it works:

The Standard Task: The AI still tries to guess the next word in a sentence (like normal).
The New Trick: The AI is also forced to guess what its own "internal thought" (its hidden state) will be in the next step, before it even sees the next word.

The Analogy: Imagine you are playing a video game.

Normal AI: You look at the screen, see a monster, and press a button to attack. Then you look at the screen again for the next move.
NextLat AI: You look at the screen, predict exactly how your character's internal status bar (health, energy, position) will change before the game actually updates. You are learning the "physics" of the game world, not just memorizing which button to press.

By forcing the AI to predict its own future "thoughts," it is forced to compress all the history it has seen into a compact, consistent summary (a Belief State). It learns a "World Model"—a mental map of how things work—rather than just a list of word associations.

What Did They Find?

The paper claims that this simple trick makes the AI much smarter in four specific areas:

1. Better Maps (World Modeling)

The Test: They trained AI on taxi rides in Manhattan.
The Result: A normal AI could predict the next street name perfectly but had a "hallucinated" map where roads went through buildings or loops didn't make sense. The NextLat AI learned a map that actually looked like Manhattan. It understood that if you turn left, you are now on a different street, and it kept track of its location logically.

2. Better Math and Logic (Reasoning)

The Test: A puzzle called "Countdown" where you have to use math to reach a target number.
The Result: Normal AIs often get stuck and make a mistake at the very end, trying to force a wrong answer to match the goal. NextLat AIs planned ahead better, avoiding those "regretful compromises" and solving the puzzle more accurately.

3. Better Navigation (Planning)

The Test: Finding a path through a complex maze (Path-Star graph).
The Result: Normal AIs often take shortcuts that look good for one step but lead to dead ends. NextLat AIs looked further ahead, understanding the whole structure of the maze, and solved it almost 100% of the time.

4. Faster Reading (Language Modeling)

The Test: Generating text.
The Result: Because the NextLat AI has a good "mental model" of the future, it can guess several words ahead of time with high confidence. This allows it to "speculate" and read/write faster. The paper claims this can make the AI 3.3 times faster at generating text compared to standard models.

Why Is This Special?

Usually, to make an AI smarter at planning, you have to change its architecture (make it slower or more complex) or train it on massive amounts of data.

NextLat is special because:

It's a "Plug-in": You don't have to rebuild the AI. You just add a small, simple "helper" task during training.
It Keeps the Speed: The AI still trains and runs as fast as before.
It's Theoretically Sound: The math proves that this method forces the AI to create a "sufficient statistic"—a perfect, compact summary of the past needed to predict the future.

Summary

The paper argues that by teaching Transformers to predict their own future "thoughts" (latent states), we force them to build a compact, logical map of the world. This makes them better at planning, reasoning, and understanding complex structures, all while making them faster to run. It's like teaching a student not just to memorize the textbook, but to understand the underlying principles so they can solve problems they've never seen before.

Technical Summary: Next-Latent Prediction Transformers Learn Compact World Models

Problem Statement

Standard Transformer architectures replace recurrence with self-attention and a memory mechanism that scales with sequence length. While this enables exceptional parallelization and predictive power, it removes the inherent inductive bias to compress history into compact latent states with consistent transition rules. Consequently, Transformers often learn complex, task-specific shortcuts that fit training data well but generalize poorly, failing to form coherent "world models" or belief states necessary for long-horizon planning and reasoning. The standard next-token prediction objective does not guarantee that the model's internal hidden states evolve as sufficient statistics for the past history, nor does it ensure that these states can recursively predict future observations.

Methodology: Next-Latent Prediction (NextLat)

The authors introduce Next-Latent Prediction (NextLat), a framework that augments standard next-token training with self-supervised predictions in the latent space. The core idea is to train the Transformer to learn latent representations (specifically, the final-layer hidden states $h_t$ ) that are predictive of the next latent state given the next token.

Theoretical Foundation

The paper establishes that if a model optimizes for both Next-Token Consistency (predicting $X_{t+1}$ from $h_t$ ) and Transition Consistency (predicting $h_{t+1}$ from $h_t$ and $X_{t+1}$ ), the hidden states $h_t$ are theoretically guaranteed to converge toward belief states. A belief state is defined as a sufficient statistic of the history $X_{1:t}$ for predicting the future $X_{t+1:T}$ . Unlike standard autoregressive training, which allows ad-hoc lookups of past tokens without compressing information, NextLat enforces a recurrent inductive bias where the latent state must contain all necessary information to transition to the next state and predict the next token.

Implementation

NextLat operates by adding a lightweight auxiliary objective to the standard Transformer training loop:

Latent Dynamics Model ( $p_\psi$ ): A simple Multi-Layer Perceptron (MLP) is trained to predict the next hidden state $\hat{h}_{t+1}$ given the current hidden state $h_t$ and the next token $X_{t+1}$ .
Loss Function: The total loss combines:
- Next-Token Loss ( $L_{next-token}$ ): Standard cross-entropy for token prediction.
- Next-Hidden State Loss ( $L_{next-h}$ ): A Smooth L1 loss regressing the predicted state $\hat{h}_{t+1}$ against the true next state $h_{t+1}$ (with stop-gradient applied to the target to prevent representational collapse).
- KL Divergence Loss ( $L_{KL}$ ): A complementary objective ensuring the token distribution predicted from the predicted latent state matches the distribution from the true latent state, acting as a soft supervision signal.
Multi-Step Supervision: The model can be trained to recursively unroll the latent dynamics for $d$ steps, providing richer gradient signals for long-range dependencies.

Crucially, NextLat leaves the Transformer's architecture, inference procedure, and parallel training efficiency unchanged. The latent dynamics model is used only during training to shape representations; during inference, the model decodes autoregressively as usual.

Key Contributions

Theoretical Guarantee: The paper proves that optimizing for next-token and next-latent consistency provably shapes transformer representations into belief states, a property not guaranteed by next-token prediction alone.
Practical Framework: NextLat introduces a computationally efficient auxiliary objective that injects recurrent inductive bias without altering the Transformer's parallel training or inference structure.
Variable-Length Self-Speculative Decoding: By learning coherent latent dynamics, NextLat enables recursive prediction in latent space. This allows for variable-length self-speculative decoding, where the model can draft tokens beyond the fixed training horizon, significantly accelerating inference.

Experimental Results

The authors evaluate NextLat across four domains: world modeling, reasoning, planning, and language modeling.

World Modeling (Manhattan Taxi Rides): NextLat learns internal maps significantly more consistent with the true world model (street topology) compared to standard GPT, MTP, and JTP baselines. It achieves higher sequence compression (0.71 vs. 0.65 for GPT) and a much lower effective latent rank (52.7 vs. 160.1 for GPT), indicating more compact and efficient representations.
Reasoning (Countdown): On the mathematical reasoning task, NextLat substantially outperforms baselines (54.8% accuracy at $d=1$ vs. 33.1% for GPT). It demonstrates superior lookahead planning, reducing "regretful compromises" where models force invalid final equations to match a target.
Planning (Path-Star Graphs): NextLat maintains near 100% solve rates across various graph topologies, whereas baselines like BST and JTP fail on larger graphs ( $G_{7,7}$ ). This suggests NextLat avoids the "Clever Hans" shortcut learning common in token-space prediction.
Language Modeling (FineWeb-Edu):
- Downstream Accuracy: NextLat preserves next-token perplexity better than MTP and JTP while achieving modest gains in downstream zero-shot accuracy (59.21% avg vs. 58.82% for GPT).
- Inference Speed: Leveraging variable-length self-speculative decoding, NextLat achieves up to 3.3× faster inference compared to standard autoregressive sampling. This speedup persists even when drafting beyond the training horizon, a capability MTP and JTP lack due to their fixed token-space horizons.

Significance and Claims

The paper positions NextLat as a simple yet effective paradigm for learning compact, predictive, and generalizable representations in Transformers. Its primary significance lies in unifying the recurrent inductive bias of compact world models with the scalability and parallelism of Transformers.

The authors claim that NextLat offers a solution to the "myopic" nature of next-token prediction, encouraging models to form structured, predictive internal states rather than relying on shallow token correlations. By enabling variable-length speculative decoding without architectural changes, NextLat provides a practical path to faster inference and better long-horizon reasoning. The paper concludes that NextLat is a step toward training objectives that endow autoregressive models with simpler, more compact, and therefore more generalizable representations of complex systems.

The authors remain modest regarding future directions, noting that they did not explore complex latent dynamics architectures (using only simple MLPs) and that the optimal design of the objective (e.g., the necessity of KL loss at larger scales) requires further systematic study. They also acknowledge that while NextLat improves representations, the underlying circuit complexity of the fixed-depth Transformer remains bounded by the TC0 class, though the co-trained latent dynamics model (RNN) appears capable of generalizing beyond these limits.