NextLat: Next-Latent Prediction Transformers Learn Compact World Models

A transformer trained to also predict its own next hidden state learns to summarize the past in a compact state.

NextLat keeps the transformer and its next-token loss and adds a small network that predicts the next hidden state from the current hidden state and the next token. The gradient of that prediction pushes each hidden state to carry everything about the past that the future depends on. The architecture and inference stay the same, and the small network can also propose draft tokens that speed up generation without changing the output.

Transformers trained on Manhattan taxi rides propose a legal next turn every time and still carry a map of the city with streets at impossible angles. With one extra training loss, the paper's models keep their next-turn accuracy, hold a more consistent map, and use about a third as many effective dimensions to do it.

The paper is Next-Latent Prediction (NextLat), from Microsoft Research, released in November 2025. The page follows the order the method needs: why attention removes the pressure to compress, what a belief state is, the three loss terms and which parameters each one trains, why the states do not collapse, what the theorem does and does not prove, and then the results on world modeling, planning, state tracking and decoding speed.

A taxi model with an incoherent map

Vafa et al. (2024) built a test where the world model a sequence model should learn can be drawn: taxi rides in Manhattan written as turn-by-turn directions, whose true "world" is the street map. Their transformers generated valid routes 96 to 99% of the time, yet when they reconstructed the map implied by each model's predictions, it had streets at impossible orientations, roads passing over other roads, and connections that do not exist.

NextLat reuses their random-walk dataset (91 million rides, 4.7 billion tokens) and trains 48-layer transformers of width 384 for six epochs. A plain next-token model, which the paper calls GPT, and the NextLat model both score 100% on the next-token test, the share of predictions that are legal turns. The other metrics separate them:

Sequence compression asks a question a correct map answers trivially. Take two different routes that arrive at the same intersection and share a destination. The remaining trip is the same problem in both cases, so a model with a coherent map should continue them identically; the metric is the fraction of such pairs where it does. GPT does so 65% of the time, NextLat 71%. Detour robustness replaces the model's top choice with a random legal turn 75% of the time and checks that the ride still reaches its destination: 85.0% for GPT, 95.0% for NextLat. Effective latent rank measures how many directions the hidden states actually spread over, and has its own figure below. All of these differences appear between models whose next-token scores are identical, and the paper concludes that next-token accuracy "cannot meaningfully assess the quality of a model's learned world model."

Why attention removes the pressure to compress

A recurrent network reads a sequence one token at a time and carries a single fixed-size vector forward. To predict token 500 it can consult only that vector, so whatever about tokens 1 to 499 matters for the future has to be packed into it. A transformer removed that constraint: self-attention lets the position at token 500 read any earlier position directly, which is what makes training parallel across positions and long-range dependencies easy to learn. It also means the hidden state at step $t$ never has to be a self-contained summary, because the raw past stays available for lookup. The paper's introduction describes the transformer's memory as growing with sequence length, with no pressure to compress history into a compact state with a consistent update rule.

Drag the query position in Figure 1. In attention mode the number of positions $\mathbf{h}_t$ reads grows with $t$; in belief-state mode each state is built from the previous state and the newest token only, and stays one box wide at any length.

Trained next-token transformers do encode belief-state structure in their activations; Shai et al. (2024) found it with linear probes. What standard training lacks is a guarantee or a pressure that every position holds a compact, updatable summary, so on structured tasks the model can fit the data with shortcuts that read the history instead. NextLat asks whether that pressure can be added back while training stays parallel.

What a belief state is

The term comes from control and reinforcement learning, where an agent receives a stream of observations and never sees the true state of the world. A belief state is a summary of the observation history that is sufficient for prediction: once you hold it, the rest of the history adds nothing. The paper's Definition 3.1 states it for token sequences. Write $X_{1:t}$ for the tokens so far and $\mathbf{b}_t = g(X_{1:t})$ for some function of them. Then $\mathbf{b}_t$ is a belief state if, for every bounded measurable function $f$ of the future tokens,

Choosing $f$ as the indicator of any particular future makes this say that the summary predicts every future event with the same probability the full history does. Equivalently, $\mathbf{b}_t$ is a sufficient statistic of the history for the future, the term the paper takes from Striebel (1965).

In the taxi task, the history is the list of turns so far. To predict the rest of the ride you need the current intersection and the destination; the turns that led there add nothing once you know those two. The pair (current intersection, destination) is a belief state, and it is far smaller than a 50-turn history. Two rides that reach the same corner bound for the same place have the same belief state, which is why a model whose hidden state is a belief state scores well on sequence compression.

