Next-Latent Prediction Transformers Learn Compact World Models

Teach a transformer to predict its own next internal state, and it compresses the past like a recurrent network.

A transformer can look up any past token on demand, so nothing pushes it to summarize the story so far. NextLat adds one self-supervised loss that does, turning the internal states into compact, predictive summaries, while leaving the architecture, parallel training, and inference untouched.

A transformer can name the correct next turn of a taxi ride every single time, and still hold a map of the city that folds streets over one another and falls apart the moment a road is closed.

That is not a thought experiment. Vafa and colleagues trained transformers on millions of turn-by-turn taxi routes through Manhattan, where the "true" answer is just the street map. The models hit near-perfect next-token accuracy and generated valid routes the vast majority of the time, and yet when you reconstruct the map they have internally, it is incoherent: streets at impossible angles, roads flying over other roads, phantom connections. Ask the model to route around a detour and it is lost. High surface accuracy sat right on top of a broken world model, and next-token accuracy, the usual metric, stayed near-perfect the whole time.

The paper this explainer is about, Next-Latent Prediction (NextLat, out of Microsoft Research), is a small, almost surgical fix for that failure. It adds one auxiliary training loss and no new architecture, and across world modeling, reasoning, planning, and language modeling the resulting transformers form internal states that are measurably more compact and more predictive of the future. It rests on a chain of ideas that are each simple on their own: why attention removes the pressure to compress, what a "belief state" is, how you train a network to predict its own future representations without it cheating, and why the result amounts to a recurrent network trained inside a transformer.

A map that passes every quiz

Start with the puzzle, because it names the target precisely. A model that has truly learned the Manhattan map should not care how you arrived at a given intersection. Two taxis that reach the same corner, both bound for the same destination, face the identical remaining problem, so a model with a coherent map would predict the identical next turn for both. A model that has instead memorized surface correlations can keep the two histories apart and answer differently for each, which is exactly what an incoherent map looks like from the inside.

The paper measures this directly with a metric called sequence compression: the fraction of cases where, given two different histories that end at the same state and head to the same place, the model predicts the same continuation for both. A perfect world model scores 1.0. A plain next-token transformer (the paper calls it GPT, since it is trained with ordinary next-token prediction like GPT-3) scores 0.65. The task NextLat sets itself is to raise that number without wrecking the near-perfect next-token accuracy that hides the problem.

Attention removes the pressure to compress

Why does a transformer end up here? Look at the mechanism it replaced. A recurrent network reads a sequence one token at a time and carries a single fixed-size state vector forward, updating it at each step. To predict token 500 it has nothing to consult but that one state, so it is forced to distill everything relevant about the first 499 tokens into a fixed-size summary. Compression is not optional; it is the only memory the model has.

A transformer threw that constraint out, on purpose, and won by doing so. Self-attention lets any position reach back and read any earlier token directly, and that random access is what makes training parallel and long-range dependencies easy. The cost is subtle. Because the model can always re-read the raw past, nothing forces the hidden state at step $t$ to be a self-contained summary of what came before. The set of things it can look up grows with the sequence; the pressure to fold the past into a compact running state is gone. Drag the query position below and watch the difference: attention's reach grows with the sequence, while a recurrent-style belief state stays one fixed-size box no matter how long the input gets.

This is not a claim that transformers cannot form compressed summaries. Trained next-token transformers do, in fact, encode a surprising amount of belief-state structure in their activations. The precise claim is narrower, and everything downstream depends on it: ordinary next-token training gives the model no guarantee and no pressure to keep a compact, updatable summary at every step, so whether it forms one is left to chance, and on structured tasks it often does not. The question the paper poses is whether you can put that pressure back without giving up the transformer's parallelism. Its answer starts by naming exactly what a good summary is.

Belief states: one compressed summary

The concept the paper borrows is the belief state. In the study of partially observed systems, if you cannot see the true state of the world and only receive a stream of observations, the belief state is the smallest thing you need to remember so that nothing else in the raw history would help you predict what comes next. It is a sufficient statistic of the past for the future: once you hold it, the original sequence is redundant.

The paper states it precisely. Writing $X_{1:t}$ for the tokens seen so far and $\mathbf{b}_t = g(X_{1:t})$ for a summary computed from them, $\mathbf{b}_t$ is a belief state when, for every bounded function $f$ of the future, conditioning on the summary gives the same expectation as conditioning on the entire history:

