NextLat: a transformer graded on its own next hidden state

A recurrent network is forced to compress. Whatever happened ten steps ago has to fit through the same fixed-width hidden state as whatever happened one step ago, so an RNN that wants to predict well has no choice but to keep only what still matters. A transformer has no such pressure — self-attention gives it ad-hoc lookup over every past token, so its "memory" grows with the sequence rather than staying fixed, and next-token cross-entropy rewards it for reading the right token off that unbounded context, not for summarizing anything. Next-Latent Prediction (NextLat), from Jayden Teoh, Manan Tomar, Kwangjun Ahn, Edward S. Hu, Tim Pearce, Pratyusha Sharma, Akshay Krishnamurthy, Riashat Islam, Alex Lamb, and John Langford (arXiv 2511.05963, code at github.com/JaydenTeoh/NextLat), adds the recurrent-style pressure back in without touching the architecture: alongside ordinary next-token prediction, train the transformer's hidden state to predict its own next value. This piece checks what "compact" and "world model" concretely mean in the paper's own tables, where the evidence for an actual world model is stronger than raw prediction accuracy, and where it's thinner than the title implies.

expand

What the objective actually is

NextLat's hidden state is not a separate encoder's output, a VQ codebook lookup, or anything with its own weights independent of the transformer — it is literally $h_t$, the transformer's own final-layer hidden state at position $t$, the same vector that already feeds the language-modeling head. There is no frozen target network and no pretraining phase to freeze afterward: the whole thing — transformer parameters $\theta$ and a small MLP dynamics model $p_\psi$ — trains jointly, in one combined loss, from initialization. $p_\psi$ takes the current hidden state and the actual next token (teacher-forced, not sampled) and predicts a distribution over the next hidden state:

Three terms make up the total loss. The first is ordinary next-token cross-entropy — nothing new here:

The second regresses the predicted hidden state toward the transformer's actual future hidden state, rolled out up to $d$ steps ahead with a Smooth L1 loss and, critically, a stop-gradient on the target:

The third is a distillation term: it decodes both the true and predicted hidden states through the (stop-gradded) language-modeling head and matches the resulting token distributions with KL divergence, so the predicted latent is graded on whether it decodes correctly, not only on whether it matches the true vector coordinate-for-coordinate:

and the full objective is a weighted sum of all three:

As pseudocode, one training step looks like this — the part worth noticing is that h_hat is produced recursively, chaining $d$ applications of the small dynamics MLP, each one conditioned on the real next token rather than its own previous guess:

# theta = transformer (produces h_t and decodes tokens); psi = small MLP dynamics model
h = transformer.encode(x[: t + 1])              # h_t, theta's own hidden state
loss_token = cross_entropy(decode(h, theta), x[t + 1])
 
h_hat = h
loss_next_h, loss_kl = 0.0, 0.0
for i in range(1, d + 1):
    h_hat = psi(h_hat, x[t + i])                 # teacher-forced next token, not sampled
    h_true = transformer.encode(x[: t + i + 1])  # theta's actual hidden state at t+i
    loss_next_h += smooth_l1(stop_grad(h_true), h_hat)
    loss_kl += kl_div(decode(stop_grad(h_true), stop_grad(theta)),
                       decode(h_hat, stop_grad(theta)))
 
loss = loss_token + lambda_h * (loss_next_h / d) + lambda_kl * (loss_kl / d)

At inference, psi never runs — the transformer decodes autoregressively exactly as it would without any of this. Nothing about the model's forward pass, its parameter count at inference, or its parallel training changes; the only thing that changes is what the training signal shapes.

Why it needs a stop-gradient at all

Readers of this site's LeVJEPA piece will recognize the shape of the problem: any objective that grades a model against its own output risks a trivial fix, where the model satisfies the loss by making every output identical rather than by predicting anything real. LeVJEPA sidesteps that risk entirely — no stop-gradient, no target network, because SIGReg's isotropic-Gaussian regularizer makes the collapsed solution structurally unreachable. NextLat takes the more conventional route: the sg[·] in $\mathcal{L}_{\text{next-h}}$ and $\mathcal{L}_{\text{KL}}$ blocks gradient from flowing into $\mathbf{h}_{t+i}$ through its role as a target, so the dynamics model has to move to match the transformer's real trajectory rather than the trajectory drifting to meet a lazy prediction. It's the same asymmetry BYOL and SimSiam rely on for image self-supervision, applied across time instead of across two augmented views. The two papers land on opposite sides of the same design question — predict in latent space to sidestep pixel detail — for a reason worth being precise about: LeVJEPA's "target" is a different augmented view of the same input, so collapsing to a constant still has to explain away real visual variation the regularizer is watching for; NextLat's target is a future timestep of the same trajectory, where a constant hidden state would still decode correctly at each step if the model routed all its real work around the latent instead of through it — a more direct collapse path, and the reason the paper reaches for the older tool.

