Toward Consistent World Models with Multi-Token Prediction and Latent Semantic Enhancement

Qimin Zhong; Hao Liao; Haiming Qin; Mingyang Zhou; Rui Mao; Wei Chen; Naipeng Chao

cs.LGApr 7, 2026

Toward Consistent World Models with Multi-Token Prediction and Latent Semantic Enhancement

Authors: Qimin Zhong, Hao Liao, Haiming Qin, Mingyang Zhou, Rui Mao, Wei Chen, Naipeng Chao

Organizations: Shenzhen University · Microsoft Research Asia

Abstract

Whether Large Language Models (LLMs) develop coherent internal world models remains a core debate. While conventional Next-Token Prediction (NTP) focuses on one-step-ahead supervision, Multi-Token Prediction (MTP) has shown promise in learning more structured representations. In this work, we provide a theoretical perspective analyzing the gradient inductive bias of MTP, supported by empirical evidence, showing that MTP promotes the convergence toward internal belief states by inducing representational contractivity via gradient coupling. However, we reveal that standard MTP often suffers from structural hallucinations, where discrete token supervision encourages illegal shortcuts in latent space that violate environmental constraints. To address this, we propose a novel method Latent Semantic Enhancement MTP (LSE-MTP), which anchors predictions to ground-truth hidden state trajectories. Experiments on synthetic graphs and real-world Manhattan Taxi Ride show that LSE-MTP effectively bridges the gap between discrete tokens and continuous state representations, enhancing representation alignment, reducing structural hallucinations, and improving robustness to perturbations.

Figures & tables

Figure 1: Overview of LSE-MTP. Given a backbone hidden state

\mathbf{h}_{n}

, horizon-specific transition layers produce multi-step predictive representations. Training combines multi-step token prediction with latent consistency and semantic anchoring losses. All transition layers are discarded at inference time.

Model	ER (Erdős–Rényi Graph)						USG (Urban Street Graph)
Model	$k=2$		$k=3$		$k=4$		$k=2$		$k=3$		$k=4$
	Sim(F)	Gain	Sim(F)	Gain	Sim(F)	Gain	Sim(F)	Gain	Sim(F)	Gain	Sim(F)	Gain
1TP	0.051	0.027	0.054	0.022	0.078	0.036	0.055	-0.005	0.082	0.018	0.072	0.005
2TP	0.232	0.210	0.102	0.074	0.094	0.062	0.264	0.214	0.126	0.066	0.112	0.048
3TP	0.229	0.195	0.194	0.167	0.136	0.107	0.249	0.197	0.244	0.186	0.148	0.083
4TP	0.223	0.176	0.201	0.162	0.204	0.171	0.230	0.178	0.235	0.180	0.222	0.163

Table 1: Representation alignment on ER and USG graphs. Sim(F) and Gain denote cosine similarity and structure gain (Sim(F) - random baseline) for

k

-step future equivalent states.

Figure 2: Two independent paths (

A\to C\to E

and

B\to D\to E

) converging at a shared future

E

Figure 3: Visualization of learned weights. Under 2TP, unobserved cross-path transitions (

A\to D

B\to C

) are strengthened relative to 1TP.

Model	$K$	ER (Erdős–Rényi Graph)				USG (Urban Street Graph)
Model	$K$	$G=,P=$	$G=,P\neq$	$G\neq,P=$	Baseline	$G=,P=$	$G=,P\neq$	$G\neq,P=$	Baseline
NTP (1TP)	1	0.29	0.11	0.09	0.01	0.22	0.09	0.10	0.03
MTP	2	0.39	0.23	0.11	0.05	0.28	0.10	0.11	0.02
MTP	3	0.43	0.28	0.14	0.07	0.30	0.11	0.10	0.02
MTP	4	0.44	0.30	0.15	0.08	0.32	0.12	0.09	0.03
LSE-MTP	2	0.40	0.25	0.12	0.05	0.34	0.13	0.16	0.05

Table 2: Belief compression on ER and USG graphs. Values report hidden-state similarity for trajectories sharing the same goal

G

and next-step position

P

under different control conditions;

=

denotes “same” and

\neq

denotes “different”.

