Do Diffusion Models Learn Semantically Meaningful and Efficient Representations?

TL;DR

The paper asks whether diffusion models learn semantically meaningful and efficient representations beyond memorization, using a controlled 2D Gaussian bump task. It trains a conditional DDPM with a UNet and analyzes internal representations via layer 4 outputs reduced with UMAP to relate latent structure to generation performance. The authors identify three latent-manifold phases (A, B, C) during training, with an ordered phase (C) yielding correct localization and high accuracy; they find that x and y representations are coupled even on imbalanced data, suggesting current vanilla diffusion models may lack fully factorized, data-efficient representations. The work highlights inductive biases as a path to enforce factorization and improve compositional generalization and data efficiency in diffusion-based generations.

Abstract

Diffusion models are capable of impressive feats of image generation with uncommon juxtapositions such as astronauts riding horses on the moon with properly placed shadows. These outputs indicate the ability to perform compositional generalization, but how do the models do so? We perform controlled experiments on conditional DDPMs learning to generate 2D spherical Gaussian bumps centered at specified $x$- and $y$-positions. Our results show that the emergence of semantically meaningful latent representations is key to achieving high performance. En route to successful performance over learning, the model traverses three distinct phases of latent representations: (phase A) no latent structure, (phase B) a 2D manifold of disordered states, and (phase C) a 2D ordered manifold. Corresponding to each of these phases, we identify qualitatively different generation behaviors: 1) multiple bumps are generated, 2) one bump is generated but at inaccurate $x$ and $y$ locations, 3) a bump is generated at the correct $x$ and y location. Furthermore, we show that even under imbalanced datasets where features ($x$- versus $y$-positions) are represented with skewed frequencies, the learning process for $x$ and $y$ is coupled rather than factorized, demonstrating that simple vanilla-flavored diffusion models cannot learn efficient representations in which localization in $x$ and $y$ are factorized into separate 1D tasks. These findings suggest the need for future work to find inductive biases that will push generative models to discover and exploit factorizable independent structures in their inputs, which will be required to vault these models into more data-efficient regimes.

Figures (9)

• Figure 1: Schematic illustration of a fully factorized solution vs. an unfactorized solution.(a) is a factorized solution where the $x$ position is generated independently from the $y$ position of the Gaussian bump by intersecting two oval Gaussian bumps localized in one dimension but not the other. (b) shows a coupled solution where a single Gaussian bump localized in both dimension is generated. One difference between the two possibilities is that a network that recognized the independence of generation in $x,y$ could learn with $\mathcal{O}(2K)$ examples, while otherwise it would take $\mathcal{O}(K^2)$ examples.
• Figure 2: Example image data.
• Figure 3: The three phases of manifold formation. The learned representations (UMAP reduced, colored by the ground truth $x$-positions) of the diffusion models undergo the three phases in increasing order of training steps as depicted in the 3D visualizations in the bottom row. In each phase, the corresponding qualitative generation behavior is demonstrated with 25 sampled images in the top row, in which the red dots mark the ground truth locations of the Gaussian bumps. (Phase A) has no particular structure in the learned representation, and the generated images either have no Gaussian bumps or multiple Gaussian bumps at the wrong locations. (Phase B) has a disordered, quasi-2D manifold with corresponding generation behavior of a single Gaussian bump at the wrong location. (Phase C) has an ordered 2D manifold with the desired generation behavior.
• Figure 4: 2D phase diagram of performance metrics as a function of increment and training steps. (a) shows the predicted label accuracy and (b) shows the R-squared averaged in predicting $x$- and $y$-positions of the Gaussian bumps from the latent representation. The models are trained with datasets of various increments from 0.1 to 1.0 and sigma of 1.0. The total number of training steps are held constant across all the models.
• Figure 5: Performance metrics of models trained using imbalanced datasets.(a) using increments of $d_x=0.1$ and $d_y=1.0$ and (b) using increments of $d_x=0.1$ and $d_y=0.5$. Models in both cases are trained with amply sufficient amount of steps to reach convergence.