Next-Latent Prediction Transformers Learn Compact World Models

• Transformers’ growing memory and flexible attention can encourage task-specific shortcuts because they lack inherent pressure to compress history into consistent latent updates.The paper motivates restoring recurrent-style compression while retaining transformer parallelization.Source: §1 · p. 2
• NextLat adds self-supervised latent-state prediction to next-token training, jointly learning compact summaries and recurrent-like latent dynamics without changing transformer architecture or parallel training.The latent dynamics model predicts the next latent state from the current latent state and next token.Source: §1 · p. 2
• NextLat theoretically shapes representations toward belief states: compact summaries of past information sufficient for predicting future observations.These representations provide the paper’s formal account of compact internal world models.Source: §1 · p. 2
• NextLat improves representation compactness, lookahead planning, and downstream accuracy across world modeling, reasoning, planning, and language modeling benchmarks.The comparison is against standard next-token prediction and other baselines.Source: §1 · p. 2
• Up to 3.3× faster language-model inference comes from variable-length self-speculative decoding enabled by recursively predicted latent states.The method avoids separate multi-token prediction heads while allowing flexible future-token speculation.Source: §1 · p. 2

Self-Supervised Learning

• Self-supervised representation learning generates supervisory signals from raw inputs, and NextLat applies self-predictive learning by modeling how latent representations evolve over sequences.The paper contrasts its fully self-supervised objective with paired-data latent prediction.Source: §2 · p. 3

Belief States

• Belief states compress history into sufficient statistics for predicting the future, whereas transformers’ sequence-growing memory does not inherently enforce this compression.NextLat is presented as more computationally efficient than BST while learning belief states without a separate transformer.Source: §2 · p. 3

World Models

• World-modeling systems commonly use latent dynamics that map a latent state and action to the next latent state, a structure NextLat adapts to transformer token sequences.NextLat treats new tokens as actions governing latent-state evolution.Source: §2 · p. 3

Beyond Next-Token Prediction

• Prior richer-supervision methods mainly predict future tokens, while NextLat supervises coherent dynamics in latent space to provide richer gradient signals.The paper positions latent-space prediction as distinct from token-space multi-token objectives.Source: §2 · p. 3

Speculative Decoding

• Unlike post-hoc latent draft models, NextLat shapes representations during pretraining and recursively composes latent dynamics for variable-length self-speculative decoding.Token-space MTP methods are typically constrained by a fixed speculative horizon determined during training.Source: §2 · pp. 3–4

3 Methodology: Next-Latent Prediction

• A belief state is a sufficient statistic of token history from which the future-token distribution can be predicted.Formally, conditioning future functions on the belief state matches conditioning on the complete history.Source: §3 · p. 4

3.1 Why Next-Latent Prediction?

• Joint optimization uses a transformer hidden state, a next-token output head, and a latent dynamics model predicting hidden-state transitions.The theoretical setup analyzes idealized optimization of both prediction objectives.Source: §3.1 · p. 4
• If next-token and transition consistency are both achieved, the hidden state must be a belief state for the preceding sequence.The transition model conditions on the current hidden state and next token, matching the transformer-induced next-state law.Source: §3.1 · p. 4
• Next-token consistency alone does not guarantee belief-state representations because self-attention can retrieve past tokens without compressing them into latent summaries.NextLat adds transition consistency to supply the missing compression pressure.Source: §3.1 · p. 5
• Latent supervision predicts the next hidden state and recursively propagates information about later states, providing denser distribution-level signals than one-hot token targets.This richer signal is presented as improving data efficiency.Source: §3.1 · p. 5

3.2 Learning to Predict Next-Latent States

• The practical objective combines next-token cross-entropy, multi-step hidden-state prediction, and a KL loss aligning predicted-state token distributions.Stop-gradients help prevent representational collapse, while multi-step supervision provides richer learning signals beyond the d = 1 theoretical guarantee.Source: §3.2 · pp. 5–6
• The latent dynamics model is needed during training to shape representations, while the transformer decodes independently at inference.Experiments use simple MLP dynamics models to isolate the core approach.Source: §3.2 · pp. 5–6

3.3 Variable-Length Self-Speculative Decoding

• Recursive latent dynamics let NextLat draft a variable number of tokens, unlike MTP’s fixed horizon, and remain coherent beyond shallow training horizons d = 1, 2.This reduces draft-verification constraints and supports faster inference.Source: §3.3 · p. 6, Figure 2 · p. 6

4.1 World Modeling

• NextLat consistently outperforms baselines across world-modeling metrics, achieving sequence compression of 0.71 and effective latent rank of 52.7, over 3× smaller than GPT’s.All models reach 100% on the next-token diagnostic, but NextLat produces a more coherent internal map and stronger OOD trajectory and detour robustness.Source: §4.1 · pp. 7–8, Table 1 · p. 7

