Technical Summary: Learning to Theorize the World from Observation

1. Problem Statement

Contemporary AI world models primarily operationalize "understanding" as the ability to accurately predict future states in latent or observation spaces. However, developmental cognitive science suggests that human understanding emerges not merely from prediction, but from constructing internal, reusable theories of how the world works, often before language acquisition. This "theory-theory" posits that cognition involves discovering structured, compositional explanations for observed phenomena.

A critical gap exists in current AI systems: they are optimized for prediction or reconstruction, often learning entangled composite transformations that capture correlations but lack the compositional structure necessary for systematic generalization. These models fail to discover explicit, reusable mechanisms that explain how and why observations are generated, making them brittle under distribution shifts or when faced with novel compositions of known factors.

The central problem addressed is: Can an artificial system learn to construct explicit explanatory theories of the world solely by observing raw, non-textual sensory inputs, without relying on language, task labels, or ground-truth program supervision?

2. Methodology: Learning-to-Theorize (L2T) and NEO

The authors propose Learning-to-Theorize (L2T), a learning paradigm where models infer executable, compositional programs (theories) from raw observation pairs $(x, y)$. A theory is defined as a sequence of reusable primitive operations that transform a source observation $x$ into a target $y$.

To instantiate L2T, the paper introduces the Neural Theorizer (NEO), a probabilistic neural model with the following core components:

2.1. Probabilistic Framework

NEO models the conditional likelihood $p_\theta(y | x)$ by introducing latent variables: a discrete program $\tau$ (a sequence of primitives) and an execution trace $s$. The model assumes a generative process where an encoder maps $x$ to an initial latent state, a "Theory Programmer" selects primitives sequentially, and a shared transition model executes these primitives to generate a final state, which is decoded to reconstruct $y$.

2.2. Learned Language of Thought

Instead of using a pre-defined symbolic language, NEO learns a Language of Thought directly from data.

• Primitives: The model induces a discrete vocabulary of operations via a Vector-Quantized VQ-VAE codebook. These primitives have no predefined semantics; their meaning is derived from how they transform latent states.
• Composition: The theory programmer $q_\phi(z_{ik} | s_k, y)$ acts as a goal-conditioned policy, selecting the next primitive to steer the latent state toward the target $y$. This allows the model to compose primitives into structured programs without explicit supervision.

2.3. Minimum Description Length (MDL) Principle

To prevent over-decomposition (creating unnecessarily long programs) and encourage the discovery of the most reusable set of primitives, NEO employs the MDL principle. The model selects an optimal program length $k^*$ that balances reconstruction accuracy with program brevity:
$$k^* = \arg \min_{k} \lambda_{MDL}^k \cdot \ell(y, \hat{y}_k)$$
This mechanism incentivizes the model to find the shortest accurate explanation, naturally favoring the discovery of fundamental primitives over memorizing entangled sequences.

2.4. State Grounding

A key innovation is State Grounding. To ensure intermediate latent states correspond to valid observations (preventing the model from learning "shortcut" transitions that drift off the data manifold), NEO enforces a decode-encode consistency loss on intermediate states. This ensures that every step in the inferred program is a meaningful transformation within the learned latent space.

2.5. Inference and Test-Time Scaling

At test time, NEO infers a theory by iteratively selecting primitives until the reconstruction error falls below a threshold. Because the model learns a compositional structure, it supports length generalization (handling programs longer than seen in training). Furthermore, the probabilistic nature of the theory programmer enables test-time search (NEO-S), where multiple candidate theories are sampled and selected via majority voting to improve accuracy on complex, unseen compositions.

3. Evaluation: Observation-to-Theory Induction Benchmark (OTIB)

The authors introduce OTIB, a benchmark designed to evaluate whether models can infer reusable theories from raw observations without program supervision.

• Setup: Training data consists of observation pairs generated by latent programs. The training set includes a fraction $\alpha$ of possible program compositions, while the rest are held out as Compositional Out-of-Distribution (OOD).
• Metrics: Performance is measured by Self-Explainability (reconstructing the target from the support pair) and Transferability (applying the inferred theory to a new query input $x'$ to predict $y'$). Transferability is the critical metric for determining if a reusable theory was learned versus an instance-specific mapping.
• Domains: The benchmark is instantiated in three domains:
  1. GridWorld: 10x10 grid with motion primitives.
  2. Arithmetic Factorization: Integer transformations via multiplication by primes.
  3. Image Editing: CIFAR-10 images transformed by operations like rotation, brightness, and masking.

4. Key Results

Experiments demonstrate that NEO significantly outperforms monolithic baselines (Discrete/Continuous VQ-VAEs and $\beta$-VAEs) across all settings:

• Compositional Generalization: While baselines achieve high in-distribution performance, they collapse to near-zero transferability on Compositional OOD tasks. NEO maintains high transferability (e.g., ~0.93 on GridWorld OOD), proving it learns reusable primitives rather than memorizing input-output mappings.
• Length Generalization: NEO successfully generalizes to programs longer than those observed during training (e.g., length 8 in GridWorld, length 6 in Arithmetic), a capability where baselines fail completely.
• Primitive Discovery: Analysis shows NEO recovers the ground-truth primitive set even when primitives are never observed in isolation during training (low $\alpha$). The model successfully decomposes entangled training examples into atomic, reusable units.
• Test-Time Scaling: The NEO-S variant, which utilizes sampling-based search at inference, further boosts performance, particularly on long-horizon tasks, confirming that the compositional structure enables effective exploration of the program space.

5. Significance and Claims

The paper positions itself as a proof of concept for a shift in AI learning objectives: moving from prediction-centric world models to explanation-driven, compositional understanding.

• Theoretical Contribution: It validates the hypothesis that structured, compositional theories can be induced from raw, non-textual observations without symbolic supervision, aligning AI learning objectives with the "baby as a scientist" perspective from developmental cognitive science.
• Methodological Contribution: It demonstrates that representing theories as executable programs with learned primitives and grounding intermediate states enables systematic generalization that monolithic latent representations cannot achieve.
• Modest Scope: The authors explicitly acknowledge limitations. The current formulation assumes small, discrete primitive sets and short program lengths, limiting scalability to highly complex, continuous, or stochastic real-world dynamics. The experiments are restricted to controlled synthetic benchmarks. The work is presented as a foundational study of learning mechanisms rather than a deployment-ready system.

In conclusion, the paper argues that by framing learning as "Learning-to-Theorize," AI systems can acquire internal explanatory structures that are reusable, compositional, and capable of generalizing to novel phenomena, thereby moving closer to a richer notion of machine understanding.