Definition 3.1 defines a sufficient statistic, not the smallest one: the raw history $X_{1:t}$ satisfies the equation trivially. So being a belief state and being compact are separate properties, and the theorem below guarantees only the first. And it is the predictive form of the idea, a summary that reproduces the distribution of future tokens, not necessarily a probability distribution over hidden world states.

How NextLat works: predicting the next hidden state

The hidden state $\mathbf{h}_t$ is the transformer's final-layer output at position $t$, the vector the output head turns into next-token logits. The output head $p_\theta$ and the transformer are trained with the usual cross-entropy:

NextLat adds a latent dynamics model $p_\psi$, a separate network with its own parameters $\psi$. It takes the current hidden state and the next token and predicts the next hidden state, the way a world model in reinforcement learning (as in World Models or MuZero) takes a latent state and an action and predicts the next latent state; here the next token plays the role of the action. In the paper it is a three-layer MLP with GELU activations, applied to the layer-normalized concatenation of $\mathbf{h}_t$ and the next token's embedding, and it outputs a change added to the current state:

The paper keeps $p_\psi$ a simple MLP so that any gain can be attributed to the objective rather than to a specialized transition architecture, and lists a richer design as future work. The residual form makes the default prediction "the state does not change", so the MLP only has to learn the update. Applied to its own output, $p_\psi$ is a recurrent cell: a state and a token in, a new state out.