4.2 Reasoning

• NextLat consistently outperforms all baselines on Countdown, exceeding same-horizon MTP and JTP by over 35.7% even with d = 1.For final-equation validity at d = 1, NextLat reaches 54.8% versus 42.3% for the next-best baseline.Source: Figure 4 · p. 8, §4.2 · p. 9, Figure 5 · p. 9

4.3 Planning

• NextLat maintains close to a 100% solve rate across all Path-Star graph topologies, unlike MTP, JTP, and BST.The result uses the more difficult fixed-sample setup with 200k examples and N = 100, and is attributed to latent-space transition consistency avoiding local shortcuts.Source: §4.3 · pp. 9–10, Figure 6 · p. 9
• NextLat provides the strongest long-range predictive representations up to 20 tokens ahead while matching GPT’s next-token performance for d ∈ {1, 8}.BST, MTP, and JTP degrade next-token or long-horizon probe performance, whereas NextLat preserves both short- and long-range prediction.Source: §4.3 · p. 10, Figure 8 · p. 10

4.4 Language Modeling

• NextLat modestly improves average zero-shot accuracy over GPT at d = 2, 59.21 versus 58.82, while preserving next-token perplexity better than MTP and JTP.The accuracy gains are not consistent across tasks, and larger models may be needed for larger improvements.Source: §4.4 · p. 11, Table 2 · p. 11
• NextLat achieves up to 3.3× inference speedup with variable-length speculative drafts, including fully valid drafts at length 10.It exceeds its training horizon and outperforms longer-horizon JTP and MTP while requiring only shallow, computationally efficient training horizons.Source: Table 3 · p. 11, §4.4 · p. 12, Figure 9 · p. 12

5.1 Comparison Against Current Approaches

• BST is over 3× slower than NextLat in FineWeb-Edu training, even after 10% prefix–suffix subsampling, while NextLat with d = 1 adds negligible GPT overhead.BST also uses two transformer encoders during training and inference.Source: Table 4 · p. 12, §5.1 · p. 13
• NextLat achieves substantially better benchmark performance than JTP across all benchmarks despite a modest sequential latent-rollout cost.At d = 8, JTP trains faster, but NextLat’s sequential computation remains far more efficient than recurrent neural networks.Source: §5.1 · p. 13
• NextLat learns belief states independently of prediction horizon d, avoiding BST’s O(T^2) gradient cost and JTP’s restrictive d ≥ k condition.Larger horizons provide richer gradient signals rather than being required for the belief-state guarantee.Source: §5.1 · p. 13

5.2 NextLat vs. Recurrent Neural Networks

• On A5, the NextLat transformer fails beyond 12 tokens, but its co-trained latent dynamics model reaches over 95% accuracy on 36-token sequences.The latent RNN uses substantially fewer parameters than the transformer, 2.62M versus 6.43M.Source: §5.2 · p. 14, Figure 10 · p. 14
• NextLat adds recurrent inductive bias without strict recurrence, using only O(d) sequential dependence during training instead of RNNs’ O(T).For d = 1, transformer and latent RNN co-training is fully parallel across time.Source: §5.2 · pp. 14–15
• NextLat’s belief-state convergence is independent of d because full transformer gradients remain intact, unlike truncated backpropagation through time.The latent dynamics model supplies outer-loop temporal consistency, while larger horizons can accelerate empirical convergence.Source: §5.2 · p. 15
• NextLat supports nonlinear latent transitions without changing transformer architecture, but its fixed-depth transformer remains bounded by TC0 expressivity.The authors note that Figure 10 raises whether the co-trained latent dynamics model shares that limitation.Source: §5.2 · pp. 14–15
• NextLat is an auxiliary objective rather than a hybrid Transformer–SSM architecture, making it compatible with diverse sequence-modeling architectures.The authors suggest it could complement hybrid models by organizing residual attention pathways.Source: §5.2 · p. 15

6 Limitations and Future Work

• All experiments use simple MLP latent dynamics, leaving the effects of more expressive architectures and hidden-layer widths unexamined.Smaller latent dimensions help on Path-Star and Countdown, but the latent bottleneck’s task-dependent capacity remains insufficiently studied.Source: §6 · p. 15
• The necessity and interaction of multi-step supervision, KL token-level supervision, stop-gradients, and Smooth L1 loss remain unclear at larger scales.The objective design is guided mainly by small-scale ablations and empirical intuition.Source: §6 · p. 15
• Evaluation omitted newer specialized MTP variants and did not fully test adaptive variable-length speculative decoding.Draft lengths were fixed per prompt and varied only statically between 2 and 10 tokens.Source: §6 · pp. 15–16
• The study does not analyze learned latent-space structure and observes optimization-sensitive loss trajectories during pretraining.The authors identify scaling and parameterizing the NextLat objective as an open direction.Source: §6 · p. 16
• Future work includes post-hoc finetuning, reinforcement-learning post-training, and higher-dimensional or hierarchical belief-state architectures.These directions aim to extend NextLat beyond its preliminary study.Source: §6 · p. 16