Model	$K$	ER (Erdős–Rényi Graph)		USG (Urban Street Graph)
Model	$K$	ISP $\downarrow$	Legal Prob $\uparrow$	ISP $\downarrow$	Legal Prob $\uparrow$
NTP (1TP)	1	$2.7\times 10^{-5}$	0.995	$2.2\times 10^{-5}$	0.998
MTP	2	$4.2\times 10^{-5}$	0.994	$4.9\times 10^{-5}$	0.996
MTP	3	$7.8\times 10^{-5}$	0.992	$7.3\times 10^{-5}$	0.994
MTP	4	1.04 $\mathbf{\times 10^{-4}}$	0.985	1.33 $\mathbf{\times 10^{-4}}$	0.989
LSE-MTP	2	$3.0\times 10^{-5}$	0.995	$4.1\times 10^{-5}$	0.997

Table 3: Next-step probability coupling on ER and USG graphs. ISP and Legal Prob report the probability of illegal shortcuts and valid actions for trajectories sharing a common future.

Model	Valid Trajectories	Current State Probe	State-wise Similarity	Compression Precision	Distinction Precision	Distinction Recall	Detour Robustness
Sample Size	1000 trials	1000 seqs	5000 trials	1000 trials	1000 trials	1000 trials	1000 trials
1TP (baseline)	0.993 (0.003)	0.926 (0.001)	0.693 (0.139)	0.108 (0.011)	0.357 (0.015)	0.210 (0.010)	0.692 (0.016)
4TP	0.997 (0.002)	0.964 (0.000)	0.722 (0.119)	0.119 (0.011)	0.298 (0.014)	0.195 (0.010)	0.708 (0.014)
8TP	0.995 (0.002)	0.964 (0.000)	0.820 (0.098)	0.114 (0.011)	0.293 (0.014)	0.182 (0.009)	0.716 (0.014)
LSE-4TP	0.997 (0.002)	0.943 (0.001)	0.791 (0.127)	0.135 (0.012)	0.327 (0.015)	0.213 (0.011)	0.727 (0.014)
LSE-8TP	0.998 (0.001)	0.967 (0.000)	0.851 (0.091)	0.143 (0.012)	0.285 (0.014)	0.201 (0.010)	0.733 (0.014)

Table 4: Evaluation results on real-world Manhattan Taxi Ride Modeling. Values are reported as mean (standard deviation).

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

Model	$K$	Hyperparams $(\lambda_{l},\lambda_{s})$	ER (Erdős–Rényi Graph)			USG (Urban Street Graph)
Model	$K$	Hyperparams $(\lambda_{l},\lambda_{s})$	Suc $\uparrow$	Disc $\downarrow$	WT $\downarrow$	Suc $\uparrow$	Disc $\downarrow$	WT $\downarrow$
NTP (1TP)	1	-	91.80	5.95	2.25	96.22	2.57	1.21
MTP		-	91.99	6.44	1.57	97.13	1.70	1.17
LSE-MTP		(0.1, 0.1)	92.68	5.43	1.89	97.98	1.31	0.71
LSE-MTP	2	(0.3, 0.3)	91.22	6.33	2.45	98.36	1.13	0.51
LSE-MTP		(0.5, 0.5)	91.05	6.66	2.30	97.92	1.33	0.75

Appendix

Table 5: Detailed navigation performance metrics across different horizons (

K

) and LSE hyperparameter configurations (

\lambda_{l},\lambda_{s}

). Suc, Disc, and WT represent Success rate, Disconnection rate, and Wrong Target rate, respectively.

Model	$K$	Hyperparams $(\lambda_{l},\lambda_{s})$	ER (Erdős–Rényi Graph)				USG (Urban Street Graph)
Model	$K$	Hyperparams $(\lambda_{l},\lambda_{s})$	$G=,P=$	$G=,P\neq$	$G\neq,P=$	Baseline	$G=,P=$	$G=,P\neq$	$G\neq,P=$	Baseline
NTP (1TP)	1	-	0.267	0.108	0.091	0.016	0.248	0.112	0.118	0.051
MTP	2	-	0.376	0.221	0.120	0.057	0.298	0.118	0.119	0.039
LSE-MTP		(0.1, 0.1)	0.396	0.263	0.129	0.069	0.346	0.152	0.176	0.078
LSE-MTP		(0.3, 0.3)	0.382	0.264	0.122	0.068	0.366	0.160	0.200	0.093
LSE-MTP		(0.5, 0.5)	0.380	0.270	0.133	0.086	0.369	0.176	0.216	0.107