For a deterministic transformer, the true next state $\mathbf{h}_{t+1}$ is a fixed vector given the tokens, so learning $p_\psi$ is regression. The paper uses a Smooth-L1 (Huber) loss, which per coordinate is squared error for gaps under 1 and absolute error beyond (PyTorch's default threshold of 1, which the code uses):

Early in training, when predictions are far off, a coordinate with a gap of 5 contributes a gradient of 1 instead of 5, so a few large gaps cannot dominate the update. The latent loss averages it over a horizon of $d$ steps:

$\mathbf{sg}[\cdot]$ is the stop-gradient operator: in the backward pass the target $\mathbf{h}_{t+i}$ is treated as a constant. For $d>1$, $\hat{\mathbf{h}}_{t+i}$ is a recursive rollout: $p_\psi$ is applied to its own previous prediction with the true tokens $X_{t+1}, \dots, X_{t+i}$ fed in, so the transition rule has to stay accurate over several steps. The theorem below needs only $d=1$; the paper uses larger $d$ as a richer training signal.

A second term checks the prediction in token space. Pass the true next state and the predicted next state through the output head, with the head's weights frozen, and require the two next-token distributions to match:

The first argument of the KL is the distribution from the true state, so the true state is the teacher and the prediction the student, as in knowledge distillation with one model playing both roles. The superscript $\mathbf{sg}$ on $p_\theta$ marks the frozen head. Where Eq. (3) asks whether $\hat{\mathbf{h}}$ lands on the right vector, Eq. (4) asks whether it predicts the same next token, which weights errors by how much they change the model's output. The full objective is

with $\lambda_{\text{next-h}}, \lambda_{\mathrm{KL}} > 0$. The code defaults are $\lambda_{\text{next-h}} = \lambda_{\mathrm{KL}} = 1$ and $d = 1$; the experiments tune them per task, as the table further down shows. Figure 2 shows Eq. (3) at work on a path of hidden states. Each teal dot is a true $\mathbf{h}_t$, each amber ring is the prediction $p_\psi$ makes for the next state, and the dashed line is the error. Drag training forward and the rings move onto the dots.

Concretely, take the paper's largest model: 22 layers, 16 heads, width $D = 2048$, a GPT-2 vocabulary of 50,257 tokens and sequences of 1,024 tokens. One forward pass gives hidden states of shape (batch, 1024, 2048). At position 500, $p_\psi$ reads $\mathbf{h}_{500}$ and the embedding of token 501, a concatenated vector of 4,096 numbers, and passes it through layers of width $4096 \to 6528 \to 6528 \to 2048$. That is about 82.7 million parameters, which is why the paper's Table 4 lists 1.40B training parameters for NextLat against 1.32B for GPT. The 2048-vector output is added to $\mathbf{h}_{500}$ to give $\hat{\mathbf{h}}_{501}$, compared coordinate by coordinate with the detached $\mathbf{h}_{501}$ for Eq. (3), and both vectors go through the frozen head to give two distributions over 50,257 tokens for Eq. (4). The same happens at every position in parallel, since all the $\mathbf{h}_t$ come from the one forward pass. After training $p_\psi$ is dropped and the model serves 1.32B parameters. One step in PyTorch-style pseudocode:

# one NextLat training step, horizon d = 1 (Algorithm 1 of the paper)
h      = transformer(x)              # hidden states, [B, T, D], one parallel pass
logits = head(h)                     # [B, T, V]
loss_tok = cross_entropy(logits[:, :-1], x[:, 1:])     # ordinary next-token loss

# predict h_{t+1} from (h_t, x_{t+1}); h_t is NOT detached
h_hat  = p_psi(h[:, :-1], embed(x[:, 1:]))    # residual MLP, [B, T-1, D]
loss_h = smooth_l1(h_hat, h[:, 1:].detach())  # Eq. (3): the target is a constant

# Eq. (4): same next-token distribution from h_hat as from the true h_{t+1}
teacher = logits[:, 1:].detach()              # head applied to the true state
student = frozen_head(h_hat)                  # head weights get no gradient here
loss_kl = kl(teacher, student)

loss = loss_tok + lam_h * loss_h + lam_kl * loss_kl   # Eq. (5)
loss.backward()   # reaches p_psi, the head (via loss_tok), and the transformer

The settings vary by task (paper Table 5 and the experiment sections). $p_\psi$ is always three layers; only its hidden width changes.

Where each gradient goes

The stop-gradient in Eq. (3) applies to the target $\mathbf{h}_{t+1}$ only. The input $\mathbf{h}_t$ that $p_\psi$ reads is not detached (in the code, line 303 of model_nextlat.py detaches h_t_next and nothing else), so the gradient of the latent loss flows through $p_\psi$ into $\mathbf{h}_t$ and from there into every transformer layer. The transformer is pushed to produce states from which one fixed rule can compute the next state, and $p_\psi$ itself is discarded after training.

Every position plays both roles. $\mathbf{h}_{501}$ is a constant target for the prediction made from $\mathbf{h}_{500}$, and a trainable input for the prediction of $\mathbf{h}_{502}$. In the KL term the head is frozen, so its gradient passes through the head's weights to $\hat{\mathbf{h}}$ without changing them, and the head is trained only by the ordinary next-token loss. Pick a loss in Figure 3 to see its backward path.

The frozen head also saves compute. A backward pass through a 2048 × 50,257 output layer is expensive, and the paper's ablation (Appendix D) found that stopping gradients on both the Smooth-L1 target and the KL head gave the best probe accuracy on TinyStories while removing that extra backward pass.

Why the states do not collapse

A model rewarded for predicting its own future states has a trivial way to score perfectly: make every hidden state the same vector. Then $p_\psi$ only has to output that vector, the latent loss is zero, and the states carry no information. This failure, representational collapse, is the standard risk of self-predictive learning, the family of methods (common in reinforcement learning) that learn representations by predicting their own future latents.

The stop-gradient blocks the main path to collapse. Without it, the gradient of the error between the prediction $p_\psi(\mathbf{h}_t)$ and the target $\mathbf{h}_{t+1}$ moves both ends of every pair toward each other, and shrinking the whole cloud to one point lowers every error at once. With the target held constant, only the input side moves. Tang et al. (2023) proved what that buys in a linear setting: if the predictor is fit faster than the representation (a small MLP next to a large transformer approximates this), the semi-gradient update keeps $\Phi^\top\Phi$ constant, where $\Phi$ is the matrix whose rows are the states, so $\Phi^\top\Phi$ holds their second moments. The states can rotate but cannot shrink or merge.

Figure 4 runs that setting. Twelve 2-D states each have a fixed successor, forming three loops of four. Every step, the linear predictor is refit exactly by least squares, then the states take one gradient step on the squared prediction error. Press Play with the stop-gradient on: the loss falls to zero and the spread stays near its starting value instead of shrinking (0.76 at step 0, 0.86 at step 300; in continuous time the second moments would be exactly constant). The states settle into three four-cornered loops, and a single linear map carries each corner to the next one around its loop. Switch the stop-gradient off and replay: the loss also reaches zero, but the spread falls toward zero as all twelve states slide to one point.

NextLat has a second guard the toy leaves out. It is still training with next-token cross-entropy, and a constant hidden state cannot produce different next-token distributions in different contexts, so collapse would cost cross-entropy. The paper's footnote 4 makes the same point: the next-token objective already grounds against collapse, and the stop-gradient is kept because it performs better. In the ablation, removing the stop-gradient on the Smooth-L1 term slightly lowered the validation losses (by about 0.02) but gave hidden states that linear probes read the future from less well, and on Countdown and Path-Star the stop-gradient was needed for high accuracy. Other self-predictive methods use the same asymmetry in other forms: SimSiam, like NextLat, puts a hard stop-gradient on targets from the same network, while I-JEPA computes its targets with a slowly updated (moving-average) copy of the encoder.

What Theorem 3.2 proves

Theorem 3.2 says that if training reaches its ideal, the hidden states are belief states. Suppose two conditions hold exactly at 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 right side of (2) is the transition the transformer's weights induce; the probabilistic notation covers stochastic transformers, and for the deterministic ones used in practice it is a point mass.) The proof is a backward induction. Split the future distribution into its first token and the rest:

By (1), the first factor equals $p_\theta(X_{t+1}\mid\mathbf{h}_t)$, a function of $\mathbf{h}_t$. By the induction hypothesis, the second factor depends on the history only through $\mathbf{h}_{t+1}$, and by (2), $\mathbf{h}_{t+1}$ is determined by $\mathbf{h}_t$ and $X_{t+1}$ through $p_\psi$. So the whole future distribution is a function of $\mathbf{h}_t$. Read as a procedure, the two maps let you generate the future from $\mathbf{h}_t$ alone, without the history:

In condition (2) the next state must be computable from the current state and the next token, with no access to the earlier tokens, so anything the future needs has to be in the current state already. The argument uses one-step transitions only, so it holds at $d=1$.

Next-token consistency alone is not enough. The Belief State Transformer paper (Hu et al., 2025, Theorem 3) constructs a distribution and a perfect next-token predictor whose head input is not a belief state. That is an existence result: next-token training does not guarantee a belief state, which is different from saying transformers cannot form one. Among the paper's baselines, GPT and multi-token prediction (MTP, extra heads predicting the next $d$ tokens) have no such guarantee. Joint-token prediction (JTP, Ahn et al., 2025) predicts each of the next $d$ tokens from the hidden state plus the true tokens before it, passed through a small attention bottleneck, so the hidden state has to carry the joint distribution of those tokens rather than each one's marginal. It has a guarantee only when $d$ is at least the observability horizon $k$ of the data-generating process: the number of future tokens whose distribution, once two histories agree on it, fixes the whole future (Definition A.1). In long text $k$ can be very large.

The theorem describes the ideal fixed point. It does not show that gradient training reaches it, and it does not say the belief state is small, since the raw history qualifies too. The abstract's phrase "provably converge towards belief states" reads stronger than this; the Provenance panel below has the details. That the learned states are compact is an experimental result, measured next.

The paper also argues that the latent target is a richer training signal than a token. A one-hot next token says which token came; $\mathbf{h}_{t+1}$ parameterizes the whole distribution over $X_{t+2}$, and because each state is trained to predict the next, it also carries information about later states. The paper links this to NextLat's results in the low-data Path-Star setting.

Manhattan: a more compact world model

In Figure 5, two routes arrive at the same intersection bound for the same destination. Toggle between the models. A model that kept the two histories apart can give them different next moves, one of them off the map; a model whose state is the belief state gives both the same move, which is what the sequence-compression metric counts (0.65 for GPT, 0.71 for NextLat).

The second compactness number describes the hidden states directly. Stack the final-layer states for a batch of 256 sequences of 256 tokens into a matrix with 65,536 rows and 384 columns (the model width). Its singular values $\sigma_1 \ge \dots \ge \sigma_{384}$ say how much the states spread along each direction. Normalize them to sum to one and exponentiate their Shannon entropy (natural log, after Roy and Vetterli, 2007):

If the states spread equally over $k$ directions and not at all over the rest, the entropy is $\ln k$ and the effective rank is exactly $k$. A spectrum that decays smoothly gets a value between 1 and 384 even though every singular value is positive, so the ordinary matrix rank stays 384. In Figure 6, switch to "k equal directions" to check the first case, then use the power-law shape and the match buttons to see how concentrated a spectrum has to be to give NextLat's 52.7 or GPT's 160.1.

The true Manhattan graph has 4,580 intersections and 9,846 edges, a small state space, so a faithful state should need only a modest number of directions. NextLat's states have an effective rank of 52.7 against GPT's 160.1, a ratio of 3.04: NextLat uses about a third as many effective directions while scoring higher on every other metric. The paper's appendix notes both values are far below the width of 384, and that a wider model did not help on this task while a deeper one did.

Planning and reasoning tasks

Path-Star. A Path-Star graph $G_{a,\ell}$ (the paper writes $G_{d,\ell}$, reusing the letter $d$) has a center node and $a$ disjoint arms of $\ell-1$ nodes each. A training example lists the edges, names a start (the center) and a goal (the tip of one arm), and gives the path. The paper follows the original setup of Bachmann and Nagarajan (2024): 200,000 fixed training graphs with node labels drawn from 1 to 100, evaluated on 20,000 held-out graphs.

Next-token training fails here because of teacher forcing, the standard practice of feeding the true path so far as the prefix during training. Once the prefix contains a node on the correct arm, the next node is the only neighbor not yet visited, which a model can copy from the edge list without knowing where the goal is. Every hop after the first is learnable this way. The first hop, choosing which arm to enter, needs the model to trace back from the goal, and a model that has learned the copy rule gets little gradient signal for it. Bachmann and Nagarajan call the shortcut the Clever Hans cheat, after the horse that answered arithmetic questions by reading its trainer's posture. The failure arises during training, when the copy rule absorbs the gradient; errors compounding at generation time play no part.

The paper tests $G_{2,10}$, $G_{5,5}$ and $G_{7,7}$ with five seeds each, and sets the NextLat horizon to $\ell-2$ so the goal lies inside the prediction horizon from the center. NextLat solves all three at close to 100%. The Belief State Transformer solves $G_{2,10}$ and $G_{5,5}$ and starts failing at $G_{7,7}$, and GPT fails throughout. The paper notes that its BST and JTP numbers are lower than the original BST and JTP papers report, because the NextLat authors used a smaller, freshly sampled setup.

Countdown. Given four numbers and a target, write three equations that combine them into the target using +, −, × and ÷. It is a generalized Game of 24, which GPT-4 solved 4% of the time with default prompting. A training example from the paper, with inputs 14, 83, 88, 91 and target 23, reads:

14,83,88,91,23 | 83-14=69, 91-88=3, 69/3=23

The setup has 500,000 training problems with targets from 10 to 100, 10% of targets held out, eight pause tokens before the answer, and 10,000 test problems, averaged over three seeds:

The gap is largest at $d=1$, where NextLat's 54.8% compares with 39.2% and 39.0% for the two multi-token methods. At $d=8$ MTP reaches 57.3%, 1.4 points behind. When the paper checks each generated equation, most errors are in the third and last one: a model that did not plan discovers at the end that its numbers cannot reach the target and writes an invalid equation that does, a behavior Ye et al. (2025) named the regretful compromise.

TinyStories. On short synthetic stories the paper freezes each trained model and fits linear probes that predict the token 1 to 20 positions ahead from a single hidden state. The extra token-level objectives of BST, MTP and JTP lower next-token probe accuracy, and MTP and JTP probes degrade quickly with distance. NextLat matches GPT one token ahead and has the most accurate probes at long range, up to 20 tokens, at both $d=1$ and $d=8$.

A₅: training a recurrent network in parallel

$A_5$ is the group of the 60 even permutations of five objects. In its word problem, each token is one such permutation, and after every token the model must output the arrangement produced by applying all the tokens so far. The task is pure state tracking: a solver that keeps the current arrangement and updates it per token handles any length, as Figure 8 shows with five lettered balls.

Transformers have trouble with it for a structural reason. The $A_5$ word problem is complete for the circuit class $\mathsf{NC}^1$ (Barrington, 1989), and constant-depth transformers with fixed-precision arithmetic can be simulated by circuits in the smaller class $\mathsf{TC}^0$ (Merrill and Sabharwal, 2023). If, as is widely conjectured but unproven, $\mathsf{TC}^0 \ne \mathsf{NC}^1$, no constant-depth transformer solves it at all lengths; a transformer needs depth growing like $\log T$ in the sequence length.

The paper trains 2-layer transformers of width 512 on one million distinct 12-token sequences, with plain next-token prediction (GPT) and with NextLat at $d=1$ using only the latent loss ($\lambda_{\mathrm{KL}} = 0$, so nothing but transition prediction shapes $p_\psi$). It tests on about 100,000 sequences of 36 tokens. For the recurrent test, the transformer only produces the first hidden state; after that $p_\psi$ steps the state forward token by token and the head reads each answer off it. That network has 2.62M parameters, against 6.43M for the transformer.

Within the 12-token training length the NextLat transformer stays near 98% (GPT slips to about 84% at length 12). Past it, both transformers drop to about 0%, as the depth bound predicts, and a GPT trained directly on 36-token sequences falls apart between lengths 10 and 13. The co-trained $p_\psi$, run as a recurrent network, stays above 95% out to length 36, with only one-step supervision and no backpropagation through time.

No complexity bound is broken. A recurrent network that updates one state per token is outside $\mathsf{TC}^0$ and can express $A_5$ tracking at any length; the open question was whether one could be learned without sequential training. Here a transformer limited to 12 steps supplied per-step targets, and a recurrent rule fit to those targets extrapolated three times further. The trained transformer itself is no more capable than before. The paper also notes that the hidden states include rotary position information (RoPE), and the transition rule generalized anyway.

Speculative decoding past the training horizon

Speculative decoding (Leviathan et al., 2022) speeds up generation with a cheap drafter. The drafter proposes several tokens, the full model scores all of them in one parallel pass, and an accept-or-resample rule keeps a prefix. The rule is built so the emitted tokens have exactly the distribution the full model's own sampling would give; only the number of full-model passes changes.

A model that can predict its own next hidden state can draft without running the transformer: decode a token from the current state with the head, step the state with $p_\psi$, decode the next token, and repeat. This is the decoding chain from the theorem, run as a loop:

# drafting with p_psi: no transformer calls inside the loop
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, embed(x))       # step the latent with the MLP
# then score the whole draft with ONE parallel transformer pass and keep
# the accepted prefix (Leviathan et al. 2022 accept/resample rule)