Read the two sides as a promise: whatever you wanted to predict about the future, the compact summary $\mathbf{b}_t$ tells you as much as the full transcript would. A worked example makes it concrete. To predict the rest of a taxi route you do not need the exact turns that got you to a corner, only which corner you are on and where you are going; that pair is a belief state, and it is far smaller than the route history. A recurrent network is built to find exactly this kind of summary and a transformer is not. Simpler summaries that still predict everything are, in the sense learning theory has always meant, the ones that generalize.

Two caveats travel with the definition, and the paper leans on both later. First, it defines a sufficient statistic, not the smallest one: the raw history $X_{1:t}$ trivially satisfies the equation too, since it obviously predicts the future as well as itself. So "is a belief state" and "is compact" are two different properties, and only the first will be guaranteed. Second, this is the predictive flavor of the idea (reproduce the distribution over future observations), not necessarily a probability distribution over hidden world-states. Keep those straight and the theorem below reads correctly.

Predicting your own next latent

Alongside ordinary next-token prediction, NextLat trains the transformer so that its own hidden states obey a consistent transition rule: from the current hidden state and the next token, a small model can predict the next hidden state. The transformer keeps producing its states in one parallel pass, as always, and a lightweight side network learns to hop from each state to the next.

First, some vocabulary the paper is careful about. The hidden state $\mathbf{h}_t$ is the final transformer layer's output at position $t$, the vector just before it is turned into token logits. The output head $p_\theta$ maps that vector to a next-token distribution, and it is trained the usual way, with cross-entropy:

The new piece is a latent dynamics model $p_\psi$. It takes the current hidden state and the next token and predicts the next hidden state. In the paper it is deliberately tiny: a three-layer MLP with GELU activations, run on the layer-normalized concatenation of $\mathbf{h}_t$ and the next token's embedding, and it predicts a delta added back to the current state, a residual update:

Why an MLP and not something bigger? Because the point is to test whether the objective shapes good representations, so the paper keeps the transition model as small and dumb as it can. And why residual? Because a one-step transition should mostly leave the state alone and nudge it, and because a residual step composes cleanly when you chain it (that composability powers the speedup later). Chained step after step, that residual MLP behaves like the recurrent cell hiding inside the transformer.

With $p_\psi$ in hand, NextLat trains its two states to agree. Its ideal target is transition consistency: the predicted next state should match the state the transformer actually produces after seeing the next token. For a deterministic transformer that real next state is a fixed vector, so this is plain regression. The paper uses a Smooth-L1 (Huber) loss, which behaves like squared error for small gaps and like absolute error for large ones, so early on, when predictions are far off, a few outliers cannot dominate the gradient:

Two symbols there need cashing out. The $\mathbf{sg}[\cdot]$ is a stop-gradient: the target hidden state is treated as a fixed constant, with no gradient flowing into it. The next section is entirely about why that matters. The horizon $d$ is how many future steps the loss covers; at $d>1$ the prediction $\hat{\mathbf{h}}_{t+d}$ is a recursive rollout, feeding $p_\psi$'s own output back in $d$ times, so the transition rule has to stay consistent over several steps rather than one. The default is $d=1$, which is already enough for the guarantee below; larger $d$ only adds richer signal.

The paper adds one more term, a KL loss, which is easiest to read as a self-consistency check in token space. Push the true next state and the predicted next state each through the (frozen) output head, and require the two next-token distributions to match:

This is the same idea as knowledge distillation, except the teacher and student are the same model: the predicted state is taught to yield the same soft next-token distribution as the real one, which supervises $p_\psi$ not just on where the next state lands in vector space but on what it would predict. The full objective is the sum, with two coefficients that both default to 1:

The figure below shows the middle term doing its work. Each teal dot is a hidden state the transformer produces as a sequence unfolds; together they trace a path through the model's latent space. From each state, $p_\psi$ predicts the next one (the amber rings), with the dashed line the Smooth-L1 gap being closed. Drag training forward and watch every ring slide onto the next dot: the predicted latents come to track the true path, which is transition consistency.

