Research map / Architectures and efficiency / Attention and transformers

Next-Latent Prediction Transformers Learn Compact World Models

NeurIPS 2026 · Architectures and efficiency › Attention and transformers · approach: Transformer expressivity and learnability · 0 citations (OpenAlex)

Read the paperCheck your idea against 108,967 papers

AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.

Proposes NextLat, an auxiliary latent next-state prediction objective for transformers that yields belief-state-like representations, improving world modeling, reasoning, planning and language modeling while enabling variable-length self-speculative decoding.

Problem

Standard next-token-trained transformers have no incentive to compress history into compact latent states with consistent transition rules, leading to shortcut solutions that generalize poorly and weak lookahead planning.

Approach

Jointly trains a transformer with a lightweight MLP latent dynamics model that predicts the next hidden state from the current hidden state and next token (with stop-gradient targets, Smooth L1 plus a KL token-space alignment), theoretically converging to belief states without changing architecture or inference.

Key results

• On FineWeb-Edu 1.3B models (100B tokens), NextLat (d=2) averages 59.21 zero-shot accuracy vs 58.82 for GPT.
• Self-speculative decoding with NextLat reaches up to 3.3x speedup (Books, d=2), versus at most ~1.9x for JTP and MTP.
• On Manhattan taxi rides, NextLat achieves the lowest effective latent rank (52.7) versus 160.1 for GPT and the highest sequence compression (0.71).

Surprising findings

• The latent dynamics MLP trained on A5 word problem generalizes to 36-token sequences (>95% accuracy) even though the transformer itself fails beyond 12 tokens.
• Smooth L1 loss increases near the end of training (LR cooldown) for NextLat pretraining, with unclear cause, while speculative acceptance keeps improving.

Limitations

• Latent dynamics model uses only simple MLPs; design space not explored.
• Ablations and hyperparameter tuning are small-scale; unclear whether multi-step supervision and KL loss matter at larger scales.
• Multiple-choice accuracy gains over GPT are small and inconsistent across tasks; did not compare against newer MTP variants.

Experimental setup

Datasets: Manhattan random walks Countdown Path-Star graphs TinyStories A5 word problem FineWeb-Edu

Benchmarks: Manhattan taxi world model metrics Countdown Path-Star graph TinyStories linear probes A5 word problem FineWeb-Edu zero-shot (LAMBADA, PIQA, HellaSwag, WinoGrande, ARC, SIQA, SciQ) Wikitext perplexity self-speculative decoding speedup

Metrics: accuracy perplexity solve rate effective latent rank sequence compression detour robustness cross-entropy probe difference speedup accepted tokens

Baselines: GPT MTP JTP BST

Builds on: Belief State Transformer multi-token prediction joint multi-token prediction self-predictive representation learning speculative sampling

Closest papers by meaning

• Reinforced Fast Weights via Next-Sequence Prediction (NeurIPS 2026 · similarity 0.88)
• Transformers Represent Belief State Geometry in their Residual Stream (NeurIPS 2024 · similarity 0.88)
• Non-asymptotic Convergence of Training Transformers for Next-token Prediction (NeurIPS 2024 · similarity 0.88)
• Towards Understanding the Universality of Transformers for Next-Token Prediction (ICLR 2025 · similarity 0.87)
• Bigram Subnetworks: Mapping to Next Tokens in Transformer Language Models (NeurIPS 2025 · similarity 0.87)
• On the Generalization Ability of Next-Token-Prediction Pretraining (ICML 2025 · similarity 0.87)