Research map / Probabilistic and scientific modeling / Molecules and proteins
When do neural networks learn world models?
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Theoretical analysis showing that, in a multi-task setting with a low-degree (Fourier-Walsh) bias, neural networks can provably recover latent data-generating variables up to simple transforms, sensitive to basis compatibility with architecture.
Problem
When and why neural networks learn world models, i.e., recover latent data-generating variables from observed data, given non-identifiability of nonlinear latent variable models.
Approach
Models Boolean data and latents, uses the Fourier-Walsh degree as a complexity measure and "realization degree" for hierarchical vs flat solutions; shows that multi-task training with low-degree task distributions makes low-degree-biased learners identify latents up to negations and permutations, and that architecture must be basis-compatible.
Key results
- Theory: under multi-task training with low-degree task distribution, the minimizer of realization degree recovers latents up to negations and permutations as n goes to infinity (Theorem on world model learning).
- Polynomial extrapolation: replacing half of ReLUs with identity/quadratic activations improves out-of-distribution extrapolation over ReLU MLPs (shown in violin plots).
- Physical law learning: the modified transformer MLPs achieve lower OOD prediction error than the standard transformer in both parabolic and elastic collision settings, while both reach near-zero training error.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace half of ReLU activations with identity or quadratic activations | MLPs for polynomial extrapolation · Polynomial extrapolation, degree-2 and degree-3 | test MSE | Helped |
| Replace transformer MLPs with modified MLPs and remaining ReLUs with GELUs | Transformer for physical-law prediction · Single-object parabolic motion and two-object elastic collision, OOD test | MSE | Helped |
| Use a quadratic activation in the MLP | MLPs trained on Boolean data · Data-generating variable identification | identification error | Mixed |
Surprising findings
- Single-task learning favors flat realizations; representation learning only emerges in the multi-task setting.
- ReLU MLPs may not exhibit a pure low-degree bias on Boolean inputs; results are very sensitive to the activation function.
Limitations
- Requires non-zero probability of sampling degree-1 tasks, which becomes exponentially small as d grows (conjectured removable).
- Experiments are small-scale proof-of-concept, mainly synthetic.
- Assumes latent variables have no specific structure (e.g., no causal structure) and that solutions adhere to low-degree bias.
Experimental setup
Datasets: synthetic Boolean data from the data generation model in Table 1 (d=m=10) synthetic polynomial extrapolation data synthetic parabolic and elastic collision motion sequences
Benchmarks: polynomial extrapolation (synthetic) physical law prediction (synthetic OOD)
Metrics: test mean square error identification error
Baselines: ReLU MLP transformer
Builds on: non-linear ICA Fourier-Walsh transform length generalization out-of-distribution setup (Hamming ball)
Closest papers by meaning
- Disentangling Representations through Multi-task Learning (ICLR 2025 · similarity 0.87)
- Evaluating the World Model Implicit in a Generative Model (NeurIPS 2024 · similarity 0.86)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.85)
- Transformers Converge to Invariant Algorithmic Cores (NeurIPS 2026 · similarity 0.85)
- Does learning the right latent variables necessarily improve in-context learning? (ICML 2025 · similarity 0.85)
- Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks? (ICML 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.