Provably a belief state — under one condition

The paper's theoretical claim (Theorem 3.2) is a backward-induction argument, not a training result: if a hidden state can (a) decode the true next-token distribution exactly, given the history, and (b) predict the true distribution over its own next value exactly, given the next token, then $h_t$ must already be a belief state — a sufficient statistic of the history, in the sense used in POMDPs and stochastic control, that carries everything needed to predict the future and nothing an observer with the full history could add. The proof composes these two guarantees recursively: decode a token from $h_t$, use it to predict $h_{t+1}$, decode the next token from that, and so on — if every step is exact, $h_t$ could not have discarded anything load-bearing. A standard transformer only gets pressure toward condition (a); NextLat is what happens when you add pressure toward (b) as well.

Two baselines in the paper's own comparisons chase the same theoretical object by different routes, and the differences are concrete rather than cosmetic. The Belief State Transformer (BST, Hu et al. 2025) gets a belief-state guarantee by training two encoders — one reading forward, one backward — and paying for it with $O(T^2)$ gradient signals per sequence; on TinyStories it trains at 0.19 iterations per second against NextLat's 3.26, over 10x slower, and needs both encoders again at inference. Joint-Token Prediction (JTP, Ahn et al. 2025) can learn belief states too, but the paper is specific about the catch: JTP's guarantee only holds once the prediction horizon $d$ is at least the data-generating process's observability horizon — how many future tokens you'd need to see to fully disambiguate the current state. For short synthetic tasks that's a small number; for open-ended language modeling the paper argues it's effectively unbounded, which makes the condition impractical to satisfy by simply cranking up $d$. NextLat's guarantee is stated independent of $d$ — larger $d$ gives a richer gradient signal, not a different guarantee.

Where "compact" gets measured

The clearest test in the paper is Manhattan Taxi Rides (Vafa et al. 2024): 91M sequences, 4.7B tokens of random 100-step traversals of Manhattan's real one-way street grid (4,580 intersections, 9,846 edges), where the model never sees the map directly — only sequences of turns — and has to infer the underlying graph from traversal statistics alone. Every model here reaches 100% next-token test accuracy, which the paper is upfront about being close to useless as a diagnostic: a model that gets every token right can still hold an internal map that's structurally wrong, and next-token accuracy alone can't tell the two apart. Four other metrics can:

expand

Effective latent rank — the exponentiated Shannon entropy of the hidden-state matrix's singular values, lower meaning more compact — is the number the title leans on: NextLat's 52.7 against GPT's 160.1 is a genuine 3x-plus reduction. But it's worth reading the full row rather than just the two ends of it: MTP's 57.7 is a much closer second than "3x more compact than everything else" would suggest, within 10% of NextLat despite MTP having no belief-state guarantee behind it at all. And a methodology note the paper states plainly in its appendix but not in the main table: effective rank isn't read from the same point in every architecture. GPT and NextLat use the final-layer hidden state; JTP uses the state immediately before self-attention in its own "Fetch" head; MTP uses the output of its next-token prediction head. These are reasonable choices, but they aren't the same object, which means the rank column is more directly comparable between GPT and NextLat than it is across all four.

