Beyond Next-Token Prediction: Building Coherent World Models Through Multi-Token Prediction and Latent Semantic Enhancement

Paper: 2604.06155 Authors: Qimin Zhong, Hao Liao, Haiming Qin, Mingyang Zhou, Rui Mao, Wei Chen, Naipeng Chao Categories: cs.LG, cs.AI, cs.CL

The Gap

Existing research has shown that LLMs can perform impressive reasoning tasks, but whether they develop coherent internal world models—structured representations of how the world works—remains contested. Next-Token Prediction (NTP) trains models to predict one step ahead, which provides weak supervision for learning long-term dependencies and causal structures. Multi-Token Prediction (MTP) emerged as a promising alternative, predicting multiple future tokens simultaneously, but prior work lacked theoretical understanding of why MTP might learn better representations.

This paper identifies a critical flaw: standard MTP, despite its theoretical advantages, suffers from structural hallucinations where the model learns illegal shortcuts in latent space that violate environmental constraints. The gap is between MTP’s theoretical promise of learning structured representations and its practical failure to respect the underlying state dynamics.

Problem: LLMs lack coherent world models
   |
   v
Observation: NTP gives weak supervision for structure
   |
   v
Hypothesis: MTP induces representational contractivity
   |                                    |
   v                                    v
Theory: Gradient coupling        Reality: Structural hallucinations
promotes belief convergence      (discrete tokens -> illegal shortcuts)
   |                                    |
   +------------------------------------+
                    |
                    v
            Solution: LSE-MTP
    (anchor to continuous state trajectories)
                    |
                    v
            Evidence: Synthetic graphs + Manhattan Taxi
    (better alignment, fewer hallucinations, robust)
                    |
                    v
        Conclusion: Bridging discrete/continuous gap
                    enables coherent world models

The Increment

One sentence: Before this paper, MTP was known to improve representation learning but lacked theoretical grounding and suffered from structural hallucinations; after this paper, we understand MTP’s gradient inductive bias toward belief states and have a method (LSE-MTP) to prevent illegal latent shortcuts by anchoring to continuous state trajectories.

Core Mechanism

LSE-MTP operates in three stages. First, during training, the model generates hidden state representations for multiple future tokens using standard transformer layers. Second, instead of only supervising the discrete token predictions, LSE-MTP extracts the ground-truth hidden state trajectory from a reference model or oracle and computes an auxiliary loss that pulls the model’s predicted hidden states toward these ground-truth trajectories. Third, the total loss combines the standard cross-entropy loss on discrete tokens with this latent semantic alignment loss, weighted by a hyperparameter.

The key innovation is the dual supervision signal. The discrete token loss teaches the model what to predict, while the latent semantic loss teaches the model how to represent the underlying state dynamics. By anchoring predictions to continuous state trajectories rather than just discrete token sequences, LSE-MTP prevents the model from learning shortcuts that produce correct tokens through invalid intermediate representations.

Input sequence
      |
      v
  [Encoder]
      |
      v
Hidden states: h_t, h_{t+1}, h_{t+2}, ...
      |
      +------------------+------------------+
      |                  |                  |
      v                  v                  v
[Token head]      [Token head]      [Token head]
      |                  |                  |
      v                  v                  v
  token_t+1         token_t+2         token_t+3
      |                  |                  |
      v                  v                  v
  [CE Loss] <-------- Discrete supervision
      
      +------------------+------------------+
      |                  |                  |
      v                  v                  v
Ground-truth:     h*_{t+1}          h*_{t+2}
      |                  |                  |
      v                  v                  v
  [MSE Loss] <-------- Continuous supervision
      
      Both losses combined -> Total gradient

Think of LSE-MTP like teaching someone to navigate a city. Standard MTP is like showing them a sequence of street names to memorize—they might learn the route but not understand the spatial layout. They could “hallucinate” shortcuts that work on paper (correct street names) but violate physical constraints (streets don’t actually connect that way). LSE-MTP is like also showing them a map with their position marked at each step. The street names (discrete tokens) tell them where to go, while the map positions (continuous states) ensure they understand the actual spatial relationships. If they try to take an impossible shortcut, the map supervision catches it because their internal position would be wrong, even if they named the right streets.

Key Concepts

• Representational Contractivity: Imagine you’re tracking a ball bouncing in a room. If you’re uncertain about its exact position, your uncertainty should shrink over time as you gather more observations—not explode into wildly different possibilities. Representational contractivity means the model’s internal uncertainty about the world state decreases as it processes more information. In MTP, gradient coupling across multiple prediction heads creates a pressure for the model’s representations to converge toward a single coherent belief state, rather than maintaining multiple inconsistent interpretations. This happens because errors in early predictions propagate through the gradient to affect later predictions, forcing the model to find representations that work consistently across all time steps.

