pith. sign in

arxiv: 2511.05963 · v3 · pith:EREDXFXXnew · submitted 2025-11-08 · 💻 cs.LG

Next-Latent Prediction Transformers Learn Compact World Models

Jayden Teoh , Manan Tomar , Kwangjun Ahn , Edward S. Hu , Tim Pearce , Pratyusha Sharma , Akshay Krishnamurthy , Riashat Islam
show 2 more authors
Alex Lamb John Langford
This is my paper

Pith reviewed 2026-05-25 07:18 UTC · model grok-4.3

classification 💻 cs.LG
keywords transformersworld modelsbelief stateslatent predictionnext-token predictioncompact representationsinductive biasself-supervised learning
0
0 comments X

The pith

Adding a next-latent prediction loss makes transformer internal states converge to belief states that compress history for future prediction.

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The paper proposes NextLat, an auxiliary objective that trains a transformer to predict its next latent representation from the current latent and the next token, in addition to standard next-token prediction. Without this term, transformers have no built-in pressure to maintain compact, consistent internal states across time steps. The authors prove that the joint objective forces the latents to converge toward belief states, the minimal information about past observations needed to predict future ones. This change is claimed to produce more coherent internal world models while preserving the original architecture, parallel training, and inference procedure. Readers would care because such representations are said to improve accuracy on tasks that require planning or reasoning over sequences.

Core claim

NextLat extends standard next-token training with self-supervised predictions in the latent space, training the model to predict its next latent state given the next token. The paper shows theoretically that the resulting latents converge to belief states, which are compressed summaries of history sufficient for predicting future observations. This auxiliary objective injects a recurrent inductive bias into transformers, encouraging formation of compact internal world models with coherent transition dynamics that standard next-token prediction does not guarantee.

What carries the argument

The next-latent prediction auxiliary loss, jointly optimized with the primary next-token loss, which drives convergence of latents to belief states.

If this is right

  • The learned latents become more compressed representations of history.
  • Downstream accuracy improves on world modeling, reasoning, planning, and language modeling tasks.
  • Variable-length self-speculative decoding becomes possible, accelerating inference up to 3.3x.
  • The transformer architecture, parallel training efficiency, and inference procedure remain unchanged.

Where Pith is reading between the lines

These are editorial extensions of the paper, not claims the author makes directly.

  • The same auxiliary objective might be added to other non-recurrent sequence models to induce similar compression.
  • Improved lookahead planning may stem directly from the enforced consistency of the learned transition dynamics.
  • Smaller models trained with NextLat could match the effective capacity of larger models trained only on next-token loss.

Load-bearing premise

Jointly optimizing the auxiliary latent-prediction loss with the main next-token loss will produce convergence to belief states without the auxiliary term dominating or destabilizing training.

What would settle it

Train a transformer under the NextLat objective and check whether its latent representations fail to compress history or produce inconsistent next-state predictions across different sequences of the same underlying process.

Figures

Figures reproduced from arXiv: 2511.05963 by Akshay Krishnamurthy, Alex Lamb, Edward S. Hu, Jayden Teoh, John Langford, Kwangjun Ahn, Manan Tomar, Pratyusha Sharma, Riashat Islam, Tim Pearce.

