How Diffusion Models Learn to Factorize and Compose

Summary

This paper investigates the capabilities of diffusion models, particularly Denoising Diffusion Probabilistic Models (DDPMs), in learning factorized representations and achieving compositional generalization. The authors aim to quantify this by analyzing mechanisms to train the model, supporting the hypothesis that the architecture of diffusion models has an inductive bias towards such factorized representations. Their results on a toy dataset suggest that, to achieve out-of-distribution compositional generalization, the training set must: i) contain at least a few compositional examples of the factors, and ii) present the factors independently of each others across the full range of their variability. If either of these conditions i) or ii) is not met, the model will fail to generalize out of distribution.

Strengths

- The manuscript is particularly well-written and presented, with clear and thorough explanations. The methodology and experiments are relevant and consistent, and they are both well-explained and well-illustrated. - The originality of this work lies in explaining the factorization capability emerging in the manifold from the perspective of percolation theory. This supports the hypothesis that the training set needs a certain level of correlation among the components of independent features to achieve a faithful representation, which in turn leads to compositional generalize.

Weaknesses

The only weakness I can think of in this work is that it is conducted in simplistic toy settings. However, I don't believe this discredits the work and interesting experiments presented in the paper. Additionally, the authors acknowledge this issue and suggest exploring more naturalistic and relevant experiments in future studies. Some typos: - Line 7: first time using the abbreviation ‘DDPMs’ in the abstract, maybe precise what it is (Denoising Diffusion Probabilistic Models) - Line 198: at the end of the sentence, there is an extra ‘not’

Questions

N/A

Rating

8

Confidence

3

Soundness

4

Presentation

4

Contribution

3

Limitations

See “Weaknesses”

Summary

This paper investigates how and when diffusion models learn factorized representations of composable features. To this end, the authors construct controllable synthetic datasets by compositionally combining 1D and 2D Gaussian data and examine the factorized representation and the compositional generalization capability of the diffusion model. Systematic analysis on the controllable dataset indicates that the diffusion model learns orthogonal but not necessarily parallel representations and is capable of compositionally generalization to OOD samples if a few compositional examples are provided. Additionally, the authors draw a connection to percolation theory, suggesting that a certain amount of correlated data is required to learn factorized representations.

Strengths

- The paper is well-written and easy to follow. - The simple yet well-controlled experiments and comprehensive analysis provide a clear understanding of factorization and compositionality in diffusion models. - Connecting the empirical findings to percolation theory improves the understanding on the emergence of factorization.

Weaknesses

- The experiments are conducted solely on a simple dataset. The 2D Gaussian Addition dataset features only basic additive compositions of 1D Gaussian sprites. It remains unclear whether the paper's claims extend to more complex compositions, such as multiplicative compositions involving scale, color, etc. - Some of the conclusions, such as those regarding the properties and requirements for compositional generalization, have already been addressed in previous work [1]. [1] provides theoretical analyses and conditions for sufficient compositional support. Therefore, the experiments and conclusions from section 3.2 provide limited additional insights. [1] Wiedemer et al., “Compositonal Generalization from First Principles”, in NeurIPS 23.

Questions

- Are the conclusions and implications specific to the diffusion model? Although the paper claims to investigate the compositionality of diffusion models, the experiments and analysis do not seem diffusion-specific. Other general generative models might exhibit similar behavior. - In the Gaussian Bump + 1D Gaussian Stripes experiment (Figure 4(g)), what is the accuracy when only 2D Gaussian Bump data is used? Comparing this value would help identify the performance boost from compositional generalization. - How did you format the inputs when training the diffusion model on 1D Gaussian stripe data? Did you input a null value (e.g., $(\mu_x, \phi)$)? - How does the model generate OOD samples in 2D Gaussian Bump Data? From the experiments on the 2D Gaussian Addition dataset, it is concluded that the model cannot interpolate well on unseen examples, as implied by the low accuracy in the intersection areas. By adding 1D Gaussian Addition data, the model learns from single 1D stripe patterns and composes that information to generate OOD samples in the test regions. However, in 2D Gaussian Bump Data + 1D Gaussian Addition data, the model never observed the 2D Gaussian Bump samples in the test regions. How could adding 1D Gaussian Addition data lead to improvement? - In Figures 4(c), 4(d), and 4(e), interestingly, the model trained only on 1D Gaussian data already predicts reasonably on $y$ values (but very low for $x$ values). Is there any particular reason for this result?