Multi-token prediction methods draft with their extra heads, one per future position, so a model trained with horizon $d$ drafts at most $d$ tokens per cycle. $p_\psi$ is one rule applied repeatedly, so the draft can run past $d$. The paper pretrains 1.3B models on 100 billion FineWeb-Edu tokens, then samples 1,024 prompts of 512 tokens per domain and generates 512-token continuations on 8 B200 GPUs, trying fixed NextLat draft lengths from 2 to 10 and reporting the fastest. At $d=2$ on Books, NextLat averages 4.86 accepted tokens per drafting step and a 3.32× speedup; MTP and JTP accept 1.83 and 1.89 tokens, for 1.72× and 1.90×. Switch the domain and the horizon in Figure 9.

At $d=1$ the multi-token methods accept under one token per step (0.91 to 0.98) for 1.38× to 1.47×, while NextLat accepts 2.66 to 3.64 for 2.29× to 2.72×. Code is the hardest domain for NextLat (2.38× at $d=2$) and Books the easiest. The 3.3× headline is NextLat's best configuration on one domain, and every emitted token is still distributed as the base model's.

Costs, scale and limits

Training cost depends on the horizon. On the 1.3B model (paper Table 4) GPT trains at 3.09 steps per second, NextLat at 3.09 for $d=1$, 2.79 for $d=2$ and 1.73 for $d=8$, because $p_\psi$ is unrolled sequentially over the horizon. The Belief State Transformer trains a second, backward transformer and accumulates gradients over all $O(T^2)$ prefix-suffix pairs; even subsampling 10% of the pairs it runs at 0.89 steps per second, over 3× slower than NextLat. NextLat's gradient terms grow as $O(Td)$.