The mechanism hinges on where the gradient flows. The target next state $\mathbf{h}_{t+1}$ is stopped, but the input state $\mathbf{h}_t$ that $p_\psi$ reads is not. So the auxiliary loss does not only train the little MLP; its gradient flows back through $\mathbf{h}_t$ into the whole transformer. The transformer is being pushed to produce states that a simple, fixed rule can step forward, and that pressure compresses the past into a belief state. The little MLP is only a vehicle for training the transformer's representation.

To make the shapes concrete, take the paper's largest model: a 1.3-billion-parameter transformer with hidden width $D = 2048$. A forward pass produces hidden states $\mathbf{h}$ of shape (batch, sequence, 2048). The dynamics model $p_\psi$ reads a 2048-vector state and a 2048-dimensional token embedding, so 4096 numbers enter a three-layer MLP ($4096 \to 4096 \to 2048$) that outputs a 2048-vector delta added back to the state. That MLP is about 42 million parameters, a few percent of the 1.3 billion, and it is discarded after training. One step of the full objective reads:

# one NextLat training step (default multi-step horizon d = 1)
h      = transformer(x)                  # hidden states, shape [B, T, D]
logits = head(h)                         # next-token distribution
loss_tok = cross_entropy(logits, next_tokens)   # the usual objective

# predict the model's OWN next hidden state from (h_t, next token)
h_hat  = p_psi(h[:, :-1], x[:, 1:])      # p_psi: 3-layer MLP, residual
target = h[:, 1:].detach()               # stop-gradient on the target
loss_lat = smooth_l1(h_hat, target)      # pull h_hat onto the true next h

# align their next-token distributions (output head frozen)
loss_kl = kl(logits[:, 1:].detach(), frozen_head(h_hat))

loss = loss_tok + loss_lat + loss_kl     # both lambdas default to 1
loss.backward()   # grad flows through p_psi AND back into the transformer

Backpropagation touches $p_\psi$ and, through the un-detached input $\mathbf{h}_t$, the entire transformer; only the target states are frozen. Nothing about the forward pass or the architecture changed, so the same trained network later decodes token by token exactly as any transformer does.

Why the states don't collapse

A reader who has seen self-supervised learning will feel an alarm here. If you reward a model for predicting its own future states, there is a way to win the game for nothing: make every state identical. Push all the $\mathbf{h}_t$ to one constant vector and the prediction is always exactly right, the loss goes to zero, and the representation now carries no information at all. This trivial solution, representational collapse, is the classic failure of methods that predict their own targets.

The stop-gradient rules that solution out by breaking the symmetry. Because the target state is detached, it is a fixed goalpost rather than a moving one: only the prediction side is free to move toward the target, and the target cannot come rushing back to meet it. The states can no longer collude to collapse together, because neither is allowed to chase the other. This is the same asymmetry that keeps self-predictive methods like SimSiam and I-JEPA from collapsing when they predict one view of an input from another. NextLat uses the simplest form of it, a hard stop-gradient on the same network, with no momentum or moving-average target.

There is a second guard, and it is specific to language. NextLat is still doing next-token prediction all along, and a constant hidden state cannot produce varying next-token distributions, so the cross-entropy loss already penalizes collapse on its own. The paper notes this, and also notes that adding the stop-gradient still helps in practice. Toggle it off in the figure and drag training forward: the prediction loss falls just as nicely, but the cloud of states rushes together to a single point and the state-spread meter crashes: a low loss bought only by throwing the representation away.

The guarantee, and its fine print

The theorem states it directly: if this training reaches its ideal, the hidden states are belief states. Concretely, suppose two conditions hold exactly for every position. The output head is a perfect next-token predictor,

and the dynamics model is a perfect transition predictor,

Then each $\mathbf{h}_t$ is a belief state for $X_{1:t}$. The proof is a short backward induction, and the intuition is a decoding chain. If (1) holds you can read the next token off $\mathbf{h}_t$; if (2) holds you can update to the next hidden state given that token; repeat, and you regenerate the entire future distribution from $\mathbf{h}_t$ alone:

Anything from which you can reconstruct the full future is a sufficient statistic of the past, which is the definition. And notice the pressure built into condition (2): the next state has to be predictable from the current state and the next token alone, without re-reading the history, and that demand is precisely what forces the state to already contain everything relevant. It holds at $d=1$; multi-step supervision is only there to teach the rule faster.

