# Improving World Models using Deep Supervision with Linear Probes

arXiv: https://arxiv.org/abs/2504.03861 · HTML: https://mlsanity.com/paper/2504.03861/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in World models for language agents, one of 11 directions in World models.
- Builds on 18 of its 30 references that are in the map.
- First posted on arXiv 4 Apr 2025.

## Placed in

- World models / [World models for language agents](https://mlsanity.com/wm/world-models-for-language-agents/)
- Mechanistic interpretability / [Circuits and mechanistic analysis](https://mlsanity.com/mechinterp/circuits-and-mechanistic-analysis/)

## Builds on (most cited first, 18 of 18)

- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [Recurrent World Models Facilitate Policy Evolution](https://mlsanity.com/paper/1809.01999/) (2018)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022)
- [DayDreamer: World Models for Physical Robot Learning](https://mlsanity.com/paper/2206.14176/) (2022)
- [Deeply-Supervised Nets](https://mlsanity.com/paper/1409.5185/) (2014)
- [Deeply-Recursive Convolutional Network for Image Super-Resolution](https://mlsanity.com/paper/1511.04491/) (2015)
- [Foundation Models in Robotics: Applications, Challenges, and the Future](https://mlsanity.com/paper/2312.07843/) (2023)
- [Contrastive Learning of Structured World Models](https://mlsanity.com/paper/1911.12247/) (2019)
- [Fast and Accurate Image Super-Resolution with Deep Laplacian Pyramid Networks](https://mlsanity.com/paper/1710.01992/) (2017)
- [Attractor and integrator networks in the brain](https://mlsanity.com/paper/2112.03978/) (2021)
- [Medical Image Segmentation with Limited Supervision: A Review of Deep Network Models](https://mlsanity.com/paper/2103.00429/) (2021)
- [Deep Supervision for Pancreatic Cyst Segmentation in Abdominal CT Scans](https://mlsanity.com/paper/1706.07346/) (2017)
- [Recurrent neural network models for working memory of continuous variables: activity manifolds, connectivity patterns, and dynamic codes](https://mlsanity.com/paper/2111.01275/) (2021)
- [Persistent learning signals and working memory without continuous attractors](https://mlsanity.com/paper/2308.12585/) (2023)
- [Simultaneous embedding of multiple attractor manifolds in a recurrent neural network using constrained gradient optimization](https://mlsanity.com/paper/2310.18708/) (2023)

## Abstract (arXiv)

Developing effective world models is crucial for creating artificial agents that can reason about and navigate complex environments. In this paper, we investigate a deep supervision technique for encouraging the development of a world model in a network trained end-to-end to predict the next observation. While deep supervision has been widely applied for task-specific learning, our focus is on improving the world models. Using an experimental environment based on the Flappy Bird game, where the agent receives only LIDAR measurements as observations, we explore the effect of adding a linear probe component to the network's loss function. This additional term encourages the network to encode a subset of the true underlying world features into its hidden state. Our experiments demonstrate that this supervision technique improves both training and test performance, enhances training stability, and results in more easily decodable world features -- even for those world features which were not included in the training. Furthermore, we observe a reduced distribution drift in networks trained with the linear probe, particularly during high-variability phases of the game (flying between successive pipe encounters). Including the world features loss component roughly corresponded to doubling the model size, suggesting that the linear probe technique is particularly beneficial in compute-limited settings or when aiming to achieve the best performance with smaller models. These findings contribute to our understanding of how to develop more robust and sophisticated world models in artificial agents, paving the way for further advancements in this field.