How Diffusion Models Learn to Factorize and Compose

We thank the reviewer for taking the time to provide us with constructive feedback. Please find our responses to the specific concerns and questions below.

Weaknesses

1. We thank the reviewer for the very positive feedback. As noted in the global response section “Regarding the toy setting,” using a simple toy setting allows us to carefully examine various effects. Future investigations on more realistic datasets are necessary but could introduce many competing effects. Despite the simplicity, our study provides valuable insights into compositionality and generalization in Diffusion models. Future research should explore why Diffusion models cannot encode continuous latent features continuously and how percolation theory of manifold formation can be applied to natural image data.
2. We thank the reviewer for pointing out our typos. We will fix them in the manuscript.

We thank the reviewer for taking the time to provide us with constructive feedback. Please find our responses to the specific concerns and questions below.

Weaknesses

1. We appreciate the reviewer’s feedback. In our paper, we explored both additive and multiplicative composition with the 2D Gaussian Addition (addition of two 1D Gaussian Stripes) and the 2D Gaussian Bump (multiplication of two 1D Gaussian Stripes) datasets. Fig. 4(g) shows that the model can generalize to new 2D Gaussian Bumps given all 1D Stripes and a few 2D Bumps. This multiplicative composition is similar to transformations like scale, color, and style, as it involves "masking" one 1D Stripe with another.
2. We thank the reviewer for bringing Ref. [1] to our attention. While it similarly concludes that gaps in the support of the training dataset hinder learning, it lacks insights from a manifold formation perspective. Their experiments use sprites with a mix of (semi-)continuous and categorical latent variables, showing general results across all types. Missing support in categorical variables makes sense, as a model cannot generate images of monkeys after only seeing cats and dogs. However, with continuous latent features like $x$- and $y$-positions, one would expect interpolation to be possible (e.g., generating a sprite at $x=16$ after seeing $x=15$ and $x=17$). Our experiments focus on the model’s ability to represent continuous latent features. In Sec. 3.1, we found that the model represents continuous features similarly to categorical variables, with some overlap. Consequently, as shown in Sec. 3.2, the model struggles with interpolation. The main takeaway is that the model excels at composition but fails at interpolation. Contrary to popular belief, Diffusion models can perform well at compositional generalization but struggle with continuous latent features, hindering interpolation and robust generalization.

Questions

1. While we believe some observations may apply to other generative models, we refrain from general claims, focusing our study on Diffusion models and their factorization and compositionality. Similar studies on vision- or language-based generative models are beyond our current paper's scope.
2. We thank the reviewer for the constructive suggestion. We refer the reviewer to the global response section on “Notes on additional figures - Additional Figure 2: Data Efficiency Scaling”.
3. When generating 1D Gaussian Stripes, we embed the data image (32 x 32) into a larger canvas (44 x 44), centering it and creating a 6-pixel border. We generate 2D Gaussian Additions centered in this extended border region, most of which (except the corner ones) will have one Gaussian Stripe partially visible in the 32 x 32 space. We then crop the central 32 x 32 pixels, keeping the label as the center of the 2D Gaussian Addition in the extended space. Thus, 1D Gaussian Stripes have one coordinate outside the 32 x 32 image space, preserving the 2D structure for both 1D and 2D data points. Details of the data generation process are in Sec. C.1 of the Appendix.
4. The experimental setup is detailed in Sec. C.1 Fig. 8 of the Appendix. The model effectively learned spatial information from the 1D Gaussian Stripes. Given a few examples of 2D Gaussian Bumps, it learned to multiplicatively compose the 1D Gaussian Stripes into Bumps. Combined with the results in response to Question 2, this demonstrates the model's ability to transfer knowledge across different forms of compositionality.
5. We appreciate the reviewer's detailed observation. The 1D models in Fig. 4(c)-(e) are trained on equal numbers of vertical and horizontal 1D Gaussian Stripes, providing the same training examples for learning $x$ and $y$. While the exact reason for the higher accuracy in generating $\mu_y$ is unclear, we observed that the model often defaults to generating 1D Stripes when it fails to generate the intersection of two Stripes, possibly skewing the accuracy distribution between $\mu_x$ and $\mu_y$.

Thank you to the authors for the detailed reponse.

I believe most of my concerns have been addressed. Although the authors employed controlled and simple toy datasets, I find their experimental setting and comprehensive analysis sufficient to investigate their hypotheses and provide insights regarding the compositionality of diffusion model. Therefore, I raise my score to weak accept.

We thank the reviewer for taking the time to provide us with constructive feedback. Please find our responses to the specific concerns and questions below.

Weaknesses

Expanding on high-level

