When do neural networks learn world models?
Authors: Tianren Zhang
Guanyu Chen
Feng Chen
Venue: ICML 2025 poster (2025)
Keywords: World Models
Representation Learning
Learning Theory
Neural Networks
Machine Learning
Primary area: theory->learning_theory
TL;DR: We show that in a multi-task setting, models with a low-degree bias provably learn world models under mild assumptions.

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.

When Do Neural Networks Learn World Models?
Tianren Zhang'
Guanyu Chen Feng Chen
Abstract
stract and interpretable features, indicating a non-trivial representation of data generation. Others, however, report marked decline in LLM performance novel tasks (Wu et al., 2023; Berglund et al., 2024; Mirzadeh et al., 2024), implying a lack of genuine world representations that enable human-level generalization in out-of-distribution settings.
Humans develop world models that capture the underlying generation process of data. Whether neural networks can learn similar world mod- els remains an open problem. In this work, we present the first theoretical results for this prob- lem, 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 recov- ery is sensitive model architecture. Our anal- ysis leverages Boolean models of task solutions via the Fourier-Walh transform and introduces new techniques for analyzing invertible Boolean transforms, which may be of independent inter- est. 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 repre- sentation hypothesis in large language models.
Despite ongoing research, the theoretical foundations of learning world models remain unclear. Notably, even the term "world model" lacks precise definition. This gives rise to several fundamental questions: what does it mean for neural networks to learn world models? When and why can they do so? More fundamentally, what constitutes bona fide world model?
The goal of this work is to address these problems by in- troducing formal framework for world model learning and presenting the first theoretical results in this area. Fol- lowing the spirit of prior work (Ha & Schmidhuber, 2018; Li et al., 2023; Gurnee & Tegmark, 2024), we first show that latent variable models (Everett, 2013) provide a natural scaffold for formulating world model learning. Specifi- cally, learning world models can be framed as achieving a non-triviaa recovery of latent data-gene variables. However, a core challenge in this formulation arises from well-known negative result showing that recovering true latents is generally impossible due to a fundamental issue of non-identifiabiiiy (Hyvärinen & Pajunen, 1999). That is, multiple solutions can fit the observed variables equally well, making true latent variables non- -on-identif from observed data alone.
1. Introduction
Humans develop internal models of the world, extracting core concepts that generate perceptual data (Ha & Schmid- huber, 2018). Can neural networks do the same? With recent advances in large language models (LLMs), this ques- tion has gamered increasing attention (Bender et al., 2021; Mitchell, 2023). Understanding if and how neural networks learn human-like world models is crucial for building AI systems that are robust, fair, and aligned with human val- ues (Hendrycks et al., 2023).
At first glance, the non-identifiability of latent variables may suggest pessimistic outlook on learning world models. However, existing results overlook an important point: solu- tions that equally fit the data are not necessarily equivalent as functions. Thus, algorithms with implicit bias in function space, such as those employed in deep learning (Kalimeris et al., 2019; Goyal & Bengio, 2020), would favor certain solutions over others. In particular, we focus towards low-complexity functions, a phenomenon widely observed in neural networks and believed to be key factor in the success of deep learning (Pérez et al., 2019; Huh et al., 2024; Goldblum et al., 2024). Yet, due to the lack of a well-established complexity measure for continuous functions, formalizing such complexity bias and analyzing its impact remains challenging problem on its own (see
Empirical findings on world model learning have been mixed. Some studies suggest that medium-sized neural networks (Mikolov et al., 2013) and LLMs (Li et al., 2023; Bricken et al., 2023; Gurnee & Tegmark, 2024) learn ab-
bias
'Dpartment of Automation, Tsinghua University, Bei- jing, China. Correspondence to: Feng Chen <chen- feng @ wai.l.siinguua.cuu.ooo
Proceedings of the 42nd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

When Do Neural Networks Learn World Models?
Section A for related work).
Benefits
In this work, we circumvent this challenge by leveraging simple yet important fact: while real-world data and la- tent variables may be continuous, all variables processed by neural networks are ultimately encoded as bit strings due to finite precision of computers. This allows us to model all variables as Boolean without loss of generality. While this may seem a subtle difference (since we only lose the information that goes beyond machine precision), as we demonstrate in later sections, it turns out to provide surprisingly powerful machinery for defining and analyzing transform of Boolean functions (O' DODonnell, 2014). Building on this foundation, we present, for the first time, a nuanced per- spective on learning world models that reveals an interplay between the cowcooplyy bias, proxy tasks, and model architecture. Our main contributions are:
Sec. 4.3
Latent variable recovery Definition
Basis compatibility
World models
Architectures
Sec.
Sec. 4.4-5
Sec. 4.1-4.2
Multi-tasking + low-degree bias
Identifiability
Figure I. A graphical summary of our framework and main results.
the complexity of solutions via the Fourier-Wail
laws (Kang al., 2025). We show that architectures inspired by our analysis outperform conventional archi- tectures such as ReLU MLPs and transformers (Vaswani et al., 2017) in these tasks.
1. In Section 2, we lay down general definitions of learning world models and discuss its core challenge posed by the non-identifiability of latent data-generating variables. This provides a foundation for future work to formally reason about learning world models and offers theoretical rigor to the recent scientific debate on this topic (Bender et al., 2021; Mitchell, 2023).
2. Formulation of Learning World Models
How to define the world model and the problem of learning world models remains debatable to date. Yet, the term "world models" has been widely referred to in the literature as models that uncover the underlying generation process of data and maintain a representation of it (Ha & Schmidhuber, 2018; Gurnee & Tegmark, 2024; Richens & Everitt, 2024). For example, pioneering works by Li ct al. (2023) and Nanda et al. (2023) define "world models" in board games as the board state that generates move sequences. This motivates formulation of world model learning under the framework of latent variable models (Everett, 2013). To this end, we first define a general data generation process.
In Section 3, we introduce our Boolean function formu- lation and the corresponding complexity measures based on notion of realization degree, offering a tractable approach for analyzing the impact of low-complexity bias on world model learning.
3. In Section 4, we present the first theoretial results on learning world models in the context training on proxy tasks using observed data. We identify two critical fac- tors for world model learning: (i) multi-task learning setting; (ii) the low-complexity bias, instantiated by low-degree bias of the model and a low-degree lask distri- bution. Together, these factors ensure the identifiability of latent data- -ata-gennr variables. Moreover, we show the provable benefits of learning world models in an out- of-distribttion generalization setting (Abbe et al., 2023) and study the impact of model architecture under a no- tion of basis compatibility (see Figure for graphical summary of our results). Technically, our analysis relies on analyzing the degree properties of Boolean functions composed with invertible transforms, which may be of independent interest.
Definition 2.1 (Data generation process). Let x E X be the observed data variables and let z E Z be the latent variables for some data space X' and latent space 2. The observed data are sampled as follows: (i) sample z ~ p(z) for some probability distribution p over Z; (ii) generate x through an invertible and non-linear function x 1b(z).
This definition resembles the data generation process used in many latent variable models such as non-linear ICA (Hyvärinen & Pajunen, 1999), invariant feature learn ing (Arjovsky et al., 2019), and causal representation learn- ing (Schölkopf et al., 2021). The main difference is that unlike these models, here we do not assume p(z) to be any structured distribution.
Given Definition 2.1, natural way to formalize "under- standing" the generation process of data is to invert it, i.e., recover the latent variables z from the observed data x. This leads to our basic formulation of learning world models.
4.
In Section 5, we illustrate the algorithmic implications of our results two representative tasks: polynomial extrapolation (Xu et al., 2021) and learning physical
Definition 2.2 (Learning world models). Let T be a set of transforms T: Z - Z. We say a representation & X + 2 learns the world model up to T if there exists T T such
Note that this differs from neural network quantization (Nagel al., 2021; Gholami et al., 2021).
2

---

When Do Neural Networks Learn World Models?
that ((x) T2)) for every = E supp(p) and æ (=)..
of Definition 2.2 is not possible with simple T by looking at the observed data alone.
For instance, if only contains the identity transform. then (x) recovers z exactly. In general, we require ' to contain only simple transform classes (e.g., linear transforms) for meaningful recovery.
Given this result, one may naturally be pessimistic about the outlook on learning world models. However, recent studies on LLMs seem to suggest otherwise: instead of learning arbitrary non-linear representations of the observed data as Lemma 2.4 implies, LLMs turn out to often learn semanti- cally meaningful, numan-interpretable representations (Li et al., 2021; Bricken et al., 2023; Marks & Tegmark, 2023; Gurnee & Tegmark, 2024). This suggests that despite having vast number of parameters and being capable of learning different task realizations (Reizinger et al., 2024), LLMs can still learn representations that are reasonably aligned with humans.
The main difficulty of this latent variable recovery problem is that z is, by definition, unobservable. To address this, we leverage the fact that models in practice are trained on some proxy tasks (e.g., next-token prediction) learn representations of implicitly. We formulate this follows.
Definition 2.3 (Task and realization). A task is defined function h: X - R. For task h, let
H(h) = ( X TR|
f(x) = h((r),
(1)
To address this puzzle, in this work we propose to incor- porate the implicit bias of neural networks into analysis instead of using task performance as the only identifiabil- ity criterion. Specifically, we explore if the bias towards low-complexity realizations, a trait of both human reasoning and deep learning models, could be used to steer the learned realization towards non-trivial recovery of true latents. Yet, the main challenge is that for continuous functions with inputs in Rd, we lack a well-established complexity mea- sure that is amenable to analysis. For example, Kolmogorov complexity (Li et al., 2008) in algorithmic learning theory offers a unified framework for defining the complexity of any object, yet being uncomputable. Fortunately, in the following sections we will show that this problem could be cireumvented by using Boolean functions.
Vz2 E supp(p), x= = [=)]]
If a function composition f(1) o.... o f(a) is in H(h), then we say it is a realization of h.
Flat and hierarchical realizations. With definitions above, our main hope is that by training on some proper tasks h, the model can learn hierarchical realization g & HH(h) with a function g Z + E and a representation & that learns the world model in the sense of Definition 2.2. How- ever, a key challenge arises: by definition, all realizations have zero training error on supp(p) and are thus indistin- guishable by their task performance. For example, every function h" E H((h) is itself flat realization of h, i.c., with- out any explicit representation learning. Thus, by looking at the observable data alone, we have no reason to expect that g o & will be favored over ha" Likewise, it is also unrea- sonable to expect that g o & should be favored over another hierarchical realization g' o op E H(hh) with '' not learning world models. Indeed, a well-known impossibility result in latent variable modeling shows that the true latent variables are non-identifiable when va is sufficiently flexible non- linear function of the latents (Hyvärinen & Pajunen, 1999; Khemakhem et al., 2020).
3. Complexity Measures
Computers encode every object, including the observed data and variables learned by neural networks, in bit strings. Motivated by this fact, we could assume without loss of gen- erality that both x and in Definition 2.1 are Boolean (after some proper encoding). Formally, in the rest of the paper we will let X' S 11, and 2 = (-1, Since & is invertible, we have m 2 d. Unless otherwise mentioned, we will also assume supp(p) = Z to ensure that all elements in 2 can be sampled with non-zero probabilities.
Lemma 2.4 (Non-identifiabiiityy (Khemakhem et al., 2020)). Let Z = Rd and z E Z be a random vector of any distribu-
tion. Then, there exists an invertüble transform T : 2 - Z such that the components of z' = T(z) are independent, standard Gaussian variables.
A direct consequence of the Boolean modeling of variables is that all functions involved are Boolean functions (see See- tion B for a brief introduction). Notably, this itself does not resolve the non-identifiabilit issue stated in Section 2, as functions with Boolean inputs and outputs can still exhibit arbitrary nonlinearity. Instead, the primary advantage of this treatment is the useful machinery it provides for defin- ing functional complexity, offered by the Fourier- transform of Boolean functions:
Lemma 2.4 indicates that we can construct new random variables ((x) z' that have the same distribution p as true latent variables & (x) = z (thus fitting the observed variables equally well) by first transforming to standard Gaussian variables, applying any orthogonal transform, and then inverting the transform (note that standard Gaussian distributions are invariant to orthogonal transforms). Ap- plying this result to our context, we conclude that without additional assumptions, learning world models in the sense
(Fourier- Walsh transform (O'Donnell, Definition 3.1
2014)). Every function f : {+1)" + can be uniquely
3

---

When Do Neural Networks Learn World Models?
expressed as a multilinear polynomial
For example, the realization degree of a flat realization h' E H(k) coincides with its degree; the realization degree of hierarchical realization h = g o + is deg(g o Q) = deg(g) + deg($). Compared to degree, realization degree takes into account the cost of implementing each function in hierarchical realizations, which better captures represen- tation learning in practice. Throughout this work, we say a model exhibits a low-degree bias if it minimizes the real- ization degree.
f(z) = E sc f(5)xs(z),
(2)
where (.n).. (=) Iles oi are parity functions, and f(S) E R are the coefficients.
The Fourier-Waat transform shows that every Boolean func- tion can be represented functions Xs that capture all non-linear relationship between inputs and outputs. It can also be shown that parity functions are a basis of the vector space Fm := [f: : {=1}" - R of n--imennsioal Boolean functions (see Section B for more details). By contrast, the space of arbitrary continuous func- tions does not have a similar, theory-friende Section C for more discussion on the use of Boolean func- tions and complexity measures.
linear combination of parity
Finally, we introduce two definitions that are useful in char- acterizing the impact of the low-degree bias on flat and hierarchical realizations.
Definition 3.5 (Min-degree solutions). For task h, we define its min- solutions Hmin (h) as the set of func- tions in (hh and that minimize the degree. We denote their degree by deg (Hmin (h)).
basis. See
Given the Fourier-Walse transform of a function, natural measure of its complexity is its degree:
Definition 3.6 (Conditional degree), For a task h and a representation @ : X + Z, let h² E Hmann (h) and assume that the set {g : 2 + R g o E H(h) is not empty. The conditional degree of the task h on the representation & is then defined as
Definition 3.2 (Degree). For every function f f: {=1]" R. its degree is defined by
deg(f) [ax f(S) 03.
deg((
+) = deg(h") max(deg(s) :oo & E H(h))-
For a Boolean function with multiple output dimensions, we define its degree by the sum of degrees of the Boolean functions mapping the input to each output coordinate.
By Definition 3.5, flat realizations with low-degree bias sat- isfy h" E Hmmia (h). Analyzing the impact of representations is more involved; the main intuition behind Definition 3.6 is that representation & only makes task h "simpler" if the conditional degree deg(h )) > O, i.e. solving the task on top of 4 yields a smaller degree compared to low-degree flat realizations. This notion will be explored further in the next section.
Remark 3.3. Note that the degree of f equals to the max- imum degree of the basis functions it uses: deg(f) = max(deg(Xe f(S) * 0}. This interpretation will be useful when we consider different basis functions as in Sec- tion 4.4.
Intuitively, the degree of a Boolean function measures how non-linear it is. It can also be viewed as an approximation of Kolmogorov complexity if we treat parity functions Xs as "function codes" with length [S]. Prior work has shown that many classes of neural networks indeed have bias towards low-degree solutions for Boolean inputs (Abbe et al.. 2023; Bhattamishra et al., 2023); in comparison, here we study whether a relevant low-complexity bias could implicitly lead to world model learning.
4. Theoretical Analysis
This section presents our main theoretical results. We first study how the low-degree bias introduced in Section 3 drives representation learning (Section 4).. We then present our main results on learning world models, showing sufficient conditions for identifying latent generatinn variables (Section 4.2). Next, we show provable benefits of learn- ing world models in an -uito--dist generalization setting (Section 4.3). We conclude this section by prelimi- nary analysis on how the model architecture impacts world model learning (Section 4.4). All proofs are deferred to Sec- tion E. Finally, Section D presents numerical experiments that substantiate our theoretical results.
Complexity measures for realizations. Based on the no- tion of degree, we next introduce complexity measures to quantify the complexity of realizations. First, note that while degree can characterize function complexity, it fails to distinguish between different realizations, as they all be- have identically when considered as whole function. To overcome this limitation, we introduce realization degree.
Definition 3.4 (Realization degree). Let h be task. For o f(g) of h, its realization degree is realization f() ....o defined as
4.1. Low-Degree Bias Drives Representation Learning
As a warm-up, we first consider a basic question: why should neural networks, such as LLMs, lear any represen- tation of teteraten variables when trained on proxy
deg o... f()) Σ deg(S())
(4)
sE/gl
4

---

When Do Neural Networks Learn World Models?
tasks? Indeed, modern neural networks can often memorize the entire dataset (Zhang al., 2017) or rely on superficial statistical patterns to solve tasks (Geirhos et al., 2019). This raises concerns about whether they truly understand the data despite generating plausible outputs (Bender et al., 2021).
2024). Indeed, modern pre-training objectives such as next-token prediction and contrastive learning can be interpreted as solving a large number of prediction tasks simultaneously (Radford et al., 2019; Arora et al., 2019; Brown et al., 2020).
Formally, this question can be modeled as a competition between flat realizations h" E Hmin (h) and hierarchical realizations g o & E Hh) in our formulation. Specifically, we seek to determine which realization minimizes the real- ization degree and is thus favored by the low-degree bias. Our first result shows that for any single task h, the flat realization is preferred.
Theorem 4.3 suggests that to facilitate the learning of representation O, proxy tasks should be chosen such that conditioning on & makes them less "complex". This provides a framework to reason about whether certain objectives in self-supervised learning (Liu et al., 2021) induce better representations than others. For instance, input reconstruction is often suboptimal as it permits a low-degree solution h (a:) = x without requiring any representation learning; masked image modeling (He et al., 2022) is likely more effective, as representation that captures image semantics could significantly reduce solution complexity by filtering out pixel-to to- pixel details.
Theorem 4.1 (Single-task learning). Let h be a task. Then, for every lt" E Hmii (l)), representation & X - Z, and g Z- + such that g o 4 E H((h), the following holds:
deeg(h") S deg(g o $).
(6)
Remark 4.2. Intuitively, this result is a consequence of learn- ing "redundant" representations: for any h, suffices to learn every parity function Xs with ((S) * 0 in its Fourier- Walsh transform. Even if there is a good universal data representation, explicitly learning it is often not the best choice since it may involve irrelevant parity functions with h(S)= = 0, resulting in a larger realization degree.
4.2. Conditions for Learning World Models
We now move on to investigate whether the low-degree bias facilitates world model learning. While Theorem 4.3 shows that training on properly defined proxy tasks drives representation learning in the presence of the low-degree bias, it does not specify which representation is ultimately learned. In this section, we further explore the multi-task setting to address the key question: can we construct proxy tasks hi...., hn to induce a representation t that learns the world models in the sense of Definition 2.2?
However, the situation changes in the multi-task setting. Suppose there are n distinct tasks hi...., hn. A flat real- ization solves each task independently by learning solutions hī E Hmin (h:) for each é E [m]. whereas a hierarchical realization can leverage the shared representation & across tasks, requiring only task-specific functions 9i for i E [m]. Our next theorem shows that in this setting, the hierarchical realization is favored if a sufficient number of tasks have positive conditional degree.
We begin by characterizing the space of all possible tasks. Given that the observed data is generated by x (z) with an invertible vb, every task h : X' + R can be equivalently defined via a latent-space function h' Z + R as h = h' a - The space of all possible tasks can thus be represented
as
Theorem 4.3 (Multi-task learning). Let hi,..., hn be n distinct tasks and let h; E Hmin (h+), Vi E [n]. Let & X + Z and g....., In satisfy that gi Q is an realization of fu, Vi E [n]. Then, the following holds for every h² =
Fd + --= = {h' o - ["F
(8)
where Frd =- (h: ( - R} denotes the set of Boolean functions on 2 = Our next theorem shows that if proxy tasks are constructed by uniformly sampling from Fd 6b_ then as n + co, all viable representations yield the same task-averaged realization complexity.
(hi..... .11)), g = (gi,..., .. and xh E (D):
deg(h*) dog(g o )*)² Z Σ vE[x]
dogg(h/ (5*- d2 (7)
Theorem 4.4 (Representational no free lunch). Let hy...., hn be tasks that are independently and uniformly sampled from Fd o o/ Then as n - oc, for any two representations $ a satisfying that there exists a bijective
Thus, if sufficient number of tasks satisfy deg(h, )) 0
and Eie(Ny dog(h; | ()) > d², then dog(h") > dog(g 4).
contrary to Theorem 4. We discuss two implications of these results:
transform T 2 - Z such that ((x) = T(z) ( for every z E supp(p) and x +f22), the following holds:
lim 1 (deg(g b) - -de(g' o +0') -0, 0, n=000
The contrast between the single-task and multi-task set- tings justifies the importance of multi- molti-taskkng in learning general-purpose purpose representations, which has been conjec- tured by prior work (Radford et al., 2019; Huh et al.,
(9)
where 9. E (Fddy satisfy that g, o 4 and gf o E' are both realizations of for every E [m).
5

---

When Do Neural Networks Learn World Models?
Remark 4.5. Note that the condition of the existence of
(i) uniformly sample a k-degree lask. Let (8" ,9" be the minimizer of the following optimization problem:
bijective transform T is a minimal requirement for $(x) containing enough information for solving all tasks. The
1 deg(g 4)
fact that all such representations have the same realization complexity suggests that the representation induced by uniformly sampling from F o & only learns the world model up to arbitrary bijective transforms.
min +
n s.t. g 1 E H((h)), ViE [n].
(10)
Then as n - x0, 46 learns the world model up to negations and permutations. i.e. there exists a permutation ".....,.e of 1, ..... d such that (2(x)) E {=21,) } for every j E [d]. with z + (x))
Theorem 4.4 can be viewed as a "no free lunch"-likk theo- rem for representation learning. The original no free lunch theorem (Wolpert, 1996) states that every learner's perfor- mance equally good when averaged over uniform distri- bution on learning problems; here we show that every viable representation is equally complex when averaged over a uni- form distribution on tasks. As we will show in Section E.4, the technical intuition of this result is that every represen- tation renders some tasks in the task space "simple" and others "complex", with the overall task-averaged complexity independent of the particular choice of the representation.
We make several remarks on this result:
The reason why k-degree tasks facilitate world model learning is that, in effect, they induce task distribution in which lower-degree tasks on true latents are drawn with larger probabilities than in the uniform setting. This overcomes the result in Theorem 4.4 by breaking the degree balance between different representations when averaged over the task distribution: representations ca- pable of solving these lower-degree tasks in "cheaper" ways would now be favored.
To overcome this result, we then move to the non-umiform case where proxy tasks are still drawn from Fr but with different weights assigned to different functions. This set- ting is of more practical interest: prior work has reported various evidence suggesting that realistic tasks are often much more structured than being purely random (Whitley & Watson, 2005: Zhang et al., 2017). In particular, real- world data tend to be highly compressible, implying that low-complexity input-output maps occur more frequently than high cogplexty ones (Dingle et al., 2018; Zhou et al., 2019; Goldblum et al., 2024). To formalize this, we define degrre tasks.
A limitation of Theorem 4.8 is that we requires non- zero probability of explicitly sampling degree-1 tasks, in which latent variables are also task outputs. However, we emphasize that this probability is exponentially small as d becomes large, and we conjecture that it can be completely removed in many settings. See Section E.6 for more discussion.
Technically, in our proof we introduce a degree analysis of k-degree Boolean functions composed with invertible transforms (Lemma E. 10), which may be of independent interest in analyzing similar problems.
Definition 4.6 (k-degree tasks). Let = (h: [=1]d
R deg(h) h)) < kj. We say task h is k-degree task if o W' E F.).
h E F 2
In other words, k-degree tasks can be solved by a function with degree not greater than is on top of true latents. One can easily verify that FV S F i for every & E [d (4-1) 1] and Fd= -d. Thus, uniform sampling from all possible tasks amounts to uniform sampling from Fy Fz k-degree tasks are also related to tasks with positive conditional degree, as shown by the following corollary.
Connection to the linear representation hypothesis. A number of recent mechanistic interpretation studies show that L.LMs often represent abstract, interpretable features as directions in their intermediate representation space (Nanda et al., 2023; Marks & Tegmark, 2023; Gurnee & Tegmark, 2024). Theorem 4.8 could be viewed as a provable, Boolean version of the emergence of such linear representations: per- mutations and negations are precisely all degree-1 Boolean functions, natural counterpart of degree- 1 real polynomials (i.e., linear functions) in the real domain. The main signif- icance of this result is that even when proxy tasks involve complex, non-linear functions over true latent variables z, we can still recover x up to very simple transforms despite the presence of such nonlinearity.
Corollary 4.7. For every task h satisfying deg(h {)!)> O, we have h E Fá_ -1
To capture the lowpccmyyyxity bias in task sampling, we assign more weights to Fe o 3 " < d). Perhaps sur- prisingly, our next theorem shows that even a slight such preference low-complexity tasks in the task distribution can induce world model learning up to simple transforms.
Theorem 4.8 (World model learning). Let PI..... .d E (0, 1) such that Eieja Pt- = 1. Let hi....,h, be n tasks that are independently sampled as follows: (i) sample a degree k E [d] according to probabilities Pr(k- = il = Pti
4.3. Benefits of Learning World Models
Up to now, we have presented sufficient conditions for learn- ing world models. As complement these results, this
6

---

When Do Neural Networks Learn World Models?
section demonstrates provable benefits of learning world models in the context of out-of-distribution generaliza- tion setting introduced by Abbe et al. (2023).
Motivated by this, in the following we analyze how different choices of basis in the Boolean function space can impact world model learning. Informally, one may also view neural networks as implementing functional bases through layer- by-layer function composition, and different neural network architectures may induce different bases in the function space (Teney et al., 2024).
Theorem 4.9 (Benefits of learning world models). Lel the latent variables during training be uniformly sampled from the Hamming ball B, : {z E {=1]² | #=) S r with r d, and let the latent variables in test be uniformly sampled from 2. Let h: {+1}m } be downstream task such that h o 1/ is parity function with (hg(h) q >
Notably, for a given input dimension n, the transform uses parity functions Xs basis of the 2 dimensional vector space J² = [ (+1} - R} (cf. Definition 3.1). To obtain a different basis, we define a basis transform U, i.e., an invertible linear transform on F such that (Uxxs) S S [n]] is a basis of F The degree of a function f [+1) - R under the new basis is then given by
k == [log2 Ei= (1) Then, if deg(h -) 2 q and deg(h o 6) > k- - q + r, the following holds: (i) the test mean square error (MSE) of any h- E Hauin(M) (h) is larger than 1; (ii) let &² be a representation that learns the world model up to negations and permutations as in Theorem 4.8 and let g" be function such that g" o T0" HH),, then the test MSE of g" Φ0 is O.
degu (f) :=))==axx)))) max (ax((((( f(S) + O), (11)
Remark 4. As also noted by Abbe et al. (2023), a prac- tical scenario reflected by sampling from B, is length gen- eralization (Anil et al., 2022; Press et al., 2022). Here we show that in this setting, hierarchical realization with the world model is provably more generalizable than any flat realization, despite that both of them achieve zero i.i.d. test error.
where U-! reflects the cost of using the new basis to repre- sent the original one. As result, the low-degree bias under dego(() (-) may deviate from that under deg(-). To capture the effect of this, we introduce the notion of basis compatibility.
Definition 4.11 (Compatibilitty. We say a basis transform U is compatible if (eg((((ss) deg(xs), VS C [m]. Vn.
It has been widely believed that learning world models leads to better generalization (Li et al., 2023; Richens & Everitt, 2024; Yildirim& & Paul, 2024). In comparison, Theorem 4.9 indicates that such benefits typically manifest when the conditional degree deg(h | 1/-1) of the downstream task h is large enough. Technically, this is because for tasks with small deg(h 2 ) solutions using world models still involve high-degree parity functions, whose learning is hampered by the restricted sampling from B,. As a practical example, semantical representations of images can make it casier to answer questions about high-level concepts (tasks with large deg(h ()), yet may make it harder to predict the intensity of a certain pixel (tasks with small or negative deg(h (h: | w-1)) Together with Theorems 4.4 and 4.8, this result suggests that low-degree task distribution is essential for both learning world models and exploiting its advantage.
In other words, a basis transform is compatible if it preserves the degrees of all basis functions. Our next result shows the impact of basis compatibility on world model learning.
Theorem 4.12 (Impact of basis compatibility). Consider the same setting as in Theorem 4.8 with n + cc. Let U be a basis transform and let the degrees of tb and g be measured under the new basis {.xs)) Then, (i) U is compatible, tb learns the world model up to negations and permutations: (ii) there exists incompatible U such that x- = T w where T is an invertible transform on Z satisfying max4e[d] deg(T, 2)
2 k.
Theorem 4. 12 implies that to facilitate world model learning, the model architecture should induce a basis that preserves the degree of the "natural" basis under which low-degree proxy tasks are drawn. Thus, we can also interpret basis compatibility as the compatibility between the model and the tasks. This explains why neural networks with differ- ent activation functions can exhibit different complexity biases (Abbe et al., 2023; Teney et al., 2024). Basis com- patibility is also related to algorithmic alignment (Xu et al., 2020), which suggests that tasks with algorithmic structures aligned with the computational structures of neural networks could be learned more sample-efficiently. In comparison, we focus on the identifiability of world models instead of sample efficiency, and we provide a unified framework for analyzing similar concepts via the lens of the degree bias.
4.4. Impact of Model Architecture
In the above analysis, we study the role of the fow-degree bias with the assumption that the task solutions perfectly adhere to it in the function space. Yet, practical training of neural networks often more involved than this abstraction. Although neural networks are known as universal function approximators (Hornik et al., 1989; Funahashi, 1989), prior work has shown that models with different architectures may represent the same function differently (Raghu et al., 2021). In particular, embedded nonlinearities such as activation functions can steer how functions are represented by neural networks (Xu et al., 2020; Ziyin et al., 2020).
7

---

When Do Neural Networks Learn World Models?
0.8
3.0
Ground Truth
Transformer
ReLU MLP
Test MSE
0.6
Test MSF
2.4
Ours
0.29
Ours
1.8
0.4
1.2
0.2
0.6
0.07
0.0 ReLU MLP
0.0
Ours
ReLU MLP
Ours
Parabolic
Collision
(b)
(c)
Figure Empirical results. (a) An example of extrapolating a degree-3 polynomial. Shaded region indicates training region. (b)
Violin plots of the test mean square enoor (MSE) of ReLU MLPs and our models in extrapolating degree-2 (left) and degree-3 (right) polynomials. (c) Results for learning physical laws. Each column indicates the task and out-of-distributton test MSE averaged over 5 runs.
5. Algorithmic Implications
basis. In practice, we replace half of ReLU functions in each MLP layer by one the two functions including the identity function (xx) = x and the quadratic function a (xx x²
In this section, we illustrate the algorithmic implications of basis compatibility through two representative tasks: poly- nomial extrapolation (Xu et al., 2021) and learning physical laws (Kang et al., 2025; Motamed et al., 2025). The first task provides an illustrative example of our results in Section 4.4, while the second is of greater practical interest for learning real-world models and is related to recent efforts in devel- oping video prediction models world simulators (Brooks et al., 2024). While our experiments are currently limited in scale and mainly serve as proof-of-concept demonstrations, we view extending the ideas discussed in this section to broader contexts as a promising future direction.
Representing (ax) = x* would be much easiler using these functions, and thus expect that a neural network with the low-complexity bias would then use them instead of the remaining ReLU despite the same expressive ability they have, resulting in ʃ ~ f. As shown in Figure 2b, this simple method indeed leads to significant improvement in extrapolation. Please see Section G.1 for more results and discussion.
One may wonder that in this task, still need to know the basis of the task a priori to achieve basis compatibility. However, we emphasize that even without prior knowledge, we can still use different activation functions simultaneously and rely on the low-complexity bias of neural networks to self-adaptively select functions that are the most compatible with the task, as empirically shown in our experiments. We thus conjecture that this approach could be quite universal. Indeed, in the next section we will show that exactly the same approach can also bring benefits in a distinct scenario.
5.1. Polynomial Extrapolation
We begin by a synthetic task in which we train multi- layer perceptrons (MLPs) to fit real polynomials (x) = Σl= a,a" (see Section F.1 for details). The space of poly- nomials has standard basis (1, 21, x2, ,...), ,....,
making it natural counterpart of the Boolean function space with the parity basis. Recovering all basis functions used by Pa could generate any polynomial by linearly combining these basis functions, which enables extrapolation.
5.2. Learning Physical Laws
We now show that the approach in Section 5.1 also benefits sequence prediction task aiming at learning physical laws, Correctly abstracting fundamental physical laws is essential for any model to be a real "world model" for the physi- cal world (Motamed et al., 2025). Notably, a model that "understands" the physical laws is expected to generalize these laws waseen distributions rather than fit them only in the training domain. Following Kang et al. (2025), we generate two types of object motion sequences reflecting basic physical laws: (i) single-object parabolic motion, and (ii) two-object elastic collision motion (see Section F.2 for more details). We train a transformer (Vaswani et al., 2017) to predict the motion of objects conditional on the first few frames in the motion sequence. The model then evaluated in out- distribbttinn generalization setup with objects
As shown by prior work (Xu et al., 2021) (see also Fig- ure 2a), common ReLU MLPs cannot extrapolate degree-k polynomials with k 1 beyond the training region. Our framework explains why this happens via the incompatibil- ity between ReLU and the polynomial basis: ReLU MLPs can learn a basis function f that approximates f(x) = * arbitrarily well in any finite training region, but the sim- plest f that fits the data is not f itself, which is expensive to represent using the composition of ReLU. As result, the actually learned f differs from and hence does not extrapolate well.
Motivated by this explanation, here we provide simple fix: replacing a portion of the ReLU activation functions with functions that are more compatible with the polynomial
8

---

When Do Neural Networks Learn World Models?
having different initial velocities and sizes. To apply our method, we simply replace every MLP in the transformer with our modified MLP in Section 5.1 (we also replace the remaining ReLUs with GELUs).
Ahuja, K., Mahajan, D., Wang, Y., and Bengio, Y. Inter- ventional causal representation learning. In International Conference on Machine Learning, 2023.
Allen-Zhu, Z. and Li, Towards understanding ensem- ble, knowledge distillation and deff-disiillation in deep learning. In International Conference on Learning Rep- resentations, 2023.
As shown in Figure 2c, our model achieves lower prediction error than transformer in both settings. Note that both mod- els achieve near zero training error. Thus, the fact that our model generalizes better out-of-distribution is not because it fits the data better, but due to it better capturing the un- derlying laws in object movements, See Section G.2 for the predicted motions from both models.
Andriushchenko, M., Varre, A. V., Pillaud- Vivien, L.., and Flammarion, N. SGD with large step sizes learns sparse features. In International Conference on Machine Learn- ing. pp. 903-925, 2023.
6. Limitations and Future Work
Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. Exploring length generalization in large language models. In Advances in Neural Information Processing Systems, volume 35, pp. 38546-38556 2022.
This work is an initial step towards formally understand- ing world model learning and the internal representations learned by neural networks, and we can see many exciting future directions. (i) Neural networks in practice learn hier- archical representations (Bengio et al., 2013), which may require a more structured modeling of the data generation model and the learned representations. (ii) We do not as- sume the latent variables to have any specific structure (e.g.. causal relations between latent variables); extending our analysis to causal representation learning by combining our results and existing results in causal inference (Schölkopf et al., 2021; Richens & Everitt, 2024) is an interesting direc- tion. (iii) We primarily focus on the low-complexity bias of neural networks; considering other implicit biases (Allen- Zhu & Li, 2023; Zhang et al., 2024) of neural networks and more fine-grained complexity measures are also promising directions for future work.
Arjovsky, M., Bottou, L.,, Gulrajani, I. and Lopez- Paz, D. Invariant risk minimization. arXiv preprint arXiv: wwXXvv1907.02880 2019.
Arora, S., Khandeparkar, H.,, Khodak, M., Plevrakis, O., and Saunshi, N. A theoretical analysis of contrastive unsupervised representation learning. In International Conference on Machine Learning, pp.
2019.
Bartlett, P. L. and Mendelson, S. Rademacher and Gaussian complexities: Risk bounds and structural results. Journal of Machine Learning Research, 3:463 482, 2002.
Bartlett, P. L., Foster, D. J., and Telgarsky, M. Spectrally- normalized margin bounds for neural networks. In Ad- vances in Neural Information Processing Systems, 2017.
Acknowledgements
This work was supported in part by the National Key Re- search and Development Program of China under STI 2030- Major Projects 202120000, and in part by the Na- tional Key Research and Development Program of China No. 2024YDLN0006, and in part by the National Natural Science Foundation of China under Grant 62176133.
Bartlett, P. L.., Long, P. M., Lugosi, G., and Tsigler, A. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 1 117744):30063-30007) 2020.
Bender, E. M., Gebru, T., AaMillan-Majon A., and Shmitchell, S. On the dangers of stochastic parrots: Can language models be too big? In Proceedings of the 2021 ACM Conference on Fairness, Accountability, and Trans- parency, pp. 610-623,11 2021.
Impact Statement
This paper presents work whose goal to advance our un- derstanding of the inner workings of deep neural networks. There are many potential socictal consequences of our work, none which we feel must be specifically highlighted here.
Bengio, Y., Courville, A., and Vincent, P. Representation learning: A review and new perspectives. IEEE Transac- tions on Pattern Analysis and Machine Intelligence, 35 (8)81798-1828, 2013.
References
Berglund, I.., Tong, M., Kaufmann, M., Balesni, M., Stick- land, A. C., Korbak, T., and Evans, O. The reversal curse: LLMs trained on "a is b" fail to learn "b is a". In International Conference on Learning Representations, 2024.
Abbe, E., Bengio, S., Lotfi, A., and Rizk, K. Generalization on the unseen, logic reasoning and degree curriculum. In International Conference on Machine Learning, 2023.

---

When Do Neural Networks Learn World Models?
Bhattamishra, S., Patel, A., Kanade, V., and Blunsom, P. Simplicity bias in transformers and their abil- ity to learn sparse boolean functions. arXiv preprint arXiv:2211.12311 2023.
Gholami, A., Kim, S., Dong, Z., Yao, Z., Mahoney, M. W., and Keutzer, K. A survey of quantization methods for efficient neural network inference. arXiv preprint arXiv:2103. 2021.
Bricken, T., Templeton, A., Batson, J.. Chen, B., Jermyn, A., Conerly, T., Turner, N., Anil, C., Denison, C., Askell, A., et al. Towards monosemanticity: Decomposing language models with dictionary learning. Transformer Circuits Thread, 2, 2023.
Goldblum, M., Finzi, M., Rowan, K., and Wilson, A. G. Po- sition: The no free lunch theorem, Kolmogorov complex- ity, and the role of inductive biases in machine learning. In International Conference on Machine Learning, 2024.
Goyal, A. and Bengio, Y. Inductive biases for deep learning of higher-level cognition. arXiv preprint arXiv:2011.15091, 2020.
Brooks, T., Peebles, B., Holmes, C., DePue, W., Guo, Y., Jing, L., Schnurr, D., Taylor, J., Luhman, T., Luhman, E., et al. Video generation models world simulators, 2024.
Gunasekar, S., Lee, J., Soudry, D., and Srebro, N. Charac- terizing implicit bias in terms of optimization geometry. International Conference on Machine Learning, pp. 1827-1836, 2018a.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert- Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler. D. M., Wu,
J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray. S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. Language models are few-shot learners. In Advances Neurul Information Processing Systems, 2020.
Gunasekar, S., Lee, D., Soudry, D., and Srebro, N. Im- plicit bias of gradient descent on linear convolutional networks. In Advances in Neural Information Processing Systems, pp. 9482 9491, 2018b.
Gurnee, W. and Tegmark, M. Language models represent space and time. In International Conference Learning Representations, 2024.
Chatterjee. and Sudijono, T. Neural networks gen- eralize on low complexity data. arXiv preprint arXiv:2409.1 12446, 2024.
Ha, D. and Schmidhuber, World models. arXiv preprint arXiv: 2018.
Chizat, I.,, Chizat, I.., and Fr, U.-P.-S. Implicit bias of gra- dient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theory, 2020.
He, K., Chen, X., Xie, S., Li, Y., Dollár, and Girshick, R. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16000-16009, 2022.
Craik, J. W. The nature of explanation, volume 445. CUP Archive, 1967.
Hendrycks, D., Mazeika, M., and Woodside, T. An overview of catastrophic Al risks. arXiv preprint N2ii:2230.12001,, 2023.
Dingle, K., Camargo, C. Q., and Louis, A. A. ouputt maps are strongly biased towards simple outputs. Nature Communications, 911,, 2018.
Hornik, K., Stinchcombe, M., and White, H. Multilayer feedforward networks are universal approximators. Neu- ral Networks, 255):359-366, 1989.
Everett, B. An introduction to latent variable models. 2013.
Friston, K., Moran, R. J., Nagai, Y., Taniguchi, T., Gomi, H., and Tenenbaum, J. World model learning and inference. Neural Networks, 144:573-590, 2021.
Huh, M., Mobahi, H., Zhang, R., Cheung, B., Agrawal, P., and Isola, P. The low-rank simplicity bias in deep networks. Transactions on Machine Learning Research, 2023.
Funahashi, K.-I. On the approximate realization of contin- uous mappings by neural networks. Neural Networks, 2 (3):183-192, 1989.
Huh, M., Cheung, B., Wang. T., and Isola, P. The Platonic representation hypothesis. In International Conference on Machine Learning. 2024.
Geirhos, R., Rubisch, P., Michaelis, C., Bethge, M., Wich- mann, F. A., and Brendel, W. ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness. In International Conference on Learning Representations, 2019.
Hyvärinen, A. and Pajunen, P. Nonlinear independent com- ponent analysis: Existence and uniqueness results. Neural Networks, 12(3):429-44 439, 1999.
10

---

When Do Neural Networks Learn World Models?
Hyvärinen, A., Sasaki, H., and Turner, R. E. Nonlinear ICA using auxiliary variables and generalized contrastive learning. In International Conference Artificial Intelli- gence and Statistics, 2019.
Liu, X., Zhang, F., Hou, Z.,, Mian, L.,, Wang, Z., Zhang, J., and Tang, J. Self-supervised learning: Generative contrastive. IEEE Transactions on Knowleadge and Data Engineering, 35(1):857-88) -577868 2021.
Jacot, A., Gabriel, F., and Hongler, C. Neural tangent ker- nel: Convergence and generalization in neural networks. In Advances in Neural Information Processing Systems, 2018.
Liu, Z., Zhong, Z., and Tegmark, M. Grokking as compres- sion: nonlinear complexity perspective. arXiv preprint arxiv:2310.05918, 2023.
Loshchilov, I. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2019.
Jiang, Z., Yang, M., Tsirlin, M., Tang, R., Dai, Y., and Lin, J. "Low-resource" text classification: parameter- free classification method with compressors. In Findings of the Association for Computational Linguisties: ACL 2023, pp. 68 10-088828, Toronto, Canada, 2023.
Lotfi, S., Finzi, M., Kapoor, S.,, Potapczynski, A., Goldblum, M., and Wilson, A. G. PAC-Bayes compression bounds so tight that they can explain generalization. In Advances in Neural Information Processing Systems, 2022.
Jin, C. and Rinard, M. Emergent representations of program semantics in language models trained on programs. In International Conference Machine Learning, 2024.
Lyu, K., Wang, R., Li, Z.., and Arora, S. Gradient descent on two-layer nets: Margin maximization and simplicity bias. In Advances in Neural Information Processing Systems, pp. 12978-12991, 2021.
Kalimeris, D., Kaplun, G., Nakkiran, P., Edelman, B., Yang. T., Barak, B., and Zhang, H. SGD neural networks learns functions of increasing complexity. In Advances in Neural Information Processing Systems, pp. 3491-3501, 2019.
Marks, S. and Tegmark, M. The geometry of truth: Emergent linear structure in large language model representations of true/false datasets. arXiv preprint arXiv:2310.06824 2023.
Kang, B., Yue, Y., Lu, R., Lin, Z., Zhao, Y.. Wang, K., Huang, G., and Feng, J. How far is video generation from world model: A physical law perspective. International Conference on Machine Learning, 2025.
Mikolov, T., Yih, W.-t, and Zweig, Linguistic regularities in continuous space word representations. In Proceedings of the 2013 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pp. 746-746674 1146-751,220 2013.
Khemakhem, I., Kingma, D. P., Monti, R. P., and Hyvärinen, A. Variational autoencoders and nonlinear ICA: A unify- ing framework. International Conference on Artificial Intelligence and Statistics, 2020.
Mirzadeh, I., Alizadeh, K., Shahrokhi, H., Tuzel, O., Bengio, S., and Farajtabar, M. GSM-symbolic: Understanding the limitations of mathematical reasoning in large language models. arXiv preprint arXiv:2410.05229, 20244,
LeCun, Y. A path towards autonomous machine intelligence version 0.9. 2, 2022-06-27. Open Review, 62(1):1-62, 2022.
Mitchell, M. Ai's challenge of understanding the world. Science, 382(6671), 2023.
Lee, J. D., Lei, Q.. Saunshi, N., and Zhuo, J. Predicting what you already know helps: Provable self-supervised learning. In Advances in Neural Information Processing Systems, volume 34, pp. 309-323, 2021.
Motamed, S., Culp, L., Swersky, K., Jaini, P., and Geirhos, R. Do generative video models learn physical principles from watching videos? arXiv preprint arXiv:2501.090. 2025.
Li, B. Z., Nye, M., and Andreas, J. Implicit representations of meaning in neural language models. arXiv preprint arXiv:2106.00737, 2021.
Nagel, M., Fournarakis, M., Amjad, R. A., Bondarenko, Y., Baalen, M. v., and Blankevoort, T. A white pa- per on neural network quantization. arXiv preprint arXiv:2 2021.
Li, K., Hopkins, A. K., and Bau, D. Emergent world rep- resentations: Exploring sequence model trained on a synthetic task. International Conference on Learning Representations, 2023.
Nanda, N., Lee, A., and Wattenberg. M. Emergent línear rep- resentations in world models of self- self-supervised sequence models. arXiv preprini wwXiv:2309.0094. 2023.
Li, M.. Vitányi, P., et al. An introduction to Kolmogorov complexity and applications, volume 3. 2008
O' D'Donelll R. Analysis of boolean functions. Cambridge University Press, 2014.

---

When Do Neural Networks Learn World Models?
Paszke, A., Gross, S., Massa, F., Lerer, A., Bradbury, J., Chanan, G., Killeen, T., Lin, Z., Gimelshein, N., Antiga, L., Desmaison, A., Kopf, A., Yang, E., DeVito, Z., Raison, M.. Tejani, A.. Chilamkurthy, S., Steiner, B., Fang. L., Bai, J., and Chintala, S. PyTorch: An imperative style, high-pens sech-eeffrrance deep learning library. In Advances in Neural Information Processing Systems, pp. 8024-8035, 2019.
Von Kügelgen, J., Sharma, Y., Gresele, L.,, Brendel, W., Schölkopf, B., Besserve, M., and Locatello, R. Self- supervised learning with data augmentations provably isolates content from style. In Advances in Neural Infor- mation Processing Systems, volume 34, pp. 16451-16467, 2021.
von Kügelgen, J., Besserve, M., Wendong, L., Gresele, L., Kekić, A., Bareinboim, E., Blei, D., and Schölkopf, B. Nonparametric identifiability of causal representations from unknown interventions. In Advances in Neural Information Processing Systems, volume 36, 2023.
Press, O., Smith, N., and Lewis, M. Train short, test long: Attention with linear biases enables input length extrapo- lation. In International Conference on Learning Repre- sentations, 2022.
Wei, C., Xie, S. M., and Ma, Why do pretrained language models help in downstream tasks? An analysis of head and prompt tuning. In Advances in Neural Information Processing Systems, volume 34, pp. 16158-16170, 2021.
Pérez, G. V., Louis, A. A., and Camargo, C. Q. Deep learn- ing generalizes because the parameter-function map is biased towards simple functions. In International Confer- ence on Learning Representations, 2019.
Whitley D. and Watson, J. P. Complexity theory and the no free lunch theorem. Search Methodologies: Intro- ductory Tutorials Optimization and Decision Support Techniques, pp. 117-339, 2005.
Radford, A., Wu, J., Child, R., Luan, D., Amodei, D., Sutskever, I., et al. Language models are unsupervised multitask learners. OpenAl Blog, 18)): 2019.
Raghu, M., Unterthiner, T., Kornblith, S., Zhang, C., and Dosovitskiy, A. Do vision transformers see like convolu- tional neural networks? Advances in Neural Information Processing Systems, 34: 12:12286,,, 2021.
Wolpert, D. H. The lack of priori distinctions between learning algorithms. Neural computation, 8(7):1341- 1390, 1996.
Reizinger, P., Ujváry, S., Mészáros, A., Kerekes, A., Bren- del, W., and Huszár, F. Position: Understanding LLMs requires more than statistical generalization. In Interna- tional Conference on Machine Learning, 2024.
Wong, I.., Grand, G., Lew, A. K., Goodman, N. D., Mans- inghka, V. K., Andreas, J., and Tenenbaum, J. B. From word models to world models: Translating from natural language to the probabilistic language of thought. arXiv preprint arXiv:2306. /2622 2023.
Richens, J. and Everitt, T. Robust agents learn causal world models. In ICLR, 2024.
Wu, Z., Qiu, L., Ross, A., Akyürek, E., Chen, B., Wang, B., Kim, N., Andreas, J., and Kim, Y. Reasoning or reciting? Exploring the capabilities and limitations of language models through counterfactual tasks. arXiv preprint arXiv:2307.024777 2023.
Schölkopf, B., Locatello, F., Bauer, S., Ke, N. R., Kalch- brenner, N., Goyal, A., and Bengio, Y. Toward causal representation learning. Proceedings of the IEEE, 109(5):
612-634, 2021.
Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. The implicit bias of gradient descent sepa- rable data. The Journal of Machine Learning Research, 19(1)2822--2878, 2018.
Xie, K., Yang, I., Gunerli, J., and Riedl, M. Making large language models into world models with precondition and effect knowledge. arXiv preprint arXiv:2409. 122777 2024.
Teney, D., Nicolicioiu, A., Hartmann, V., and Abbasnejad, E. Neural redshift: Random networks are not random functions. arXiv preprint arXiv:2403.022411 2024.
Xu, K., J., Zhang, M., Du, S. S., Kawarabayashi, K.-i., and Jegelka, S. What can neural networks reason about? In International Conference on Learning Representations, 2020.
Tosh, C., Krishnamurthy, A., and Hsu, D. Contrastive learn- ing, multi-view redundancy, and linear models. In Algo- rithmic Learning Theory, pp. 1179-1206, 1206, 2021.
Xu, K., Li, J.. Zhang, M., Du, S. S.,, Kawarabayashi, K.-i., and Jegelka, S. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference on Learning Representationss 2021.
Vapnik, V. The nature of statistical learning theory. 1999.
Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, L.,, and Polosukhin, I. Atten- tion is all you need. In Advances in Neural Information Processing Systems, pp. 5998 -0088, 2017.
Xu, Z.-Q. J., Zhang, Y., Luo, T., Xiao, Y., and Ma, Z. Fre- quency principle: Fourier analysis sheds light on deep neural networks. arXiv preprint arXiv: /90/ 0056666,,, 2019.
12

---

When Do Neural Networks Learn World Models?
Yildirim, and Paul, L. From task structures to world models: What do LLMs know? Trends in Cognitive Sciences, 2024.
Zhang, C., Bengio, S., Hardt, M., Recht, B., and Vinyals, O. Understanding deep learning requires rethinking gen- eralization. In International Conference Learning Representations, 2017.
Zhang, T., Zhao, C., Chen, G., Jiang, Y., and Chen, Fea- ture contamination: Neural networks learn uncorrelated features and fail to generalize. In International Confer- on Machine Learning, 2024.
Zhao, C., Zhang. T., Chen, G., Jiang, Y., and Chen, M Identifying and exploiting view bias of prompt learning. Transactions Machine Learning Research, 2024.
Zhou, W., Veitch, V., Austern, M., Adams, R. P., and Orbanz, P. Non-vacuous generalization bounds at the ImageNet scale: A PAC-Bayesian compression approach. In Inter- national Conference on Learning Representations, 2019.
Zimmermann, R. S., Sharma, Y., Schneider, S., Bethge, M., and Brendel, W. Contrastive learning inverts the data generating process. In International Conference on Machine Learning, pp. 12979-12990, 2021.
Ziyin, L., Hartwig, T., and Ueda, M. Neural networks fail to learn periodic functions and how to fix it. In Advances in Neural Information Processing Systems, 2020.
13

---

When Do Neural Networks Learn World Models?
A. Related Work
World models. The term "world model" in the machine learning context originates from Ha & Schmidhuber (2018), who describe it as a human-likk "mental model of the world" that learns an abstract representation of information flow and can be used predict future events. This definition closely aligns with the concept of "mental model" in cognitive science, i.., internal representation of external reality (Craik, 1967), making it naturally connected to the field of representation learning in machine learning (Bengio et al., 2013).
Recently, the remarkable capabilities of large language models (LLMs) have sparked a scientific debate on whether these models merely exploit superficial statistical patterns to generate predictions without genuine "understanding" of natural language (Bender et al., 2021; Mitchell, 2023), or whether they develop models that serve as compact and interpretable representations of the underlying data generation process. A series of studies have demonstrated the presence of internal representations in language models trained on synthetic tasks (Li et al., 2023; Nanda et al., 2023; Jin & Rinard, 2024). For real-world LLMs, research in mechanistic interpretability suggests that these models learn compact, interpretable, and causal features within their intermediate layers (Li et al., 2021; Bricken et al., 2023; Marks & Tegmark, 2023; Gurnee & Tegmark, 2024). At the same time, many studies report a significant decline in LLM performance on tasks that are assumed to be underrepresented in their pre-training distribution (Wu et al., 2023; Berglund et al., 2024; Mirzadeh et al., 2024).
Beyond sequence models, world models have also gained attention in reinforcement learning (Ha & Schmidhuber, 2018; LeCun, 2022: Xie et al., 2024), probabilistic learning (Friston et al., 2021; Wong et al., 2023), and causal discovery (Richens & Everitt, 2024). However, despite the growing number of empirical studies, the theoretical foundations of world model learning remain largely unexplored.
Latent variable recovery. Our definition of world models falls within a broad class of latent variable recovery prob- lems (Everett, 2013), where observable data is generated by latent variables through an unknown generation function. It is well established that, without additional assumptions, recovering the true latent variables from observed data is generally impossible if the generation function is nonlinear (Hyvärinen & Pajunen, 1999; Khemakhem et al., 2020), a fundamental result in non-linear independent component analysis (non- (non-linear ICA).
To overcome this impossibility, subsequent research has explored various structural assumptions on latent variables, such as conditional independence between the latent variables and an observable auxiliary variable (Hyvärinen et al., 2019; Khemakhem al., 2020; Lee et al., 202 1), distributional constraints (Zimmermann et al., 2021; Wei et al., 2021), and causal interventions (Von Kügelgen et al., 2021; Ahuja et al., 2023; von Kügelgen et al., 2023). Some studies have also linked these assumptions to contrastive learning (Hyvärinen et al., 2019; Tosh et al., 2021: Zimmermann et al.. 2021). However, incorporating such structural assumptions often leads to complex and less scalable training paradigms compared to the pre-training framework of modern LLMs (i.e., next-token prediction). As a result, these studies do not directly address the central question of our work: Can the ongoing paradigm of LLMs learn world models? Meanwhile, the fact that LLMs already acquire non-trivial representations (Li et al., 2021; Bricken et al., 2023; Marks & Tegmark, 2023; Gurnee & Tegmark, 2024) suggests that they must leverage some form of implicit bias rather than explicit structural assumptions on input data, which motivates our study.
Implicit bias of neural networks. Overparameterized neural networks have been shown to possess the capacity to memorize entire training datasets (Zhang et al., 2017). However, their ability to generalize well in many settings suggests that they exhibit implicit preferences for certain as linear models, random feature models, and two-layer neural networks, a body of theoretical work demonstrates that (stochastic) gradient descent imposes specific forms of implicit regularization on the learned solutions (Soudry et al., 2018; Gunasekar et al.. 201 Ba;;; Bartlett et al., 2020; Chizat et al., 2020; Lyu et al., 2021: Allen-Zhu & Li, 2023; Andriushchenko et al., 2023; Abbe et al., 2023; Zhang et al., 2024).
combonll referred to as implicit bias. In simple models, such
Empirical studies further suggest that practical neural networks extend many of these implicit regularization effects through a form of simplicity bias, favoring "simpler" solutions over more complex ones (Pérez et al., 2019; Kalimeris et al., 2019; Xu et al., 2019; Bhattamishra et al., 2023; Huh et al., 2023; Zhao et al., 2024). However, the notion of "simplicity" is often defined empirically and varies across studies. In this work, we formalize the concept of simplicity within our theoretical framework and analyze its relationship to learning world models.
14

---

When Do Neural Networks Learn World Models?
Complexity measures. Kolmogorov complexity (Li et al., 2008) in algorithmic learning theory provides unified framework for quantifying the complexity of any object, including functions. However, it is not computable in general. In machine learning, conventional statistical learning theory typically employs the VC dimension to bound the generalization eπтоr of models (Vapnik, 1999). More recent complexity measures include Rademacher and Gaussian complexities (Bartlett & Mendelson, 2002). However, these measures primarily assess the complexity of function classes rather than individual funciions.
complexity measure
growing body of research introduces complexity measures tailored for neural networks trained via (stochastic) gradient descent (Bartlett et al., 2017: Jacot et al., 2018: Zhou et al., 2019; Lotfi et al., 2022: Chatterjee & Sudijono, 2024), often leveraging them for generalization analysis. However, these measures inherently depend on the neural network parameterization, rather than capturing function complexity independently of specific parameterizations. Some studies propose alternative complexity metrics inspired by Kolmogorov complexity in the machine learning context (Xu et al., 2020; Liu et al., 2023), but these metrics are typically problem-specifie and do not enable a direct complexity analysis in the function space, unlike the approach we take in this work. See Section C for more discussion on some commonly-usee complexity measures.
B. Preliminaries on Boolean Function Analysis
This section introduces basic definitions and properties of Boolean functions for readers who are not familiar with Boolean function analysis. For a more detailed introduction, we recommend the first few chapters of the book by O'Donnell (2014).
Following the convention in Boolcan function analysis, throughout this work we use the term Boolean functionss to refer to functions with the form
f: [=1-1,1}" +0,
(12)
where Y could be any subspace of Rd for an arbitrary integer d, such R. {-1,1]3, etc.
There are different ways to represent the input bits in the above definiton. A natural way is to use 0 and as elements of the field 1F2. In this way, a (single-output) Boolean function with n input coordinates (bits) is represented by f: {0 1)th + 13. It is also convenient to use -1 and 1, thought as real numbers, and define f as function from f-1, to 11. The latter representation can be easily transformed from the former one by the mapping b - (-1)² over 13. .. In our analysis, we will mostly use the latter representation as is more compatible with the Fourier- Fousier-Walsh transform of Boolean functions.
Fourier- Walsh transform. As in Definition 3.1, every function {-1,1}" IR can be expressed as multilinear polynomial, i.e., we have
f(x) = Σ f(S)xs(x),
(13)
8gIm]
where f(S) Fourier- Walsh coefficients and (x) Iles x; monomials, also called parity functions. The name parity function is due to the fact that it computes the logical parity or XOR of the bits. As an example, consider = max2, i.c., the maximum function on 2 bits, and its fourier-Walss transform is f(11, = 1/2 + f + _ - =
The existence of Fourier-VV Walsh transform can be proved by construction, using an interpolation method similar to constructing
Lagrange polynomials. For each point = (a1,..., an) E {[1, 1)", define
define
1 1 0:
1 (x(x)=
II
(14)
2
<E(n)
It is easy to verify that (x)) takes value 1 when = a and 0 otherwise. We then have
f(x) = Σ
f(a)la(x)
(15)
-E(-1,13
and arrive at a polynomial representation of f. Since any factors of 27, Vi E [n]] can be replaced by 1, we further know that this polynomial must be multilinear.
15

---

When Do Neural Networks Learn World Models?
For functions f F2 - R, we can also define their Fourier- transforms by extending the Xs notation using the mapping b- (-
Xs(x) (-1)
(16)
where E F2 We can thus write the Fourier-Wand Walsh transform of f :F2 IR in the same form equation (13).
Parity functions, orthogonality, and Parseval's Theorem. R and g : (-1,1)" IR by
We define the inner product (, of functions f (-1,1)"
(.g) -B=-)(((-Lay
(17)
A key fact about parity functions is that they orthogonal under the above definition: for every S, [n], we have
1, S S" (xs Xs')
(18)
0, otherwise
Consequently, we consider the vector space V containing all functions f: (-1,1) -> IR, then parity functions form an orthonormal basis of V. can be verified that the Fourier- Walsh coefficients in equation (13) satisfy
f(S) (f.xs)
(19)
for every S - [m].
Purseval's 's Theorem shows that for every : (-1,13 + R.
(f.f) Σ f(S)?
(20)
SCIm
In particular, if f is Boolean-valued, then (f. 1. The uniqueness of the Sourier-Walske transform of Boolean functions can also be proved using the Parseval's ' Theorem (O'Donnell, 2014).
C. Discussion on Complexity Measures
Besides the Boolean complexity measures considered by this work, there are also continuous approximations of Kolmogorov complexity, e.g., continuous complexity measures based on standard Fourier transform, the order of approximation polynomials, or compression (Xu al., 2019; Pérez al., 2019; Jiang al., 2023; Teney et al., 2024). However, an important downside of these continuous complexity measures is that they do not apply to arbitrary continuous functions, neither constituting a basis of the continuous function space, which makes them less amenable to theoretical analysis.
For example, the linear function (xx ar + b, E R does not have standard Fourier transform since it is not absolutely integrable. Even if we constrain it to some finite interval [-L, L] and periodically extend it across the real axis, the resulting Fourier transform would involve frequencies up to infinity, suggesting infinite complexity if we use the highest frequency as the complexity measure. This is contrast to the empirically observed simplicity bias of gradient descent, where linear models are often learned first (see e.g., Kalimeris et (2019)). The measure of polynomial order has a similar problem since not all continuous functions can be linearly represented by set of polynomials. In fact, the degree of Boolean functions could be viewed as a "discrete" version of polynomial order, but it is more principled since the Boolean function space indeed has polynomial basis given by the Fourier- transform. Finally, we found other compression-bast approximations of Kolmogorov complexity hard to analyze theoretically since they rely on specific compression algorithms.
D. Numerical Experiments
This section presents numerical experiments that corroborate our theoretical results. In all experiments, we train MLPs using stochastic gradient descent (SGD) and binary data from {m generated by the data generation model in Table 1 with d = 10. A summarization of results is as follows:
Low-degree bias of neural networks: Figure 3 shows that MLPs indeed have a low-degree bias when solutions with different degrees exist for the same training distribution.
16

---

When Do Neural Networks Learn World Models?
Impact of multi-task training on identifying eata-generating variables: Figure 4 shows that the identification of -lta-generating variables significantly decreases as more training tasks are drawn, while single-task learning does not lead to similar identification.
Identifying data-generating variables benefits out-of-distribution (OOD) generalization: Figure 5 shows that in Hamming sampling setting as in Theorem 4.9, MLPs with a representation & that identifies true data-generating variables achieve superior OOD generalization compared MLPs without such representation.
Please refer to the figures for more details and discussion.
It is worth noting that in all experiments, we use MLPs with quadratic activation (x)) = x to ensure the compatibility between the model and the Fourier- Walsh transform (i.e., the polynomial expression) of Boolean functions. Our preliminary experiments showed that the low-degree bias and the identification error are very sensitive to the activation function used, which is consistent with our basis compatibility analysis in Section 4.4. This also explains the phenomenon reported by Abbe al. (2023) that MLPs with ReLU activation do not exhibit "pure" low-degree bias for Boolean inputs.
Table 1. The data generation model æ (2) used in numerical experiments.
4/2
373
444
35
11
21
21.22
21 2g
21 Pg
25
24
2824282627
24
28
26
29
Error
1.0
1.2
1.0-
0.8
Full
Full
0.8
0.6
MSE
0.6
0.4
0.4
0.2
0.2
200
600
800
1000
200
800
1000
Iterations
Iterations
(a) Task 1: f(x) = 25445
Error
OOD Error
1.4
1.2
Full
1.0
Full
1.0-
MSE
0.8
0.6
0.6-
0.4
0.2
0.2
200
600
800
1000
200
800
1000
Iterations
Iterations
(b) Task 2: f(x) = ur + 0.5
Figure 3. Low-degree bias of neural networks: Training latents are uniformly sampled from the Hamming hall B, = (z (+1)4) #-1(z) r): the out-of-distriiuttin (OOD) test latents uniformly sampled from (+1)* (d 10). For small r, low-degree realizations of tasks (target functions) exist in the training distribution. A model that lears these low-degree realizations would have low in distribution (ID) error yet high OOD error due SGD-trained MIPss learn ground-truth target functions when 10 (denoted by "Full" in the legends), they indeed fit low-degree realizations for training data with r and 3, resulting in high OOD error. separately train a three-layer MLP each task for each r.
loss high-degree components in target functions. observe that while
17

---

When Do Neural Networks Learn World Models?
identification Enror # Tasks
Identification Error # Tasks
Identification Error vs # sasks
Identfiiation Eneer (MSE)
0.25
0-20-
0.23
Ideetificaaton Error
0.20
Identification Errer
0.15-
0.15-
0.10
n
64
0.85
1020
2000
2000
4000
5000
1000
2000
5000
4000
5600
1000
2000
3000
4000
5000
Iterations
Iterations
Iterations
(a) 2-degree tasks
(b) 3-degrec tasks
(c) 4-degree tasks
Figure 4. Impact of multi-task training identifying date-gnngg variables: The identification error of dala-generating variables decreases as the task number increases, measured by the MSE of the linear probe trained to fit true latent variables z given the learned representation & (x)). Each task a randomly sampled parity function of with degree
larger than le for k-degree tasks. Compared to single-task task training (n = 1). multi-task training with 64 significantly reduces the identification error for all k-degree training task distributions with
3, 4, corroborating Theorem 4.3 and Theorem 4.8 in the main text. Both the representation & and each task-specific function gi, i E [ni] are instantiated three-layer MLP.
ID Error
OOD Error
1.0
W/o WM
1.2
W/O WM
0.8-
0.6
MSE
MSE
0.4
0.2
0.2
0.0
0.0
200
600
800
1000
200
400
600
800
1000
Iterations
Iterations
(a) Task 1: f(x) 25.5545
ID Error
OOD Error
1.4
w/o WM
w/o WM
1.2
1.0
1.0
MSE
0.8
0.8
0.6
0.4
0.2
0.2
0.0
0.0
200
400
600
800
1000
200
400
600
800
1000
Iterations
Iterations
(b) Task 2: f(x) 3133 0.5x11550
Figure 5. Identifying data-generating variables benefits out-of-distribution (OOD) generalization: For the same tasks considered in
Figure 3 with 2. SGD-trained MLPs top oracle representation that identifies true data- seatteegnertnnnggg variables (denoted by "w/ WM" in the legends) achieve OOD generalization. By contrast, SGD-trained MLPs without the representation (denoted by "w/o WM" in the legends) fail. This corroborates Theorem 4.9 in the main text. We separately train MI.P for each task and for the setting with/without the oracle representation.
18

---

When Do Neural Networks Learn World Models?
E. Proofs
This section provides complete proofs of all theorems the main text, organized as follows.
In Section 1, we introduce some definitions and technical lemmas.
In Section E.2, we provide the proof of Theorem 4.1.
In Section E.3, we provide the proof of Theorem 4.3.
In Section E.4, we provide the proof of Theorem 4.4.
In Section E.5, we provide the proof of Corollary 4.7.
In Section E.6, we provide the proof of Theorem 4.8.
In Section E.7, we provide the proof of Theorem 4.9.
In Section E.8, we provide the proof of Theorem 12.
Notation and conventions. We use [n] to denote the set {1,...,n] n} for positive integers n. For a set S, we denote its cardinality by (S]. For a probability distribution p over some set S, we denote by supp(p) := {a E p(s) 0} denote the multi-output function satisfying that f(x) (f:) ......... f (z)); conversely, for function f with output dimensions, we use fi to denote the function mapping the inputs of to its i-th output dimension for i E [ni]. As defined in the main text,
its support. For n functions fi...., we use f (fi....,
the notation F = (f [+1] R and Fh = (ʃ (+1] > R deg(f) S k- for positive integers n. Note that although both F and F are infinite, only finite functions in them are implementable by computers due to bounded precision. Therefore, in our proofs we will treat them as finite yet exponentially large sets. This enables us to use, e.g., |F and Σver" dog(h') in our proofs and helps avoid non-essential technical nuances.
E.1. Technical Lemmas
This section presents additional definitions and technical lemmas that may come in handy in our proofs. We begin by introducing lemma that upper-bounds the degree of min-degree solutions for every task with d-dimensional latent variables. Lemma E.1. Suppose that the input data variable x E C (-1, 1)th is generated as in Definition 2.1 by a d-dimensional lantent variable = 1-,,1 d S m. Then, for every task h X + {-1,1}, we have
deg( Hmin (h)) s d.
(21)
Proof. The case of d m is trivial, in what follows consider the of d >> We first prove the following lemma: Lemma E.2. For eververy S = (i1,..., S [m] with |.|| = d d+1, there exist by.... bati with by E (-1,1), [d+1] such that Isseja + by) = holds for every x e X.
Proof of Lemma E.2. We can prove Lemma E.2 by contradiction: assume that it does not hold, i.e., for some S, there exists a E X such that IIyeja+1j + b,) # 0, then we must have r;,- = by. Vj E [d+ 1]. Since applying this argument to every (b1, (bi,..., bat1) E 114,1 we can find 2d+1 elements in X that differ in at least one coordinate in S, which yields
ba+1 are arbitrary. by
| > 2d+1, On the other hand, recall that we assume supp(z) Z. Due to the invertibility of f, we have |2'| = | = 2d, which contradicts | 2 2d+11 Hence, the initial assumption is false and Lemma E.2 holds.
Applying Lemma E.2, we have that for every degree-d + subset S = -fi......a id+1 C [m].
II
(Ii, + b,)
II jES
Σ
bg II 24 0, Vo X
(22)
sE[d+1]
KES'
holds for some by....,bu44 where by E {-1,1} for every S' C S. This means that we can thus replace every degree-
monomial (s( = Π 2j by degree-d polynomial Σ Srees by I Ty. By iteratively using this
19

---

When Do Neural Networks Learn World Models?
replacement in the Fourier-Wasd Walsh transform of h (Definition 3.1), one can eventually obtain polynomial with degree d or less without changing the value of the function on every X. This completes the proof.
Remark E.3. Without the latent structure, trivial upper bound of the min-degree solution of any task h is deg((uui) (h)) -
m.. Hence, Lemma E.1 is an example of how the min-degree bias can exploit the latent structure of data by favoring low-degree solutions with degree independent of the data dimension m.
Before presenting our next lemma, we first introduce the concept of influence for Boolean functions. Definition E.4 (Influence). Let ʃ 1- + {-1, be Boolean function. Then, the influence of coordinate i, E on f is defined as
[v]]
Inf;()) = Pr Pra=-((-11)))// + f(x
(23)
where x f (I1 Ti-1, Diy Fi+1 ,,.,,,,,....,))
By Parseval's Theorem, we have a formula between influence and the Fourier- Walsh coefficients (O' Donnnell, 2014).
Lemma E.S. For f {-1,1)" + { and every i E [nl, the following holds:
[nf;(f)= = E
f(5))
(24)
SE(S'CIa| |
where f(S) is the Fourier- Walsh coefficients of f as in Definition 3.1.
With the above definitions, our next lemma introduces a necessary and sufficient condition of bijective Boolean functions being degree-1, 1, based on the restricted influence of all input coordinates.
Lemma E.6. Let f: {-1,11" + 1-1,1}" be a bijective function and let 1 S k S n - 1 be an integer. Then, we have (eg(ft) = vi E [m] and only for every S C [n]] with |S| = k, there exists T C [n]] with [T] = k such that:
for every j T. Inf,(f)) > for at least one i E S;
for every j [m]] TT, (ff) for every i E S.
Proof. Note that deg(f:) 1, vi E [m]] together with the fact that f is bijective indicates the existence of a permutation ii.....in in of I,..., n such that (x) = Ii, or (xx = -I,, for every j E [n], which trivially gives the result. In the following we prove the other direction. Note that the case of k = 1 is trivial. Hence, to prove that deg(ft) = 1, it suffices to prove (*): for every i E [ni], there exists j E [n] such that (f,) > O and Inf = O, Vm + j. m E [m]]. for every 2 < S n-
We first prove that for every S. I' is unique, by contradiction. Suppose that for some S, T is not unique, i.e. there exists T' + T C [m] such that T' satisfies the condition. Then, by Definition E.4, changing the values of the coordinates z;, i S can change only the values of { j E T but not the values of other fz(x),,j E [M] \ (T U T'). This results in [f(x),, x E X31 S 2n-18| 2 2 2--b+|/mn 2", contradicting the bijectivity of f. Therefore, the assumption is false and T is unique. This allows to define mapping ø S
Let S = (S C [n] |S| = k} be the set of all subsets of [ri] with cardinality k. We then prove that Φ : S - T is bijective on Sg. To this end, it suffices to show that every S + S" E Sz are mapped to different T. This can be similarly proved by contradiction as in proving the uniqueness of T: if is false, then there exists a subset S" = S S and T such that Inf,(S,) = 0, Vi E 5", E \ T. Then, by Definition E.4, changing the values of the coordinates 1, E 5" can change only the values of (x).j j E T but not the values of other fz()) & T. This results in [ff()) x E x3 S 2 15-1 = 2 1s"1 < 2", contradicting the bijectivity of f. Therefore, the assumption is false and every S # S" E Sy are mapped to different T.
Given that : 5 + T is bijective, we know that for every subset S C Sy, there exists unique subset T = (T Φ(S) S E 8) C Sk such that [S] = [T]. We then move on to prove the proposition (*) by contradiction: suppose it is false, i... for some i E [n]. there exists U C [r] with [U] 2 such that [f,) > 0 for every j E U and [f,) = 0 for every j E [n] \ U. If [0] > k, then for any S such that i E S, there does not exist feasible T, which is a contradiction. 2 < [U] & k. consider S = {S c Sx E $} with |S| = c+1 By the definition of o, we know that
T T (T c S& U T) which gives TT| s - |0
< ck-1 1 = (S). Thus, the assumption is false and proposition (*) is
true. This completes the proof.
20

---

When Do Neural Networks Learn World Models?
Our next lemma shows that any bijective transform 1 can induce a bijective transform F*.
Lemma E.7. Let T (- - [-1,1] be a bijective transform. Then, F T Fd.
Proof. It suffices to show that the mapping It' + ht' T for h' F bijective. On one hand, it is obvious that h of Fd for every h' E Fr. On the other hand, for each h Frd there exists ht = h o T such that ht o T h². This completes the proof.
We then present lemma from Abbe et al. (2023) that guarantees the uniqueness of low-degree solutions when the training data sampled from a Hamming ball.
Lemma E.8 (Abbe et al. (2023), Theorem 5.1). Consider a Boolean fincttio f {= + R. Then, there exists a unique function fr : {+1]" deg(f+) S.. r.
#-1(2) S r), we have fr(z) f(z) and such that for every z E B, {z {=1}*
Our next lemma shows that a bijection on F² that preserves the degree all parity functions preserves the degree of all functions.
Lemma E.9. Let U : : + F be an invertible linear transform. If (fde(U((xs) = deg(xs) every S C [n], then we have
deg(U(()) deg(f), V.
(25)
Proof. By the linearity of U, have
U()= = Σ
f(S)xs)
Σ [(S)U(xs)-
(26)
SCIu]
SCIM
It then follows from Definition 3.2 that
(e)(U(()) max (deg(U(xs))= f(S) #0} max {deg(xs) f(8) # +0} deg(s).
(27)
This completes the proof.
E.2. Proof of Theorem 4.1
Proof. Since g o i is an realization of h, we have g o & E H(h). This gives
deg(h") (hh deg(h") ") deg (h)) I deg(g @).
(28)
Thus, it suffices to prove that deg(g )) < deg(g + for every Z + R and tb X + 2. Let the Fourier- Walsh transform of g and 4,, E [d] be
g(z) = Σ 3(G) II
2
(29)
GCI4
IGG
and
(,(1)) Σ (S))) II
(30)
SC mm)
jES
respectively. Plugging (30) into (29) with 2, (x) gives
())
Σ
O(G) II
Σ 8Clm)
(c(S)
II aFk
(31)
GCId
iEG
Σ
O(G)
II IEG
Σ
(
II Zb
(32)
GE(
(
KES
21

---

When Do Neural Networks Learn World Models?
We thus have
deg(g Φ) v
max
max
[S]
(33)
(ld)" (0) iEG
(5))
max
max
[S]
(34)
G d
IEG (S))
max ()
[S]
(35)
i
dc
(36)
On the other hand, Definition 3.4 gives
deg(g b) deg(g) + (e) deg(g) +
Σ
deg( (2))
(37)
ie/d
deg( P,).
(38)
Therefore, we have deg(g o +) S deg(g e $). This completes the proof.
E.3. Proof of Theorem 4.3
Proof. By Lemma E. we can upper-bound the degree of each Φ5 E [d] by
deg (1) S
(39)
Expanding the LHS of (7) then gives
deg(h") - deg(go x0)
Σ deg(h? hī})
deg(g:)
Σ
deg(D) Φ7)
(40)
sEim]
EE |n]
SE(d]
deg
hh))
deg(g:)
d.a.
(41)
sEin]
E [m]
Note that h E Hmin (h) and o + is an realization of hy for every E [n]. By Definition 3.6, we have
deg(h, 4*) = deg(h?) deg(g:), Vi e [n].
(42)
Plugging equation (42) into (41) completes the proof.
E.4. Proof of Theorem 4.4
Proof. Note that for every task h E Fdd o o we can write h = H - for some H/ E Fdd Thus, for every g: 2 + R
and db X + Z such that & E H((h), we have (g o (x) = h(x) = (h' o (z)) Vx E X, which amounts to
g(z)= = (h' o ( Vz E
For every hy..... hn, g : 2 + (-1, 1)", and 4 : X + such that g, o E H(h.)) for every E [n], we have
1 Σ
Σ deg(gi) ):
lim deg(g o 2) lim n + n
deg(
()) +
(43)
10000
n jE
Ta iE(m)
Applying Lemma E. we have deg()); dj E [d]. This gives
lim 12+000 m
Σ
deg( ()) 0.
(44)
jE(d]
22

---

When Do Neural Networks Learn World Models?
Meanwhile, since hg X R, [ni]
independently and uniformly sampled from F o + have
lim n 70000 7 EE/n]
deg(g)) = En-U(rrop s)deg(h for
(45)
=Ew-eecsadeeg(hi o 1
(46)
Σ MERR
deg(h' o oT
(47)
[Fq
By Lemma E.7, we have Fd T = Fd.
This gives
deg(h' T-
Σ
deg(h").
(48)
MEF
MEFa
Plugging equations (44), (47), and (48) into equation (43) gives
lim
deg(g +)
deg(h'),
(49)
Fdd
N'EF
which is a constant independent of T (and thus independent of 4). Therefore, for any two viable representations 4, /'
and g' E (Fd)² with D. c and o & Vi E [n]. we must have lima-roc 1 (deg(g b) deg(g' o x0)) 0. This completes the proof.
E.5. Proof of Corollary 4.7
Proof. By Definition 3.6, we have
deg(h | |サ-1 > 0 + deg(Hminn (h)) > deg(g)
(50)
for g o a E Hh). By Lemma E.I, have degg(mminn(h (h)) for every h X IR. Thus, for deg(Hmiin (h)) deg(g) to hold, we must have deg(h 16) deg(g) s d 1. By Definition 4.6, this gives h Fd 1 - completing the proof.
E.6. Proof of Theorem 4.8
Proof. We first prove the following lemma that characterizes the averaged degree change for Boolean function in F when composed with invertible transforms.
Lemma E.10. For every integer 1 < k S d and every bijection T: --1, we have
Σ
deg(h' (h' T) 2 Σ deg(h').
(51)
MEFI
MEFI
In particular, when k 1, the equality holds if and only if deg(TT)) for every i E [d].
Proof of Lemma E. 10. For every g C Fd, let G o T = {h' o T N E g}. By Lemma E.7, we know that the mapping h + o T is bijective on Fd, We thus have ( T = F for every k. For Ft. a T, there are two possibilities:
1. F o OT== =F. This immediately gives
Σ
deg(h' oT)= = Σ
deg(h') =
Σ deg(h').
(52)
WEF
MEFFOT
WEFI
2. Fd. T + Fd. We can then decompose ENERY deg(h' hh' T) as follows:
Σ MEAL
deg(h' o T)
Σ
(h') +
Σ
deg(h')
(53)
ME(TOTYYNTE
Σ deg(h')
Σ
deg(h')-
Σ
deg(h')
(54)
MEFI
(EFFN(F(OT)
23

---

When Do Neural Networks Learn World Models?
Note that [ T| = [ gives (Fr (F)) deg(h() S k for every h' E Fy and deg(h') k for every h' Fu Ft. Taking these two facts together, we have
(Ft o T) Meanwhile, by Definition 4.6, we have
E
deg(h')
E
deg(h') 0.
(55)
(FF
(F
Plugging equation (55) into equation (54) then gives Enery deg(h' o TT) > deg(h))))(h)
Combining the above two cases, we conclude that EwerY deg(h' T) ENEFY deg(h') for every S k S d. Note that the above analysis also gives a necessary and sufficient condition for the equality to hold: Fi o To Fty
In particular, when k = 1. ENEFY deg(h' T) = ENEFY deg(h') holds only for T sstisffiig that F o T = F( Note that for every non-constant function h' E FY, there exists E [d] such that (h' o (x) E z E {-1,1] Due to the arbitrariness of h/, we must have deg(T;) for every i [d].
for every
We now move on to prove Theorem 4.8. Our aim is to prove that the minimizer (Ø',g") gg") of the optimization problem (10) learns the world model by negations and permutations when the number of tasks n - oc. Due to equations (43) and (44), the original problem equals to
min
lim
deg(gs)
d: -..x=22.EEPP
VEIN]
(56)
s.L.
9 no + H(h)), Vi E [n].
Due to the constraint gi & H(hs), Vi E [n]. we know that there must exist a bijection T: 2 - Z such that for every X. (xx) T(z), with e (xx) (x) being the true latent variable. Therefore, it remains to prove the existence of bijection T: : E + 2 with (T;(T)) 1, vi E [d] such that for every E X, (x) T(z).
For every gi o + E H(hh)). we have h; = gi o - 9: To We then have
lim n-000 iE(n]
deg(g;) = R ... pa) B PA-U(EOO sydeg(h o fo 1-1)
(57)
Ex . Pa) Ey U o 11
(58)
Σ kEld [4]
deg(h' o 1)
(59)
WEF!
By Lemma E.10, we have
deg(h' 1 2 Σ deg(h')
(60)
MEFE
MEFL
for every &s [d]. Plugging (60) into equation (59) gives
lim
deg(g:)
Σ kE|d]
Ph:
deg(h'),
(61)
70000
1E|m)
WEFI
where the equality holds only if deg(T:) 1 for every i E [d]. This completes the proof.
Remark E. A limitation of Theorem 4.8 is that we requires non-zero probability of explicitly sampling degree-1 tasks (i.e., Pi > 0), in which latent variables are essentially observed as task outputs. If O, we can still prove equation (61) by applying Lemma E.10; in other words, we can still prove that every representation $ that learns the world model up to negations and permutations is a minimizer of the optimization problem (10). However, such minimizers may not be unique, because Lemma E. 10 only proves the equivalence between the equality and deg(T;) = 1, Vi E [d] when k = but not k S In fact, we can construct hard examples showing that in some cases, there indeed exists other that minimizes ENEFY deg(h' T) for every k E [d] {1).
24

---

When Do Neural Networks Learn World Models?
Example E.12. Let d and lei T : [11,11 + {-1,1} be a bijective transform defined as
(z) 21, [=(x) = 21 Tz(z) = 2123.
(62)
One can easily verify that for every lo E (2.3), every parity function X Xs (T(2)) = -Le (z)) with |S| S k satisfy deg(xs(T(z))) I k. By the Fourier- Walsh transform, this amounts to F T Fi for k {2,3). which gives
Ewerg deg(h' T) En deg(h") by the proof of Lemma E.10. Thus, in this case require p1 > 0 to ensure that the representation & satisfying ((x)= = (z) for every E X' and = + (x) is not minimizer of (10).
Nevertheless, we do believe that cases like Example E.12 are rare. This is because by the proof of Lemma E. such examples must construct a bijection T such that F o T = F for every k E [d] \ [1], which is increasingly difficult when d becomes large. For example, if we increase the dimension of Z from 3 to 4 in Example E. and keep Ti. T2 and T3 as is, it could be verified that there does not exist a T4 : {1, 1)s - {-1, satisfying that F o T F for every & {2, 3]. We believe that this intuition could be rigorously proved using e.g., Lemma E.6 other techniques and leave it as future work.
E.7. Proof of Theorem 4.9
Proof. We first prove the following lemma:
Lemma E.13. Assume that the latent variables are uniformly sampled from the Hamming ball B, {= E {=1]"
#2) S r} with < d. Then, for every task h, we have
d (
deg(H (h)) S
log2
Σ 1=0
(63)
Proof of Lemma 13 The main idea of the proof is similar to that of Lemma E. and Lemma E.2.
Let l: = [log2 Ex= (-]]- Due to the invertibility of 4/s we know that [{ p( > 0}| = |B,| = Eiz (9)
Therefore, for any S = { ix } S [d] such that k' > k. there must exist bi...., ........ by 2--,, 1)a² such that
II (oi, + b,) JE|M]
II jES
2 +
Σ
bs II 2x ===0
(64)
S'CS
KES'
for every x E X' := {x p( (x)) > 0}. where by. E {-1, for every S' C S-if this does not hold, then we have
|x" 2 28 > 2* > Ei n (), which contradicts |X'| Ei (9) By equation (64), we can replace every degree-k'
monomial (s)x) = IL 2y by a degree-k' - 1 polynomial -sscss S1SK by Ilkes: Ty. Iteratively using this
replacement in the Fourier- Walsh transform of h gives the desired result.
We can now prove Theorem 4.9.
Proof of (i). By Lemma E.13, for every ht* Hminnh), we have deg(hi) k - -llog Ei (9). The test MSE of any h* thus satisfies
h(x)){1 err(h*) = E2- (h"(x) (x)))xx))) = (1,00 + -(x)==(2**** 2hh
(65)
+ B2-U( (1 (x)²² 2
2 (M(
(66)
1 2(h² w. h o v))
(67)
We then prove that k 2 r + 1. To see this, recall that d > r and one can verify
loga Σ 1=0 r
(1++)] = [loga (2r+1 - -1] +
k: == loga
2
(68)
1=0
Since deg(h | 1) 2 q - r, by Definition 3.6 we have deg(h") dog(g) 2 q - r for every g satisfying g 1--1 h² This gives deg(g) S deg(1) (h) -q k q + r, i... dhg(h" o 0b))= < k q + r. Recalling that h o is parity function xs with deg(h v) |S] k k-q+r and applying equation (19), have
(h" o vh h o v) = ho (S)
(69)
25

---

When Do Neural Networks Learn World Models?
Plugging equation (69) into (67) completes the proof.
Proof of (ii). Since deg(h & () 9 and & learns the world model up to negations and permutations, have deg(h Φ*) deg(h ل 2 r and hence deg(g") S By Lemma E.8, g" is unique. We thus have h = g t
which gives the desired result.
E.8. Proof of Theorem 4.12
Proof. We first prove the following lemma:
Lemma E.14. U is compatible, then (J) deg(s) holds every F
Proof of Lemma E. By Defintion 4.11, have (U deg(xs) for every compatible U. Applying Lemma E.9, further have (xg(U(xs) (xs) (0g(UU((xx)) This gives
degu (f) = max { deg(U- f(S) # o} = max {deg(xs): f(S) #0} deg(f),
(70)
which completes the proof.
We now ready prove Theorem 4.12. Note that under the new basis (U(s)) the original optimization problem (10) becomes (also applying the equivalence between (10) and (56)):
lim Σ degu (9t)
min
VE(N] 9 go H(h,), Vi E [n].
(71)
s.t.
Proof of (i). If is compatible, then by Lemma E. 14 we have dege(() (f) = deg(J) J for any Boolean function f. This immediately gives the equivalence between (71) and (56) and hence their minimizers. We thus have that 4* learns the world model up negations and permutations as Theorem 4.8.
Proof of (ii). If U is not compatible, then due to the invertibility of U, there exists at least one h" F such that degu(h' (h" oT-1) > deg(h") for some k < d. Since composing Boolean functions with degree- transforms does not change the degree of functions, we have deg (h" T ) > deg(ht" and hence ΣNeTY degu (h' o TT²)) > Every deg(h").
In particular, let (xi....,Xed be the set of all parity functions with d-dimensional inputs. For every k [d], we construct U such that:
1. UX22)) is permutation of x X2d;
2. For every E [d], we have U(x(a)) = Xs for some C [d] with [S] = and U(xs) = X(i).
Recall that degy (M/ o T-1) deg(U ( Therefore, to ensure that Σv degy (N' oT ENErY deg(h)
for k 1. we thus must have 177 Xs for some S S [d] and [d]. with |S| This means
max deg 1I;-1) 2 deg (T = k. (E(d)
(72)
which gives the desired result.
F. Experimental Details
This section presents additional experimental details, All of our experiments were conducted using Py Torch (Paszke et al., 2019) on NVIDIA V100/A 100 GPUs.
26

---

When Do Neural Networks Learn World Models?
(a) An example from the training distribution (sampled 5 frames with uniform spacing).
(b) An example from the test distribution with larger radius a larger initial velocity (sampled frames with uniform spacing).
Figure 6. Two visualized examples of the parabolic motion.
E.1 Polynomial Extrapolation
Dataset. We consider fitting and extrapolating degree-n polynomials with the form (a - a,z" Given an input a R. the label is given by y = (x). our experiments, we consider three families of polynomials with degree 1, 2, and In each family, every coefficient a, {0, 1, ..... n) uniformly sampled from [0. 1). In our experiments, we sample 50 polynomials each family for the violin plots. Other data parameters are follows:
Training, validation, and test data uniformly sampled from [-1, [-1,1), and (-, respectively.
For each polynomial instance, we sample 50, 000 training data, 1, 000 validation data, and 10, 000 test data.
Model and hyperparameters. We consider MIPs with the following architecture:
(z) w (wid-i) (...o (w x + +bw)) + -bld-1) + B(a)
(73)
where a is the activation function and for every i E [d] ( and ((() are weights and bias of the i-th layer, respectively. For ReLU MLPs, all activation functions are set to ReL.U; for our architecture, we replace half of ReLUS in every layer by the identity function (x) = æ and the quadratic function (zx) = z², with the number of identity functions and quadratic functions being the same (i.e., both functions constitute 25% activation functions, while the remaining 50% are still ReLUs).
We search the following hyperparameters for MLPs:
Number of layers d is set
Width of each ( from (128, 256, 512].
We train all MLPs with the mean square error (MSE) loss with the AdamW optimizer (Loshchilov & Hutter, 2019), Training hyperparameters are as follows:
Initial learning rate from {le 3, le le 5J. We a cosine learning rate scheduler.
Weight decay is set to 0.1.
Batch size is set 512.
Number of epochs set to 400.
Evaluation metric. We evaluate all models using MSE test data.
27

---

When Do Neural Networks Learn World Models?
(a) An example from the training distribution (sampled 5 frames with uniform spacing).
(b) An example from the test distribution with larger initial velocities (sampled 5 frames with uniform spacing).
Figure 7. Two visualized examples of the collision motion.
F22 Learning Physical Laws
Dataset. Inspired by Kang et al. (2025), create training and test sequences representing ball-shaped object movements that adhere to two basic physical laws: (i) single-object parabolic motion (reflecting Newton's 's second law of motion), and (ii) two-object elastic collision (reflecting the convervation energy and momentum). In both settings, we consider a 2-dimensional environment which each object is encoded by dirre-dimensional tuple (Ft, zt: Ms) at every time step t, where ra represents the radius of the ball and (zt, yt) stands for the 2-dimensional coordinates of the ball. To encode the timestamp information, also include dimension for the current timestamp f. This results in 4-dimensional inputs (i.e., (Te, Ft, Mr, f)) in the parabolic motion setting and 7-dimensional inputs (i.e., (rl #t. Mi rr, # M² L) in the collision motion setting. Each motion sequence consists of 32 frames with a timestep of 0.1. Details of both settings are as follows.
1. Parabolic motion. This motion describes the process where ball with initial horizontal (i.e., along the axis) velocity falls due to fixed gravity g 9.8 (along the y axis). We use the following training and test parameters:
Radius is uniformly sampled from [0.7, 1.5) in training and from [1.5, 2.0] in test.
Initial velocity is uniformly sampled from [1, 4) in training and from [4.5, 6.0) in test.
2. Collision motion. This motion describes the process where two balls with different sizes and different initial velocities
move horizontally towards each other and collide. We assume that the collision is perfectly elastic and all balls
with the same density, the velocities of both balls can be inferred from their radii and initial velocities. We use the following training and test parameters:
Radius each object is uniformly sampled from [0.7, 1.5] in training and from 1.05 test.
Initial velocity is uniformly sampled from [2, 4] in training and from [4.5, 6.0] in test.
The horizontal distance between two objects uniformly sampled from [5, 15] both training and test.
See Figure 6 and Figure 7 for visualized examples of the parabolic motion and examples of the collision motion.
For both settings, we sample IM training sequence and 50, 000 test sequence.
Model and hy pepraaaraeeers. We train a decoder-only transformer (Vaswani al., 2017) conditioned on the first frames to predict the remaining frames (expect for the timestamp dimension). For our method, simply replace every MI.P in the original transformer with modified MLP as in Section F.1 (replacing ReLU by GELU), We use teacher-foreing in training that is similar to next-token prediction, i.e., the model only needs predict the next frame (starting from the 4-th frame to the last 32-th frame) given all ground-truth frames prior to it. We use the following model hyperparameters:
Number of layers of transformer is set to 4.
Number of heads of transformer is set to 4.
Width of transformer is set to 512.
28

---

When Do Neural Networks Learn World Models?
0.05
0.04
Test MSE
0.02
0.01
0.00
ReLU MLP
Ours
Figure 8. Violin plots of the test MSE of the ReLU MLP and our model in extrapolating degree-1 polynomials.
We train all models using the MSE loss with the Adam W optimizer. Training hyperparameters are as follows:
Initial learning rate randomly sampled from [le 6, le 33. 3]. We use cosine learning rate scheduler.
Weight decay is set to le - 4.
Batch size is set to 1024.
Number of epochs is set to 300.
Evaluation metric. For test, we iteratively use the trained model to predict all missing frames given the first 3 frames. The predicted frames will be used together with the given frames for the model to predict the next frame. We evaluate all models using MSE on test data, averaged over all predicted 29 frames for each sequence.
G. Additional Results and Discussion
This section presents additional empirical results and discussion.
G.1. Polynomial Extrapolation
In the main text, we report extrapolation results degree-2 and degree-3 polynomials in Figure 2b; for completeness, here we also report extrapolation results degree-1 polynomials, i.e., linear functions. As shown by (Xu et al., 2021), ReLU MLPs can also extrapolate well in this setting. The violin plots of the test MSE of both the ReLU MLP and our model are in Figure 8. We can see that while both models achieve a much smaller extrapolation eπrоr compared to those in extrapolating higher degree polynomials, our model still outperforms the ReLU MLP. We also provide more examples for extrapolating degree-1, 2 and degree-3 polynomials Figure 9, Figure 10, and Figure 11, respectively.
Comparison with Xu al. (2021). While Xu et al. (2021) also show that 2-layer MLPs with quadratic activation functions extrapolate quadratic functions better than ReLU MLPs, we emphasize that there are two key differences between results and theirs:
1. Xu et al. (2021) replaces all ReLU functions with quadratic activation functions, while we only replace half of ReLU functions with quadratic activation functions and identity activation functions. This difference is important in scenarios where we do not know the exact structural form of target functions- note that our method can viewed as an "ensemble" of different activation functions, which enables the neural network to adaptively select activation functions that are the most compatible with the task as shown by our empirical results.
29

---

When Do Neural Networks Learn World Models?
2. Xu et al. (2021) only considers 2-layer MLPs for learning quadratic functions, while consider using 4-layer MLPs for learning degree-2 and degree-3 polynomials. This difference enables verify that neural networks can learn bases that require functiin compositions. For example, degree-3 polynomials need basis function y = x which cannot be composed using a 2-layer MLP but is composable using MLPs with more than two layers and with quadratic and identity activation functions. We also note that the inclusion of identity functions is important since MLPs with only quadratic activations can only represent basis functions y = x² with even degrees
G.2. Learning Physical Laws
We present visualization results of the transformer baseline and our model in Figure 12 in test collision motion example. We can see that while both models yield accurate predictions before the collision, our model outperforms the baseline on the predictions of the object velocities after the collision.
30

---

When Do Neural Networks Learn World Models?
Truhh
Ground
Gruns
(b)
(c)
Figure 9. Selected examples for degree-1 polynomial extrapolation. Shaded regions indicate training regions.
Truhh NLP
Rebi
Cura
Grouns Thtt Rec MLP
(a)
(b)
(c)
Figure 10. Selected examples for degree-2 polynomial extrapolation. Shaded regions indicate training regions.
Ground Tuth Nuw RLP
Ground Trl Rock vur Cur
Grouns Trts
nou YIP
(a)
(b)
(c)
Figure 11. Selected examples for degree-3 polynomial extrapolation. Shaded regions indicate training regions.
31

---

When Do Neural Networks Learn World Models?
(a) Ground truth.
(b) Prediction results of our model (MSE = 0.0039).
(c) Prediction results of the transformer baselinc (MSE 0.5090).
Figure 12, Visualization results a test collision motion example. All three rows seleet the same frames with uniform spacing.
32


Reviews:
[{"review_id": "FP0xJLa0XV", "reviewer": "Reviewer_PBCW", "summary": "The authors explore the theoretical and empirical conditions under which neural networks could learn world models. To do so, they abstract the possible task that a model would need to learn as a set of Boolean functions which can be approximated using polynomials over binary inputs. The authors claim that this is a suitable abstraction since numerical representations in computers are represented with limited precision anyway and it allows to use simple complexity metrics. They then proceed to prove several theorems on how models would behave: will they learn shallow representations or deeper ones? The end by showing how their approach (fitting polynomials) performs against Transformers.\n\n\n## update after rebuttal\n\nI've raise my score to a week accept. I like the general direction of the paper and their responses (and my own independent analysis) have cleared some things up but I feel like it is really hard to read at times. The authors try to be very precise with their approach which I understand, but it makes following the paper harder than it should be, in my opinion.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful review. We provide point-by-point responses to your concerns in the following.\n\n**Definition of the world model:** Our abstraction of the world model is motivated by a line of research in the machine learning literature. In brief, the fundamental aspect of our world model definition is _\"a model that learns a representation of high-level data-generating variables\"_ (Definition 2.2). We emphasize that many works in the existing literature share a similar definition as ours (e.g., [1-4] and many others; see our responses to _Reviewer L9Q8_ for further elaboration on this point) and thus feel that our definition is well-motivated and captures the main spirit of learning world models _in the context of machine learning_.\n\n  - We did _not_ claim that our formulation precisely captures how the world model works in the brain. As you have mentioned in the review, the main point of our abstraction is to model how high-level representations/abstractions of data-generating variables can be learned from low-level inputs---note that this is also the main motivation of studying world models in the machine learning context [1]. On the other hand, we agree that how cognitive-level representations in the brain emerge might be a much harder question and for now, remains a mystery.\n\n  - Our world model definition does _not_ require Boolean functions: Note that Definition 2.2 (and the whole Section 2 that presents our formulation) does not depend on Boolean functions at all. Boolean functions are simply used as a _mathematical tool_ to characterize the impact of _low-complexity bias_  on identifying data-generating variables, which is otherwise difficult to analyze in the real domain (note that this tool has also been used by other machine learning works [5, 6]). That being said, we believe that this treatment still captures many interesting machine learning settings and applications, e.g., representation learning, out-of-distribution generalization, self-supervised learning, etc.\n\n  - Your review mentions that \"If the input signals are so low-level and it is not clear if higher-level abstractions emerge from them...\". We are not sure if we understood this correctly, but our formulation does require that high-level variables _can_ be obtained by inverting the data-generating function, which we believe is a minimal assumption to ensure that high-level variables are learnable from the inputs.\n\n**Detailed comparisons with ICA and CRL:** Thank you for your question. ICA assumes that different coordinates of the latent data-generating vector $\\mathbf{z}$ are independent random variables, while CRL assumes that different coordinates of $\\mathbf{z}$ are governed by a causal graph and aims to identify the causal relations between different coordinates. By contrast, Definition 2.1 does not impose any distributional constraints on the coordinates of $\\mathbf{z}$, nor do our results rely on any of these assumptions. We will also add this discussion to the next version of our paper.\n\n**Multi-task learning:** The interpretation that pre-training objectives such as next-token prediction could be viewed as multi-task learning has been well-studied by prior work. Notably, the GPT-2 paper \"Language models are unsupervised _multitask_ learners\" [7] formulates next-token prediction as modeling the conditional distribution $p(output|input, task)$ with natural language as _task conditioning_. Task conditioning can also be implemented at an architectural level, i.e., using task-specific decoders as in our formulation.\n\n**Please let us know if you still have any questions or concerns. We are happy to engage in further discussion.**\n\n---\n\n[1] Ha and Schmidhuber. World models. NeurIPS, 2018.\n\n[2] Li et al. Emergent world representations: Exploring a sequence model trained on a synthetic task. ICLR, 2023.\n\n[3] Richens and Everitt. Robust agents learn causal world models. ICLR, 2024.\n\n[4] Gurnee and Tegmark. Language models represent space and time. ICLR, 2024.\n\n[5] Barak et al. Hidden progress in deep learning: SGD learns parities near the computational limit. NeurIPS, 2022.\n\n[6] Emmanuel et al. Generalization on the unseen, logic reasoning and degree curriculum. ICML, 2023.\n\n[7] Radford et al. Language models are unsupervised multitask learners. 2019.", "claims_and_evidence": "There are several claims which do not convince me: that their representation abstraction (using polynomials over parity functions) is a reasonable proxy for the kind of representations that world models would operate on; that their tests correctly mimic what a world models would need to process.\n\nThe first is not at all clear to me because while it is true that computers operate with binary functions, human cognition does not. Sure brains operate with electrical signals, but the representations at the cognitive level are more abstract. Indeed the point of world models is that such abstractions emerge from this more basic substrate. It is not clear to me how such an equivalent process happens in this setting.\n\nThe second is related to this. If the input signals are so low-level and it is not clear if higher level abstractions emerge from them, that how do we know that complex functions as defined here mimic what a world model needs to learn?", "methods_and_evaluation_criteria": "See above. The evaluations and methods make sense if we are evaluating the capability of networks to model functions using polynomials over parity functions, but I feel like equating these to world models is a stretch.", "theoretical_claims": "No.", "experimental_designs_or_analyses": "I checked them, they are fine.", "relation_to_broader_scientific_literature": "As expressed above, I am not sure the findings are as relevant as the authors claim to the problem of learning world models.", "essential_references_not_discussed": "No that I can recall.", "other_strengths_and_weaknesses": "The paper tends to be well written, though it can get bogged down in very strict formality. On the other hand as a I have said it is not clear to me that this is a good abstraction of what world models must solve and the experimental evaluation is a bit limited.", "other_comments_or_suggestions": "I would say that one way is to reframe this a problem of learning Boolean functions and not claim anything about world models.", "questions_for_authors": "It is not clear when defining the problem setting how is this any different from the standard ICA, CRL, etc settings. And the authors clarification that \"we do not assume that $p(z)$ is any structured distribution\" is just vague. Can you provide a better definition?\n\nIn lines 264-268 you say \"Indeed, modern pre-training objectives such as next-token prediction and contrastive learning can be interpreted as solving a large number of prediction tasks simultaneously (Radford et al., 2019; Arora et al., 2019).\" But is this really true? Is it not stretching the definition of task over its implied meaning? For me, these are just ways to learn a distribution, which may or may not correspond to different tasks. What separates different tasks from different instances of a task in this framework?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "igj2CA9myz", "reviewer": "Reviewer_i1Ec", "summary": "This paper provides a theoretical framework to reveal whether neural networks are capable of learning world models that capture the underlying data generation process. The framework connects the latent variable models with the world model learning. The theoretical results provide insights for future research.\n\n## update after rebuttal\nThanks to the authors for providing the rebuttal. I've read the author's response and comments from other reviewers. I have no further questions at this time. I will keep my original positive rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and your approval of our work. The following are our responses to your questions.\n\n**Applications:** Thank you for your questions on potential algorithmic applications of our results. Yes, we believe that our results may inspire future algorithmic research in the following ways.\n - The main message of Theorem 4.3 is that the conditional degree of proxy tasks is important for representation learning. This may inspire the development of self-supervised representation learning objectives by maximizing their conditional degree (e.g., pixel-level prediction may not be a good objective since it potentially has a small conditional degree), which could be applied to designing auxiliary/proxy tasks in model-based RL, training video generation models, and language modeling. We are now exploring this idea further in the context of training language models by augmenting next-token prediction with auxiliary training objectives.\n - The notion of basis compatibility shows that accommodating the inductive biases of the proxy tasks in the network architecture can benefit world model learning, as demonstrated in our proof-of-concept experiments. This may inspire future model design for learning and extrapolating physical laws in a larger context. For example, for video generation models to robustly capture periodicity in real-world physics, we may need to explicitly embed modules with a natural periodicity behavior, rather than merely hoping that models with standard architectures could learn such behaviors from data.\n- More broadly, an important takeaway of our work is that \"while neural networks typically exhibit 'simplicity biases', the underlying meaning of 'simplicity' is not fixed and is _dependent_ on the network architecture\". To our knowledge, this is an underexplored aspect in the existing literature on the simplicity bias of neural networks. This could also explain why the simplicity bias tends to be sometimes beneficial and sometimes harmful in different tasks. Moving forward, we believe that designing network architectures that could _adaptively_ fit such simplicity biases to the \"right\" inductive biases of training tasks is a promising direction. We are currently working on a project on this topic.\n\nWe will also include more discussion in the next version of our paper.\n\n**Please let us know if you still have any questions or concerns. We are happy to engage in further discussion.**", "claims_and_evidence": "The theoretical framework and results support the paper's claims.", "methods_and_evaluation_criteria": "This paper mainly proposes a theoretical framework and experiments on polynomial extrapolation and  learning physical\nlaws.", "theoretical_claims": "I have checked the main theoretical claims, but not in every detail.", "experimental_designs_or_analyses": "The experimental results support the claims.", "relation_to_broader_scientific_literature": "This paper is among the first to propose a theoretical framework to study the problem of when neural networks learn world models.", "essential_references_not_discussed": "To the best of my knowledge, the references are sufficiently covered.", "other_strengths_and_weaknesses": "[Strengths]\n1. The overall paper is solid and well-structured.\n2. The theoretical results are comprehensive and insightful.\n3. In addition, the experimental results are provided in this paper to further verify the claims of this paper and help readers understand the proposed concepts more easily.\n\nOverall, this paper is solid and provides insightful observations/proof for future research in the areas of physical AI, robotics, autonomous driving, etc. Despite some questions listed in the following Questions For Authors section, no major technical concerns are found at this time.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1. Whether the derived insights could be applied to model-based RL tasks?\n2. Considering that video generation-based world models play crucial roles in robotics and autonomous driving, I'm wondering whether the derived insights could benefit the design of large video world models by allowing them to capture real-world physics from data.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BElr237nnw", "reviewer": "Reviewer_GtGz", "summary": "this paper aims to answer the question \"when do neural nets learn world models\". This is an ambitious goal that has been studied in many papers. the paper proposes a theoretical results for this problem. then the paper also shows the algorithmic implications of the results.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review. We are happy to engage in further discussion if you have additional questions or concerns in the author-reviewer discussion stage.\n\nOn the point of causality, we agree that causality is also an important aspect of learning world models. Moreover, we would like to emphasize that a well-known difficulty in causal discovery is the identification of causal latent variables when the observed data is generated from them in a _non-linear_ fashion (see e.g., [1]), for which we show that such identification could be hoped if the model has a low-complexity bias. Therefore, combining our results and the results in causal discovery is interesting future work.\n\n---\n\n[1] Schölkopf et al. Toward causal representation learning. Proceedings of the IEEE, 2021.", "claims_and_evidence": "yes", "methods_and_evaluation_criteria": "yes", "theoretical_claims": "yes", "experimental_designs_or_analyses": "yes", "relation_to_broader_scientific_literature": "i believe this paper will have some impact to a sizable audience in the ML community", "essential_references_not_discussed": "references are adequate", "other_strengths_and_weaknesses": "i believe causality is important in world models, but is not explored in this paper.\n\nI'm not very familiar with the content in this paper so I'll reevaluate the paper after reading the comments from the other reviewers.", "other_comments_or_suggestions": "n/a", "questions_for_authors": "please see the comments above", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DzoRGG3tHU", "reviewer": "Reviewer_L9Q8", "summary": "The paper explores whether neural networks can develop internal world models akin to those humans create to understand the underlying data generation processes. The research presents the first theoretical results in this area, demonstrating that in a multi-task setting, models with low-degree bias can recover latent data-generating variables under mild assumptions. This recovery, however, is sensitive to model architecture. Using Boolean models of task solutions via the Fourier-Walsh transform, the researchers introduce new techniques to analyze invertible Boolean transforms. They show that the implicit bias in deep learning towards low-complexity functions can favor certain solutions over others. Their findings have implications for self-supervised learning, out-of-distribution generalization, and the linear representation hypothesis in large language models. The paper provides a formal framework for world model learning and presents theoretical results, identifying key factors such as multi-task settings and low-complexity bias necessary for the successful learning of world models. The research also illustrates algorithmic implications through tasks like polynomial extrapolation and learning physical laws, showing that architectures informed by their analysis can outperform conventional ones.", "questions": "", "limitations": "", "rating": 5, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your detailed review and your approval of our work. We provide point-by-point responses to your questions in the following:\n\n**Synthetic experiments:** Thank you for these questions; we believe that they greatly help to improve our work. Inspired by them, we conducted additional synthetic experiments using Boolean inputs and MLPs to further support our theoretical claims. _Please find our experimental results in this anonymous link: https://anonymous.4open.science/r/rebuttal_figures-8C6A/rebuttal_figures.pdf._ (Backup links: https://ibb.co/mrdwJt9T and https://ibb.co/gNX4FYT if the former link is not available). A brief introduction on how these results answer your questions is as follows:\n\n  - **Low-degree bias of neural networks (Q1):** _Figure 1_ shows that MLPs indeed have a low-degree bias when solutions with different degrees exist for the same training distribution.\n\n  - **Multi-task learning and the identification of data-generating variables (Q2 & Q3, Theorem 4.3 & Theorem 4.8):** _Figure 2_ shows that the identification error of data-generating variables significantly decreases as more training tasks are drawn, while single-task learning does not lead to similar identification.\n\n  - **Identifying data-generating variables benefits OOD generalization (Q3, Theorem 4.9):** _Figure 3_ shows that in a Hamming sampling setting as in Theorem 4.9, MLPs with a representation $\\Phi$ that identifies true data-generating variables achieve superior OOD generalization compared to MLPs without such a representation.\n\nWe will add these experiments to the next version of our paper.\n\n**Definition of the world model:** We believe that defining the world model as a model that learns data-generating variables, although not formally established as a standard term, is well-motivated by the existing literature:\n\n  - In the first work that introduces the term \"world model\" in the machine learning context [1], the authors describe the world model as learning \"an abstract representation of both spatial and temporal aspects\" of the information flow, which naturally motivates a representation learning formulation. The authors of [1] also experimentally demonstrate this concept by training an RNN-based _generation_ model on top of a learned _latent vector_, which involves latent data generation models. Thus, our definition can be viewed as a natural abstraction of their model.\n\n  - Many existing works explicitly use very similar definitions as ours: For example, the authors of [2] wrote \"An alternative hypothesis is that LLMs, in the course of compressing the data, learn more compact, coherent, and interpretable models of the _generative process_ underlying the training data, i.e., a world model\", while [3] defines a \"causal world model\" as \"a causal model of the _data generating process_\".\n\n  - A series of empirical works (e.g., [4-8]) also share a similar spirit: [4-6] treat world models as \"understandable models of the process producing the sequences they (language models) are trained on\" [4]; empirically, they test if transformers learn world models/world representations by connecting the models' internal representations to environment states that generate the training sequence. [7] adopts a similar approach to the task of generating program traces. [8] describes \"linear world models\" as models with high-level features linearly decodable from their latent spaces (see also our discussion \"Connection to the linear representation hypothesis\" in Section 4.2).\n\nGiven the above discussion, we feel that our current title is consistent with the related literature. We will also include more discussion in the next version of our paper.\n\n**Please let us know if you still have any questions or concerns. We are happy to engage in further discussion.**\n\n---\n\n[1] Ha and Schmidhuber. World models. NeurIPS, 2018.\n\n[2] Gurnee and Tegmark. Language models represent space and time. ICLR, 2024.\n\n[3] Richens and Everitt. Robust agents learn causal world models. ICLR, 2024.\n\n[4] Li et al. Emergent world representations: Exploring a sequence model trained on a synthetic task. ICLR, 2023.\n\n[5] Nanda et al. Emergent linear representations in world models of self-supervised sequence models. EMNLP, 2023.\n\n[6] Toshniwal et al. Learning chess blindfolded: Evaluating language models on state tracking. AAAI, 2022.\n\n[7] Jin and Rinard. Emergent representations of program semantics in language models trained on programs. ICML, 2024.\n\n[8] Marks and Tegmark. The geometry of truth: Emergent linear structure in large language model representations of true/false datasets. COLM, 2024.", "claims_and_evidence": "The claims are well supported by the theorems. They could be better supported if the authors added some experiments.", "methods_and_evaluation_criteria": "No proposed method", "theoretical_claims": "I haven't checked every detail of the proof due to time constraints, but I believe the logic of the proof is correct.", "experimental_designs_or_analyses": "I think additional experiments should be added to better support the main claim.", "relation_to_broader_scientific_literature": "This is pioneering work in studying the data generation model when learned by neural networks. Most assumptions are realistic and facilitate the proof of the results. Theorems 4.3 and 4.8 are novel to me and are practically important. Therefore, I believe this paper will make a significant impact in this field in the following ways: the assumptions can be reused for other theoretical studies, and the conclusions can inspire future empirical work.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "**Strengths:**\n\n1. This is pioneering work in studying the data generation model when learned by neural networks.\n\n2. Most assumptions are realistic and facilitate the proof of the results.\n\n3. Theorems 4.3 and 4.8 are novel to me and are of practical importance.\n\n**Weaknesses:**\n\n1. The major weakness is that the experiments are not strongly related to the theorems. I would improve the score to \"strong accept\" if this point is addressed. (Please refer to Questions 1, 2, 3.)\n\n2. \"World model\" in the title is defined as the data generation model. This does not seem to be a widely accepted term. The authors should support this claim with more existing works or consider changing the title.", "other_comments_or_suggestions": "No", "questions_for_authors": "1. A major assumption is that the neural network will learn the function with minimal (realization) degree (e.g., (ii) in Theorem 4.8). The authors support this assumption with literature on \"simplicity bias,\" but I believe that additional support from synthetic experiments is important.\n\n2. Theorems 4.1 and 4.3 state that the data generation model $\\Phi$ emerges if we adopt a multi-task learning scheme. Could the authors support this claim with synthetic experiments?\n\n3. Could Theorems 4.8 and 4.9 be supported with synthetic experiments?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper presents a novel theoretical framework for understanding when neural networks can learn "world models"—internal representations of latent data-generating variables. By combining a low-degree bias in Boolean transforms with multi-task learning, the authors prove identifiability results under mild assumptions. The work is well-motivated, connects to broader themes like out-of-distribution generalization and representation learning, and includes proof-of-concept experiments that demonstrate architectural implications.

While one reviewer questioned the abstraction’s connection to practical world models, the authors addressed this satisfactorily in the rebuttal. The reviewers generally appreciated the paper's theoretical depth and relevance. Overall, the paper provides foundational insights and is a valuable contribution to the theory of representation learning.
Decision: Accept (poster)