SAI
← All ICML 2026 orals

Learning to Theorize the World from Observation

Doojin Baek, Gyubin Lee, Junyeob Baek, Hosung Lee, Sungjin Ahn

OralReplication not startedPaper PDFOpenReview

Learning to Theorize the World from Observation

SAI paper + code review · Referee report

Summary

The paper proposes Learning-to-Theorize (L2T) — a paradigm that frames the goal of world modeling not as accurate future prediction but as inferring an executable programmatic theory that explains how a source observation transforms into a target one. The instantiation, Neural Theorizer (NEO), factors this map through a pretrained encoder/decoder, a VQ-VAE "theory programmer" that iteratively selects discrete primitives from a learned codebook, a shared transition model that executes them, an MDL-style rule that adapts program length per instance, and a decode–encode "state grounding" cycle that keeps intermediate latents on the encoder manifold. The Observation-to-Theory Induction Benchmark (OTIB) — GridWorld motions, arithmetic factorization, and CIFAR-10 image edits, each with an α\alpha-controlled compositional split and a length-OOD split — evaluates whether the induced program transfers: a theory extracted from a support pair must generalize to a fresh query pair drawn from the same latent program.

The conceptual move is worth taking seriously. Positioning developmental theory-theory arguments and neural program induction next to observation-driven world modeling is a clean idea, and the support/query transfer protocol cleanly separates instance memorization from acquisition of a reusable mechanism. Empirically, the compositional-vs-monolithic gap is convincing: monolithic latent-action baselines transfer to compositional OOD at near-zero accuracy while NEO retains most of its in-distribution transfer; the code–primitive alignment analyses show that NEO recovers primitives (single motions in GridWorld, single multiplications in arithmetic, missing image edits) that are never observed in isolation; and the grounding and λMDL\lambda_{\mathrm{MDL}} ablations tie the mechanism to a specific design choice.

Where the paper is unevenly supported is on its headline claims. Plain NEO's Arithmetic length-OOD transfer is 0.019–0.045 across α\alpha — essentially failure — and the strong 0.5–0.7 numbers come from NEO-S with B=1024B=1024 test-time samples, a variant introduced almost in passing. Plain NEO also trails Disc-Mono in-distribution. The encoder/decoder are pretrained rather than learned end-to-end, softening the "from raw observations" framing. The formal test-time inference (an argmax over the ground-truth test set), an inverted branch and an off-by-one in the pseudocode, a λMDL\lambda_{\mathrm{MDL}} sign convention that contradicts the "exponential penalty" prose, and several missing displayed equations mean that a careful reader cannot reconstruct the exact procedure from the text. All benchmarks are synthetic and no code repository is announced. The contribution is real, but the framing and the technical presentation both need tightening.

Strengths

  • Conceptual contribution. Reframing world modeling as inference of a transferable executable theory (support→query program transfer) rather than next-frame prediction is a clean and defensible objective that connects developmental theory-theory to neural program induction.
  • Compositional-vs-monolithic gap is convincing. On compositional OOD in every domain, monolithic latent-action baselines collapse to near-zero transfer, while NEO retains most of its in-distribution performance, giving a clear signal that the compositional inductive bias matters.
  • Primitive recovery from entangled compositions. The code–primitive alignment analyses show NEO recovers primitives that never appear in isolation at training time — e.g. single L/R/D moves at α=0.33\alpha=0.33 in GridWorld, single ×3\times 3 / ×5\times 5 in arithmetic, and multiple missing image edits at α=0.33\alpha=0.33/0.660.66.
  • Mechanism-tied ablations. Removing state grounding causes primitiveness to collapse to 0.002 and every transfer score to zero; sweeping λMDL\lambda_{\mathrm{MDL}} tracks program length as Figure 8 shows. Both results tie the empirical behavior to a specific design choice rather than a generic architectural win.
  • Test-time scaling story. NEO-S demonstrates that the compositional structure enables usable search — accuracy monotonically improves with BB on GridWorld and lifts Arithmetic length-OOD from 0.02 to 0.7 — while monolithic baselines remain flat.
  • Benchmark honesty. OTIB deliberately holds out compositions and enforces support/query separation, so the reported gaps cannot be trivially explained by memorization of individual (x,y)(x, y) pairs.