At 1.3B parameters the language-modeling gains are small. Zero-shot accuracy averaged over eight benchmarks is 59.21 for NextLat at $d=2$ against 58.82 for GPT, 58.69 for MTP and 58.63 for JTP, and the paper says the gains are not consistent across tasks. FineWeb-Edu perplexity is 10.52 for GPT, 10.88 for NextLat, 11.00 for MTP and 11.18 for JTP: NextLat costs less next-token quality than the multi-token methods but still some. The clear wins are on the structured tasks above. The paper does not test frontier scale.

The paper lists its own open questions. $p_\psi$ was only ever an MLP, and its hidden width, which acts as a bottleneck on how much the belief state can hold, was not studied (smaller widths helped on Path-Star and Countdown). The loss design comes from small ablations, and it is unclear whether $d>1$ or the KL term matter at larger scale. The speculative decoder used a fixed draft length per run, not an adaptive one.

NextLat changes the representations, not the computational class. A fixed-depth transformer trained with NextLat is still a fixed-depth transformer, bounded by $\mathsf{TC}^0$; the extra capability in the $A_5$ experiment belongs to $p_\psi$, a recurrent network the training produced on the side. The method differs from hybrid transformer and state-space architectures (such as Mamba layers mixed with attention): it adds no layers and works as a training loss on an unmodified transformer.