7 Conclusion

• NextLat augments next-token training with self-supervised latent-space prediction, enabling transformers to learn belief-state representations.The objective encourages concise latent summaries sufficient for predicting the future.Source: §7 · p. 16
• Across world modeling, reasoning, planning, and language modeling, NextLat yields more compact, predictive, and generalizable representations.These gains occur without changing transformer architecture, parallel training efficiency, or inference procedure.Source: §7 · p. 16
• NextLat combines recurrent models’ bias toward compact, temporally consistent representations with transformers’ scalability and parallelism.The paper presents this combination as broadly applicable to autoregressive sequence modeling.Source: §7 · p. 16

A Belief States in Sequence Modeling

• Belief states compress sequence history into representations sufficient for predicting the distribution of future tokens.They preserve all information needed for future prediction while discarding irrelevant historical detail.Source: A Belief States in Sequence Modeling · p. 23
• BST shapes hidden representations into belief states by jointly training next-token and previous-token predictors across prefix–suffix decompositions.Its bidirectional objective uses separate transformers for prefix and suffix representations.Source: A Belief States in Sequence Modeling · p. 23
• JTP learns belief states in k-observable systems when its joint next-k prediction recovers the true next-k conditional distribution.Under this condition, the hidden state determines the full-horizon future conditional.Source: A Belief States in Sequence Modeling · p. 23

B Formal Proof of Theorem 3.2

• NextLat’s proof uses next-token and transition consistency to establish measurable maps for recursively decoding future tokens from latent states.The formal argument proceeds by backward induction, showing that if h_k+1 is a belief state then h_k is also a belief state.Source: B Formal Proof of Theorem 3.2 · p. 24
• The induction concludes that earlier latent states are belief states because they generate the next token and transition to a later belief state with the correct conditional dynamics.The base case follows from next-token prediction at the final step.Source: B Formal Proof of Theorem 3.2 · p. 24

C More Details on NextLat Implementation

• The latent transition model is a three-layer GELU MLP that predicts a residual update from the normalized current hidden state and next-token embedding.The predicted delta is added to the current hidden state to form the next latent estimate.Source: C More Details on NextLat Implementation · p. 24
• Training masks next-token and KL losses on context tokens but retains next-hidden supervision there, while sequence packing masks latent losses across document boundaries.This preserves belief-state development during context processing without mixing separate documents.Source: C More Details on NextLat Implementation · p. 24
• The implementation supports recursive multi-step latent prediction by repeatedly applying the transition model to shifted hidden states and future tokens.The default multi-step horizon is one, but the loop permits longer recursive rollouts.Source: C More Details on NextLat Implementation · p. 25
• The implementation uses detached hidden-state targets and a detached output head for auxiliary losses, reducing unwanted gradient paths through the transformer and head.The pseudocode applies stop-gradient before the Smooth L1 and KL computations.Source: C More Details on NextLat Implementation · p. 25
• NextLat’s loss combines next-token cross-entropy, Smooth L1 next-hidden prediction, and KL alignment between output distributions from predicted and target hidden states.The overall objective weights the latent and KL terms alongside the standard next-token loss.Source: C More Details on NextLat Implementation · p. 25

D Ablation Studies

• At d = 8, KL + Smooth L1 gives the best probe performance, while Smooth L1 alone is strongest at d = 1.The combined objective also achieves lower validation KL and Smooth L1 losses than either loss alone, indicating complementary supervision.Source: D Ablation Studies · p. 26, Figure 12 · p. 26
• KL matching yields more predictive representations than CE at shallow multi-step prediction horizons, but this advantage diminishes at larger horizons.KL incurs higher memory overhead because it materializes full logits and softmax distributions over the vocabulary, unlike fused CE implementations.Source: D Ablation Studies · p. 27
• Applying stop-gradients to both the Smooth L1 target and KL output head yields the best probing performance across token offsets.Removing the Smooth L1 stop-gradient slightly improves validation losses by about 0.02 but does not improve representation quality.Source: D Ablation Studies · p. 27, Figure 14 · p. 27
• Stop-gradients improve efficiency by eliminating an additional backward pass through the vocabulary-sized output head.The paper also reports that stop-gradient, especially for Smooth L1, is essentially needed for high accuracy on Countdown and Path-Star graph.Source: D Ablation Studies · p. 27
• Overall, the ablations attribute NextLat’s effectiveness to the combination of Smooth L1, KL, and carefully placed stop-gradients.Together these choices support predictive belief-state representations rather than optimizing any single auxiliary objective in isolation.Source: D Ablation Studies · p. 27