Weaknesses

  • Length generalization is a property of NEO-S, not NEO. On Arithmetic length-OOD, plain NEO reaches 0.019–0.045 across α\alpha — essentially failure — and only NEO-S (test-time sampling with B=1024B=1024) attains the 0.5–0.7 range. The introduction lists length generalization as a headline contribution of NEO; the empirical evidence supports this only in GridWorld.
  • NEO underperforms baselines in-distribution. At α=0.33\alpha=0.33 NEO scores 0.914/0.911 while Disc-Mono is 0.988/0.983 and Cont-Mono-Opt is 0.994. The paper does not discuss why compositional structure trails a monolithic VQ-VAE on the training distribution or why NEO-S closes the gap.
  • Formal test-time inference is circular. Section 2 defines τ^=argmaxτTtestp(τx,y)\hat\tau = \arg\max_{\tau \in \mathcal{T}_{\mathrm{test}}} p(\tau\mid x, y) with τTtrain\tau \notin \mathcal{T}_{\mathrm{train}}, restricting the argmax to the held-out test set — an object the model has no access to. This does not match Section 3.5's iterative-primitive-plus-threshold procedure.
  • Algorithm 3 branching is inverted. Line 15 defines SS as successful support-reconstructions; the algorithm then runs majority voting when S=S = \emptyset (arg max over the empty set) and returns failure when SS \neq \emptyset. This is inverted and directly contradicts the prose in Section 5.6.
  • MDL length-selection indexing differs between algorithms. Algorithm 1 scores at D_θ(s_k) while Algorithms 2 and 3 score at D_θ(s_{k+1}). Since kk transitions produce sk+1s_{k+1}, one of these is wrong.
  • λ_MDL sign convention contradicts the prose. Section 3.3 says λMDL>1\lambda_{\mathrm{MDL}} > 1 penalizes longer programs exponentially; GridWorld tables use λMDL=0.95\lambda_{\mathrm{MDL}} = 0.95 and 1.001.00. Under the arg-min-of-λk\lambda^k \cdot \ell rule, λ<1\lambda < 1 makes the length factor decrease in kk, i.e. rewards length among equally-fitting prefixes — the opposite of "prefer shorter correct programs."
  • q_ϕ conflated between categorical and continuous. Section 3.2 defines qϕq_\phi as a categorical distribution over MM' primitive categories, while Algorithm 3 line 6 treats it as producing a continuous vector whose L2 distance to codebook entries forms logits.
  • Encoder/decoder pretraining softens "raw observations." The perceptual grounding is fixed before NEO sees anything; the transition model learns primitives inside a pretrained latent geometry. Whether primitive recovery depends on that geometry is untested.
  • Baselines are all monolithic strawmen. Related work correctly positions L2T against NPI, LEAPS, HPRL, LPN, DreamCoder, and differentiable interpreters, but no such baseline is included, so the reported gap tests only the "no program structure" hypothesis.
  • Identifiability asserted, not established. "Training compositions alone contain sufficient evidence, in principle, to recover the full primitive set via such decomposition" is a load-bearing assumption with no condition or proof.
  • Metric direction of d flips silently. Section 2's dobsd_\text{obs} is an error (lower better); Tables 1 and 2 bold the highest cells (higher better); Figure 3 reverts to lower-better 1\ell_1. The same word ("self-explainability", "transferability") therefore refers to opposite-direction quantities across domains.
  • Missing displayed equations. The Cont-Mono ELBO, the Cont-Mono-Opt update, the MDL selection rule (Eq. 5), and the state grounding loss are all announced but not rendered.
  • Only synthetic benchmarks. Four grid motions, four multiplicative operations, and eight CIFAR-10 edits, all with hand-designed primitives and short compositions. The Limitations section is honest about this, but the framing invites over-generalization.
  • Statistical rigor. All main tables report only 3-seed means without per-cell variance; NEO-S's sampling temperature for the reported cells is not tabulated, even though Figure 20 shows strong temperature sensitivity on Arithmetic.
  • Figure 1 uses primitives (Rotate, Left, Down, Paint) that appear in no benchmark. A reader is likely to conflate the figure with GridWorld and expect a Paint operator.
  • Presentation glitches. "NEOS" vs "NEO-S" appears in the same sentence group; Figure 6(b) caption and body text disagree on which colour is argmax vs sampled; "the 3 repeated-multiplication sequences" is followed by a list of four.

Reproducibility & code

Veritas inspected the released code/data package without executing anything. The manuscript does not include an explicit code-availability statement, and the inspected package does not contain the training or evaluation scripts needed to reproduce Tables 1, 2, 11, Figures 3, 6, 7, 8, 20, or the ablations in Table 3 — so on the current release, reproduction has to work from the text alone. The appendix is unusually thorough (Tables 18–27 give per-domain hyperparameters; Appendix C.3/C.5/C.6 enumerate anchor and OOD compositions), but several parameters that materially affect the reported cells are missing.

  • Pseudocode contains a load-bearing bug. Algorithm 3's success/failure branch is inverted; a re-implementer would silently reproduce broken NEO-S selection unless they notice.
  • Off-by-one in MDL length selection. Algorithm 1 uses D_θ(s_k); Algorithms 2/3 use D_θ(s_{k+1}). The correct convention needs to be settled before Table 1/2/3 numbers can be matched.
  • Central equations missing. The MDL rule (Eq. 5), the grounding loss, the Cont-Mono ELBO, and the Cont-Mono-Opt update are all referenced by their surrounding "where …" clauses but not displayed in the rendered text.
  • Seeds not documented. Every table reports "mean across three runs" without seed values.
  • NEO-S sampling temperature for reported cells not fixed. Figure 20 shows dramatic temperature sensitivity on Arithmetic length-OOD, but the temperature used for the tabled NEO-S cells (Tables 1, 2) is not stated, nor is it clear whether it was tuned on the OOD split.
  • Image-editing pixel-difference threshold unspecified. This threshold controls which samples survive dataset construction and therefore changes every dataset size in Tables 15–17, the identity-baseline L1 values, and every reported NEO L1 number.
  • MDL scheduling only partially specified. The Arithmetic and Image-Editing schedules give start/end and a "scheduling ratio", but the horizon (fraction of steps or epochs) is ambiguous, and Image Editing α=1.0\alpha=1.0 disables scheduling with no rationale — leaving readers unsure whether the λ\lambda crosses 1 during training as the "first stabilize, then compose" mechanism requires.
  • CNN encoder/decoder topology missing. Only hidden and feed-forward dimensions are documented; layer count, kernel sizes, strides, and non-linearities are not, which matters because primitive recovery is measured inside the pretrained latent geometry.

