Compositional Generalization via Forced Rendering of Disentangled Latents
Authors: Qiyao Liang
Daoyuan Qian
Liu Ziyin
Ila R Fiete
Venue: ICML 2025 poster (2025)
Keywords: Factorization
Compositionality
Compositional Generalization
Data Efficiency
Primary area: deep_learning->everything_else
TL;DR: Standard networks given separate x/y cues still memorize and stitch examples rather than combine them. Forcing each cue into the output space—via low-rank embeddings or simple stripe data—enables true compositional generalization.

Abstract: Composition—the ability to generate myriad variations from finite means—is believed to underlie powerful generalization. However, compositional generalization remains a key challenge for deep learning. A widely held assumption is that learning disentangled (factorized) representations naturally supports this kind of extrapolation. Yet, empirical results are mixed, with many generative models failing to recognize and compose factors to generate out-of-distribution (OOD) samples. In this work, we investigate a controlled 2D Gaussian "bump" generation task with fully disentangled $(x,y)$ inputs, demonstrating that standard generative architectures still fail in OOD regions when training with partial data, by re-entangling latent representations  in subsequent layers. By examining the model's learned kernels and manifold geometry, we show that this failure reflects a "memorization" strategy for generation via data superposition rather than via composition of the true factorized features. We show that when models are forced—through architectural modifications with regularization or curated training data—to render the disentangled latents into the full-dimensional representational (pixel) space, they can be highly data-efficient and effective at composing in OOD regions. These findings underscore that disentangled latents in an abstract representation are insufficient and show that if models can represent disentangled factors directly in the output representational space, it can achieve robust compositional generalization.

Compositional Generalization via Forced Rendering of Disentangled Latents
Qiyao Liang* Daoyuan Qian*
Liu Ziyin 13
Ila Fiete'
Abstract
1. Introduction
A core motivation for disentangled representation learning is the belief that factoring an environment into independent latent dimensions, or "concepts", should lead powerful combinatorial generalization: if model can learn the rep- resentation of each factor separately, then in principle it can synthesize new combinations of those factors without extensive retraining. This property is often referred to as compositional generalization and is considered crucial for data efficiency and robust extrapolation high-dimensional tasks, such as language understanding and vision (Lake & Baroni, 2018).
Composition-thhe ability to generate myriad vari- ations from finite icans-is believed to underlie powerful generalization. However, compositional generalization remains key challenge for deep learning. A widely held assumption is that learn- ing disentangled (factorized) representations nat- urally supports this kind of extrapolation. Yet, empirical results are mixed, with many genera- tive models failing recognize and compose fac- tors to generate out-of-distribution (OOD) sam- ples. In this work, investigate controlled 2D Gaussian "bump" generation task with fully dis- entangled (I, y) inputs, demonstrating that stan- dard generative architectures still fail in OOD regions when training with partial data, by re- entangling latent representations in subsequent layers. By examining the model's learned kernels and manifold geometry, we show that this failure reflects "memorization" strategy for generation via data superposition rather than via composition of the true factorized features. We show that when models are forced foroed-through architectural modi- fications with regularization or curated training data to render the disentangled latents into the full-dimensional representational (pixel) space, they can be highly data-efficient and effective composing in OOD regions. These findings un- derscore that disentangled latents in an abstract representation are insufficient and show that if models can represent disentangled factors directly in the output representational space, can achieve robust compositional generalization.
Despite many attempts, disentangled representation learning has shown mixed results when it comes to compositional generalization. Some studies report clear data-efficiency gains and OOD benefits (Higgins et al., 2017a; Chen et al., 2018), while others find no significant advantage or even detrimental effects (Locatello et al., 2019; Ganin al., 2017). This tension motivates deeper theoretical explo- ration:
a disentangledfluctorieeed representation alone sufficient for compositional generalization? If not, why, and what other ingredients or constraints are necessary?
In this paper, we develop mechanistic perspective on why even fully disentangling the intermediate (or input) rep- resentations is often insufficient enable robust out-of- distribution (OOD) generalization. Concretely, synthetic 2D Gaussian "bump" generation task, where a network learns to decode given (æ, y) coordinates into a spatial image. We provide a detailed empirical investigation into why disentangled representations often fail to achieve robust compositional generalization. Specifically, we show that:
focus
From representational manifold viewpoint, even if the input or bottleneck (latent) layer are fully disentangled, subsequent layers usually "warp" and remix factors, lead- ing to poor out-of-distributtoo (OOD) extrapolation.
'Equal contribution 'Massachusetts Institute of Technology, Cambridge MA. USA 02139 Uniieersity of Cambridge, Cam- bridge CB2 IEW, U.K "NT Research. Correspondence to: Qiyao Liang <qiyao@ mit.edu>.
From a kernel-based perspective, networks failing to
Proceedings of the 42nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
We use the terms "disentanglemen"" and "factorization" inter- changeably. See Section A.2 for discussion of the definitions and distinctions.

---

