Compositional Generalization via Forced Rendering of Disentangled Latents

Architecture Choice (Size & Type)

• 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.

• 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.

• 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.

Factorization vs. Mere Reversibility

• 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.
• 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.
• 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.

Although 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.

Applicability to More Complex Datasets & Real Data

See our discussion for Reviewer 3 (iydz).

Novelty & Positioning Relative to Prior Work

See our response for Reviewer 3 (iydz).

Discussion on Loss, Architecture Details, & Additional Points

• 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.
• 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.
• 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.
• Minor Stylistic Issues: We will address all writing and visualization issues raised by the reviewer.

We 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.

Novelty and Positioning Relative to Prior Work

We 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. In 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.

Applicability to More Complex Datasets and Real-World Data

We 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).

In 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.

Relationship to Curriculum Learning and Cognitive Science

We 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.

Regularization Techniques and Architectural Constraints

We 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.

Regarding 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.

Terminology and Clarity

We 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.

We 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.

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.

Definition of Compositionality and OOD Generalization

We 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.

In 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.

Experimental Details and Reproducibility

We 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:

We will include this summary table in the revised manuscript to clearly communicate the architectural details and facilitate reproducibility.

Clarification of Equation (8)

We 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.

Minor comments

We will revise the manuscript to get rid of quotation marks around superposition accordingly.

In 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.

Input Encoding Types

We 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.

Shape of the OOD Region

We 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.

Quantitative Averaging Over Multiple Runs

We 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.

Fixed Coordinate Inputs

We 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.

Potential Issues with Stripe Augmentation

We 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.

Generalizing Data Augmentation Strategy

Please refer to our response to Reviewer 3 (iydz) regarding the same point.

Fixed Coordinate Bumps Experiment

We 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.

Clarification on Figure 1e and Section 3.3

We 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.

Other References

We thank the reviewer for highlighting important references and will ensure their inclusion and proper discussion in the revised manuscript.

Regarding Novelty and Main Contribution

We 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.

Additional Comments and Suggestions

We will revise the manuscript accordingly.

Once again, we sincerely thank the reviewer for these valuable suggestions, which will undoubtedly strengthen our paper's clarity and impact.