Compositionality and Systematic Generalization › Implementation: Factorized Versus Monolithic

We will now build both models in PyTorch, train them on identical data from the six training worlds, and compare fixed-action autoregressive rollouts on all nine. The monolithic model receives the world's identity as a single index and uses one shared prediction head. The factorized model receives the two mechanism indices separately and computes its prediction as the sum of two mechanism-specific heads. The implemented models do not have matched parameter counts: the monolithic model has 5,001 trainable parameters and the factorized model has 9,906. The experiment therefore illustrates the two constructions under matched data and optimization, with model capacity left as a confound; it is not a capacity-controlled causal estimate of structure alone.

Before reading the code, it helps to know exactly what the comparison is and is not. It is a single-seed, single-split pedagogical illustration under matched data and optimization but unmatched parameter count. It does not satisfy the minimum serious-evaluation protocol stated above: it has no repeated seeds or splits, no planner that selects actions, and no adaptation study. It shows how this particular untrained world-index representation behaves on cells the data never touched; capacity-matched controls and the full evaluation suite would be required to isolate structure's causal contribution or make a deployment claim.

First, import the remaining libraries and define the ground-truth world. The mechanism tables store the parameters for each variant, and a single rollout function composes them exactly as the equation above specifies.

Note how the simulator is written: the mechanism tables are looked up independently and their contributions are summed. The ground-truth generator is itself factorized, which means the swap test could be run exactly by exchanging rows of the two tables. This is one of the advantages of a synthetic study: every claim about mechanism identity is verifiable by construction, whereas in a real domain the same claims would require the diagnostics from the previous section.

Now the dataset. For each of the nine worlds we generate trajectories driven by random actions. The training split is six worlds; the anti-diagonal is held out. Each factor value still appears in multiple training worlds, which is exactly the primitive-coverage condition from the previous section.

The training pool holds 4,800 transitions (six worlds, eight trajectories each, 100 steps per trajectory), and the held-out worlds contribute 1,200 more. Every mechanism value remains abundantly observed in training, which is the primitive-coverage condition the split was designed to satisfy, so the split tests novelty of composition rather than novelty of primitives. The two models share their trunk but differ in one decisive place. The monolithic model embeds the world index (a single integer from 0 to 8) and has no way to relate world $(1,1)$ to worlds $(1,0)$ and $(2,1)$: its embedding for index 4 is simply never trained. The factorized model embeds each mechanism separately and combines them with two additive heads, which mirrors the additive structure of the true dynamics.

It is worth pausing on what "never trained" means for the categorical-ID baseline, because it explains the coming divergence precisely. The embedding for world index 4 (the pair $(1,1)$) was initialized randomly and receives no direct gradient, since no training transition carries that index. Its contribution is therefore an untrained random vector, although the shared trunk is updated by seen cells and can change the missing-cell output indirectly. The baseline's held-out failure is consistent with an arbitrary world label that supplies no sharing rule; the unmatched parameter counts mean this run does not exclude capacity as a contributor to the performance gap.

The factorized forward pass is the additive composition rule from the missing-cell discussion made literal: pred_a + pred_b. Each head sees the state and action along with only its own mechanism's embedding, so the model is structurally incapable of learning a mechanism-A behavior that depends on mechanism B's identity. On the true additive dynamics, the sum can represent the conditional mean. The individual heads are not uniquely identifiable, however: because both receive $(x,u)$, any shared function $c(x,u)$ can be added to one head and subtracted from the other without changing the prediction. The experiment therefore tests recombination of the summed predictor, not whether either learned head recovers the semantic ground-truth mechanism. Recovering that stronger interpretation would require additional restrictions or normalization constraints.

Both models are trained with the same objective and budget: mean squared error on one-step prediction, an empirical expectation over transitions drawn from the six training worlds, optimized with Adam for 4000 minibatch steps. The loss never directly evaluates either model on a held-out combination. Observed held-out behavior therefore reflects the architectures' inductive couplings together with finite data, initialization, optimization, implicit regularization, shared-weight updates, and the unmatched capacity—not inductive bias in isolation.

Evaluation happens in two modes. The one-step test uses fresh trajectories from the held-out worlds with the true state always as input, isolating compositional accuracy from compounding. The autoregressive test feeds each model's own predictions back as inputs for 120 steps while replaying the same fixed exogenous action sequence for the true system and both models. This measures multi-step predictive fidelity under state feedback. It is not a closed control loop, because no policy or planner chooses actions from the predicted state, and it does not by itself measure deployment decision usefulness.

The held-out comparison supports the intended mechanism, but the training cells expose an additional stability issue. Four of the six training worlds have rollout MSE below $0.13$ for both models. The monolithic model alone is below $0.13$ on five: for the factorized model, $(A{=}2,B{=}1)$ is about $0.38$, and the seen world $(A{=}2,B{=}2)$ is unstable for both at about $5.46$ and $6.12$. On the three held-out worlds, the factorized errors are about $0.55$, $0.016$, and $0.47$, whereas the monolithic errors are about $9.93$, $107.54$, and $86.73$. The untrained monolithic embeddings account for a direct missing-cell mechanism, but the capacity confound and single run preclude a complete causal attribution; factorization alone does not guarantee stable autoregressive dynamics on every seen world.

Read the table with the taxonomy in mind. The held-out cells are pure novel recombinations: every mechanism value in them was seen in training. In this run, the factorized model's held-out MSEs are far below the categorical baseline's but two are roughly one to two orders of magnitude above most stable training cells; the experiment does not identify the cause of that gap or establish how it changes with more data. In particular, the additive heads have gauge freedom and need not correspond individually to the named mechanisms. The monolithic model's held-out embedding receives no direct gradient, so adding more transitions only from the same six seen worlds would not train that embedding; changing the representation or adding the missing world indices would.

The heatmaps separate two effects. The monolithic panel has large errors on all three outlined held-out cells, consistent with their untrained embeddings, but the seen $(2,2)$ cell is also unstable. The factorized model is much better on the held-out cells, yet its field is not uniform: its MSE ranges from about $0.0048$ to $6.12$, with the maximum again at the seen $(2,2)$ cell. This single run illustrates useful recombination and an autoregressive stability failure that affects both architectures; it is not a capacity-controlled causal estimate.

The panels also show why representational scaling and empirical error scaling must be kept separate. A factorized representation stores components along the mechanism axes rather than a separate model for every grid cell, but this two-mechanism experiment does not establish how its error profile scales to ten mechanisms. That stronger conclusion would require higher-dimensional runs or a bound covering estimation error, interaction misspecification, and autoregressive stability.

Reading the trajectory itself on a symlog scale, the factorized prediction qualitatively follows the true state over this rollout horizon, while the categorical-ID prediction later develops a large, nonmonotone transient excursion. The observation is consistent with the intended difference between an untrained world embedding and shared mechanism embeddings, but the single seed and unmatched capacities do not isolate that representation choice as the sole cause.

The symlog scale is worth a second look. On this trajectory, the monolithic state reaches magnitude about $42$, compared with about $8.3$ for the true system and $8.8$ for the factorized model: a large runaway excursion, but not a multiple-orders-of-magnitude difference in state. The orders-of-magnitude separation appears in the per-step squared error below, whose maxima are about $1.62\times10^3$ and $0.22$. Feeding each wrong prediction into the next step creates an autoregressive prediction failure that one-step evaluation misses; a planning claim would additionally require a policy or planner to choose actions.

The horizon-resolved view completes the diagnosis: error is plotted against prediction depth on the held-out world.