• Structural Hallucinations: These aren’t the usual “making up facts” hallucinations. Structural hallucinations occur when a model produces correct surface outputs (right tokens) through invalid internal reasoning (wrong latent states). Imagine a student who gets the right answer on a math test by accidentally canceling two mistakes. In MTP, the model might learn to predict “turn left” → “Main Street” → “arrive” by encoding an impossible spatial trajectory in its hidden states—like teleporting through buildings. The discrete token supervision only cares that the words are right, not that the implied path is physically valid. This is why standard MTP can score well on token accuracy while failing to learn coherent world models.

• Latent Semantic Enhancement: This is the core fix. Instead of only checking if the model outputs the right words, we also check if its internal representations match the ground-truth state trajectory. Think of it as the difference between grading only the final answer versus also checking the work shown. The “latent semantic” part means we’re supervising the meaning encoded in hidden states, not just surface tokens. The “enhancement” means we’re adding this supervision on top of standard token prediction. Practically, this requires access to ground-truth state trajectories during training—either from a simulator (synthetic tasks) or a pre-trained reference model (real-world tasks).

Framework Shift

Before (Standard MTP):                After (LSE-MTP):

Token sequence                        Token sequence
  |   |   |                             |   |   |
  v   v   v                             v   v   v
[Predict] [Predict] [Predict]         [Predict] [Predict] [Predict]
  |   |   |                             |   |   |
  v   v   v                             v   v   v
Loss on tokens only                   Loss on tokens
                                        |   |   |
Hidden states float free               v   v   v
(can violate constraints)             [Align] [Align] [Align]
                                        |   |   |
                                        v   v   v
                                      Ground-truth states
                                      (enforces constraints)

From supervising only discrete outputs to supervising both discrete outputs and continuous internal states, the core shift is anchoring representations to the actual dynamics of the world being modeled.

Expert Assessment

Problem choice: This is a real gap. The question of whether LLMs develop world models is central to understanding their capabilities and limitations. The observation that MTP suffers from structural hallucinations is novel and well-motivated—it explains why MTP sometimes fails to improve reasoning despite better perplexity. The problem sits at the intersection of representation learning and world modeling, which is currently hot in the field.

Method maturity: The theoretical analysis of gradient coupling is elegant, but the solution feels somewhat brute-force. LSE-MTP requires access to ground-truth state trajectories, which limits applicability. For synthetic tasks, this is fine—you have the simulator. For real-world tasks, they use a pre-trained model as the oracle, which is circular (you need a good model to train a good model). A more mature approach would learn to infer state trajectories without external supervision, perhaps through self-consistency constraints or contrastive learning. That said, the method is a solid first step and the empirical gains are convincing.

Experimental integrity: The baselines are fair—they compare against standard NTP and MTP. The synthetic graph navigation task is well-designed to isolate the structural hallucination problem. The Manhattan Taxi dataset provides real-world validation. However, I’d like to see more analysis of failure cases. When does LSE-MTP still hallucinate? How sensitive is it to the quality of the ground-truth trajectories? The ablation studies are decent but could be deeper. The paper also doesn’t discuss computational cost—LSE-MTP requires computing and storing ground-truth trajectories, which isn’t free.

Writing quality: The paper is generally well-written, but the theoretical section (Section 3) is dense and could benefit from more intuitive explanations before diving into the math. The connection between gradient coupling and representational contractivity is crucial but gets buried in notation. If the authors rewrote Section 3 to lead with intuition and simple examples, then formalized it, the paper would be much more accessible. The experimental section is clear and well-organized.

Verdict: weak accept — The theoretical insight about MTP’s gradient inductive bias is valuable, and LSE-MTP demonstrates a practical way to reduce structural hallucinations, but the method’s reliance on ground-truth trajectories limits its generality and the experimental analysis could be more thorough.

Takeaways

Practitioners working on sequence modeling can steal the idea of dual supervision: don’t just supervise outputs, also supervise intermediate representations. This applies beyond LLMs—any model that generates sequences with underlying state dynamics (video prediction, robotics, time series forecasting) could benefit from latent semantic alignment. The specific technique: if you have access to ground-truth states (from a simulator or oracle), add an auxiliary loss that pulls your model’s hidden states toward those ground-truth trajectories. Weight it carefully—too much and you overfit to the oracle, too little and you get structural hallucinations.

The broader insight: discrete token supervision alone is insufficient for learning coherent world models because it allows illegal shortcuts in latent space. If your task has environmental constraints (spatial, physical, logical), you need to explicitly supervise the model’s internal representations to respect those constraints. This is especially relevant for reasoning tasks where getting the right answer through wrong reasoning is common.