When Do Neural Networks Learn World Models?

The central question

Does a model understand the world—or merely fit the data?

Prediction is not yet a world model.

Many internal representations can explain the same training data. We therefore define world-model learning as a problem of representation identifiability and derive conditions under which neural networks recover the variables that actually generated their observations.

Formal definitionIdentifiability theoryAlgorithmic implications

Observed data generated from latent variables, with world-model learning shown as identification of those variables

A formal definition

Learning the world means inverting its generation process.

x = ψ (z) ⟹ Φ (x) = T (z)

Latent variables $z$ are rendered into observations $x$ through an invertible generation function $ψ$ . A representation $Φ$ learns a world model when it recovers $z$ up to a simple invertible transformation $T$ .

A world model is an inverse of reality.

The obstacle

The same observations can hide many different worlds.

Without structural assumptions, latent variables are not identifiable: multiple incompatible worlds can generate exactly the same observations. Rather than imposing a narrowly specified data model, we ask whether the learner's own inductive bias can remove this ambiguity.

Two distinct latent variables and generation functions producing the same observations — Different latent explanations, indistinguishable observations.

Main result

Three forces make the hidden world recoverable.

Identifiability emerges from an interaction between how the model learns, how many tasks it sees, and which tasks the world presents.

Multi-task learning

Shared structure becomes cheaper than solving every task with an isolated shortcut.

Simplicity bias

A low-degree preference filters out plausible but unnecessarily complex latent worlds.

Aligned task distribution

Real tasks slightly favor low-degree functions of the true latent variables.

Sufficient condition

Φ^{*} (x) = P D z

The lowest-complexity representation recovers the true latent variables up to permutation and sign flips—simple transformations that preserve the structure of the world.

Single taskShortcuts win

A flat predictor is always the simplest realization.

Many tasksSharing pays off

A common representation amortizes complexity across tasks.

Uniform tasksNo free lunch

Every information-preserving representation is equally good.

Low-degree tasksThe world emerges

The true latents become the simplest shared solution.

From theory to algorithms

Change the basis. Change what the model learns.

Architecture determines which functions feel simple. Guided by the theory, we modify standard MLP and Transformer blocks so the target low-degree structure is easier to represent.

Polynomial extrapolation and physical-law prediction results comparing the proposed method with standard neural networks — Basis-aware architectures improve polynomial extrapolation and prediction of physical dynamics beyond the training distribution.

Parabolic motion · OOD MSE0.07standard Transformer: 0.16

Elastic collision · OOD MSE0.13standard Transformer: 0.29

Theoretical benefitOODprovably lower error under suitable shifts

Controlled recovery

More aligned tasks reveal the latent world.

Proof-of-concept experiments track identification error as the number and degree of tasks change. The empirical phase transition mirrors the theory: enough appropriately distributed tasks drive recovery error toward zero.

Identification error across task number and task degree

Citation

@inproceedings{zhang2025world,
  title     = {When Do Neural Networks Learn World Models?},
  author    = {Zhang, Tianren and Chen, Guanyu and Chen, Feng},
  booktitle = {International Conference on Machine Learning},
  year      = {2025}
}