When Do Neural Networks Learn World Models?

TL;DR

This paper provides the first theoretical analysis showing that neural networks with certain biases can recover latent data-generating variables in multi-task settings, highlighting the importance of architecture and offering new analytical techniques.

Contribution

It introduces a theoretical framework demonstrating conditions under which neural networks can learn world models, using Boolean Fourier analysis and addressing architecture sensitivity.

Findings

Low-degree biased models recover latent variables in multi-task learning.

Recovery is sensitive to neural network architecture.

New techniques for analyzing Boolean transforms are developed.

Abstract

Humans develop world models that capture the underlying generation process of data. Whether neural networks can learn similar world models remains an open problem. In this work, we present the first theoretical results for this problem, showing that in a multi-task setting, models with a low-degree bias provably recover latent data-generating variables under mild assumptions--even if proxy tasks involve complex, non-linear functions of the latents. However, such recovery is sensitive to model architecture. Our analysis leverages Boolean models of task solutions via the Fourier-Walsh transform and introduces new techniques for analyzing invertible Boolean transforms, which may be of independent interest. We illustrate the algorithmic implications of our results and connect them to related research areas, including self-supervised learning, out-of-distribution generalization, and the linear representation hypothesis in large language models.