Recommended Changes

Essential

  • Reframe the length-generalization claim. In the introduction and abstract, attribute Arithmetic length OOD to NEO-S (compositional structure plus search), not to NEO alone. Restrict "length generalization" claims for plain NEO to GridWorld, and either add a plain-NEO length-OOD number to the abstract or drop the plain-NEO framing there.
  • Rewrite the formal test-time inference. Replace τ^=argmaxτTtestp(τx,y)\hat\tau = \arg\max_{\tau \in \mathcal{T}_{\mathrm{test}}} p(\tau\mid x, y) with the actual procedure (iterative primitive selection over the space of learned compositions, terminated by the ε\varepsilon threshold or MDL rule), and remove the hard τTtrain\tau \notin \mathcal{T}_{\mathrm{train}} constraint — that is a property of the evaluation set, not of inference.
  • Fix Algorithm 3. Reverse the branching so majority voting runs on SS \neq \emptyset and failure returns when S=S = \emptyset. Also make the notation for qϕq_\phi inside the algorithm consistent with Section 3.2 (either qϕq_\phi outputs a categorical distribution and the softmax over codebook distances is a separate step, or reformulate Section 3.2 to describe qϕq_\phi as producing an embedding).
  • Fix the MDL-index off-by-one. Pick one convention (D_θ(s_{k+1})) and use it in Algorithm 1, Algorithm 2, Algorithm 3, and the Section 3.3 prose.
  • Reconcile the λMDL\lambda_{\mathrm{MDL}} sign convention. Either state that the length term is λMDLk\lambda_{\mathrm{MDL}}^k under a minimization so λ<1\lambda < 1 rewards length (matching Fig. 8 and Table 19), or invert the operational values. Remove "λMDL>1\lambda_{\mathrm{MDL}} > 1 penalizes longer programs exponentially" until it matches the tables.
  • Release code, seeds, and dataset generators. Include NEO-S sampling temperatures for every reported cell, the image-editing pixel-difference threshold, the exact MDL schedule endpoints and horizons, and full CNN topology. Provide a runnable reproduction of Tables 1, 2, 3, 11.
  • Add per-cell variance (std or bootstrap CI) across the three seeds in Tables 1, 2, 3, 11 so borderline comparisons (e.g. NEO-S transfer across α\alpha) are interpretable.
  • Restore missing equations. Display the MDL selection rule, the state grounding loss, the Cont-Mono ELBO, and the Cont-Mono-Opt gradient-ascent update in the main text.

Suggested

  • Add a compositional / program-induction baseline (e.g. LEAPS, LPN, or a simplified DreamCoder-style search) so the reported gap is not solely against monolithic latent-action models.
  • Add an ablation with an untrained (or jointly-trained) encoder/decoder to test whether the primitive-recovery result depends on the pretrained latent geometry, since the "from raw observations" framing implicitly promises this.
  • Address identifiability. State the condition on α\alpha-sampling under which primitives are guaranteed to be recoverable from entangled compositions, or reframe the "in principle" claim as an empirical observation supported by the primitiveness experiments.
  • Fix "the 3 repeated-multiplication sequences" to "the 4", since four are listed and both Table 6 and the primitive-observability logic depend on all four being anchors.
  • Rename the tabled metric to something unambiguously higher-is-better (e.g. "success rate") or reserve d for the error case; reconcile with the 1\ell_1 direction used in Figure 3.
  • Justify or unify the train-time MDL rule and the test-time ε\varepsilon threshold (they currently use different criteria).
  • Redraw Figure 1 with primitives that appear in an actual benchmark, or state clearly that it is schematic.
  • Fix Figure 6(b) colour legend so text and caption agree, and correct "selctions" → "selections".
  • Replace all "NEOS" with "NEO-S".
  • Note in the main text that all experiments use the deterministic special case of the transition operator, since the general stochastic formulation is never tested.
  • Add a short mechanistic analysis of state grounding (e.g. what intermediate s_k look like without grounding) so the ablation is not just a "flip this off and it collapses" result.
  • Extend at least one benchmark to a non-toy perceptual manifold (natural video with distractors, or a stochastic-dynamics variant) as a first step toward the real-world extensions the framing suggests.