Rating

6

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

The conclusion section describes the limitation regarding the restricted scope and simplicity of the toy dataset. While the study provides valuable insights, the findings are based on synthetic datasets with simple structures, which may not fully capture the complexities of real-world data.

Summary

This paper investigates, on a very simple toy dataset, how conditional diffusion models learn factorized representations of the data, and the extent to which they can compositionally generalize out of distribution. Additionally, the authors make a connection to percolation theory in physics.

Strengths

The motivation and research questions are very interesting and relevant for the community. They are outlined in a compelling way in the introduction. The experiments are carefully designed and rather interesting.

Weaknesses

Main high-level weaknesses: - This paper takes a promising approach to a crucial research question, but to me it does not deliver. Although using toy data allows for broader experimentation, the evaluation is overall quite limited. For this reason, though the authors make an effort to clarify the relevance of their results for realistic settings (Sec. 4.1), this still sounds unconvincing. - Clarity, especially in the presentation and discussion of the results. I found the results and conclusions significantly more difficult to parse than would be expected from toy experiments. --- Expanding on the above: - In general, the design and motivation for the experimental study is a bit lacking. I appreciate toy experiments, but they should reflect more realistic cases as much as possible, and in this case we are assuming that the model observes basically all information necessary to reproduce the data (as opposed to typical cases where the conditioning signal has significantly less information than the data itself). The authors should provide a solid justification for this choice, and more in general argue how such a toy scenario may be informative for more complex settings. The most natural next step would be to include a dataset that comes significantly closer to realistic settings, although of course some investigations will probably be impossible in that case. The trade-off between relevance and controllability is a hard one, and the current paper seems to be heavily on the latter side. - Even sticking to the current toy data, a broader evaluation would be possible and useful. There are several degrees of freedom that can be explored further, e.g., the UNet input noise level at evaluation time, the mutual information between condition and data (which could range from perfect - the current unrealistic case - to zero - the unconditional case), different compositionality patterns (as done in some references in the paper), the layer of the UNet at which representations are extracted. - The abstract states "paving the way for future research aimed at enhancing factorization and compositional generalization". This may be an overstatement, given my point above. What are concretely actionable insights from these experiments? - The experiments here investigate the representations in the last layer of a UNet. Why this choice? Representations in diffusion models could also be considered to be the activations at different layers, especially the bottleneck, which has been investigated in the literature. Another representation can be the latent variable deterministically corresponding to the data using the probability flow ODE. - The bottleneck idea is mentioned only in Appendix A.1, and in a negative way. This reinforces my belief that this toy scenario is too far from realistic settings, where the bottleneck is widely used as representation in diffusion models. However, let me still point out that I find the toy scenario a very interesting and promising direction. - As far as I can tell, there is no mention of the noise schedule for training DDPM. - When evaluating the representations in the UNet, what is the noise level in the input? I would expect this to significantly affect the representation (especially since you're using the last UNet layer and the UNet is trained to predict noise -- but this is just a hunch). - The models here are basically trained to convert the $(\mu_x,\mu_y)$ conditioning pair to an image that is deterministically determined from such a pair. I would be a lot more interested in investigating the representations of an *unsupervised* model trained on such data, where the labels are inaccessible but can be used for evaluation -- similarly to how e.g. disentanglement is evaluated. - Alternatively, to mimic real-world conditional (e.g. text-to-image) generation, there could be some stochasticity involved, such that the observed data is not trivially obtained from the conditioning signal. - To make the experiments and results fully understable and accessible to a wider range of researchers, I would strongly recommend including a quick introduction to the relevant concepts from geometry/topology. - In addition, the individual results subsections seem to lack a clear structure, which makes them not particularly easy to follow. Some more intuitive explanation, as well as highlighting the main takeaways from each experiment, might help. A few minor or more detailed comments: - In the contributions: "differing values of the same feature are also treated similarly" - what does this mean? - In Section 3.1 there's mention of $x$ and $y$: are these actually $\mu_x$ and $\mu_y$, since these are the ground-truth generative factors of the data? - At the beginning of the results section, the dataset is modified to have a torus topology. Why not define the dataset like this in the first place? - Line 124. "If the model were to parameterize x and y independently with two ring manifolds, we expect to see a Clifford torus in the neural activations, rather than a different geometry such as the 3D torus." What does it mean exactly that the model *parameterizes* x and y? My interpretation is that, when generating data conditional on $x$ and $y$ (which I take to be the means), we can observe how the activations in the pre-determined UNet layer change, as we change $x$ and $y$. - Line 127 and following: "we first confirm that the model indeed learns a torus representation of the dataset". What would the model alternatively learn? Since we're considering such simple toy datasets, I think it's expected that we can get a full intuition of what is going on. In my opinion, this is not the case. The following lines also involve technical terms from topology that are not properly introduced. - It's unclear what exactly effective dimensionality is, why it is important here, and what we can learn from it. - The conclusions drawn in the last paragraph of page 4 are not very clear. For example: "These results suggest that x and y are independently encoded in pairwise orthogonal subspaces, but different values of x’s and values of y’s respectively are not encoded in the same way, i.e. in parallel subspaces". - line 170: "compositionally generalizing out of the training distribution as well as its ability to spatially interpolate in a single variable alone". What is the precise difference between these two scenarios? In general, interpolation and composition don't seem to be accurately defined (although they can be implicitly inferred e.g. by lines 172-179).

Questions

See weaknesses.

Rating

5

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

The limitations are mentioned in the discussion section. Unfortunately, I believe some of them may be too large to ignore.

Summary

In this work, the authors investigate how diffusion models achieve compositional generalization. Through controlled experiments on conditional DDPMs with 2D Gaussian data, the authors find that these models learn semantically meaningful, factorized manifold representations of composable features. These representations are orthogonal for independent feature variations but not aligned for different values of the same feature, resulting in superior compositionality but limited interpolation over unseen feature values. The study reveals that a small amount of compositional examples can enhance this capability and links the formation of these representations to percolation theory in physics. This work provides insights into the mechanism of compositionality in diffusion models, guiding future research to improve their factorization and generalization for real-world applications.

Strengths

1. The work investigate the factorization and compositionality of diffusion model, which is an interesting and valuable problem.

Weaknesses

1. The major concern is that the work only considers a highly reduced setting, which makes it hard to validate wether the conclusions generalize to real world applications. Since the setting is oversimplified, it's also unclear how the conclusions can be extended to popular applications like conditional text2image generation. 2. The performance of diffusion is mostly evaluated with customized metrics without comparison to other models. It's hard to get a sense of how well/bad the model performs in term of the metrics.

Questions

1. Given the dataset contains only ~100K synthetic images, the model can simply memorize the data. This makes it hard to argue the generalization behavior of the model can be extended to real image diffusion models. 2. Why the output of layer 4 is used to investigate the internal representations? Do the authors try representations from other layers?

Rating

5

Confidence

3

Soundness

2

Presentation

2

Contribution

2

Limitations

The work sufficiently addressed the limitations.

Response to rebuttal

I thank authors for answering my questions. Though as also pointed by other reviewers, the paper only considers toy setting and lacks experiment on larger and more practical settings. After reading the discussions between authors and reviewers (including me), I believe the paper does provide some interesting observations of the underlying representation structures learned by diffusion models. And it could potentially guide the development of more efficient and powerful diffusion models. Therefore I have raised my score to 5 meaning I won't go against accepting the paper.