Figure 1
Figure 1. Figure 1: Reconstructed maps from sequences generated by transformers trained on Manhattan taxi rides [ [PITH_FULL_IMAGE:figures/full_fig_p001_1.png] view at source ↗
Figure 2
Figure 2. Figure 2: Illustration of different predictive mechanisms at time step t = 3. Other methods supervise only the token-level emissions, leaving intermediate latent representations implicit. In contrast, NextLat explicitly learns latent dynamics that predicts hidden state hˆ t+1 from (ht, xt+1). Token-level supervision is then applied to the hˆ t+1. Therefore, accurate multi-token predictions (beyond the next token) em… view at source ↗
Figure 3
Figure 3. Figure 3: Reconstructed maps from transformers trained on Manhattan taxi rides using different objectives. [PITH_FULL_IMAGE:figures/full_fig_p007_3.png] view at source ↗
Figure 4
Figure 4. Figure 4: Performance on Countdown. Best result is bolded, and second best is underlined. Eq. 1 Eq. 2 Eq. 3 0 20 40 60 80 100 Validity (%) GPT BST MTP JTP NextLat [PITH_FULL_IMAGE:figures/full_fig_p008_4.png] view at source ↗
Figure 5
Figure 5. Figure 5: Validity of equations (i.e., LHS = RHS) generated on Countdown. All models in this plot use d = 1. Setup. Following Gandhi et al. [2024], we generate 500k training problems with target numbers ranging from 10 to 100 and reserve 10% of the targets for out-of-distribution evaluation. During both training and testing, we insert eight ‘pause tokens’ [Goyal et al., 2023] after the target number, allowing models… view at source ↗
Figure 7
Figure 7. Figure 7: Illustration of a G5,5 Path-Star graph [Bach￾mann and Nagarajan, 2024]. 8 [PITH_FULL_IMAGE:figures/full_fig_p008_7.png] view at source ↗
Figure 8
Figure 8. Figure 8: Cross-entropy loss difference relative to GPT, obtained from linear probes trained on frozen hidden states [PITH_FULL_IMAGE:figures/full_fig_p009_8.png] view at source ↗
Figure 9
Figure 9. Figure 9: Illustration of the latent transition model pψ. We parameterize the latent transition model pψ with a three-layer MLP using GELU [Hendrycks and Gimpel, 2023] activations. The latent transition model takes as input the layer-normalized [Ba et al., 2016] concatenation of the current hidden state ht and next-token embedding Xt+1, and outputs a delta update applied via residual connection: hˆ t+1 = pψ(ht, Xt+1… view at source ↗
Figure 10
Figure 10. Figure 10: Full plot version of [PITH_FULL_IMAGE:figures/full_fig_p023_10.png] view at source ↗
read the original abstract

Transformers replace recurrence with a memory that grows with sequence length and self-attention that enables ad-hoc lookups over past tokens. Consequently, they lack an inherent incentive to compress history into compact latent states with consistent transition rules. This often leads to learning solutions that generalize poorly. We introduce Next-Latent Prediction (NextLat), which extends standard next-token training with self-supervised predictions in the latent space. Specifically, NextLat trains a transformer to learn latent representations that are predictive of its next latent state given the next token. Theoretically, we show that these latents provably converge towards belief states, compressed information about the history necessary to predict the future. This simple auxiliary objective injects a recurrent inductive bias into transformers while leaving their architecture, parallel training efficiency, and inference unchanged. NextLat effectively encourages transformers to form compact internal world models with coherent belief states and transition dynamics -- crucial properties not guaranteed by standard next-token prediction alone. Empirically, across benchmarks in world modeling, reasoning, planning, and language modeling, NextLat demonstrates significant gains over standard next-token prediction and other baselines in downstream accuracy, representation compression, and lookahead planning. Furthermore, NextLat enables variable-length self-speculative decoding, accelerating inference by up to 3.3x in language modeling. NextLat offers a simple yet effective paradigm for learning compact, predictive representations in transformers that generalize better. Our code is available at https://github.com/microsoft/NextLat.

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, simulated authors' rebuttal, and a circularity audit. Tearing a paper down is the easy half of reading it; the pith above is the substance, this is the friction.

Referee Report

2 major / 2 minor

Summary. The paper proposes Next-Latent Prediction (NextLat), which augments standard next-token cross-entropy training of transformers with an auxiliary self-supervised loss that predicts the next latent representation given the next token. It claims that the resulting latents provably converge to belief states (compressed sufficient statistics of history for future prediction), inject a recurrent inductive bias without altering architecture or inference, and yield empirical gains in world modeling, reasoning, planning, language modeling, representation compression, and up to 3.3x faster self-speculative decoding.

Significance. If the convergence result holds under the joint objective actually used in training and the reported gains are robust to controls, the method would offer a lightweight way to encourage compact, predictive internal world models in transformers while preserving their parallel training advantages. The combination of a theoretical fixed-point argument with downstream improvements in lookahead planning and variable-length decoding would be a notable contribution to representation learning for sequential decision-making tasks.

major comments (2)
  1. [Abstract / theoretical analysis section] The theoretical claim that latents converge to belief states is derived only for the auxiliary latent-prediction objective (abstract and the paragraph beginning 'Theoretically, we show...'). No argument is supplied that the stationary points of the combined loss (next-token cross-entropy plus weighted auxiliary term) remain belief states; the next-token term could in principle select representations that are locally predictive of the immediate token but violate the belief-state fixed-point condition. This gap is load-bearing for the central 'provably converge' claim.
  2. [Experiments section] The empirical sections report gains on world-modeling, reasoning, and planning benchmarks, but the manuscript does not include an ablation that isolates whether the auxiliary loss alone (without the next-token term) already produces the claimed belief-state convergence, nor does it verify that the learned latents satisfy the belief-state property on the actual trained models.
minor comments (2)
  1. [Abstract] The abstract states 'Our code is available at https://github.com/microsoft/NextLat' but does not specify the commit or tag corresponding to the reported experiments.
  2. [Theoretical analysis] Notation for the latent transition and belief-state definitions should be introduced with explicit equations rather than prose descriptions to allow direct comparison with the derived fixed-point condition.

Simulated Author's Rebuttal

2 responses · 0 unresolved

We are grateful to the referee for the constructive feedback on our manuscript. We address the two major comments below and plan to make revisions to strengthen the theoretical claims and experimental validation.

read point-by-point responses

  1. Referee: [Abstract / theoretical analysis section] The theoretical claim that latents converge to belief states is derived only for the auxiliary latent-prediction objective (abstract and the paragraph beginning 'Theoretically, we show...'). No argument is supplied that the stationary points of the combined loss (next-token cross-entropy plus weighted auxiliary term) remain belief states; the next-token term could in principle select representations that are locally predictive of the immediate token but violate the belief-state fixed-point condition. This gap is load-bearing for the central 'provably converge' claim.

    Authors: The referee correctly identifies that our theoretical analysis derives the convergence to belief states specifically for the auxiliary latent-prediction objective. We did not provide a proof that the stationary points of the joint loss are belief states. In the revised manuscript, we will revise the abstract and the theoretical section to accurately reflect that the auxiliary objective has belief states as fixed points, and that the combined training is intended to encourage this property while maintaining next-token prediction performance. We will also add a discussion on why the next-token loss is not expected to violate the fixed-point condition under suitable hyperparameter choices for the auxiliary weight. revision: yes

  2. Referee: [Experiments section] The empirical sections report gains on world-modeling, reasoning, and planning benchmarks, but the manuscript does not include an ablation that isolates whether the auxiliary loss alone (without the next-token term) already produces the claimed belief-state convergence, nor does it verify that the learned latents satisfy the belief-state property on the actual trained models.

    Authors: We agree that the manuscript would benefit from an ablation isolating the auxiliary loss and direct verification of the belief-state property on the trained models. In the revision, we will add such an ablation study where possible (noting that training solely with the auxiliary loss may require adjustments for stability) and include metrics or checks to verify that the learned latents act as sufficient statistics for future predictions. This will help confirm the empirical realization of the theoretical property. revision: yes

Circularity Check

0 steps flagged

No circularity; theoretical claim presented as independent proof

full rationale

The paper states a theoretical result that latents converge to belief states under the auxiliary latent-prediction objective. No quoted equations or self-citations reduce this claim by construction to fitted inputs, renamed empirical patterns, or load-bearing prior work by the same authors. The joint next-token + auxiliary loss is acknowledged as the actual training objective, but the provided text frames the convergence as a separate proof rather than a statistical consequence of the fit itself. This is the normal case of a self-contained derivation; the skeptic concern is a potential applicability gap, not a circular reduction.

Axiom & Free-Parameter Ledger

0 free parameters · 0 axioms · 0 invented entities

Abstract-only review; no explicit free parameters, axioms, or invented entities are described.

pith-pipeline@v0.9.0 · 5821 in / 995 out tokens · 54763 ms · 2026-05-25T07:18:57.466192+00:00 · methodology

discussion (0)

Sign in with ORCID, Apple, or X to comment. Anyone can read and Pith papers without signing in.

Lean theorems connected to this paper

Citations machine-checked in the Pith Canon. Every link opens the source theorem in the public Lean library.

  • IndisputableMonolith/Cost/FunctionalEquation.lean washburn_uniqueness_aczel echoes
    ?
    echoes

    ECHOES: this paper passage has the same mathematical shape or conceptual pattern as the Recognition theorem, but is not a direct formal dependency.

    Theoretically, we show that these latents provably converge towards belief states, compressed information about the history necessary to predict the future. ... For these consistency objectives to be satisfied, ht must converge to a belief state

  • IndisputableMonolith/Foundation/RealityFromDistinction.lean reality_from_one_distinction echoes
    ?
    echoes

    ECHOES: this paper passage has the same mathematical shape or conceptual pattern as the Recognition theorem, but is not a direct formal dependency.

    Optimizing for next-token consistency ... and transition consistency ... ensures existence of measurable maps ... ht must jointly optimize toward a belief state

What do these tags mean?
matches
The paper's claim is directly supported by a theorem in the formal canon.
supports
The theorem supports part of the paper's argument, but the paper may add assumptions or extra steps.
extends
The paper goes beyond the formal theorem; the theorem is a base layer rather than the whole result.
uses
The paper appears to rely on the theorem as machinery.
contradicts
The paper's claim conflicts with a theorem or certificate in the canon.
unclear
Pith found a possible connection, but the passage is too broad, indirect, or ambiguous to say the theorem truly supports the claim.

Forward citations

Cited by 4 Pith papers

Reviewed papers in the Pith corpus that reference this work. Sorted by Pith novelty score.

  1. Semantic Step Prediction: Multi-Step Latent Forecasting in LLM Reasoning Trajectories via Step Sampling

    cs.LG 2026-04 unverdicted novelty 7.0

    Applying STP at consecutive semantic reasoning steps achieves 168x more accurate multi-step latent prediction on ProcessBench than frozen baselines, with trajectories forming smooth curves best captured by non-linear ...

  2. Improving Sampling for Masked Diffusion Models via Information Gain

    cs.CL 2026-02 unverdicted novelty 7.0

    Info-Gain Sampler improves MDM decoding by using bidirectional information gain to reduce cumulative uncertainty, outperforming greedy samplers on reasoning accuracy and creative writing tasks.

  3. Exploring the Potential of Probabilistic Transformer for Time Series Modeling: A Report on the ST-PT Framework

    cs.LG 2026-04 unverdicted novelty 6.0

    ST-PT turns transformers into explicit factor graphs for time series, enabling structural injection of symbolic priors, per-sample conditional generation, and principled latent autoregressive forecasting via MFVI iterations.

  4. The Depth Ceiling: On the Limits of Large Language Models in Discovering Latent Planning

    cs.LG 2026-04 unverdicted novelty 6.0

    LLMs discover latent planning strategies up to five steps during training and execute them up to eight steps at test time, with larger models reaching seven under few-shot prompting, revealing a dissociation between d...