Better World Models Can Lead to Better Post-Training Performance

Prakhar Gupta, Henry Conklin, Sarah-Jane Leslie, Andrew Lee


Abstract
We study how explicit world-modeling objectives affect the internal representations and downstream capability of Transformers, using Rubik’s Cubes as our training domain. We ask: (1) how does explicitly pretraining a world model affect a model’s latent representations, (2) how does world-model quality affect post-training performance, and (3) how should a finite data budget be split between pretraining and task-specific fine-tuning? We compare standard action-prediction fine-tuning with two strategies that add explicit state-prediction supervision: state pretraining followed by fine-tuning, and joint action and state training. We measure task accuracy after Group Relative Policy Optimization (GRPO). We further find that explicit world-modeling yields better representations in terms of higher probing accuracy and steerability of the model, and that better representations yield larger gains from GRPO, especially on harder cube states. Finally, when the fine-tuning budget is held fixed, probe accuracy strongly predicts the GRPO improvement. Under a fixed total data budget, accuracy is maximized by allocating only a small fraction to pretraining.
Anthology ID:
2026.findings-emnlp.1012
Volume:
Findings of the Association for Computational Linguistics: EMNLP 2026
Month:
October
Year:
2026
Address:
Budapest, Hungary
Editors:
Hung-Yi Lee, Sunipa Dev, Jessy Li, Desmond Elliott
Venue:
Findings
SIG:
SIGDAT
Publisher:
Association for Computational Linguistics
Note:
Pages:
21387–21399
Language:
URL:
https://aclanthology.org/2026.findings-emnlp.1012/
DOI:
Bibkey:
Cite (ACL):
Prakhar Gupta, Henry Conklin, Sarah-Jane Leslie, and Andrew Lee. 2026. Better World Models Can Lead to Better Post-Training Performance. In Findings of the Association for Computational Linguistics: EMNLP 2026, pages 21387–21399, Budapest, Hungary. Association for Computational Linguistics.
Cite (Informal):
Better World Models Can Lead to Better Post-Training Performance (Gupta et al., Findings 2026)
Copy Citation:
PDF:
https://aclanthology.org/2026.findings-emnlp.1012.pdf