This matters because next-token consistency by itself does not get you a belief state. The Belief State Transformer paper constructs a distribution and a perfect next-token predictor whose hidden state is provably not a belief state. A model can predict every next token correctly and still fail to summarize the past, which is the Manhattan failure restated as a theorem. Condition (2) is the ingredient standard training was missing.

Read the theorem for exactly what it claims, because the abstract's phrasing invites more. It is a statement about the ideal fixed point: if both objectives are driven to exact consistency, then the states are belief states. It does not prove that training reaches that fixed point, and it does not say the belief state is the compact one, since the raw history qualifies too. That the learned states come out genuinely small is not a corollary of the math; it is an experimental result, and measuring it is next.

A more compact world model

Back to Manhattan, now with the training changed. Toggle the two models in the figure below. A next-token transformer, arriving at a shared corner by two different routes, can still pick two different next turns, one of them heading off the map, because it kept the two histories apart. NextLat folds both histories into one belief state, so the continuation is shared and legal. The sequence-compression metric counts exactly that behavior, same state implies same future, and NextLat lifts it from GPT's 0.65 to 0.71.

The sharper evidence is a number about the representation itself: effective latent rank. This is not the ordinary matrix rank. You take the singular values of the hidden-state matrix, normalize them to sum to one, and compute

which is the exponential of the entropy of the normalized singular values: an effective number of directions the representation actually spreads across. Two matrices can share the same true rank while one packs nearly all its variance into a handful of directions and the other smears it across many. On the taxi data the true map has only 4,580 intersections and 9,846 streets, so a faithful state needs only a modest number of directions. NextLat's hidden states come in at an effective rank of 52.7 against GPT's 160.1, roughly a third as many, a representation that has concentrated itself onto what the task actually requires: compact and accurate at once, on the exact benchmark that exposed the incoherent maps.

Reasoning, planning, and A₅

The same pressure that compresses a map also helps the model plan, and the cleanest demonstration is a task built to punish shortcuts. In the Path-Startask you get a graph with a center node and several arms, a start, and a goal at the tip of one arm, and you must output the path. It looks trivial and it is notoriously hard for next-token training, for a specific reason. During training the true path is fed back in as the prefix (teacher forcing), so once you are on an arm, every next node is only the neighbor of the revealed current node, a copy anyone can do. The only step that requires actually looking ahead to the goal is the very first hop: which arm to enter. A model that learns the easy copy rule never learns to choose the arm, so it guesses, and its accuracy falls as the number of arms grows. The paper calls this the Clever Hans cheat, after the horse that read its trainer instead of doing arithmetic.

NextLat solves the Path-Star task at close to 100% across every graph size the paper tries, from two arms up to seven; the Belief State Transformer manages the small graphs but starts failing on the largest, and plain next-token training fails throughout. The same pressure that fixed the map fixes this. Choosing the arm requires the state to carry a belief about where the goal sits, and next-latent prediction builds exactly that belief, leaving a copy shortcut nowhere to hide.

The pattern repeats on Countdown, an arithmetic-search puzzle (a generalization of the game of 24) that frontier models find hard. Here the gain is largest at the shortest supervision horizon: trained with a one-step horizon, NextLat reaches 54.8% against 39.2% and 39.0% for the multi-token-prediction baselines at the same horizon, and above the 33.1% of plain next-token training and the 42.3% of the Belief State Transformer. Give it a longer horizon and it climbs to 58.7%. When the paper inspects the failures, the losing models tend to botch the last equation, forcing an invalid final step because the shortfall only surfaces at the end, when the running total cannot reach the target, a lack of lookahead by another name.

Take the word problem for the group $A_5$, the even permutations of five objects: you are handed a sequence of shuffles and must output the single permutation they compose to. This is a pure state-tracking task, and it is known to sit in a complexity class ($\mathsf{NC}^1$) that a fixed-depth transformer is believed unable to reach, since such transformers live in the smaller class $\mathsf{TC}^0$, and $\mathsf{TC}^0 \subsetneq \mathsf{NC}^1$ is a standard (still unproven) conjecture. The paper trains tiny two-layer transformers on length-12 sequences and tests them on length 36. Both the plain and the NextLat transformer fail past 12, as the theory predicts, and so does a plain transformer trained directly on length 36. But the little latent dynamics model $p_\psi$, co-trained inside NextLat, generalizes to length 36 at over 95%.