Appendix

Table 6: Detailed representation similarity metrics across different horizons (

K

) and LSE hyperparameter configurations (

\lambda_{l},\lambda_{s}

G

and

P

indicate Goal and current Position, with

=

and

\neq

representing identical or different conditions. Baseline refers to the

G\neq,P\neq

condition.

Model	$K$	Hyperparams $(\lambda_{l},\lambda_{s})$	ER (Erdős–Rényi Graph)		USG (Urban Street Graph)
Model	$K$	Hyperparams $(\lambda_{l},\lambda_{s})$	ISP $\downarrow$	Legal Prob $\uparrow$	ISP $\downarrow$	Legal Prob $\uparrow$
NTP (1TP)	1	-	$7.9\times 10^{-5}$	0.989	$2.2\times 10^{-5}$	0.999
MTP	2	-	$1.22\times 10^{-4}$	0.988	$5.7\times 10^{-5}$	0.998
LSE-MTP		(0.1, 0.1)	$\mathbf{4.6\times 10^{-5}}$	0.990	$\mathbf{2.5\times 10^{-5}}$	0.998
LSE-MTP		(0.3, 0.3)	$5.2\times 10^{-5}$	0.990	$3.9\times 10^{-5}$	0.998
LSE-MTP		(0.5, 0.5)	$6.2\times 10^{-5}$	0.989	$3.1\times 10^{-5}$	0.998

Appendix

Table 7: Detailed next-step probability coupling and structural hallucinations across different horizons (

K

) and LSE hyperparameter configurations (

\lambda_{l},\lambda_{s}

). ISP and Legal Prob denote illegal shortcut and valid action probabilities, respectively.

Explore similar work

LightMTP: Lightweight Latent Multi-Token Prediction
Oct 5, 2026Tamara Czinczoll, Julie Kallini, Gerard de Melo +1Language Model Pretraining Multi-Token Prediction
Hasso Plattner Institute / University of Potsdam · Stanford University · Tel-Aviv University
NITP: Next Implicit Token Prediction for LLM Pre-training
May 24, 2026Xiangdong Zhang, Debing Zhang, Shaofeng Zhang +3Language Model Pretraining Representation Geometry
School of AI, Shanghai Jiao Tong University · Dots Studio, Xiaohongshu Inc. · University of Science and Technology of China +1
Hierarchical Latent Prediction for Language Models
Aug 6, 2026Chang Shi, Tim Pearce, Manan Tomar +2Language Model Pretraining Multi-Token Prediction
University of Texas at Austin · Microsoft Research

Oct 5, 2026J/K move · Enter open · S save

LightMTP: Lightweight Latent Multi-Token Prediction

Tamara Czinczoll, Julie Kallini, Gerard de Melo, Chen Shani

Hasso Plattner Institute / University of Potsdam · Stanford University · Tel-Aviv University

Language Model Pretraining Multi-Token Prediction

Open paper AskPDF arXiv

Next-token prediction (NTP) is the standard pretraining objective for large language models, yet it provides an explicit training signal only for the immediate next token, which can lead models to exploit local patterns instead of capturing longer-range structure and ideas. Multi-token prediction (MTP) addresses this by training models to predict several future tokens. However, existing MTP methods often introduce a large number of new parameters with limited improvements in downstream performance. Latent MTP approaches address this efficiency issue by encoding future tokens into a vector representation. However, these approaches usually rely on external helper models for future token encoding. We propose LightMTP, a lightweight, i.e., parameter-efficient, latent MTP approach that bootstraps the future token representations from the model's own hidden states. Our two LightMTP variants extend supervision to more future tokens without requiring the additional computational overhead of conventional MTP nor the external supervision latent MTP normally relies on. LightMTP adds at most 1% extra parameters, retains better performance on general language modeling benchmarks, and achieves similar gains in planning, coding, and reasoning.

Abstract

Explore similar work

LightMTP: Lightweight Latent Multi-Token Prediction

NITP: Next Implicit Token Prediction for LLM Pre-training

Hierarchical Latent Prediction for Language Models