E NextLat Pretraining Quirks

• Muon and AdamW produce substantially different Smooth L1 trajectories under the same learning-rate schedule, while AdamW achieves better KL optimization.The authors suggest Muon’s orthogonalized updates may shape latent representations more efficiently, but this explanation is not established.Source: E NextLat Pretraining Quirks · p. 27, Figure 15 · p. 28
• Both optimizers show rising Smooth L1 loss during learning-rate cooldown, with the increase substantially more pronounced for AdamW.The authors explored several loss, gradient, coefficient, decay, and weight-decay changes without resolving the phenomenon.Source: E NextLat Pretraining Quirks · p. 28
• Speculative-decoding acceptance over a five-token lookahead continues improving while AdamW’s Smooth L1 loss rises.This suggests the loss increase may reflect latent-state scale or geometry changes rather than degraded token-level transition coherence.Source: E NextLat Pretraining Quirks · p. 28, Figure 15 · p. 28
• Gradient-norm spikes during MTP training with Muon caused instability and limited the pretraining analysis to AdamW.This constrains direct optimizer comparisons in the pretraining study.Source: E NextLat Pretraining Quirks · p. 28

F Extra Experiment Details

• The supplementary hyperparameter search varies λnext-h and λKL, finding λKL = 1.0 robust across tasks and λnext-h = 1.0 effective in most settings.For FineWeb-Edu, all compared methods use uniform loss weights because extensive tuning is computationally prohibitive.Source: F Extra Experiment Details · pp. 28–29

F.1 Manhattan Taxi Rides

• Six epochs and sequence packing support Manhattan Taxi Rides training, while evaluation is limited to paths of up to 50 steps to avoid unsupported long-horizon planning.The 48-layer, 384-dimensional models use sequence packing for efficiency; the evaluation adjustment aims to separate world-model incoherence from planning beyond the training distribution.Source: §F.1 · p. 29

F.2 Countdown

• The Countdown setup uses atomic integer tokens from 1 to 10,000 and eight pipe-symbol pause tokens before solution generation.Prompt losses are masked, and operators and delimiters receive separate token indices.Source: §F.2 · p. 30

F.3 Path-Star

• Path-Star follows a fixed, challenging graph distribution and increases weight decay to 0.1, with the task probing whether latent-space prediction supports generalizable solutions under limited data.The experiment targets superficial teacher-forced strategies and evaluates correct-arm generation on 20k held-out instances.Source: §F.3 · p. 30

F.4 TinyStories

• TinyStories freezes pretrained models and trains one-layer linear probes for 20k steps using the same dataset and pretraining hyperparameters.Probe hidden-state choices vary by model, with NextLat using the same final-layer states as in the effective-rank analysis.Source: §F.4 · pp. 30–31

F.5 A5 State Tracking

• A5 state tracking trains on 1 million 12-token sequences and tests length generalization on approximately 100k 36-token sequences using one-step latent supervision alone.Rotary position embeddings are retained because removing positional embeddings degraded performance, yet the learned latent dynamics still generalize beyond training length.Source: §F.5 · p. 31

F.6 FineWeb-Edu Pretraining

• FineWeb-Edu pretraining uses a 4e−4 peak learning rate, 0.1 weight decay, 500M-token global batches, 1024-token sequences, and a warmup-stable-decay schedule.Evaluation covers nine zero-shot language-modeling benchmarks and self-speculative decoding across Wikipedia, books, code, and math domains.Source: §F.6 · pp. 31–32

F.6.1 Results with MTP/JTP (d=4)

• Training JTP and MTP with d = 4 is substantially more expensive than training NextLat with d = 1 or d = 2.The comparison attributes this difference to the increased training cost of larger-horizon multi-token prediction methods.Source: §F.6.1 · p. 32
• NextLat (d = 2) achieves the best average accuracy overall, while increasing MTP and JTP to d = 4 yields no meaningful multiple-choice improvement.This comparison is reported in Table 6 for 1.3B-parameter models trained on 100B FineWeb-Edu tokens.Source: §F.6.1 · p. 32, Table 6 · p. 32
• NextLat (d = 2) achieves the highest overall speedup and continues improving with longer draft lengths, whereas d = 4 MTP and JTP curves saturate earlier and remain lower.Across diverse domains, d = 4 methods generally fail to surpass NextLat except in Code.Source: §F.6.1 · p. 32, Figure 17 · p. 32, Table 7 · p. 32