No complexity bound was broken. The $A_5$ problem is a plain finite-state computation, and a sequential network that walks the sequence updating one running state, an ordinary recurrent network, can do it at any length; $p_\psi$ stepped recursively is such a network. The surprise is about learnability, not expressivity. Training entirely in parallel, with no recurrence unrolled through time, still produced a recurrent rule that extrapolates far past its training length, and past what its own parent transformer can do. NextLat at $d=1$, viewed this way, is a transformer and a recurrent network trained together, the transformer supplying targets and the recurrence learning to march.

A lossless decoding speedup

The dynamics model was introduced to shape representations during training, and for ordinary generation it can be thrown away: the trained transformer decodes exactly as any transformer does, token by token, and $p_\psi$ is not called. But there is a bonus in keeping it, because a model that can predict its own next hidden state can predict several, cheaply, without running the transformer at all.

Speculative decoding needs exactly that: a cheap drafter proposes a run of tokens, the full model checks them all in one parallel pass, and an accept-or-resample rule keeps the longest correct prefix. The rule is designed so the tokens it emits are distributed exactly as the full model's own sampling would produce, down to floating-point. The speedup is free of quality cost: it changes how fast you decode, never what you decode.

NextLat's drafter is $p_\psi$ composed with itself. Take a real hidden state, decode a token through the frozen head, step the latent forward with the MLP, decode the next token, and repeat. Because the transition is a single reusable rule rather than a fixed set of prediction heads, the draft can run as long as the latents stay coherent, well past any fixed training horizon.

# inference-time rollout: draft ahead in latent space, no transformer
h = last_hidden_state          # from one real transformer pass
draft = []
for _ in range(gamma):         # gamma can exceed the training horizon d
    x = sample(head(h))        # decode a token from the current latent
    draft.append(x)
    h = p_psi(h, x)            # step the latent forward with the tiny MLP
# then verify the whole draft in ONE parallel transformer pass

Here NextLat pulls away from the multi-token-prediction methods it is compared against. Those train a fixed number $d$ of extra output heads, so they can draft at most $d$ tokens per cycle. NextLat, trained with a horizon of just two, still has verifiers accept close to five drafted tokens per cycle, because the recursive latent rollout keeps proposing good tokens past the horizon it was trained on. The figure shows the comparison: the multi-token methods pile up against the $d=2$ wall at roughly 1.7 to 1.9 times faster, while NextLat's accepted run reaches past it for up to a 3.3-times speedup, and every token it emits is still the base model's exact output.

What it changes, and what it can't

The costs are modest and the paper is candid about them. At a one-step horizon the extra work is a single small MLP and one more forward pass through it, negligible next to the transformer, and the parameter count barely moves (on its 1.3-billion-parameter model the dynamics MLP is a few percent). That is a different bargain from the Belief State Transformer, which reaches belief states by training two transformers and accumulating gradient over all prefix-and-suffix pairs, at more than three times NextLat's training cost. The gradient signals here grow only with the horizon, as $O(Td)$, not with the square of the sequence length.

At the largest scale the story is more measured, and the paper says so. Pretraining 1.3-billion-parameter models on 100 billion tokens, NextLat's edge on multiple-choice accuracy is small and uneven (a 59.21 average against 58.82 for plain next-token training), though it preserves next-token perplexity better than the multi-token baselines and delivers the decoding speedup. The clean wins are on the structured tasks, world modeling, reasoning, planning, where compact belief states obviously matter; whether they translate into large gains at frontier scale is left open.

NextLat reshapes the transformer's representations; it does not change what the transformer can compute. A fixed-depth transformer trained with NextLat is still a fixed-depth transformer, bounded by the same circuit class. The extra power in the A₅ experiment lives in the recurrent $p_\psi$, a different kind of machine, that NextLat happened to train in parallel as a side effect. What the paper offers is a way to install a recurrent inductive bias, compression into a running summary, into a transformer without paying the recurrent network's sequential training cost, and the evidence is that on the tasks where that bias matters, the transformer generalizes, plans, and compresses better than next-token prediction alone ever delivers.