Sequence compression is the metric that comes closest to testing state abstraction rather than raw accuracy, and it's the one most likely to get skipped past: it's the percentage of cases where two different routes that happen to arrive at the same intersection, heading to the same destination, produce identical continuations from that point on. A model that has genuinely learned "you are here, heading there" as a state — rather than memorizing route-specific continuations — should answer identically regardless of how it got there. NextLat wins this one by more than its rank lead would predict (71% vs. GPT's 65%), and JTP's collapse here is the sharpest number in the whole table: 32%, worse than plain GPT. JTP's extra token-level supervision didn't just fail to help compression — it measurably hurt it.

Discards nuisance detail, or changes the optimization?

This is the question the paper doesn't cleanly separate. There are two live explanations for why next-latent prediction helps, and they point in different directions. One is representational: the latent target is compressible in a way pixels or discrete tokens aren't, so the objective structurally permits (even rewards) discarding whatever doesn't help predict the future — the belief-state argument above. The other is about the shape of the optimization landscape: Section 5.1 of the paper argues that plain token-level supervision is myopic — early training on next-token objectives tends to resemble $n$-gram modeling, which the paper cites prior work as showing can delay or trap models in local minima that undermine long-horizon planning, independent of anything about representation size. NextLat's own explanation for its Countdown and Path-Star results leans on this second story as much as the first.

The paper does not run the ablation that would cleanly separate the two. There's no experiment here that forces representational compression through some other mechanism — a narrow bottleneck layer, say — without the next-latent prediction objective, to see whether compression alone reproduces the planning gains; nor is there one that keeps the multi-step gradient signal but removes the pressure toward a self-consistent latent, to see whether the optimization-dynamics story survives on its own. What the paper does establish is that JTP and MTP — which also add multi-step gradient signal on top of the same token-level structure, without NextLat's latent bottleneck — get smaller and less consistent gains, and in JTP's compression score, an outright regression. That's evidence the specific routing through a latent target matters, not just the presence of extra gradient signal, but it's evidence by comparison across methods, not a controlled decomposition within one.

Is this actually a world model, or just better prediction?

The paper's title makes a strong claim, and the honest answer is that some of its evidence supports it more directly than others. Downstream accuracy — solving Countdown, planning Path-Star routes — is consistent with a good learned model, but consistent with a lot of things; a model can get better at a task through reasons that have nothing to do with internal coherence. Three pieces of evidence here are stronger, because they're designed to fail if the model is merely predicting well without an underlying consistent structure:

Sequence compression, described above, is close to a direct probe for state abstraction — it asks whether the model treats two different histories that reach the same state as the same state, which is closer to "does an internal representation of location exist" than any accuracy number could be.

Detour robustness is the closest thing to an intervention in this paper: on out-of-distribution pickup-dropoff pairs, the evaluation overrides the model's own top-1 prediction with a random (but legal) detour 75% of the time, then checks whether the resulting trajectory still reaches a valid state. This is testing something accuracy can't: whether the model can recover a coherent continuation after being forced off the path it would have chosen — which requires the model to have something like a map it can re-plan from, not just a policy tuned to its own typical trajectories. NextLat ties MTP here at 95.0%, both well ahead of GPT's 85.0% and JTP's 87.0%.

Linear probing on frozen hidden states (TinyStories) tests something orthogonal to both: whether information about tokens far in the future is linearly recoverable from the current hidden state at all, independent of whether the model ever needs to use it during ordinary decoding.

expand

The pattern here is the most direct evidence for "compact, predictive representation" in the whole paper: BST, MTP, and JTP all degrade next-token probe accuracy relative to plain GPT — the extra supervision at future offsets measurably hurts the thing the model is supposed to be best at — and JTP and MTP's advantage at longer offsets shrinks back toward zero as the offset grows. NextLat is the only method that matches GPT's next-token probe accuracy and keeps a real advantage out to 20 tokens ahead. That's the strongest single piece of evidence that something durable, rather than a training artifact, is encoded in the hidden state.

None of this is proof of a full world model in the sense of, say, being able to simulate arbitrary counterfactual rollouts — the paper doesn't run latent-space interventions (perturbing $\hat{h}$ directly and checking whether the decoded continuation changes in the way a real state change would) or test transfer to a structurally different map. The evidence is real and multi-pronged, but it is all downstream-task and probing evidence, not a demonstration that the latent supports arbitrary planning queries the way, say, an explicit transition model would.

Reasoning and planning, briefly

On Countdown (combine four numbers via arithmetic to hit a target, following Gandhi et al. 2024), NextLat beats MTP and JTP at the same shallow horizon ($d=1$) by more than 38%. The paper's more specific finding is about where errors happen: most invalid equations occur in the final step of a solution, which the paper — borrowing a term from Ye et al. (2025) — calls "the regretful compromise": a model realizes only at the last step that its plan doesn't work, and is forced into an invalid final equation to match the target anyway, unable to revise earlier choices. NextLat gets the final equation right 54.2% of the time at $d=1$, against 42.3% for the next-best baseline — evidence read as better lookahead, not just better arithmetic.

On Path-Star graphs (Bachmann and Nagarajan 2024) — a center node with disjoint arms, where the model must generate the correct arm from start to end — NextLat holds close to 100% solve rate across all three tested topologies ($G_{2,10}$, $G_{5,5}$, $G_{7,7}$), while BST solves the two smaller graphs but "begins to fail at the larger graph $G_{7,7}$," in the paper's own words. Worth flagging: the paper's own setup here is deliberately harder than BST's and JTP's original papers (a fixed 200k-sample training set and node values up to $N=100$, versus a smaller $N=50$ with fresh graphs generated every batch), which the paper states plainly rather than hiding — a fair replication difference to know about before comparing these numbers to the original BST or JTP papers directly.

What this costs, and who it's compared against

NextLat's inference parameter count matches GPT's exactly — 57M, the same number, because $p_\psi$ never runs after training — while training is only slightly slower than plain GPT (3.26 vs. 3.72 iterations/second at $d=1$) and far cheaper than BST's dual-encoder setup. All models here are small: this table's largest model is BST at 114M training parameters, and the Manhattan GPT/NextLat models — the ones with the compactness numbers above — are 89M-parameter, 48-layer transformers, deliberately made deep rather than wide, since the paper found depth mattered for the state-tracking demands of the task and width didn't (MTP's Manhattan variant runs larger still, per the paper's own appendix, though it doesn't state the exact count).