1. We kindly refer the reviewer to the first two sections of the global response.
2. We appreciate the reviewer's suggestions. To summarize your suggestions: a) Input noise level for UNet, b) Imperfect conditioning label, c) Different compositionality patterns, d) Different layers of UNet. We have explored a), c), and d). For a), we have studied the bottleneck and layer 4 at various noise levels (diffusion timesteps) and found minimal differences between representations at different timesteps, with the bottleneck having diminishing signals. Thus, we used the final timestep output from layer 4 for our analysis. For c), as shown in Fig. 4, the model learned both multiplicative and additive compositionality with the full range of 1D Gaussian Strips and a few compositional examples. For d), we investigated all UNet layers and found layer 4 better reflects the latent structure of the dataset than other layers, so we used it for our analysis. We have not considered b). While interesting, this doesn't allow precise prompting of the model to output a specific Gaussian at a desired location, complicating evaluation. As a reminder, our focus is on whether the model can even, given perfect label information, organize learned data into an efficient and meaningful representation for generalization. The pursuit of b) falls outside of our research scope.
3. We thank the reviewer for this feedback. We will modify this sentence in the abstract.
4. Although Diffusion model bottlenecks have been shown to encode semantic information, in our case, layer 4 acts as the bottleneck due to the skip connection. This does not affect our main point, as we are focused on the factorization of the model's learned representation.
5. We kindly refer the reviewer to the global response section on “Notes on additional figures - Additional Figure 1: Impact of Annealed Noise Schedule on Learning Rates”.
6. See response to “Expanding on high-level” 2.
7. Investigating representations learned by unsupervised models is intriguing. Previous work has focused on disentanglement in unconditional Diffusion models, but evaluating compositional generalization without explicit conditional input is challenging. In terms of stochasticity, our conditional training uses classifier-free guidance, dropping labels 10% of the time to co-train conditional and unconditional models. Moreover, we know that the model isn't solely relying on conditional input, as it sometimes fails to learn even with perfect conditions (e.g., Fig. 4(c)-(e)). Future work should nevertheless evaluate unconditional models or conditional models with imperfect input for their factorization and compositionality. 8-9. We will refine our manuscript to reflect the suggestions.

Minor

1. By the original sentence, we meant that different values of the same feature (e.g., $x = 13, 14, 15$) are encoded similarly to how different features ($x = 13$, $y = 13$) are encoded, treating $x$ and $y$ more like categorical than continuous variables. We will clarify this in the original sentence.
2. We refer to $x$ and $y$ as features and $\mu_x$ and $\mu_y$ as the center locations of the Gaussian bumps in a data image. The latter is intrinsic to the dataset, while the prior is learned by the model.
3. Although periodic boundary conditions could be enforced in all datasets, the resulting torus manifold is nonlinear and less straightforward to analyze using linear methods. Otherwise, there's no fundamental difference between datasets with and without these conditions.
4. By “parameterize $x$ and $y$,” we mean how the model represents the two features. Jointly, the model may learn a 2D look-up table (3D torus), while independently, it may learn $x$ and $y$ as 1D rings (4D Clifford torus). Our geometry/topology test distinguishes between these scenarios.
5. Our tests detect whether the learned representation is a 3D or Clifford torus. This requires the object to be a torus, which isn't always obvious during training. The persistent homology results in Fig. 2(b) help clarify this. We will refine the manuscript to be more precise.
6. Effective dimension, computed by the participation ratio (lines 134-135), measures the intrinsic dimensionality of learned representations. Plotting it over training epochs helps us understand if the model learns a Clifford torus (dimensionality of 4) and whether it does so from first learning a 3D torus (dimensionality of 3). Fig. 2(c) shows higher than 4-dimensional effective dimensionality, indicating independent learning of $x$ and $y$. Fig. 2(d) eigenspectra and Fig. 2(g) PCA projections suggest the learned representation isn't a perfect Clifford torus but higher-dimensional and cone-shaped. We will add more details on this in the manuscript.
7. The model treats $x$ and $y$ as categorical variables with non-zero overlaps. This means $x=16$ is a separate category from $x=17$, not neighboring values of a continuous variable. Technically, if $x$ and $y$ were continuous, different values should be encoded in parallel subspaces, like a Clifford torus. The model learns a hyper-factorized version, representing $x$ and $y$ as categorical. We will clarify this in the manuscript.
8. "Compositional generalization" means combining components in ways not seen in the training set (e.g., given ${(a, b), (c, d)}$, output ${(a, d), (c, b)}$). Interpolation refers to combining values (e.g., output $x=16$ given $x=15$ and $x=17$). These are distinct forms of generalization, which we probe separately in Sec. 3.2. We will define these clearly at the beginning of Sec. 3.2.

We thank the reviewer for taking the time to provide us with constructive feedback. Please find our responses to the specific concerns and questions below.

Weaknesses

1. We thank the reviewer for the feedback, we kindly refer the reviewer to the global response Section “Regarding the toy setting” addressed to all reviewers.

2. Due to the simplicity of our toy setting, we design the custom metrics such that it is tailored to our objective of investigating the factorization and compositionality in Diffusion models. Towards this goal, we have specifically designed the task and the metrics accordingly such that they target the model’s ability in recovering the spatial location of the Gaussian bumps.

Questions

1. Naively, when a model resorts to memorization, it typically suffers from poor performance in the out-of-distribution task sets. We have carefully selected our task and model such that the task is not too trivial for the model and that the model can just memorize all the training data. This is indeed shown through the out-of-distribution model performance evaluation in, for example, Fig. 5. We note that even given a subset of the original dataset, the model has similar in-distribution and out-of-distribution performance, which means that the model is not simply memorizing all in-distribution data but rather trying to learn the correct representation.

2. The reason why we have chosen layer 4 of the UNet as our internal representation is because it better reflects the latent structure of the dataset than other layers. Specifically, we have consistently noticed that the bottleneck layer of the UNet gives diminishing signal due to the utilization of the skip connection. As a result, we chose the layer immediately following the bottleneck layer as our internal representation. Moreover, we have investigated all the layers of the UNet and found that layer 4 gives the strongest semantic signal.