Compositional Generalization via Foreed Rendering of Disentangled Latents
compose have overwritten their disentangled inputs by "memorizing" training data, rather than combining inde- pendent factors; they simply superpose memorized states when attempting OOD generalization.
of mixed discrete and continuous features. However, Xu et al. (2022) introduced evaluation protocol to assess compositional generalization in unsupervised representa- tion learning, discovering that even well-disentangled rep- resentations do not guarantee improved out-of-distribution (OOD) behavior. Similarly, Montero et al. (2021; 2022) concluded that a model's capacity to compose novel factor combinations can be largely decoupled from its degree of disentanglement; and (Lippl & Stachenfeld, 2024) derived and demonstrated numerous failure modes in compositional generalization with disentangled latents. These observa- tions temper the once-common assumption that factorized latents automatically yield systematic compositional gen- eralization. Consequently, gap remains in understanding how-and what extent-models that appear to disentan- gle latent factors can also robustly compose them under large distribution shifts, and how to encourage such general- ization. Our work aims to rigorously illuminate this gap in the context of a controlled task that enables characterization of computational generalization performance starting from exactly disentangled latents. Generalization is tested under large distributional shifts, and we investigate the detailed mechanisms of compositionality by the deep neural net- work decoders, providing complementary perspective to prior disentanglement benchmarks and broader generative modeling approaches.
Using Jacobian-basee metric tensor to study the repre- sentational manifold, there is a layer-by-laye by-yer erosion of network disentanglement in standard CNN decoders. Finally, forcing disentangled inputs to render into a rep- resentation that matches the output (pixel) space via architectural constraints or curated datasets-enables the model to overcome memorization strategies and learn genuinely compositional rules. This approach leads to data-efficient, antor--distribution generalization in cases where factorized latents alone fail.
In addition, we provide in Appendix A general frame- work for interpreting compositionality and factorization that further clarifies these observations. Overall, these results caution against the assumption that factorized bottleneck automatically confers compositional extrapolation, and they point toward more comprehensive design strategies, such as disentangled processing and data curriculum that enforce the rendering of disentangled latents, to achieve genuinely compositional neural models." 2
2. Related Work
Disentangled representation learning. There is an ex- tensive body of work on disentanglement, ranging from early theoretical proposals (Bengio et al., 2013) to various VAE-based approaches (e.g. 3-VAE (Higgins et al., 2017a), FactorVAE (Kim& & Mnih, 2018)) and GAN-based methods (Chen et al., 2016). Empirical metrics to quantify disentan- glement include VactorVAE score, MIG score, and others (Eastwood & Williams, 2018). Yet, these often focus on axis-aligned or linear independence in the latent space, with- out explicitly evaluating compositional out-of-distribution performance. Indeed, (Locatello et al., 2019) show that un- supervised disentanglement is sensitive to inductive biases and data assumptions, raising questions about the connec- tion to systematic compositional generalization.
3. Why Disentangled Latents Often Fail Compositional Generalization
We present a toy 2D Gaussian "bump" generation task to highlight and analyze a core phenomenon: even with explic- itly disentangled inputs or bottleneck, standard feedforward networks can fail to generalize compositionally.
3.1. Toy example: 2D Gaussian bump generation
Task setup. We consider a generative model that must output an N x N grayscale image containing a Gaussian "bump" at a specified 2D (x, y) location within some bound- ing box (e.g. [0, 2N)²). The training set covers square-donut shaped region of (z, y) with large OOD region in the cen- ter of the training distribution (Fig. 1(a)); alternatively, the OOD region is an equivalent area in the bottom left corner of image. In both cases, the model sees all x and y values in training but many combinations are held out.
Compositional generalization & disentanglement. Fac- torization and compositional generalization have been widely investigated in various architectures (Zhao et al., 2018; Higgins et al., 2017b; Burgess et al., 2018; Montero et al., 2021; Xu et al., 2022; Okawa et al., 2023; Wiede- mer et al., 2023; Lippl & Stachenfeld, 2024; Liang et al., 2024b). While many works study how 3-VAEs and related methods can learn disentangled representations (Burgess et al., 2018; Higgins et al., 2017a), they typically do not explore whether these representations solve the problem of compositional
To distinguish the behavior of the decoder from the encoding process that must learn (and may fail to fully disentangle) its latent representation, we explicitly construct fully disentangled latent representation and ask whether the downstream decoder can leverage that factor- ization for compositional generalization. Concretely, we focus on a CNN-based "decoder-only" architecture that maps disentangled latent inputs to image outputs. We also
Architecture.
particuclly in the presence
"Code available at ggthhh..ongyyoollagnoennnnennnnn
2

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Disettanlled Encodiing
Population Encodings
M(y
NXN
Ramp based Encoding
Ratoo
(b)
2.2064
Reconstructed Reconstructed
OOD
(e)
OOD coard index
Binary Factorized Kame
Position
B 2048
200
9,0092
100
0.02
9.0016
00D coord index
Overlap
0.5008
Position
MSF
200
0.0072
0,0094
100
30
20
0,0000
100
200
300
400
Posttion
DD coord index
Figure I. Various disentangled input/latent encodings fail to support compositional generalization. (a) shows experimental setup for the 2D Gaussian bump toy experiment, where the scalar (zz y) coordinate pair is encoded into disentangled representation (population- based vs. ramp-based encodings), which single 2D Gaussian bump centered that corresponding (x, location. The training dataset excludes all images that contain 2D Gaussian bumps centered within the red-shaded OOD region in the center the image field. (b-c) shows the MSE error contour plots and sample generated ID/OOD images of bump-based and ramp-based encoding of the and y imput, the compositional OOD region is marked with by the red-dashed bounding box and the ground truth bump location is marked by red cross. For conp-oopoiiio network trained with bump-encoded inputs, (d) demonstrates that learns to "superpose" generalize, and (c) shows the agreement between a theoretical binary factorized kernel with the similarity matrix (computed based on the pixel overlap) between the model's ID and OOD generated samples.
to decoder-only architecture to generate a N grayscale image a
ID training data when asked compositionally
experiment with MLP decoders and provide those results in Appendix F.3, observing qualitatively similar findings.
learn to generate the correct outputs beyond the ID region, as seen the MSE error landscapes of generated images and from the images themselves Fig. 1(b-c). Networks trained with different input encodings, bumps (æ: and y in- embedded as the centers of ID Gaussian bumps in two vectors of length N, Fig. 1(b), right) or ramps (z and represented as analog ramping rates in two units, Fig. (cc), right) demonstrate different degrees of "local" generaliza- tion: some input encodings can support bumps partially inside the OOD region they overlap with the ID region). However, all networks largely fail in the OOD region, and all consistently generate multiple bumps when prompted with an OOD (z, y) input.
Disentangled input encodings. Previous work suggests that vector-based representations may be key to compo- sitional generalization (Yang et al., 2023). Inspired by stimulus coding within brains (Georgopoulos et al., 1986; O' Keefe & Nadel, 1978; Shadlen & Newsome, 1998), we test various input encodings:
Population-basse coding, e.g. 1-hot, local-bump positional encoding of and y.
Ramp coding, an analog ramping code for (I, y) in 2
units.
Hence, simply providing factorized inputs, regardless of the form of encodings, does not suffice. The deeper lay- ers appear to re-entangle z and y or rely on memorized local features, leading to poor compositional generalization. While not shown, models trained with disentangled 1-hot and positional input encodings also fail to compositionally generalize. These findings confirm that disentangled latents, across a variety of encoding formats, are not sufficient for compositional generalization.
In each case, the input representation is disentangled w.r.t. a and (meaning æ and are independently encoded by separate neurons). These inputs drive an image-generating CNN.
3.2. Model fails to compositionally generalize across various disentangled encodings.
Despite receiving disentangled inputs that correctly specify OOD (x, y) coordinate combinations, the model does not
3

---

Compositional Generalization via Forced Rendering of Disentangled Latents
For the remainder of the manuscript, use inputs encoded as local bumps or ramps.
memorizing the ID data and "superposing" activations cor- responding all or some of the seen data of a and y when asked to generalize to (z, V)ODD.
3.3. Model "superposes" seen data when asked to generalize.
In Fig. I(d), we show that the model's output on a novel input pair ((oo)) is the combination of all of the model outputs with the "superposed" in- puts, xspperposed = (I, EMM y /MM) and similarly y superposed 1E MID x(6)) /MID. y), i.e. D((z, v)OOD) = D (1auperposed) + D(y (ppauperposed)) Here MID is the number of ID data. To be more concrete, in the language of kernel, we can define a binary factorized kernel,
The form of failure in OOD generation is suggestive: In many runs, the CNN appears to memorize the training distri- bution and when prompted with an OOD (z, y) combination, seems to perform a lookup: it finds an image or images in the training data with similar values of x and the closest values of y, and superposes these to generate a new image, Fig. 1(b-c), right.
Kbinary ((z2, v.). (zz,, 91)) = Ke (x1, 8OR ky (v., 91)
The consistency of these patterns across input encoding for- mats and architectures indicate a shared underlying mecha- nism of generalization that we further characterize here.
(3)
= (x, + Ky Uy) - Krz (1., Ky (yi UU), (4) where H2 (1), = Ky (yis VU) = 8ú is given by the Kro- necker delta function. With normalization, the Gram matrix element is given by Kig- = Es Kiney (t.3")
Kernel-based approach to characterizing OOD gener- alization. One approach to characterizing model's per- formance on unseen data is to characterize a kernel K : 2 x 2 R, where Z- = X x Y' is the cartesian product input domain of æ and y coordinates of the Gaussian bumps in our toy example. Here we will focus our discussion to the 2-feature scenario for 2D Gaussian bump generation. A generalized introduction to kernel-basee approach in charac- terizing factorization and OOD generalization is given in the sym-
Fig. I(e), we visualize the similarity between representations of OOD- generated samples (sorted by coordinates) and an idealized binary kernel, where similarity equals if coordinates match exactly and 0 otherwise. Specifically, Fig. compares the ideal binary kernel (top panel) to the actual similarity ma- trix derived from model-generated Gaussian bump images (bottom panel). The agreement between these two matrices quantifies how closely the model's 's learned similarity struc- ture aligns with this ideal binary factorized kernel. Fig. ((e) corresponds specifically to the top left sections of the full matrices shown in Fig.8 in Appendix F.1, which provide a comprehensive mapping of similarities between all ID and OOD generated samples. These results indicate that the disentangled input structure enables the model indepen- dently map the x and y coordinates in-distribution (ID) samples. a form of novel OOD generation that is not the same as OOD compositional generalization.
K Kuney (t.5)
Appendix A.4. We note that any positive-semidefing metric, and square-integrable kernel K can be decomposed into a Schmidt decomposition (for the 2-feature case):
K (z)
Σ Xr N1P (v.v'), (1)
where each term 1:5 (x, x²) koy" (v,3²) y') is rank-1 with re- spect to the (x)) and (y) factors, and { are nonnegative eigenvalues. The number of terms with nonzero 2 deter- mines how mixed the kernel is. In practice, we can construct the discrete approximation, the Gram matrix K on a finite set of data inputs (z((), outputs:
In a set of follow-up experiments on MNIST image ro- tation, we observe that this same superposition of memo- rized activation patterns drawn from related in-distribution data-emerges as a general OOD generalization strategy in neural networks that have learned memorize ID data. Crucially, this phenomenon applies not only to composi- tional tasks but also to single-dimession interpolation and extrapolation. We detail these results in Appendix F.2.
()())))))))))))) ) and their corresponding
K (1,1) (((x00,00)) (()))
Non-compositional models learn a binary factorized ker- nel. In the kernel language, we would expect a compo- sitional model to break down novel input combination (z. v)OOD into factorized kernel, each independently map- ping x and y to the seen components via K2 and Ky- From the generated OOD outputs of the networks shown in Fig. 1(b-c), we observe that the generated bumps seem to align with the ground truth x and y coordinates of the OOD bump. Empirically, we found that this is due to the model
Similarly, in CNN-based guided diffusion models tasked with Gaussian bump image generation, (Liang et al., 2024a) reported that the model generated a superposition of mul- tiple ID bumps when guided to generate an OOD bump. Similar to both these findings, (Kamb & Ganguli, 2024) showed that guided diffusion models generate new images by gluing together local patches from the closest relevant images in the training data. Together with (Lippl & Stachen- feld, 2024), these observations suggest an emerging consen-
4

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Ramp Enzoding Bump Enccoing
(c)
Input: Bump Encoding
(d)
Input: Ramp Encoding
2.02
Depthe2
Inpuc
R²
Tecoortettoon Mectee
Depth
Depthwd
Fortortzation Motro
0.95
Depthee
0.90
3.98
Depthwb
000000
0.50
Inpult
R²
2.98
Desthee
900095
npptt
Hoput
[e)
Burg Enzzeing
Layer 0
Lagerirred Layer
Layer 2
Layer
Layer Index
Layer 4
Layer 5
Layer Index
Layer
Layer
Layer 1
Layer 2
Layer
Layer 4
Layer 5
Layer
Rate Encoding
Figure Cause failure compositional generalization despite disentangled inputs/latents: input memorization by decoder undoes factorization. (a-b) show the linear probe metries of the learned representation as function of x and for bump- and ramp- input-encodings. and (c-d) show the factorization score (Eq. 7), and (e-f) show the volume metric 1og (dv)). as function of layer depth, for models trained with bump-based vs. ramp-based input encodings. The linear probe metric is defined as the R2 scores of fitting two linear classifiers with respect to and y.
sus view of how various generative deep networks function in OOD generation - by combining literal patches of their training data, rather than combining factors from the data, even when provided with information what these factors should be in the form of disentangled inputs.
alently) in the generation process. Let the network be a composition of layers
Ty o ..... o Ti. with cach Ti : 2 - 2,. The input measure so (assumed factor- ized) on Zo is "pushed forward" through these layers to match the in-distribution measure 3L on the final space 2,,. Formally, 1 (Ti ..... o Ti) No for (= 1, I.....,, so μı T.Ho should approximate 2 If pL differs topo- logically or geometrically from (e.g. having an OOD "hole"), each layer Ti may need to warp coordinates signifi- cantly. These warping steps can easily re-entangle originally factorized inputs, leading failures in OOD regions.
3.4. Manifold warping through dimension expansion ruins factorization in latent space.
Why is disentangled input insufficient for learning disentan- gled factors of variation? Here seek deeper mechanistic insight into model's 's learned representations along the net- work layers. Specifically, we track how strongly æ and remain disentangled across layers. While the input layer is disentangled (by construction). subsequent convolutional layers produce mixed features with large cross-terms, partic- ularly near the final decoding layer. This "manifold warping" effectively destroys compositional structure, leaving the net- work reliant on patterns that do not extend to unseen (xr, y) combinations.
Characterization of manifold warping and factorization. To qualitatively investigate the warping and its effect the compute factorization of the layer-wise representations, the Jacobian-bati metric tensor inspired by the approach of Ref. (Zavatone- Veth et al., 2023). The generalized defini- tion can be found in Appendix B. In short, the metric tensor is defined by
g( = J(x, v) J(I,
A transport perspective. From transport perspective, the model layers effectively learn to interpolate between the input (source) distribution and the output (target) dis- tribution. Here the source distribution corresponds to the disentangled input representation manifold, while the output corresponds to the ID data representation manifold, which has an OOD "hole" in the middle. If there is nontrivial topological or geometrical deformation to the source or the target distribution representations, the model will have a hard time maintaining the integrity of the factorization, especially through dimension expansion (or reduction equiv-
where J is the Jacobian matrix w.r.t. the different input feature dimension, which corresponds to and in our synthetic toy setting. We then visualize the volume metric
do(x, v) = del v)) dx dy
(6)
for representation output at each of the network layer. The metric visualizations log1o(dv) are compared between networks with bump-based vs. ramp-based encodings in Fig. 2(c) and (f). Visually, we observe significant distor- tion as a function network depth, especially in the OOD
5

---

Compositional Generalization via Forced Rendering of Disentangled Latents
region. To accompany the distortion analysis, we plot the linear probe metric (linear classifier Re w.r. x: and y) as well as the factorization metric proposed based g in Eq. (7) as a function of network depth in Fig. 2(a-d). Here the factorization metric is defined as follows
(n) (() composition, ( are scalar singular values, and of are orthogonal basis vectors along the row and column di- mensions, respectively. This formulation enables to di- rectly regularize the singular value spectrum and the spatial smoothness of basis vectors, promoting low-rank structure in each embedding matrix.
Iay xr 8)
Factorization(g(zz.)) v))
9an (z. 8)
+
Iyn (z, 8)
This regularization is particularly effective because cach input dimension is associated with dedicated 2D embed- ding matrix, which constrains the model to learn spatially structured factors aligned with individual input dimensions. In practice, we observe that this encourages the model to discover localized or "stripe-like" patterns in the learned embeddings (see Fig. 3(b)), facilitating compositional gen- cralization. To further stabilize learning and reduce redun- dancy, we additionally penalize the entropy and variance of the singular values, which promotes sparsity and consistent usage of the factorized components.
(7)
where Dij's 's are the matrix elements of the 2 x 2 matrix g- We note that both metrics decay relatively gracefully over the layer depth for both networks with bump-based and ramp-based input encodings. Intriguingly, shallow networks with ramp-based encoding experiences a sharper drop in factorization. This is likely due the higher demand for dimension expansion with the ramp-based 2-neuron input, which leads to higher distortion when allowed fewer net- work layers to accommodate the expansion. The volume metric visualization as well as the linear probe metric as a function of layer depth for different network depths are shown in Fig. 11 and Fig. 12 (Fig. 13 and Fig. 14 for MLPs) as well as generated OOD images in Fig. 10 in Appendix F.3.
The squared magnitudes [2(1) define probability distri- bution across n:
[201)2
7(m)
(9)
E (m))
4. Encouraging Compositional Generalization
which we regularize using an entropy penalty:
Given the above model failures, we ask: Which architec- tural or training strategies can preserve factorization and yield robust compositional generalization? We provide two promising avenues next.
Len = m Σ X In X(m)
(10)
encouraging dominance of only a few modes to enforce low rank, Additionally, to discourage memorization of individ- ual indices, we penalize variance in (a()) and (");
4.1. Architecture regularization for low-rank embeddings
Leaar 7/2 Σ Var(u/()) + Var(u("))].
(11)
One approach we adopt is to architecturally encourage the expansion of abstract latent representations into structured 2D embedding filters, while applying explicit rank regular- ization during training (Fig. 3(a).. This design promotes interpretability and compositional disentanglement in the learned features. For example, Cichocki et al. (2009) demon- strate that enforcing low-rank tensor factorization can miti- gate spurious interactions between factors. Inspired by this, we apply matrix (or tensor) factorization penalty to the set of 2D embedding filters associated with each input dimen- sion, encouraging them to remain approximately low-rank and spatially consistent.
The total loss function combines these terms with standard reconstruction loss:
Livial = LMSE + Lent + Lvar-
(12)
Applying these regularizations to the embedding matrices of both x- and y-factors results in more structured and fac- torized representations, improving generalization to unseen coordinate pairs. Fig. 3(b) illustrates that, without regular- ization, the model struggles to preserve signal integrity in OOD regions. With regularization, the embeddings exhibit structured banding aligned with coordinate axes, indicating a discovered factorized representation beneficial for com- positional extrapolation. In addition, Fig. 3(c) shows that our regularization works even for a corner OOD region, in- dicating stronger form of (extrapolative) compositional generalization. With regularization, the learned represen- tation is significantly smoother than without, as evidenced in Fig. 3(d), despite having the same factorized embedding architecture. These results demonstrate the effectiveness of our proposed regularization technique.
To formalize this, in the 1-hot input encoding setup, we generate embedding matrices of shape N x N for each token x and separately for y (Fig. 3()). There are total of 2N such matrices. Denoting one of them as Wig, we decompose this matrix using an SVD-like form:
W7,
Σ (n)
u (m)
(n)
(8)
where r is a chosen upper bound on the rank of the de-
6

---

Compositional Generalization via Forced Rendering of Disentangled Latents
(a) Disentangled H-Hot input
Disentangled processing
........
(b)
No Rogulantzannn rippulaitzationn
y-14
(c)
Extract Stacking
encodings
embedding
N'X
y/
Regulorization
(d)
Layen
Layer
Layer2
(e)
Regulrizztion No Regularizaton
25
log10 (dv)
20
20
1.
Layer D
Leyer
Layer
Layer 5
loggo (dv)
Varfarce Only
Entropy-Vee
Figure 3. Inducing compositional generalization through architectural rendering and regularization constraints. (a) Schematic of architecturally forced rendering of the initially disentangled representations (with 1-hot input encodings) into space matching the output space (disentangled processing). (b) Sampled embedding activations corresponding to x 14 and y for networks trained without and with regularization, respectively. (c) Generated images when the OOD region is on the top right corner, for the non ronullaulrd (top) regularised networks (bottom) respectively. (d) Volume metric comparison between networks trained without and with regularization function of layer depth, respectively. (e) OOD vs. ID MSE plot for various ablation studies over many runs.
Ablation. To assess the impact of our regularization strat- egy, compare four configurations: (1) No Reg (baseline MSE loss), (2) Entropy Only (applying Eq. (10)), (3) Vari- ance Only (applying Eq. (11)), and (4) Entropy+ Var (com- bining both). Fig. 3 (c) plots MSE in distribution (ID) versus out-of-distribution (OOD). The condition No Reg yields the lowest error but the highest OOD error, reflecting poor generalization. Entropy Only marginally improves OOD performance, while Variance Only significantly enhances extrapolation at the cost of higher ID loss. Combining both constraints (Entropy+ Var) achieves the best OOD perfor- mance while maintaining low ID error, striking the optimal balance for robust factorization.
4.2. Dataset augmentation
Another approach that successfully enabled compositional generalization in the toy setting is dataset curation. Asking the network to generate independent factors of variation in the form of separate x and "stripes" help the network discover a representation that enables compositionality.
Concretely, we generate ID vertical and horizontal stripes by fixing single coordinate (either æ or y) and applying ID Gaussian in the orthogonal direction (samples shown in Fig. 4(a), full set shown in Fig. 7(a)). Note that here the 2D input structure to the model is still maintained via setting the coordinate to the fixed dimension to -1. These stripes in the full output representational space act as build- ing blocks that allow the network to learn each factor of variation independently in way that abstract factorized inputs unable Optionally, additionally include small number of the 2D bumps (outside the OOD region, as before) in the training set, so the model also sees small number of 2D examples. Consequently, the network ac- quires strong compositional ability: it generalizes to novel (r, y) combinattio outside the original training distribution (even with zero 2D bumps). Empirically, demonstrated in Fig. 4, this approach markedly improves compositional generalization performance.
Intuitively, the regularization leverages simple, low- dimensional structure in the learned embedding matrices. By limiting the effective rank, it forces the network to rep- resent each factor (for example, a coordinate dimension) through small set of shared directions rather than sepa- rately memorizing every possible input combination. Con- sequently, once the model learns an appropriate low-rank basis for in-distribution data, it can more easily compose those basis elements to handle new, unseen factor combi- nations in the OOD region. In other words, restricting the embedding to a few dominant modes discourages the net- work from relying on specific entries for every (x, y). thus reducing overfitting and promoting a genuinely factorized representation that better generalizes to novel coordinates.
Data efficiency scaling, Remarkably, we observed that models trained on a dataset consisting of just the stripes

---

Compositional Generalization via Forced Rendering of Disentangled Latents
(a)
Vertical Stripe Ox
Horizonia Stripe =y
(b)
Bumps
Burpss
50 Bumps
100 Bumps
(c)
(d)
Data Count (log scale)
lurii only
Layen
Layer
Layer?
Layer
Citic teterrce 18?)
Ion Fdv
Strpes Bampe
231
0
1.2
5.4
3x
4x30" Grid Size llog sca e)
5
i
Neuron (c=1,h=2,
Layer 0
Mean over
Mean over
(e)
8.5
h=0
0.4
Avg Act
0.4
0.3
02:
0.1
0.1
Layer
Neuron 0c-0.0- 00,0000-1-11111 13,w111)
0.6
0.5
0.3
h=6
0.
0.81
10 20 0
10
10
==0
Figure 4. Inducing generalizable composition by training single disentangled factors to render (data curriculum). (a) Sample grayscale images of Gaussian "stripes" at = and y = 14. (h) Generated OOD output as a function number of Gaussian bumps (all outside the upper right OOD region) included in the training dataset. (c) Data/sample efficiency: data scaling of the stripes-only (~> N) vs. bumps-only (~> N datasets to reach 90% accuracy of and y generation as a function of image size (accuracy assessed based the location of the darkest pixel). Here the stripes bumps dataset consists of 7N bumps 2.N stripes; learning breaks the curse dimensionality due to the ability to compositionally generalize zero-shot. (d) Volume metric as a function of layer depth of the network trained dataset of stripes bumps. (e) 2D neuron activations across different channels in the first layer (layer 0) of network trained trained on a dataset of stripes + 50 bumps. (f) Neural tuning curves of two sample neurons at layer O and layer 3 of the same network in (c).
are able compositionally generalize to generate additive stripe conjunctions in the OOD region without having seen any compositional examples (inside outside the OOD region) during training, as shown in Fig. 4(b) when 0 2D Gaussian bump images are included in the training dataset. Similar to Fig. 3(c), we show here an example of corerr OOD region to demonstrate that our method is robust for extrapolative compositional generalization. The remaining panels of Fig. 4(b) show that as we increase the number of 2D Gaussian bump images included in the training dataset, the model Icarns to generate proper bumps in the OOD re- gion (also see Fig. 7(b)). We further characterize the data efficiency of this approach to learning to compositionally generate bumps in the OOD region (learning with stripes bumps) versus a traditional approach to learning with just bumps. The total data count scaling with respect the im- age size N is shown in Fig. 4(c). Here benchmark the compared models based on the minimal threshold amount of data needed to reach target accuracy of 90% in gener- ating the correct & and coordinates of the stripe conjunc-
tions/bumps. To our surprise, we found that models trained datasets consisting of purely bumps have data threshold that scales N and hence the data requirements grow as the cube of the size of the problem. On the other hand,
due to the zero-shot generalization capability of models trained on the stripe dataset, the scaling is linear with N, since the data set consists of 2NV stripes in total, N for æ and N y (together with or a small number of bumps). Fur- thermore, with the increase of stripes training with bumps to obtain bump outputs rather than plus glus-shaped ones, the minimum required number of training bumps scales only linearly with N, meaning that the total training data remains linear in N, indicative of compositional generalization.
For comparison with the earlier results, we again plotted the volume metrics for the model trained on the stripes + bumps dataset, function of model layer depth, Fig. 4(d). The representations remain relatively smooth and uniform even in the OOD region, contrasting the models trained with just the bumps shown in Fig. 2. In addition, we show in Fig. 4(e) the 2D activations of neurons in the first layer of a
8

---

Compositional Generalization via Forced Rendering of Disentangled Latents
network trained on a dataset of stripes + 50 bumps. The ac- tivations show consistent "stripe" patterns that resemble the embeddings shown in Fig. 3(b) with regularization. More- over, we show neural tuning curves w.r.t. to z and for two sample neurons in layer 0 and layer 3 of the same network in Fig. 4(D), which exhibits "stripes" in early layers, which localizes to become a "bump" in downstream layers (more sample neural tuning curve can be found in Fig. 15).
domain-centric data curation (training with independent factors), the network generates pixel-level embeddings of the individual æ, y latents to achieve strong composi- tional generalization.
Limitations & Future Directions. While our investiga- tion provides a clean, mechanistic insight into why disentan- gled representations often fail on a synthetic 2D bump task, its scope inherently limited by the toy nature of the setting, Extending these insights to large-scale models (e.g., diffu- sion models, autoregressive architectures, Transformers) and high-dimensional data domains (e.g., vision, language) is an important next step. Modern architectures like vision transformers or large language models may already em- bed implicit factorization thruuh attention mecha- nisms or prompt-based modularity-bui rigorously assess- ing whether these biases genuinely support compositional generalization remains open challenge.
Intuitively, the early-layer stripes allow the network to build "scaffold" in its representations that are straightforward for it to use to perform additive composition. Later lay- ers in the model then leverage this scaffold to handle the multiplicative compositional task- - to generate stripe in- tersections instead of unions- - thereby achieving efficient bump generalization in the OOD region. This highlights the value of data-centric strategies that emphasize the gen- erative factors in isolation, consistent with the finding that dataset augmentation with stripes also improves compo- sitional generalization and data-efficiency scaling for 2D bump generation in diffusion models in Ref. (Liang et al., 2024b). While the prior work concluded that stripe training generated disentangled latent representations, the present work shows that supplying or learning disentangled latents is not sufficient for compositional generalization, and that stripe training even with disentangled latents helps form pixel-space latents representations that are necessary for OOD compositional generalization.
We also acknowledge that translating insights from our syn- thetic study to more complex, real-world datasets is nontriv- ial but essential. As highlighted by Montero et al. (2022), interactive compositionality is unlikely to be addressed by a one mne-size-fits-ll solution. Accordingly, we do not claim our approach is universally applicable, but rather that it reveals two promising and generalizable directions: (1) training modular, output-level embedding filters dedicated to each disentangled input dimension, and (2) leveraging dataset augmentation strategies that isolate factors of variation. Our simplified setting was deliberately chosen to cleanly expose the underlying failure modes, such as the model memo- rizing ID configurations for OOD generalization, without confounding interactions.
Summary. The commonality between the two disparate approaches that enable OOD generalization is that they both encourage the network to form disentangled representations that are encoded in the space of the 2D pixel outputs. By contrast, even directly supplying the network with factor- ized representations does not work if those representations are abstract and in a different latent space than the 2D im- age space. Factorization in a latent space may not survive subsequent dimension expansion by a downstream decoder. By ensuring that the network learns to separate each under- lying factor (e.g. and y) at the pixel level, we preserve the necessary compositional structure for strong extrapolation to novel conditions.
In future work, we plan to explore how these principles scale to standard disentanglement datasets and more natu- ralistic settings, explicitly investigating whether modular embedding and low-rank factorization constraints retain their utility. Finally, we emphasize the need for new met- rics and frameworks characterize partial and hierarchical compositional structures, which are far more representative of real-world cognition. We provide preliminary attempt at capturing such structures in Appendix A. Developing prin- cipled tools to measure subtler forms of compositionality may better guide both architectural and algorithmic choices in the pursuit of robust generalization in complex domains.
5. Discussion & Conclusion
Our results highlight two overarching principles:
1. Disentangled latents are not enough. Factorizing sin- gle latent layer (e.g. the bottleneck) into a disentangled representation does not guarantee compositional struc- ture in deeper layers; manifold warping can re-entangle factors, leading to failure in OOD regions.
Ultimately, bridging the gap between disentangled represen- tations and compositional generalization requires holistic view of how a network's entire forward pass maintains or destroys factorization. We hope this work provides step toward that goal, clarifying both the pitfalls of standard bottleneck- sauleneck-based approaches and the promise of structural or regularized solutions.
2. Rendering disentangled latents into the same rep- resentational space as the outputs is key. Through a specific architecture with low-rank regularization or

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Acknowledgements
3-8, 2018, Montréal, Canada, pp. 2615-2625, 2018. URL http:///accccdnnsssserrri.cccccp 8/hash/ be3dcck806623333379972222224
The authors would like to thank Alan (Junzhe) Zhou, Zim- ing Liu, Federico Claudi, Megan Tjandrasuwita, Mitchell Ostrow, Sarthak Chandra, Ling Liang Dong, Hao Zheng, Tomaso Poggio, Cheng Tang for helpful discussions and feedback at various stages of this work.
Chen, X., Duan, Y., Houthooft, R., Schulman, J., Sutskever,
I., and Abbeel, P. Infogan: Interpretable representation learning by information maximizing generative adver- sarial nets. In Lee, D. D., Sugiyama, M., von Luxburg, U., Guyon, I., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 29: Anmual Conference on Neural Information Processing Systems 2016, Decem- ber 5-10, 2016, Barcelona, Spain, pp. 2172-2180, 2016. URL https///ppcceeiinggsnnnniiippppp 6/hash/ (ecpeeeaaddaaaaabbbbbbeeeebb--e Abstrbbtbbbttttthtll
Impact Statement
Our work develops theoretical and empirical insights into why disentangled or factorized inputs do not necessarily yield compositional generalization, and proposes practical strategies (data augmentation, low-rank regularization) to preserve factorization within the network. By illuminating how and why factorization can break down-or be main- tais research can guide future model designs that are more transparent, data-efficient, and capable of robust extrapolation beyond the training distribution.
Cichocki, A., Zdunek, R., Phan, A. H., and Amari, S. Nonnegative Matrix and Tensor Factorizations Appli-- cations to Exploratory Multi-way Data Analysis and Blind Source Separation. Wiley, 2009. ISBN 978- 0-470-74666-0.. doi: 10.1002/97884407777778 URL http://dbboog//11002229777777770 100777777777777777777
From an ethical societal standpoint, do not foresee immediate risks arising directly from these factorization techniques: the methods and analyses we present are primar- ily at the conceptual and architectural level. However, we note that improvements in compositional generalization can enable models to be more adaptable, including in sensitive real-world domains such as medical imaging or autonomous systems. While better extrapolation is desirable, it also entails that models could behave in unintended ways un- der distribution shifts if they mix factors in novel, untested combinations. Ultimately, we believe our framework, by clarifying the conditions needed for truly factorized repre- sentations, will enhance the interpretability and reliability of neural networks, net positive for the broader machine learning community.
DeMoss, B., Sapora, S., Foerster, J. N., Hawes, N., and Pos- ner, I. The complexity dynamics of grokking. CoRR, abs/2412.09810, 2024. doi: 10.48550/ARXIV22412. 09810. URL https://doi.rr// htwp:///000.55555002222222209990
Eastwood, C. and Williams, C. K. I. A framework for the quantitative evaluation of disentangled representations. In 6th International Conference on Learning Representa- tions, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL http://openeeeveeeeee///oornaiie AZ.
Ganin, Y., Ustinova, E., Ajakan, H., Germain, P., Larochelle, H., Laviolette, F., Marchand, M., and Lempitsky, V. S. Domain-adversarial training of neural networks. In Csurka, G. (ed.), Domain Adaptation in Computer Vi- sion Applications, Advances in Computer Vision and Pattern Recognition, pp. 189-209. Springer, 2017. doi: 1 10. URL https:///oooo.rg 10.1007797833331195834411-
References
Bengio, Y., Courville, A. C., and Vincent, P. Representation learning: A review and new perspectives. IEEE Trans. Pattern Anal. Mach. Intell., 35(8):1798- 1828, 2013. doi: 10. 111100TTAM.220...5 URL https://doi.or//1..1100 201ML2003.50.
Georgopoulos, A. P., Schwartz, A. B., and Kettner, R. E. Neuronal population coding of movement direction. Sci- ence, 233(4771) 1986. doi: 10. 10.1126/science. 3749889.
Burgess, C. P., Higgins, I., Pal, A., Matthey, L., Watters, N., Desjardins, G., and Lerchner, A. Understanding disen- tangling in f-vae. CoRR, hbs//8804003599 2018. URL http://rriivoor/abb//88004
Higgins, I., Matthey, L., Pal, A., Burgess, C. P., Glo- rot, X., Botvinick, M. M., Mohamed, S., and Lerch- ner, A. beta-vae: Learning basic visual concepts with a constrained variational framework. In 5th Interna- tional Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings. OpenReview.net, 2017a. URL https: huppenrreiee.n////oum////ds///us
Chen, T. Q.. Li, X., Grosse, R. B., and Duvenaud, D. Isolating sources of disentanglement in variational autoencoders. In Bengio, S., Wallach, H. M., Larochelle, H., Grauman, K., Cesa-Bianchi, N., and Garnttt, R. (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018 NeurIPS 2018 December
10

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Higgins, I., Matthey, L.,, Pal, A., Burgess, C. P., Glo- rot, X., Botvinick, M. M., Mohamed, S., and Lerch- ner, A. beta-vae: Learning basic visual concepts with constrained variational framework. In 5th Interna- tional Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings, OpenReview.net, 2017b. URL https: fupeancciionn///ooumt/iiiiyyyyy
Beach, California, USA, volume 97 of Proceedings of Ma- chine Learning Research, pp- 4124, PMLR, 2019. URL http://mwweeednnggssmmrrsrss////
Montero, M., Bowers, J., Ponte Costa, R., Ludwig, C., and Malhotra, G. Lost in latent space: Examining failures of disentangled models at combinatorial gen- eralisation. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems, volume 35, pp. 10136-10149, Curran Associates, Inc., 2022. URL http:://noceeeeiigggeeeeeeeeccpp hieelccckaabbbbccc77777999900000 Conference.pdl:
Kamb, M. and Ganguli, S. An analytic theory of creativity in convolutional diffusion models. CoRR, abs/2412.20292,, 2024. doi: 10.46550/ARXIV2241222222 URL https: 0tt.or//00005555000/222.2..00.00
Kim, H. and Mnih, A. Disentangling by factorising. In Dy. J. G. and Krause, A. (cds.). Proceedings of the 35th International Conference on Machine Learning. ICML 2018, Stockholmsmüsssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp.
Montero, M. L., Ludwig, C. J. H., Costa, R. P., Malho- tra, G., and Bowers, J. S. The role of disentanglement in generalisation. In 9th International Conference on Learning Representations. ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. URL
PMLR, 2018. URL http:// 8b.html.
Okawa, M., Lubana, E. S., Dick, R. P., and Tanaka, H. Com- positional abilities emerge multiplicatively: Exploring dif- fusion models on a synthetic task. CoRR, /////3300.9336, 2023. doi: 10.488550/AXXI.23300000 2ooo..org/1.445550aaXx..2.22
Lake, B. M. and Baroni, M. Generalization without sys- tematicity: On the compositional skills of sequence-to- sequence recurrent networks. In Dy, J. G. and Krause, A. (eds.), Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsnüsssa, Stockholm, Sweden, July 10-15, 2018, volume 80 of Pro- ceedings of Machine Learning Research, pp. 2879-2888. PMLR, 2018. URL http://wwweeddnnnn.m.rrrrsseeeee lake hta.
URL hitps:
O' Keefe, J. and Nadel, L. The Hippocampus as a Cognitive Map. Oxford University Press, Oxford, UK 1978.
Shadlen, M. N. and Newsome, W. T. The variable discharge of cortical neurons: implications for connectivity, com- putations, and coding. The Journal of Neuroxcience, 18 (10):3870-3866, 1998. doi: 10.1523/NEEROOSCI.S 10-03870.1998.
Liang, Q., Liu, Z., and Fiete, I. R. Do diffusion models learn semantically meaningful and efficient representations? In ICLR 2024 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2024a. URL https:
Wiedemer, T., Mayilvahanan, P., Bethge, M., and Brendel, W. Compositional generalization from first principles. In Oh, A., Naumann, T., Globerson, A., Saenko, K., Hardt, M., and Levine, S. (eds.), Advances in Neural Information Processing Systems 36: Annual Conference on Neural Information Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16. 2023, 2023. URL http://paperr.nips.ccoppe paper/2023/hashi httee/1f6aalll00 1191118888887555559999993380 Abstract- Chiraat-Coofeeeec..hmm..
Liang, Q., Liu, Z., Ostrow, M., and Fiete, I. How diffusion models learn to factorize and compose. In 38th Confer- ence on Neural Information Processing Systems, NeurIPS 2024, 2024b.
Lippl, S. and Stachenfeld, K. When does compositional structure yield compositional generalization? A kernel theory. CoRR, hass/2005.1639) 2024. doi: 10.48550/ ARXIV.2405.1639. URL https://doi.rr/100455555 1609.
Xu, Z., Niethammer, M., and Raffel, C. Composi- tional generalization in unsupervised compositional representation learning: A study on disentanglement and emergent language. In Koyejo, S., Mohamed, S., Agarwal, A., Belgrave, D., Cho, K., and Oh, A. (eds.), Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans. IA, USA, November 28 December 2022, 2022.
Locatello, F., Bauer, S., Lucic, M., Rätsch, G., Gelly, S., Schölkopf, B., and Bachem, O. Challenging common assumptions in the unsupervised learning of disentangled representations. In Chaudhuri, K. and Salakhutdinov, R. (eds.), Proceedings of the 36th International Conference on Machine Learning, ICML. 2019, 9-15 June 2019, Long

---

Compositional Generalization via Forced Rendering of Disentangled Latents
URL http://mmeessssiiicccccheeeeehap
hat1999oooff0002222//////aaaa2aa Conference.html.
Yang. T., Wang, Y., Lan, C., Lu, Y., and Zheng, N.
Vector-based representation is the key: A study on dis- entanglement and compositional generalization. CoRR, abs/2305. 18062, 2023. doi: 10.4850/ARXIVV23005 18063. URL
Kavatone-Veth. J. A., Yang, S., Rubinfien, J. A., and Pehle- van, C. Neural networks learn to magnify areas decision boundaries. CoRR, abs/2301.11373 2023. doi: 1044850/ARXXY::301113.5.55 URL https://doi.org/10. 48550%arXiw.2200111335.
Zhao, S., Ren, H., Yuan, A., Song, J., Goodman, N. D. and Ermon, S. Bias and generalization in deep generative models: An empirical study. In Bengio, S., Wallach, H. M., Larochelle, H., Grau- man, K., Cesa-Bianchi, N., and Garnett, R. (eds.), Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Pro- cessing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 10815-10824, 2018. URL https///nneeeeinnss..rrrr.ccccpp 8/hash/ 15116699911887115666068844449999 -bbttracttttttmm
12

---

Compositional Generalization via Forced Rendering of Disentangled Latents
A. A Framework for Compositionality, Factorization, and Modularity
Existing literature is divided on whether disentangled representation learning truly fosters compositional generalization and data efficiency. Despite extensive empirical studies and numerous proposed metrics, there is still no unified framework that clearly defines factorization, disentanglement, and compositionality or their interrelationships. A key challenge lies in the lack of a systematic "language" to characterize the various types and degrees of these concepts. Furthermore, real- woold datasets often exhibit multiple, overlapping forms of compositionality, each with its own failure modes that can obscure one another. In response, we aim to develop a more general framework that makes these distinctions explicit.
This section aims to establish a framework for unifying the concepts of compositionality. factorization, disentanglement, and modularity. Specifically, we introduce two core components of our framework. First, we define the compositional complexity of a target function or distribution via the minimal number of rank-1 product terms needed represent it. Then, we formalize factorization (disentanglement) for the intermediate activations of neural network, allowing for arbitrary invertible transformations.
A.1. Definition of compositionality and compositional complexity
Here take on a function learning perspective of defining factorization and compositionality. Specifically, for learning a compositional solution, expect network to first learn how to "break down" (factorize) component factors and then flexibly "re-combine" them for compositional generation. For a neural network learning generic N N-way composite target function we can write the tensor decomposition of this function as
f X ... x AN + RM of N features
R
(aN) IN) = A (.) ((w) (ZN).
(21)....
(13)
where or (144) (1)) (xi): :2 R and A, E are lumped factors.
Note that this decomposition not unique, meaning that a given function f might permit multiple degenerate decomposition with the same rank different ranks. This means that the network not guaranteed to learn a compositional decomposition that aligns with the canonical axes of variation, required by some of the disentanglement metrics in existing literature.
Probability Distributions as a Sum of Product Factors. In some settings, rather than learning deterministic function f: X1 x...Xx XXN + RR,, we aim to model a joint probability distribution p(r1,..., .N) factorization can be written as:
An analogous Ssum-of-produce
R PN (EN), with Σ X = 1, X 2 0.
p( IN) = Σ X Pi (1a)
(14)
Here, each (r) (a;) denotes a univariate (or lower-dimensional) distribution on X; and X, are nonnegative mixture weights. In discrete domains, (r) can be categorical distributions; in continuous domains, they can be densities (e.g. Gaussian components or neural-network-baasd PDFs). Intuitively, we are representing p as mixture of "fully factorized" product distributions, each contributing a rank-1 1 tensor pi x ... x PN in the joint space.
Defining Compositional Complexity. There are some compositional functions or probability distributions that are inherently more difficult to learn than others. Here we formalize the "compositional complexity" of a target function or distribution via the minimal number of rank-1 (factorized) terms required in sum-of-product decomposition. Let N be the number of factors (or dimensions), and let {ol ((())) or {p!") (*) denote per-factor functions or distributions.
Case 1: Deterministic Function. For f X2 X ... x X' + RM, assume M 1 for simplicity, or treat each output dimension independently. sum-of-product (CP) representation
n N
f(e
Σ II
(())
(z+),
(up scalar factor A, if desired).
(15)
Define the exact rank of f as:
R N =
R*( = min R
f(II...,,NN
Σ II
(mx)
(16)
13

---

Compositional Generalization via Forced Rendering of Disentangled Latents
If cannot be exactly decomposed with finite R. consider an e-approximation in some norm measure ||
R
N
RE min R
f
Σ1
II
of
2x))
(zx)
S eJ.
(17)
A smaller R (f) (or RC(() indicates f is more "compositional" factorized.
Case 2: Probability Distribution. For a joint probability p(r1,.... IN),, a sum-of-products corresponds to a finite
mixture of product distributions:
N
p IN) =
X
Pi
(z:);
Σ A = 4.2. 0.
(18)
Its exact rank is:
R
N
(m))
(p) = min
R
p
=
Σ X
1
(19)
If no finite R yields exact factorization, allow s-approximation w.r.t suitable divergence D(.-) (e.g. KL divergence total variation distance):
f
(p) min
R
D p
Σ
E
(20)
Again, smaller ranks correspond to "low compositional complexity" (fewer factorized terms), whereas large R indicates more intricate cross-factor entanglements.
A.2. Definition of FactoriaationDieenaaageemee of Neural Networks
Having introduced the notion intrinsic compositionality for target function or distribution, we now shift to defining fuctorization (a.k.a. disentanglement) in the context neural network's activation space. This concept applies to an intermediate layer or the entire network.
Defining factorization vs. disentanglement. Despite the frequent usage of these terms, there is no universal consensus on their precise definitions in the literature. Some works treat them as strictly synonymous, while others reserve "factorization" " for representations whose coordinates directly correspond to physically or semantically independent factors (e.g. each latent dimension controlling one ground-trute variable). By contrast, "disentanglement" can refer to a broader class of methods criteria (e.g. axis-aligned independence, mutual information scores, partial correlations). In this paper, do not distinguish sharply between these nuances. Wherever refer to "disentangled" or "factorized" representations, mean that model's 's latent (or intermediate) dimensions reflect independently varying components of the data. The specific implementation whether formalized via a 3-VAE objective, mutual-information measures, or other constraints-are secondary to the overall idea that each latent dimension ideally captures a different underlying factor. Hence, we use the two terms interchangeably throughout the manuscript.
Activation Matrix Setup. Let
X1
XN
RD
(21)
NN). x,.. the output (I,I,..,,,N) E RD. Over NiD training samples (21,,,...,, NID we assemble an activation matrix
represent the activations of a given network layer. For an input (æI,
R ND
where row is aZ
(22)
Optionally, we may project a to a lower dimension < via W E RXD (e.g. PCA), yiclding
Ofeat (21, .N) = a(z IN) T E R
(23)
14

---

Compositional Generalization via Forced Rendering of Disentangled Latents
General Definition of Factorization. A purely linear test might require an invertible matrix U E RR such that
Ofeaat (11,........ NN) u U [z1(2)). .... (xN)))
(24)
where each block of coordinates depends only a single factor zj. However, factorization may emerge only under nonlinear reparametrization. Hence, we relax the requirement to any invertible map
T:
RD (e.g. a diffeomorphism or normalizing flow).
(25)
Definition (General Factorization). We say a: &1 .... NN R is nonlinearly factorized w.r.t. (21,...;22)) 2NN if there exist:
1. An invertible transformation T R + RD,
2. A partition of {I,..., D} into N disjoint subsets U...U IN.
such that defining
= Na(Ei,,... E ERD
(26)
we have, for every coordinate E J,
A(E Ba(xz) depends only on zz--
(27)
Concretely, in these transformed coordinates (B,,..., BN). each subset of coordinates implements the feature(s) of single factor Tyz, with no cross-dependence among different factors.
Modularity via Factorization. A representation is functionally modular if, under some invertible transform T, we can rearrange the coordinates into sub-blocks that each depend on exactly one factor Zy. This directly matches the factorization criterion above. In contrast, a structurally modular network places each factor's neurons in physically separate submodules (e.g. block-diagonal connectivity). making the separation visible without any transform. Hence, functional modularity (factorization) is more general and only requires suitable reparametrization in RDD whereas structural modularity is stricter design property at the raw parameter level.
A.3. Partial Factorization, Entanglement, and Reversibility
Sections A.2 gave definition of factorization vs. entanglement: if any invertible map T can reorder the coordinates into blocks that depend purely on purely on y. we call the representation factorized; otherwise, it is entangled. In practice, however, real-world data and deep networks often produce partial factorization, where some features mix æ and y only mildly or separate them only in certain regions. Further, certain layers introduce irreversible steps (e.g. pooling, saturations). which can destroy factor structure. We claborate on these points below.
Partial Factorization. A representation need not split perfectly into two blocks (1z and I,, ) to convey some compositional benefits. For example, consider a multi-factor domain (x, y. 2) in which the representation merges two of the factors (z, y) but leaves z separate. One might say the network has partially factorized the input. More formally, partial factorization arises if some invertible map T can partition the D coordinates into sub-blocks, each depending on subset (not necessarily size 1) of the factors. This differs from pure factorization in that each block may mix more than one factor, yet remain independent of other blocks. Concretely:
Taaaz, 2)) = [B,y(E, v), {=)),, with minimal cross-dependence among blocks.
(28)
This partial structure still yields some compositional benefits (e.g. reusing the (x. y) sub-block in new contexts of 2). But may not qualify as fully "disentangled" in the strict sense of Section A.2. Indeed, such partial factorization commonly appears if x and y are correlated in training data if the network's 's architecture merges them early.
Entanglement via Irreversible Mixing. Even a representation starts factorized in an early layer, the network may subsequently entangle factors by applying irreversible operations such as:
Pooling or Striding: Combining multiple spatial positions (or multiple channels) into one, discarding the ability to invert them individually.
15

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Nonlinear Saturation: A ReLU or sigmoid can "clip" signals, so the original amplitude differences from vs. y can longer be fully recovered.
Dimension Reduction: Fully connected layers that collapse from higher dimension D to smaller one D' D can forcibly merge features.
Because definition of factorization requires invertible map 7 to unscramble the coordinates, any inreversible merge that loses information effectively prevents us from re-separating and y. Thus, irreversible mixing is a key mechanism by which networks destroy factorization.
Reversible Mixing Need Not Break Factorization. By contrast, a reversible or bijective layer can combine factors in intermediate channels without permanently entangling them, so long as there exists an inverse transform. For instance, a normalizing-Ni block may shuffle or mix (z, y) channels in a fully invertible manner; the overall system can still preserve factorization if the subsequent training does not "pinch" those degrees of freedom. In principle, any invertible map that merges coordinates can be "undone" later, so the final representation might remain factorized.
A.4. A Kernel-Based Perspective Factorization
To operationalize the above definitions of fictorization/dissttaaglcmmm concretely, we here propose kernel-based perspective and define a factorization metric that assess the factorization of the kernel learned by the neural network. The goal here is to characterize such kernel, given its potentially factorized form, understand the different regimes in which it compositionally generalize (or not).
We now consider an N-partite domain X1 X ... X Ax and a positive-semideffnite (PSD) kernel
K (X ... X(v) (X ... x xNN) R.
(29)
Under suitable conditions (e.g, an analogue of Mercer's theorem), K often admits sum-of-produch (canonical polyadic) expansion:
R
N
K k( (((1.,...,, IN), (21 xxx))
Σ
X II
(z),
(30)
where each term A, 1 (r) xz)) is rank-1 across the N factors, and Re may be finite or infinite.
Interpretation. If R 1 in (30), we say K is purely factorized (rank-1) kemel, effectively separating each factor xy. Larger Rº > 1 signals partial or full entanglement across these factors in K.
Rank-Based Measure of Factorization. We define the CP rank of K by
rank(K) K)) = min f R (30) holds exactly),
(31)
in analogy to the Schmidt rank in bipartite systems. A small rank implies that K is closer "fully separable" structure across (NN). while large rank means cross-term mixing.
A.5. Polynomial Expansions and Generalized Factorization Metric
Suppose K admits a polynomial-lic expansion enumerating single-factor, pairwise, and higher-order cross terms:
K = Σ 02 Ke,
(32)
where each K₁ might be n-factor product kernel on subset of coordinates, and 0g is its coefficient. We assign each 0g a "weight" w(rr), for instance (0gl or f8. A generalized factorization metric then measures the fraction of K's "weight" contributed by chosen subset of "lower-order" terms. Formally, if S denotes the subset of indices corresponding to single-factor (or low-order) terms, we define
Etesw(()) E, w(e)
(FFC
(33)
By varying S (e.g. only single-factor vs. single-plus-pairwise). we obtain family of "explained fraction" values values-analogous to cumulative explained variance in PCA. A higher fraction indicates that K is dominantly factorized by low-order expansions (less entangled), whereas large higher-order terms imply more cross- variable mixing.
16

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Operational Regimes of NNs in light of this decomposition. Based on the defined rank-based measure of factorization, imagine two operational regime limits of the neural networks, namely
Pure memorization regime: in which case the network resorts to tabular learning (memorization), i.c. learning single N-factor kernel term
Pure factorization regime: in which case the network learns a completely compositional kernel (assuming each of the factors are independent).
These are obviously theoretical limits, and in reality, the network probably operates regime that interpolates between the two, i.e. learning partially/imperfeetly factorized solutions.
How to estimate this metric in practice. In real-world scenarios, we rarely have direct access to the continuous kernel K: (X₂ X ... x Xx)² R. Instead, we can discretize K by selecting a finite set of data points ((Z1,).... [N, and forming the M x M Gram matrix:
K (i,) = (( (II......,, .N.)).
(34)
We can interpret this matrix as discrete approximation to K. To apply our generalized factorization metric (SA.5), we then:
1. Approximate expansions. Attempt to factor (or partially factor) K into a sum-of sum-of-product form among the N factors, e.g. through low-rank tensor methods or suitable polynomial expansions on the sampled coordinates.
2. Coefficient extraction. If we assume polynomial-lit like expansion, K x Ee Og Kr, we obtain finite- dintesiiinniional coefficients } by fitting or projecting onto chosen basis.
3. Compute the factorization ratio. Restrict attention to the "low-order" terms (e.g. single-factor or pairwise) to form a subset S. Define
Eles w(2))
PF (S)
(35)
w0,)
just as in the continuous setting, but now l indexes finite set of basis or fitted components.
This procedure yields a practical, data-driven measure how "factorized" the empirical kernel is on the sampled domain. Notably, if the sampling sparse or unrepresentative, the resulting K may fail to capture subtle cross-factor structure. In that sense, one can view PF_ (S)ef(((S) (K) as an approximate or lower-dimensioonl proxy for PF_Coeff(S) (K) in the continuous limit.
B. Metric Tensor and Factorization Metric for N-Dimensional Inputs
In this appendix, we outline a generic procedure to quantify the local geometry of a neural representation defined on an N-dimensional input space, with the metric tensor g from Ref. (Zavatone- Veth et al., 2023), and then specialize to the 2D case with an explicit factorization metric that we propose based on g.
B.I. General N N-Dimensional Definition Representation and Jacobian.
Let
f: EN - RRD
(36)
where f(x (fi...., fo) for each input vector x (I1, (UI...,, IN). We define the Jacobian J/ of f a point by
0fi/0z
0fi/āz
J(x) =
E D
(37)
0fp/0x1
@fp/Oan
17

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Metric Tensor.
The induced metric lensor g(x) is then N N matrix,
g(x) J(x J(x),
(38)
where (x) /0 Geometrically, 9 captures how infinitesimal distances in input space map to distances in the representation space IR By computing g, we gain insights into how the native space becomes wraped in the embedding space (Fig. 5).
native space
embedding space
[(
Figure The action of the neural net can imagined as wrapping of (originally flat) native space (left) into embedding space (right).
Volume Element.
The local volume element is given by
(x(x) (e(( dæ ... dæN.
(39)
This indicates how the network "stretches" or "distorts" volumes near each x.
Factorization. While there is no single universal scalar for "factorization" in the N N-dimensional case, one can investigate off-diagonal elements of Jur sub-blocks of the Jacobian) to see how strongly certain input dimensions (zz, mix with others (st,)). For instance, if 0f /rx approximately orthogonal 0f/0:x 0rrr for all μ A the representation preserves independence among coordinates.
B.2. Specialization to the 2D Case
When . 2, let the input be (Ir v). Then the metric is 2 x matrix:
0f
0f
g
(40)
By
Its volume element is
du(z, det(g(z, v)) dz dy.
(41)
A Simple Factorization Metric in 2D. If desires single scalar capturing how "uncoupped"" and y remain, define:
1 (x,, v)
Factorization(((r,,
(42)
9 (2,, w)1 + (
where 92z = 90,0- 9yv 91.1, and Juy = 90.1 91.0- In words:
Ive and Oyy measure how much the representation changes when moving purely in the y-direction.
Jay captures cross-direction entanglement.
The closer Factorization(g) is to 1, the more orthogonal (uncoupled) these directions are in D
By averaging this score over grid (a, 3/ points, one obtains a global measure of whether the representation remains factorized becomes entangled.
18

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Without Regularization
With Regularization
Ze0
1420
Figure 6. Visualization of and y embedding matrices without and with regularization.
C. Low-rank Regularization for Compositional Generalization
In our generative approach, each input factor (e.g. the one-hot code for coordinate or y) is mapped to a two-dimensional embedding matrix. Concretely, suppose we have factor "y (which can æ or y), and we define a matrix
(w,( E Iddd
(43)
where the indices (i,j) correspond row and column in this 2D embedding. The matrix W()) is then passed, possibly along with other embeddings, into subsequent layers (e.g. small convolutional network) to produce the final output image.
Low-Rank Decomposition. To encourage each embedding matrix Waj to remain factorized, represent it via an SVD-like decomposition:
Wis Σ (n) A (m)) (m)))
(44)
where is an upper bound on the rank, (m) scalar coefficients, and 0s (n) (m) for the n-th mode. We then regularize these parameters encourage low-rank and spatially consistent solution.
uj are the row- and column-basi vectors
Entropy and Variance Penalties. First, interpret the squared magnitudes [A(m)]2 as probability distribution across
[1mn)22
(m)
(45)
2
and add an entropy penalty (DeMoss et al., 2024)
Cent
1 (()) In [m))
(46)
Minimizing Lent reduces the distribution's entropy, making only few modes dominant (i.e. low rank).
Second, we penalize large row- or column-wise variance of the vectors (n)), (n), Concretely, let Var(u()) denote the sample variance (n) across indices Then we define
Lvnr 1/2 Σ Var(ulul) + Var(uin))]
(47)
19

---

Compositional Generalization via Forced Rendering of Disentangled Latents
(a)
(b)
pHON
p-208
p-40%
p-bo%
p-Bo%
p- 100%
Figure 7. Augmenting 2D Gaussian bump dataset with 1D Gaussian stripes enables compositional generalization. (a) shows the ID Gaussian stripe dataset for N 28. (b) shows sample generated OOD images for models trained on ID Gaussian stripes p% ID 2D Gaussian bumps subsampled.
Minimizing Lvar discourages the model from "memorizing" individual entries (i,j), thus promoting spatial invariance. Summing both these terms along with the usual reconstruction loss (e.g. MSF) yields
Liotal
LMSE
+
Lent + Loaa-
(48)
By applying (46) to the embedding matrices for both and y-factors, the network learns factorized, compo- sitional representation that generalizes better to unseen coordinate combinations. Fig. 6 demonstrate the 2D embedding matrices corresponding to different values of x and y in the cases without and with the entropy + variance regularization. In the case of no regularization, the model fails to preserve the integrity of the signal across the OOD region. In contrast, with regularization, each embedding matrix exhibits more pronounced vertical or horizontal banding that aligns cleanly with the chosen or g-coordinate. This indicates the network has discovered a more factorized structure, making it easier to recombine and y generate 2D Gaussian bump in unseen OOD configurations.
D. Experimental Details
In this section, we describe the neural network architectures used for generating bump functions, including exact layer sizes, activation functions, optimization algorithms, hyperparameters, and training protocols. Code to reproduce the experiments can be found this Github repository. We explored various input formats and model depths; below are detailed descriptions of representative CNN and MLP models using bump-encoding inputs:
Property
CNN Architecture
MLP Architecture
Input Dimensions
56 (reshaped to 56 x 1 1)
56
Output Dimensions
x 28 28
x 28 x 28
Layers
TranTranspose2d layers (upsampling from x to 28 x 28) 64 channels per hidden layer
fully-connected linear layers
Hidden Layer Size
256 units per hidden layer 272K
Parameter Count
~315K
Activation Function
ReLU
ReLU Adam W (Learning Rate: 1 x 10-3)
Optimizer
AdamW (Learning Rate: 1 x 10-3)
Table Comparison of CNN and MLP architectures
These details clearly communicate the architectural specifies and ensure reproducibility.
20

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Cosine similarity
L2 distance
Soft overlap
1.000
2.5
0.005
OOD
0.999
OOD
OOD
Bump Encoding
600
600
600
2,0
0.004
500
500
500
400
1.5
400
0.003
0.997
300
ID
300
ID
1.0
300
ID
0.002
200
0,996
200
200
0.5
0.001
100
100
100
0.0
0.000
200
400
600
200
400
600
200
400
600
Cosine similarity
L2 distance
Soft overlap
1.000
700
700
700
OOD
OOD
00D
Rate Encoding
600
600
2.0
600
500
500
500
0.003
0,998
0.002
300
300
1.0
300
ID
0.997
ID
ID
200
200
200
0.001
100
0.0
0.000
600
200
400
Figure 8. Full similariy/distanecovrrap matrices visualization across the entire ID-OOD data distribution, compared between networks with bump (top) and rate (bottom) input encodings.
E. Dataset Augmentation
One way of improving compositional generalization our synthetic toy example to augment the 2D Gaussian bump dataset with ID Gaussian stripes, thereby encouraging more factorized representation the pixel domain. Concretely, we generate vertical stripes by fixing a particular xzindex and applying 1D Gaussian along the y-axis, and horizontal stripes by fixing y-index and applying ID Gaussian along the r-axis. Fig. 6(a) shows these stripes for N 28, with each panel corresponding to specific row or column coordinate. We then combine all stripes with fraction p% of the original 2D bump dataset, so that examples of the ID) patterns are always present even when the 2D samples are subsampled.
This augmentation strategy allows the network to learn separate "axes" variation (along x and y) more explicitly, Empirically, as shown in Fig. 6(b), models trained on these augmented datasets exhibit improved extrapolation out- of-distribution (OOD) regions, even with limited 2D bump data. The presence of ID stripes encourages form of compositional inductive bias, wherein the model learns to additively combine horizontal and vertical factors in the pixel domain. Consequently, we gain the insight that guiding models toward simpler, factorized building blocks can significantly enhance their ability to recombine learned components in novel settings, an essential capability for robust compositional generalization.
F. Supplementary Experimental Results F.1. Full Kernel Characterization
In Sec. 3.3, we characterized the model's OOD generalization behavior from a kernel-based perspective. In Fig. I(c). showed the agreement between the similarity (overlap) matrix with the theorized binary factorized kernel, which showed that the model independently "processes" novel input combinations of and y. Here, we provide the full similarity/distanceovveia matrices across the entire generated dataset, ID and OOD. We note that Fig. I(c) essentially shows a slice of what are shown in Fig. 8. We compare between generated image by models trained with bump-encodin and ramp-ence annp-ennoding inputs, and the matrices are constructed based on three metrics: 1) cosine similarity, 2) L2 distance, and 3) pixel overlap. We notice that the kernel matrices of bump-encoded model exhibit stronger "block" structure as compared
21

---

Compositional Generalization via Forced Rendering of Disentangled Latents
(a)
Recordtruction MSE Error Rotation
160 0M
162.1
164 2'
166.3
168 4'
170.5"
172.6'
174,7
176.8'
178.9
Out ctan
Dicributtonn Meae: 0.607772
A_
A
dip Moon 031297
0.64
0.03
101 2"
183.2
185 3
187,6*
191.6" 6*
193,7
195.8"
197,9*
200,00
0.02
0.03
(d)
160
162.1
164 2'
166.3'
168 4'
170.5
172.6'
176.7
176.8'
178.9'
135
180
325
915
362
X
(b)
Eirbobling
Aodl theg v00)
182.3 3"
183.2
185 3"
187,4
189-5
191.6"
193.7'
195.8"
157.9"
200.0
20
Figure 9. MNIST digit "4" image rotation experiment with OOD range of 160" 200 (a) shows the reconstruction MSE as a
function of rotation angle. (b) shows the 2D UMAP embeddings for representations at different layers of the network. (c) and (d) shows the generated rotated images in the OOD range vs. the pixel-wise interpolated images hetween 160 and 200".
to those of the ramp-encoded model. Here we note that the matrices have interesting diagonal structure from the non-finite width of the Gaussian bumps.
F2. MNIST Digit Image Rotation Experiments
In Sec. 3.3, we mentioned that the phenomenon of activation superposition corresponding memorized ID data seem be generically applicable beyond the compositional setting. In Fig. 9, we present simple toy experiment rotating single MNIST image data of number "4." We generate a dataset of this rotated image "4" for various rotation angles 0 E [0, 360"] while leaving out an OOD range gOod E (160", 200%). We then train CNN on the ID data, with input (cos sin 0). and test if the model can interpolate/extrapolaat to the good range. Unsurprisingly, the model fails to generalize to the unseen rotation angle range, as shown by the MSE plot in Fig. 9()). Nevertheless, the generated OOD images reveal that the model performs linear interpolation between the edge-most ID samples, namely the images rotated at 160" and 200". In Fig. (d), we show the linear interpolation between the ground truth rotated images at 160% and 200", which strongly resembles the generated OOD images shown in Fig. (c).. Intriguingly, in Fig. 9(b) we show that the UMAP-embedded latent representations across the network layers remain smooth without significant deformation, unlike the case of 2D Gaussian bump generation presented in the main text. Nonetheless, the OOD outputs of the model as linear interpolation of ID images confirms that the model is indeed memorizing all ID images and superposing the corresponding activation patterns, which resonates with our findings for the 2D Gaussian bump generation experiment.
F.3. Volume Metric Plots and Generated OOD Images
In Sec. 3.4, we characterized manifold warping using the Jacobian based volume metric, as well as factorization and linear probe metrics as a function of layer depth. Here we show the same metrics for networks of different depths, for both CNNs (Fig. 11 and Fig. 12) and MLPs (Fig. 13 and Fig. 14). We also show the generated images as a function of network depth for CNNs in Fig. 10. An interesting phenomenon here is that the MLP models seems to have the opposite "warping" effect in the OOD region compared the CNN models. Moreover, the MLPs seem to have smoother but more drastic degradation of factorization function of layer depth. Nonetheless, expect similar conclusions regarding disentanglement and factorization to hold for CNNs and MI.ss alike.
E.4. Neural Tuning Curves
In Sec. 4.2, we discussed the potential benefits of training models with augmented datasets that contain isolated factors of variation. Specifically, we showed that models trained with stripes+bumps are able to compositionally generalize in data efficient manner. Mechanistically, the stripes provide the model a shortcut bias that leads the model to form additive composition in early layers of the network, which manifests as "stripe conjunctions," which become "bumps" in downstream layers. In Fig. 15, we visualize four layers of a CNN trained with stripes only (left) and with stripes+ 100 bumps (right). Each sub-panel shows 2D activation map for selected neuron, along with its mean activation when collapsed over y
22

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Bump Encoding
Rate Encoding
5
10
20
10
20
10
20
10
20
10
20
20
20
20
20
D
10
10
10
10
10
10
10
20
10
20
10
20
10
20
10
20
10
20
10
20
10
20
Dapile3
20
20
20
20
70
70
70
Ove
10
10
10
1.0
10
10
10
10
0.
10
20
10
20
10
20
10
20
Depihwe
Deptheu
10
10
10
20
10
20
Degthas
X
X
10
10
20
10
20
10
20
10
20
10
10
20
20-
20
20
20
Depthwc
10-
10
10-
10-
10
20
10
20
10
20
0-
10
20
10
20
10
20
0
10
20
Figure 10. Comparison of generated OOD images for different network depths between networks with bump-based (left) and ramp-based (right) input encodings.
(middle plot) or (right plot). Across the layers, we observe that early-layer neurons primarily capture stripe-like features, with broader activations that vary more strongly along one dimension than the other. In deeper layers, these "stripes" begin to intersect more tightly, forming increasingly localized "bump" patterns. This indicates that the model successfully learning to capture the complex compositional structures present in the augmented data, reflecting the additive- to-multiplicative transition its internal representations. Furthermore, as bumps accumulate, the activations demonstrate clear shift toward more focused and precise tuning, which supports the idea that the neural network generalizing bump-like structures efficiently through additive composition early in the architecture and refining this understanding as the network deepens.
23

---

Compositional Generalization via Forced Rendering of Disentangled Latents
1.00
Depthe2
Factorizatian Metric
Depth-3 Depthwd
0.99
Dopth-5
Depthm6
0.98
Layard
Layer
Lover
11
-25
0.97
Input
Layer Index
20
Liver
Leeer
Unver
Lawer
-19
Lacro
Lggar1
Loyar3
Layer )
Lewer
Lawerd
348
2200
Placs
Laver
Leart
lever3
Iasr4
Layer 5
1000
26
Figure Evolution of manifold density and linear probe as a function layer depth for different CNN network depths (hump encoding).
1ape:
Tape:
1.00
Factorization Metricc
Soure
0.95
Depth-2
Depthe3 Dopth
130
Daptiws
Depti-6
-12
13
Layer
Layrr
Lswr1
-10
61 wny
12pm
Dsyr1
Pate
08
-13
Lager
Userr1
Liver 2
Liver
Livers
Laee
4.2
64599
-10
-1.0
-0009
-2
Lave BACK
Figure 12. Evolution of manifold density and linear probe as a function of layer depth for different CNN network depths (rate encoding).
24

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Depthm
Factorization Mecric
Depthed
Depth-
9.90
Dopthw6
Depth-?
80
0.75
Input
Layer Index
27
6.30
0.500
0.0
0.00
0.7
Figure Evolution of manifold density and linear probe as function layer depth for different MLP network depths (hump encoding).
1.00
Depth=3
0.95
Depth-4
Dopth=5
0.90
Depthm6
Dapth-7
0.85
03
0.30
Pary
0.75
0.70
Input
Layer Index
110000
Luwer
-16
Lser1
Lieert
12y0r
Figure 14. Evolution of manifold density and linear probe as function of layer depth for different MLP network depths (bump encoding).
25

---

Compositional Generalization via Forced Rendering of Disentangled Latents
Layer Neuron [c=1.he2 mm
Layer
Nean overy
Mean over
20
Avg AC
25 2.
3.45
2.4
Nozn overy
Meon over
Neuron IcmS, (mm)
Near over
Moan
hun 6.w55)
Dey Ket
Stripes 100 Bumps
-0.50
Stripes Only
1.0
Neuron (C-E.h----55
Moan overy
Mean over
Neuren 3.=3,01m8 Wwdd)
Mean every
Nean over.
Mct
Avg Kr
St
24
25
Nouron (cwó he 3. Welll 131
Mean overyy
Mean over
Neuron (c=0) 20m0) 200.)222... mm
Layer
Mean overlyy
Mean over
0.23
152
21-
Figure 15. Neural tuning curves for sample neurons across different layers of 4-layer CNN trained dataset of only stripes (left) and stripes + 100 bumps (right).
26


Reviews:
[{"review_id": "uLnvn1g0l2", "reviewer": "Reviewer_vKC8", "summary": "The authors investigate the role of disentangled representations in compositional generalization, which remains unclear in the literature. The authors observe that while inputs may be disentangled, this disentanglement can \"erode\" through subsequent layers, such that the model overall is not able to generalize OOD. Looking deeper, the authors find that this \"erosion\" usually stems from models memorizing training instances. Finally, the \"erosion\" can be mitigated via architectural constraints or curated datasets. Overall, disentangled representations seem to be a necessary but not sufficient condition for compositional generalization.", "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": "**Input Encoding Types**\n\nWe appreciate the reviewer's suggestion regarding the comprehensiveness of input encoding formats. Indeed, we have included both rate-based (scalar) and population-based encodings (e.g., 1-hot, Gaussian bumps, and positional encoding), as depicted in Fig. 1. These variations produced qualitatively similar outcomes. We will explicitly clarify this range of encodings and add a table summarizing ID/OOD MSE comparisons for clarity in the revision.\n\n**Shape of the OOD Region**\n\nWe agree the shape of the ID/OOD region could influence results. To address this concern, we previously conducted experiments comparing \"circular\" and \"square\" OOD regions, which showed no qualitative differences. Regarding correlated factors (e.g., diagonal strips), our primary focus was on the simpler case of independent factors, as even this simpler scenario already presented significant challenges for compositional generalization. We will clarify this reasoning explicitly in our revised manuscript.\n\n**Quantitative Averaging Over Multiple Runs**\n\nWe confirm that quantitative results shown in Figures 3 and 4 are indeed averaged over multiple experimental runs. To enhance transparency and reader confidence, we will include a clearly labeled table of these averaged ID/OOD performance metrics in the revised manuscript.\n\n**Fixed Coordinate Inputs**\n\nWe appreciate the opportunity to clarify the input setup. The \"fixed\" coordinate for the stripes was provided explicitly as -1 (off-canvas), ensuring clear separation between stripe and bump datasets. This detail will be clearly stated in the revised manuscript.\n\n**Potential Issues with Stripe Augmentation**\n\nWe appreciate the reviewer’s insightful concern about the stripe augmentation potentially promoting superposition-based memorization. The key significance of our stripe augmentation experiment lies precisely in demonstrating compositional generalization: the model generates correct 2D bumps in OOD regions where it never previously observed bumps during training. Even if stripes are memorized individually, the model successfully leverages this information to construct completely novel, unseen bump configurations. Thus, the stripes serve as a spatial scaffold, facilitating pixel-level signals necessary for learning factorized variations within entirely unseen regions. We will highlight this clearly in our manuscript.\n\n**Generalizing Data Augmentation Strategy**\n\nPlease refer to our response to Reviewer 3 (iydz) regarding the same point. \n\n**Fixed Coordinate Bumps Experiment**\n\nWe appreciate the reviewer’s insightful suggestion. While practically feasible, this approach might not encourage compositional generalization effectively because each Gaussian bump inherently includes both coordinates, regardless of the dataset's internal variation. The key challenge for the model is still establishing the correspondence between different input dimensions and pixel-level output variations. We will explicitly discuss this consideration in our manuscript revision, clarifying the challenge presented by compositional product structures.\n\n**Clarification on Figure 1e and Section 3.3**\n\nWe apologize for any ambiguity regarding Fig. 1e. To clarify, Fig. 1e visualizes similarity between representations of OOD-generated samples (sorted by coordinates) and an idealized binary kernel (similarity=1 if coordinates match exactly, 0 otherwise). The \"agreement\" quantifies how closely the learned similarity matrix aligns with this ideal binary factorized kernel. Specifically, Fig. 1e illustrates overlap between Gaussian bump images centered at coordinates (x,y) and (x′,y’), comparing the ideal binary kernel (top panel) to actual model outputs (bottom panel). Fig. 1e corresponds to the top-left sections within matrices shown in Fig. 8. We will clearly state equations defining \"agreement,\" explicitly describe the data subset (OOD samples), and provide a detailed caption to eliminate confusion.\n\n**Other References**\n\nWe thank the reviewer for highlighting important references and will ensure their inclusion and proper discussion in the revised manuscript.\n\n**Regarding Novelty and Main Contribution**\n\nWe agree with the reviewer’s insightful point that the observation \"disentanglement alone is insufficient for compositional generalization\" is not novel. Our primary contribution is identifying and explaining the failure mechanism (superposition and manifold warping), providing mechanistic insights into why factorized latent representations fail to generalize compositionally. Please refer to our response to Reviewer 3 (iydz) regarding the same point. To clearly communicate this, we will revise our manuscript title and abstract accordingly.\n\n**Additional Comments and Suggestions**\n\nWe will revise the manuscript accordingly. \n\nOnce again, we sincerely thank the reviewer for these valuable suggestions, which will undoubtedly strengthen our paper's clarity and impact.", "claims_and_evidence": "1. Factorization alone, independent of input encoding formats, is not sufficient for comp. gen.\n    - This claim is mostly well supported by the data, especially Fig. 1b-c.\n    - More comprehensive numbers (e.g., MSE ID/OOD for different encodings in a table) would have been nice to see.\n    - The different input encodings are not exhaustive, e.g., what about simple (normalized) scalar inputs instead of 1-hot encoding?\n    - Does the shape of the ID/OOD region play any role? E.g., if the ID region was a diagonal strip as in [3] (see below), correlation of factors in the training data might be an additional issue.\n2. Failures are due to memorization and \"superposition\" of seen data points.\n    - This claim is somewhat supported by the plots in Fig. 1b-d.\n    - The kernel perspective from Sec. 3.3 and Fig. 1e should give a more comprehensive perspective, but it is somewhat unclear to me what Fig. 1e depicts, see below.\n3. Manifold warping ruins factorization\n    - This claim is clear and well supported by the evidence in Fig. 2.\n4. Architecture/regularization can encourage compositional generalization.\n    - This claim is mostly supported by the visualizations in Fig. 2 and especially the ablation in Fig. 2d.\n    - However, the visualizations in Fig. 2b,c are for one specific model only. A more comprehensive comparison, e.g., in terms of average ID/OOD performance over multiple models with/without regularization in a table could give increased certainty that the results hold in general.\n5. Dataset augmentations can encourage comp. gen.\n    - As with 4, this claim is mostly supported by the results, but more general results averaged over multiple models would increase confidence in the visualizations.", "methods_and_evaluation_criteria": "The proposed evaluations make sense for the most part. However, I have some questions about the evaluation setup in Sec. 4.2.:\n\n- Was the fixed coordinate still an input to the network?\n- I find the stripes somewhat problematic since, by definition, their superposition results in the Gaussian blob. If the model simply memorizes and superimposes factors as before, this would mostly \"solve\" the OOD case, as we can see in Fig. 7b with $p=0$%.\n- What could this augmentation look like for other types of compositional data? I'm having a hard time imagining an analogy to the 1d stripes for factors such as \"color\", \"shape\", \"size\" in, e.g., a sprites setting.\n- How does this compare to model performance when restricting the ID set to x/y combinations with 1 fixed coordinate, but where the output is still a bump (i.e., the dataset consists solely of bumps on the left with varying y, or bumps on the bottom with varying x)? I find this to be there more widely applicable data augmentation that also provides the model with independent influence of each factor.", "theoretical_claims": "The paper makes no theoretical claims.", "experimental_designs_or_analyses": "The analyses are mostly clear, except for:\n\n- Sec. 3.3/Fig. 1e: What does \"agreement between the binary factorized kernel and the similarity matrix between the OOD and ID generated samples\" mean? It would be helpful if this could be formalized in an equation. What regions of the complete plots in Fig. 8 does Fig. 1e correspond to? What regions in Fig. 1e correspond to ID/OOD? From the y-axis label it seems like the entire plot corresponds to OOD samples?", "relation_to_broader_scientific_literature": "In my understanding, the main and titular observation that disentangled representations are not sufficient for comp. gen. is well known in the literature and has been shown in multiple prior works, including Wiedemer et al., 2023 (cited in the paper) as well as [1,3,4] (see below).\n\nWiedemer et al., 2023, and [3,4] specifically show theoretically that compositional inputs _alone_ are not enough and additional assumptions in the form of regularizations, architectural constraints, or conditions on the training data are required.\n\nThat said, the additional kernel and transport perspectives on how disentangled representations might be \"diluted\" throughout a model, and the (to my understanding) novel regularization scheme in Sec. 4.1. are still interesting, but the main claim should be adapted to properly reflect prior work, which might also have to be reflected in the title.", "essential_references_not_discussed": "- [1] _Montero et al., 2022, Lost in Latent Space: Examining failures of disentangled models at combinatorial generalisation_. This is a follow up to Montero et al., 2021, which the authors cited. It investigates whether failure to generalize compositionally is largely due to a model's inability to disentangle inputs (the encoder), or its inability to generate new compositions (the decoder) and finds decoder errors to be most prominent.\n- [2] _Schott et al., 2022, Visual Representation Learning Does Not Generalize Strongly Within the Same Domain_. This paper is similar to Montero et al., 2021, but focuses on the encoder rather than the decoder\n- [3] _Wiedemer et al., 2023, Provable Compositional Generalization for Object-Centric Learning_. Similar to [1], this paper shows decoder errors to be responsible for failures to generate compositionally, which can be overcome by additional architectural constraints. The paper discusses how models that disentangle latent factors can be guaranteed to robustly compose them, albeit with a focus on object-centric learning methods.\n- [4] _Lachapelle et al., 2023, Additive decoders for latent variables identification and cartesian-product extrapolation_. Like [3], the authors explore compositional generalization in object-centric learning.\n\nSpecifically, [1], [3], and [4] hint at the finding that factorization must be actively maintained, e.g., via additional architectural constraints.", "other_strengths_and_weaknesses": "As outlined above, I believe the main contribution (in its current phrasing) is not novel, however, additional insights provided in Secs. 3.3-4.2 are still interesting.", "other_comments_or_suggestions": "- I would suggest picking either \"disentangled\" _or_ \"factorized\" and using it consistently throughout\n- unclear how the OOD coordinates in Fig. 1e come to be\n- Fig. 2a-b, e-f and Fig. 4e are too small to read. There is white space to both sides, a different arrangement of plots might be possible.\n- Page 6, left, first paragraph ends in incomplete sentence", "questions_for_authors": "Please refer to the questions 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": "LjjjjXeiie", "reviewer": "Reviewer_g6yu", "summary": "The paper investigates conditions under which a neural network learns to generalize \"compositionally\". The setting involves learning to generate a 2D \"bump function\". A key result is that a \"disentangled\" representation is not sufficient to ensure compositional generalization. The authors then describe data curation and regularization strategies that appear, in their synthetic setting, to be sufficient for compositionality.\n\n## update after rebuttal: As described in my comment later in the thread, I am keeping my initial 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": "We sincerely thank the reviewer for their thoughtful and encouraging feedback on our work. We greatly appreciate the insightful comments and constructive questions raised, and we respond to each point concisely below.\n\n**Definition of Compositionality and OOD Generalization**\n\nWe thank the reviewer for this accurate observation. Indeed, compositional generalization, or “OOD generalization” in general is a very broad term. For a task that involves extrapolating far from the training regime (as studied in our work), a central challenge is the lack of well-defined “ground truth.” Thus, we adopted a setting where OOD generalization explicitly involves combining two independently varying factors, representing perhaps the simplest compositional scenario.  Interestingly, even in this minimal setup, our results illustrate that neural networks struggle without additional interventions, highlighting broader challenges in compositional generalization. There are indeed many more different forms of compositionality, see our response to Reviewer 3 (iydz) regarding novelty for a discussion of interactive vs. non-interactive compositionality. \n\nIn terms of the model's superposition strategy, while the model indeed learned a highly nontrivial, compositional solution, it is at the activation level, rather than the pixel level, which led to failure of the desired generalization performance, revealing the fact that it is memorizing the ID data rather than breaking them down into composing factors. Our definition of compositionality and our performance metrics aim at assessing the model’s ability to construct novel compositions from factors of independent variation. In the superposition case, the model fails at constructing novel compositions correctly.\n \n**Experimental Details and Reproducibility**\n\nWe acknowledge the importance of reproducibility and will provide precise architectural and training details in the revised manuscript. Specifically, we will clarify the architecture of the neural networks used for generating bump functions, including exact layer sizes, activations, optimization algorithms, hyperparameters, and training protocols. Additionally, we will release the code to ensure complete reproducibility and facilitate further exploration of our findings. We have explored a plethora of different input formats and model depths. The architecture details for a sample CNN and an MLP taking bump-encoding inputs are given below:\n\n| Architecture | Input Dimensions          | Output Dimensions | Layers                                           | Hidden Layer Size | Parameter Count (example) | Activation | Optimizer & LR |\n|--------------|---------------------------|-------------------|--------------------------------------------------|-------------------|---------------------------|------------|----------------|\n| **CNN**      | 56 (reshaped to 56×1×1)   | 1×28×28           | ConvTranspose2d (upsampling from 7×7 → 28×28)   | 64 channels       | ~315K (4 hidden layers)   | ReLU       | AdamW (1e-3)   |\n| **MLP**      | 56                        | 1×28×28           | Fully-connected Linear layers                    | 256 units         | ~272K (3 hidden layers)   | ReLU       | AdamW (1e-3)   |\n\nWe will include this summary table in the revised manuscript to clearly communicate the architectural details and facilitate reproducibility.\n\n**Clarification of Equation (8)**\n\nWe thank the reviewer for pointing out this potential source of confusion. The vectors and eigenvalues referenced in Equation (8) are specifically used to construct the weight matrix within our network. This formulation enables targeted regularization of different components of the weight matrix more efficiently, promoting factorization and preventing representational warping. This is especially helpful since a separate 2D embedding matrix is dedicated to each dimension in the factorized input, which encourages the model to learn for pixel-level factors. Indeed, the low-rank regularization of these embedding matrices encouraged the model to find “stripe-like” factors as shown in Fig. 3b. In the revised manuscript, we will elaborate further on Equation (8), explicitly explaining the mathematical intuition behind this approach and clarifying how it aids compositional generalization.\n\n**Minor comments**\n\nWe will revise the manuscript to get rid of quotation marks around superposition accordingly. \n\nIn summary, we sincerely thank the reviewer again for these valuable comments, which significantly help enhance our manuscript. We are confident that incorporating these suggestions will improve the clarity, rigor, and accessibility of our paper.", "claims_and_evidence": "The claims seem reasonably supported, and I appreciated the extensive use of visualizations to analyze the results. The setting is simple enough that the results seem fairly transparent—that's a plus, and kudos to the authors for creating such an easy-to-analyze model.\n\nIf I have any reservations, it's that the specific definition of \"compositional generalization\" (as described briefly in the body of the paper, and more extensively in Appendix A) seems potentially too narrow. When I look at the function learned in Figure 1d, for example, it actually seems like a fairly clever OOD generalization: the network appears to have learned that (1) center pixels are empty, (2) there must be bumps at the given x- and y-coordinates; (3) the centroid of the bumps appears at the given x- and y-coordinates. And this generalization appears \"compositional\" in the sense that it is superposing (adding) separate solutions for x- and y-coordinates.", "methods_and_evaluation_criteria": "The methods seem reasonable for the given toy problem. One issue is that I didn't immediately see exact details on some of the networks used (e.g. for the bump functions). It's not clear to me there are sufficient details to reproduce the authors' experiments precisely.", "theoretical_claims": "Generally the claims made sense to me. However, I have to admit I don't completely follow what's going on with equation 8, and would like to see more detail.", "experimental_designs_or_analyses": "The experimental designs made sense, but note the question above about experimental detail.", "relation_to_broader_scientific_literature": "The related work sections seems good: in particular, I appreciate that it was focused and did not bring in irrelevant details.", "essential_references_not_discussed": "n/a", "other_strengths_and_weaknesses": "Strengths:\n- The results from the initial toy model are not particularly surprising in retrospect, but are a very nice illustration that generalization doesn't always happen in the way that one would expect. I don't know of this example in the literature, and it seems like a good addition.\n- The regularization described in section 4 appears to be a simple way to encourage compositional generalization; I'm not 100% sure I understood it, but if I did, it seems potentially useful.\n\nWeaknesses:\n- I would have liked a more careful description of why this particular definition of \"compositional\" was chosen\n- Not clear there's enough detail to replicate the experiments exactly.", "other_comments_or_suggestions": "I don't think you need to keep adding quotes around the word \"superpose\".", "questions_for_authors": "Can you say more about equation 8?\n\nIs code available for this? If so, it would help in reproducing the results; apologies if I missed this in the text.", "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": "NzrKxP3okV", "reviewer": "Reviewer_iydz", "summary": "The authors investigate why disentanglement is not sufficient for compositional generalization. First they observe that models are unable to reconstruct simple bumps in unseen locations in visual space from fully factorized latents. They use this as evidence that disentanglement alone is not enough for compositional generalization. They then proceed to show how different forms of regularization and data augmentation can help models to achieve better generalization. These include penalizing the entropy and variance of the singular value of filters in the transformation from representation to input and modifying how data is presented so that models learn about each factor independently.", "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": "**Novelty and Positioning Relative to Prior Work**\n\nWe fully agree that claiming the insufficiency of disentanglement is not novel (also see response to Reviewer 1, vKC8). Our primary contribution lies instead in providing a detailed mechanistic explanation of why disentanglement fails—specifically, through manifold warping (due to topological deformation) and superposition-based memorization. We thank the reviewer for highlighting Montero et al. (2022), who primarily attribute compositional failures to encoder limitations rather than decoder issues. In contrast, we focus exclusively on decoder failures. Montero et al. observe that models succeed in non-interactive but struggle with interactive composition, consistent with our findings that the default composition mode is superposition-based, suitable only for non-interactive scenarios. They suggest interactive composition requires learning causal interactions between factors; notably, our proposed methods (embedding regularization and data augmentation) explicitly encourage learning these causal interactions, which we demonstrate to be effective empirically.\nIn our revision, we will amend the title and abstract to clearly highlight our contribution as providing deeper diagnostic insights into compositional generalization failures, explicitly acknowledging prior work and positioning our study accordingly.\n\n**Applicability to More Complex Datasets and Real-World Data**\n\nWe agree that translating insights from our synthetic study to more complex, real-world datasets is an important future direction. Indeed, as previously emphasized by Montero et al., interactive compositionality cannot be addressed by a one-size-fits-all solution, and accordingly, we do not claim to propose a universal remedy. While our specific approach may not directly generalize to complex datasets, it highlights two valuable exploratory directions: (1) training modular, output-level embedding filters dedicated to each disentangled input dimension, and (2) dataset augmentation with isolated factors of variation. Our simplified setting was specifically chosen to clearly illustrate underlying mechanisms of compositional failure (such as models resorting to memorized ID data for OOD generalization—a nontrivial yet unsuccessful mode).\n\nIn our revision, we will explicitly acknowledge that generalizing our diagnostic insights to more complex datasets remains open and valuable. We appreciate the reviewer’s suggestion of explicitly exploring common disentanglement datasets and will discuss this as a key direction for future research.\n\n**Relationship to Curriculum Learning and Cognitive Science**\n\nWe greatly appreciate this insightful connection. Indeed, our data curation method—introducing each factor independently—does resonate closely with curriculum learning strategies in cognitive science. In the revised manuscript, we will explicitly acknowledge this parallel, highlighting that isolating concepts before combining them is a recognized effective strategy both in human learning and potentially in artificial neural networks. We will cite relevant cognitive science literature as suggested, discussing this interesting link.\n\n**Regularization Techniques and Architectural Constraints**\n\nWe thank the reviewer for this important caution. Our regularization method differs from prior work by explicitly targeting the singular-value structure of the decoder’s weight matrix to mitigate topological deformation (warping). Nevertheless, we recognize and appreciate the reviewer’s caution about prior regularization strategies’ limited success in broader domains. We will clarify this distinction explicitly in our revision, emphasizing that while our targeted regularization approach is effective in our controlled setting, its broader applicability is an open and important question.\n\nRegarding architectural constraints, we strongly agree that identifying general architectural motifs applicable across modalities and concepts is a promising direction. We will explicitly discuss this in the revised manuscript, noting that the identified failure modes and solutions (such as factor-specific architectural constraints) could be generalized across other data modalities.\n\n**Terminology and Clarity**\n\nWe thank the reviewer for this valuable suggestion. We will comprehensively revise the manuscript to adopt a more precise and appropriate tone, clearly positioning our paper as providing mechanistic diagnostic insights rather than fundamental novelty in identifying the insufficiency of disentanglement. We agree this will significantly enhance clarity and ensure accurate positioning within the existing literature.\n\nWe sincerely thank the reviewer again for their detailed comments, constructive criticisms, and thoughtful suggestions. These revisions will substantially improve the clarity and positioning of our paper, and we appreciate the reviewer’s help in guiding this important improvement.", "claims_and_evidence": "Yes, the claims made in the article are well supported by empirical simulations.", "methods_and_evaluation_criteria": "They do, though as is often the case in ML/DL, there is a hope to see insights in toy datasets translated to more complex datasets. In this case the datasets used are very simple (just gaussian bumps), so it is unclear if these insights translate to other settings.", "theoretical_claims": "There were no proofs in the manuscript.", "experimental_designs_or_analyses": "The experimental design and analysis is sound.", "relation_to_broader_scientific_literature": "They are, but also not very novel. My feeling is that the authors put too much emphasis on what they have found (which is not novel) when they should put more emphasis on their analysis. Specifically, the idea that just having fully disentangled representations does not help compositional generalization was already pointed out in Montero et al., 2021 (which they discuss, making this omission a bit puzzling. See the last section in that article.). Additionally, their design where a central pattern is excluded from reconstruction was already explored in [1], though admittedly it appears in the appendix. Finally, the insight that enforcing some factorization in output space is required to perform compositional generalization was already hypothesized in [2] (final section). Thus the article sits in a weird position where it retreads some points that have already been made, reaching similar conclusions even if they do so via different methods and perspectives (which I believe is still valuable).\n\n\n- [1] Watters, N., Matthey, L., Burgess, C. P., & Lerchner, A. (2019). Spatial broadcast decoder: A simple architecture for learning disentangled representations in vaes. arXiv preprint arXiv:1901.07017.\n- [2] Montero, M., Bowers, J., Ponte Costa, R., Ludwig, C., & Malhotra, G. (2022). Lost in Latent Space: Examining failures of disentangled models at combinatorial generalisation. Advances in Neural Information Processing Systems, 35, 10136-10149.", "essential_references_not_discussed": "See above.", "other_strengths_and_weaknesses": "The main strength of the article is that it provides a more thorough account of some findings/insights that have been previously found. Specifically, I believe that viewing compositional generalization failures in generative models from the perspective of transport operators is a valuable one. But since both the findings and insights are not novel it is unclear how useful this will ultimately be.", "other_comments_or_suggestions": "I think my main suggestion is to shift the tone of the article from one where the authors claim to make this fundamental discovery (which for better or worse has already been done), to one where they characterize it based on transport perspective. Specifically, they can ask how layer depth and the like affect this issue and then discuss their regularization techniques/data augmentation. I would actually say that this last part feels more like curriculum learning, which has some interesting links to cognitive science and how we learn about concepts in isolation before learning how the interact with each other. Then I would try to translate their insights to more sophisticated datasets (We don't need ImageNet, but there are plenty of datasets that are used to explore disentangled representations out there). Be aware though, this idea of regularization has been tried before and in general does not lead to good results (see references), so the authors need to make a clear case as to why their approach is fundamentally different instead of just mathematically. Finally, I agree that architectural constraints are good, as shown in [4]. The question would be, are the general architectural motifs that apply across modalities/concepts?\n\n\n[1] Kim, H., & Mnih, A. (2018, July). Disentangling by factorising. In International conference on machine learning (pp. 2649-2658). PMLR.\n[2] Burgess, C. P., Higgins, I., Pal, A., Matthey, L., Watters, N., Desjardins, G., & Lerchner, A. (2018). Understanding disentangling in $\\beta $-VAE. arXiv preprint arXiv:1804.03599.\n[3] Zhu, X., Xu, C., & Tao, D. (2021, July). Commutative lie group vae for disentanglement learning. In International Conference on Machine Learning (pp. 12924-12934). PMLR.\n[4] Montero, M. L., Bowers, J. S., & Malhotra, G. (2024). Successes and Limitations of Object-centric Models at Compositional Generalisation. arXiv preprint arXiv:2412.18743.", "questions_for_authors": "I have no questions apart from the ones 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": "5f9NfRemox", "reviewer": "Reviewer_M163", "summary": "The paper develops theoretical and empirical results as to why a disentangled representation in\ninput does not necessarily lead to OOD compositional generalization.\nFirst, the paper demonstrates the failure of common generative architectures (decoder-only CNN / MLP) to\nperform compositional generalization, even when given disentangled input. This is demonstrated on a task of\ngenerating synthetic images of a 2D Gaussian bump, centered at the (x, y) input coordinates. They then suggest\nand test an explanation for this failure - that the input becomes re-entangled in deeper layers. This is shown\nquantitatively using kernels and manifold. The paper suggests that instead, the models memorize training data.\nFinally, two methods are suggested to improve the compositional generalization: architectural modifications with\nregularization, and additional training on curated data.", "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": "**Architecture Choice (Size & Type)**\n\n- **Data–Model Size and Memorization:** In our experiments, the network does not memorize when we alter the data (e.g., 1D stripes with few conjunctions), even with significantly fewer samples. Thus, we do not see evidence that a “small dataset vs. large model” alone must lead to memorization. Instead, we find that memorization arises when factorization is not encouraged (via e.g. regularization, data augmentation), and it exhibits the “superposition” strategy—activating all relevant seen patterns simultaneously.\n\n- **CNNs/MLPs as Foundational Backbones:** We tested both CNNs and MLPs (the backbone of many generative models, from VAEs, UNet-based Diffusion models, to transformers). Both show manifold distortion and OOD failure, suggesting that “warping” in the presence of incomplete data coverage is fairly universal. We also tried positional encoding with CNN/MLP (omitted for brevity) and observed similar OOD failures.\n\n- **Transformers & Spatially-aware Models:** When we experimented with Transformer-based architectures for the 2D bump task, we found a notable difference: standard Transformer encoder layers with self-attention across pixel tokens performed worse in out-of-distribution bump placements compared to a simpler stack of per-token linear layers and non-linearities. Since each pixel token already has direct access to (x,y) context via positional embeddings, pooling or mixing tokens through self-attention seems unnecessary and can inadvertently entangle factors, degrading OOD generalization. Conversely, a purely feed-forward per-token approach preserves the simple, distance-like mappings needed for forming a 2D bump, yielding better compositional generalization.\n\n**Factorization vs. Mere Reversibility**\n\n- **Why Not Just Invertibility?** We appreciate the suggestion that invertibility might ensure compositional generalization. However, we emphasize compositionality at the manifold level, as single-sample invertibility alone may still allow factors to become entangled across the manifold, hindering new combinations. For instance, invertible models like normalizing flows don't inherently show better compositionality, suggesting that invertibility by itself is insufficient.\n- **Jacobian‐Based Manifold Factorization:** Our Jacobian‐based metric checks how well the manifold locally preserves directions ∂/∂x vs. ∂/∂y If these directions become “cross‐wired,” the network can no longer compose new factor pairs effectively—despite potentially having a globally invertible map. In short, local entanglement disrupts factor reuse.\n- **Partial Rather Than Strict Factorization:** We do not require each neuron or layer to exclusively encode one factor. Rather, we need that the overall manifold maintain approximately independent axes for each factor. Empirically, when the representation heavily mixes them, the model reverts to memorizing ID combinations rather than systematically composing new ones.\n\nAlthough an invertible transformation can theoretically map individual samples in or out, robust compositional generalization requires more than local invertibility—it requires the activations' distribution to avoid entangling distinct factors, thus enabling reuse in novel combinations. Our Jacobian-based metric tracks this manifold-level factorization, emphasizing sufficient (rather than strict) factorization across the manifold for OOD generalization. We will clarify this distinction to avoid confusion.\n\n**Applicability to More Complex Datasets & Real Data**\n\nSee our discussion for Reviewer 3 (iydz). \n\n**Novelty & Positioning Relative to Prior Work**\n\nSee our response for Reviewer 3 (iydz). \n\n**Discussion on Loss, Architecture Details, & Additional Points**\n\n- **Loss Functions:** We appreciate the note on Earth‐Mover distance. While it may be more spatially appropriate for “bump” images, MSE is simpler and general. Off‐manifold images, which need not look like bumps, might not benefit from EMD.\n- **CNN/MLP vs. VAEs:** Our “decoder‐only” approach effectively isolates the portion of generative models responsible for mapping factorized latents to pixel space, which also occurs in VAEs and Diffusions. We wanted to remove as many confounding factors as possible to highlight the re‐entanglement phenomenon.\n- **Why This Specific Bump Task?** We wanted a minimal environment with known factorization. This clarity allows us to discover the superposition and warping phenomena in an unambiguous way that might otherwise be obscured in complex tasks.\n- **Minor Stylistic Issues:** We will address all writing and visualization issues raised by the reviewer.\n\nWe appreciate the reviewer’s feedback and will clarify this in our revision. Our findings underscore that factorized inputs alone don't guarantee compositional generalization unless factor separation is maintained throughout the network’s forward pass. We hope these insights inform robust model designs and data strategies more broadly.", "claims_and_evidence": "Strengths:\n* Studying the impact of factorization is a very important task, since a breakthrough in compositional embedding can break many limits in current learnable approaches.\n* The toy example is analyzed extensively and thoroughly, in terms of memorization, manifold analysis, and augmentation.\n* The paper’s demonstration of the memorization of the in-distribution training data is presented clearly and is\njustified using the binary kernel factorization. Likewise, the presentation of the re-entanglement using the Jacobian\ntensor is valid.\n* The low-rank approach is an interesting way to force factorized structure and prevent memorization. In addition,\naugmenting the training data allows improving compositional generalization of current models without the need\nto modify architecture.\n\n\nWeaknesses:\n* It appears problematic to generalize into holistic and general insights out of a single and limited toy example (it is indeed mentioned in the Limitation section, however this may not be a simple limitation but a biased *general* insight). Experiments may be biased due to various reasons, such as: 1) the relationship between the size of the network and the amount/variability of the learning data may extremely bias the network into memorization instead of generalization. 2) relying on CNN based network for a task containing locality, and hence spatial aware network is preferable, may also be problematic. It can shift the network into memorization, since the architecture is less adequate to solve the problem. For instance, it may be that more complicated networks, specifically spatial aware ones (like transformer containing positional encoding phase) may lead to different conclusions.\n* The authors show that subsequent layers re-entangle the representation, however it is not proven to be the cause\nof failure for OOD generation. The paper does not test the less constrained hypothesis, that it is sufficient for the\nrepresentation to be merely invertible into the original factors. The requirement of strict factorization seems too\nconstraining.\n* The generality and applicability of the suggest methods are unclear without tests on more complicated datasets.\nIn addition, generating additional curated data is domain-specific and can be very expensive and demanding \nfor real dataset.\n* The paper’s claim - that standard architectures fail to achieve compositional generalization despite being pro-\nvided with explicitly factorized inputs — implicitly assumes that a factorized input should be sufficient for such\ngeneralization. However, prior work (e.g. Montero 2021) has already shown that compositional generalization depends on more than\njust disentanglement at the input level. Consequently, the paper refutes a weaker claim than what is typically\nasserted.", "methods_and_evaluation_criteria": "* Nowadays VAEs fall short in terms of generative capabilities, way beyond diffusion models. I think that theoretically it is important to understand compositional generalization in general, but investigate it on CNNs is reducing the impact of the findings. (Also here I noticed that it mentioned in the Limitation section as well, I agree with this limitation and mention it here as a weakness that it is not \"cancel\" the insights but tone them down).\n* The ablation study is done on various combinations of the loss terms but all on the same architecture. If the tuning is only on the learned loss it can easily be applied to other architectures (which is a pro in fact), then I would expect this type of ablation study.", "theoretical_claims": "The methodolgy and reasoning here seems not to take into account all possible cases for their final conclusions. see above.", "experimental_designs_or_analyses": "Insufficient, toy, single architecture that is not optimal for the problem, see above.", "relation_to_broader_scientific_literature": "OK", "essential_references_not_discussed": "Related Work should contain a section on loss terms for compositionality.", "other_strengths_and_weaknesses": "Minor weaknesses:\n* The paper does not list the architectures used for the CNN and MLP used in section 3.1.\n* VAE play a central role in generation in general and compositionality research in particular, and hence the paper\nis lacking by not including it in its tested architectures.\n* The visualizations contain too small font size and information. An extreme example is in (e) in Fig. 4.\n* Line 309 right column, a sentence started but not finished, \"both results show\"\n* Not clear that MSE loss for this task is preferable. What if the network required to set a bump in (x,y) and set it in (x+d,y+d). The same MSE loss will be reached regardless the distance from the required. I think that earthmovers metric (optimal transport loss in general) would be more appropriate here.\n* Why do the authors decide on this specific task? Why generative task, and under generative ones, why this specific bump task? I did not find intuitive explanation in the paper why this task is preferable for some reason.\n* Inconsistencies in capitalizing. For example subsection 3.4 contains capital letter only for the first word, while subsections 4.1, 4.2 and more capitalize each of the words.", "other_comments_or_suggestions": "see above", "questions_for_authors": "see 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": {}}]

Author rebuttal:

Meta-review: All reviewers agree that the paper provides insights for better understanding  why disentangled representation does not guarantee compositional generalization, the claims made in the article are well supported by empirical simulations,  and the extensive visualization and analysis help to explain the main points. Although the reviewers raised the concerns on toy examples which might lead to bias, affecting the insights' generalization to other settings, the authors' rebuttal help to address the concerns. Based on the consistent positive reviews and the rebuttal, I recommend weak acceptance.
Decision: Accept (poster)