Toy domains at submission, real scale in a later revision

Every experiment above — Manhattan, Countdown, Path-Star, TinyStories — is synthetic or semi-synthetic, and every model involved is small (114M parameters at the largest, BST's dual-encoder TinyStories setup). That's the honest scope of the paper as it was submitted (arXiv v1, November 8, 2025), and the figures and tables in this piece are drawn from that version specifically. As of this writing, the current revision on arXiv adds something the original didn't have: a 1.3B-parameter language model, pretrained on 100B tokens of FineWeb-Edu, evaluated on zero-shot multiple-choice benchmarks and — the more novel addition — self-speculative decoding, where NextLat's latent dynamics model drafts multiple tokens by recursively chaining $\hat{h}_{t+1} \to \hat{h}_{t+2} \to \dots$, each drafted token verified in parallel against the base transformer using standard speculative sampling. Because the draft comes from latent-space rollout rather than a fixed number of token-prediction heads, the draft length isn't capped by the training horizon $d$ the way MTP's or JTP's is:

Inference speedup relative to standard autoregressive decoding, measured on 8x NVIDIA B200 GPUs. This is a real and useful result — it's the first evidence in the paper that NextLat's latent dynamics stay coherent well past the training horizon, since a model trained only to predict one or two steps ahead ($d{=}1,2$) is drafting sequences several tokens longer than that and still getting them accepted. But it's worth being precise about what it does and doesn't extend: this is a decoding-speed result at 1.3B scale, not a compactness or probing result at that scale. Whether effective latent rank still drops 3x, or whether the sequence-compression and detour-robustness gap holds up, at a billion-plus parameters and on real language rather than a synthetic taxi grid, is not something the current revision demonstrates — it's a real gap between what's been shown small and what's been shown large, and it's the single most useful piece of due diligence a reader should carry forward.

Where this sits next to the rest of the latent-prediction landscape

NextLat and LeVJEPA agree on the core move — predict in a representation you get to shape, rather than in raw observation or token space, precisely to avoid spending capacity on nuisance detail — but they're solving different collapse problems in different geometries: LeVJEPA predicts across space, one view of a clip against another, with no target network and no stop-gradient because SIGReg's Gaussian-matching constraint forbids the collapsed solution outright; NextLat predicts across time, one hidden state against its own future self, and reaches for the more classical stop-gradient fix because a constant hidden state genuinely can satisfy its loss if nothing stops it. Multi-token prediction is the most direct point of comparison in NextLat's own tables — MTP is close on effective rank and clearly behind on sequence compression and probing, which is the paper's own case for why routing supervision through a latent bottleneck beats simply adding more token-level heads. And LOTUS makes a related bet from the reasoning side — that computation belongs in hidden states rather than in an emitted token stream — worth reading against NextLat's belief-state framing, since both are arguments that a transformer's internal state, not just its output, is where the interesting representational work should happen.

Checked, in one table

The take

NextLat's actual contribution is narrower and more checkable than "transformers learn world models" suggests on its own: one auxiliary loss, one stop-gradient, one small MLP discarded before inference, and a specific theoretical guarantee — belief-state convergence, independent of prediction horizon — that its closest rivals either pay far more to get (BST) or only get conditionally (JTP). The paper's own tables hold up to scrutiny: MTP really is a closer competitor on raw compactness than the headline suggests, JTP really does make its own internal map less consistent despite added supervision, and the strongest evidence for "world model" over "good prediction" is a metric — sequence compression — that's easy to read past in favor of the flashier effective-rank number. The evidence earns the claim more than most papers with "world model" in the title manage, largely because the paper reaches for probing and intervention-adjacent tests rather than resting on downstream accuracy alone. What it hasn't yet shown is that any of this survives the jump past 89M parameters and synthetic domains — a gap this piece's later revision starts to close on inference speed, and hasn't yet closed on compactness.

Sources: NextLat (arXiv 2511.05963v1), read via its arXiv HTML rendering, for every table, equation, and figure in this piece; the current arXiv revision, read the same way, for the self-speculative decoding results and Table 3 (Section 3.3, Section 4.4); the NextLat GitHub repository; Belief State Transformers (Hu et al., 2025); Vafa et al. (2024) for the Manhattan Taxi Rides benchmark and its reconstruction algorithm. Figures 1-3 are the paper's own, fetched from its v1 arXiv HTML rendering and shown for commentary. The objective-comparison and compactness-explorer diagrams are original, built from the sources above.