How Diffusion Models Learn to Factorize and Compose

In this section, we aim to understand the factorization of the model’s learned representation via explicit inspection of its topology and geometry. Given a conditional DDPM trained on the 2D Gaussian bump dataset described above, we first investigate whether the model learns a coupled or factorized representation. Naïvely, a 2D Gaussian dataset with two independent features of variation, $x$ and $y$, has a 2D plane-like latent representation. Unfortunately, simply inspecting the learned representation would not allow us to easily differentiate between a coupled versus a factorized representation since the Cartesian product of two lines $L^{1}\times L^{1}\subset\mathbb{R}\times\mathbb{R}=\mathbb{R}^{2}$ is topologically and geometrically equivalent to a plane $P^{2}\subset\mathbb{R}^{2}$. To overcome this issue, we perform a simple modification to the Gaussian dataset. We impose periodic boundary conditions in the image space to connect the left-right and top-bottom boundaries of the image such that the latent representation of the dataset forms a torus. Mathematically, a torus is defined by the Cartesian product of two circles $S^{1}\times S^{1}$. A Clifford Torus is an example of a factorized representation of the torus, with each circle independently embedded in $\mathbb{R}^{2}$, resulting in a 4-dimensional object. However, the most efficient representation in terms of extrinsic dimensionality is the regular torus $T^{2}$, which is embedded in $\mathbb{R}^{3}$, albeit unfactorized. Various 2D projections of the 3D torus and the Clifford (4D) torus are shown in Fig. 2(a). Due to geometric differences between the Clifford and the 3D torus, we can now distinguish whether the model learns a factorized representation. If the model were to represent $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.

To apply the geometry tests for differentiating between the 3D and the Clifford torus, we first need to confirm that the model indeed learns a torus representation of the dataset by computing the topological features of the learned representation via persistent homology. In Fig. 2(b), we compare the persistence diagrams of a standard torus (left) with the final learned representation of the model (right). Both diagrams exhibit the similar topological features, rank-1 $H_{0}$ group, rank-2 $H_{1}$ group (with overlapping orange points on top in both diagrams), and rank-1 $H_{2}$ group, signaling that the model indeed learns a torus representation.

We then investigate when this torus representation emerges during training. In Fig. 2(c), we show the model’s performance as measured by the accuracy of its generated images (top) and the corresponding effective dimension (bottom) of the learned representations across training, as measured by the participation ratio ( given by $(\sum_{i}\lambda_{i})^{2}/\sum_{i}\lambda_{i}^{2}$, where $\lambda_{i}$ is the eigenvalue of the $i$-th principal component Gao et al. (2017)). Intuitively, the participation ratio counts the dominant eigen-components, providing a measure of the effective dimension of the geometric object of interest. A detailed inspection of the effective dimension of the learned representations over the training duration informs us of whether and how the model arrives at a factorized, 4D representation. We see in Fig. 2(c) that as training progresses, the model’s internal representation first undergoes a dimension increase then a decrease, eventually converging to around 7 dimensions after 200 epochs. While the dimensionality converged to a higher dimension than 4, the top 4 eigenvectors became notably more prominent as the training converges, as indicated in the eigenspectra in Fig. 2(d). This signals that a 4D, rather than 3D, torus representation is eventually learned. Various PCA projections of the learned representations along the training process are shown in Fig. 2(g)-(i), where the projections in the last epoch resembles those of the Clifford torus shown in (a).

To confirm whether the model learns a Clifford torus, we employ the orthogonality and parallelism tests of tori Cueva et al. (2021b) (details found in Appendix B). In an ideal Clifford torus, the rings along the poloidal direction should be parallel with each other, and similarly for those along the toroidal direction. Moreover, the rings along the poloidal direction should be orthogonal with respect to the rings in the toroidal direction. The orthogonality and parallelism tests measure the orthogonality and parallelism of subspaces spanned by various rings on the torus. If the model learns a Clifford torus representation, the subspaces that code for $x$ and $y$ should be orthogonal, and the subspaces that code for $x$ ($y$) for any given $y$ ($x$) should be equivalent (parallel). Given a pair of $x$ and $y$, we construct 2D subspaces of rings $S(\text{fix }x)$ and $S(\text{fix }y)$ spanned by fixing $x$ while varying $y$ and fixing $y$ while varying $x$, respectively. We perform 3 sets of tests: 1) $x$-on-$y$ tests test orthogonality and parallelism between surfaces $S(\text{fix }x)$ and $S(\text{fix }y)$ for all combinations of $x$ and $y$; 2) $x$-on-$x$ tests test orthogonality and parallelism between $S(\text{fix }x_{1})$ and $S(\text{fix }x_{2})$ for all combinations of different $x$’s; 3) $y$-on-$y$ tests test orthogonality and parallelism between $S(\text{fix }y_{1})$ and $S(\text{fix }y_{2})$ for all combinations of different $y$’s.

In Fig. 3, we compare the orthogonality and parallelism test statistics between an ideal 3D torus (left), the model’s learned representation at the final epoch (center), and an ideal Clifford torus (right). Comparing test 1 statistics (top row) between the model representation and the standard 3D and Clifford tori, we see that the model encodes $x$ and $y$ independently in mostly orthogonal subspaces, similar to a Clifford torus. The test 2 and 3 statistics (middle and bottom row) of the model representation, however, does not match those of either the 3D or the Clifford torus. In fact, the statistics from tests 2 and 3 of the model representation closely resemble those from test 1 of the 3D torus. This means that the model is encoding different values of $x$ in an almost orthogonal, as opposed to parallel, fashion, similarly for different values of $y$. These observations signify that the models “hyper-factorizes” each individual factor of variation by treating them more similarly to categorical variables with nonzero overlaps between neighboring categories rather than fully-continuous variables, which explains why the effective dimension of the model’s learned representation is higher than 4 after convergence in Fig. 2(c). Although these observations are made from a model trained on a torus dataset, we expect the model to have similar encoding schemes across all variants of our Gaussian bump datasets.

In this section, we examine the model’s ability to compositionally generalize. Specifically, we train the model on incomplete datasets of 2D Gaussian SOSs, in which we leave out all Gaussian SOSs centered in the red-shaded test regions (Fig. 4(f), sampling distributions shown in Fig. 4(a), (b)). We then assess the performance of the models in generating 2D Gaussian SOSs centered within and outside of the test regions. Here we choose the width of the cuts to be around 6 pixels wide, which roughly corresponds to the width of the Gaussian stripes of $\sigma=1.0$. We design the lesions such that we can probe the model’s ability at compositionally generalizing out of the training distribution as well as its ability to spatially interpolate in a single variable alone. There are 4 possible outcomes based on the model’s ability to compose and interpolate, and we give predictions for the model’s generalization performance in each case: 1) the model cannot compose or interpolate: here, we would see low performance across the test regions; 2) the model interpolates but cannot compose: here, we would see high performance across the test regions; 3) the model composes but cannot interpolate: here, we would see high performance in one dimension but not the other in the non-intersecting part of the test regions, low performance in the intersection; 4) the model can compose and interpolate: here, we would see higher performance across the test regions. We note that case (2) and (4) are indistinguishable via the behavior of the model.

In our first experiment, we simply train our model on a 2D Gaussian SOS dataset where all data centered in the test regions are left out, i.e. excluding all $\mu_{x},\mu_{y}\in[13,19]$ as shown in Fig. 4(a). We call this model the 2D model since the model is trained on only 2D Gaussian SOS data. We then examine the terminal accuracy of the 2D model in generating Gaussian SOSs at the correct $\mu_{x}$ and $\mu_{y}$ in various parts of the test regions, which we section into a horizontal part, a vertical part, both excluding the area of their intersection, and the intersection itself (shown schematically in Fig. 4(f)). We note that the 2D model achieves high accuracy in generating $\mu_{y}$ while suffers low accuracy in generating $\mu_{x}$ in the vertical section. Similarly, the 2D model achieves high accuracy in generating $\mu_{x}$ while suffers low accuracy in generating $\mu_{y}$ in the horizontal section. The model suffers low accuracy in generating both $\mu_{x}$ and $\mu_{y}$ in the intersection region. These observations have two implications: i) the model is factorized and compositional since it is able to generate the correct $\mu_{x}$ or $\mu_{y}$ irrespective of the other; ii) the model has limited ability to spatially interpolate, which suggests that it does not learn a fully continuous manifold in its activation space (see Fig. 10(c)). These observations meet our expectation for outcome case (3), when the model can compose but not interpolate and resonates our conclusion from Sec. 3.1 that the model has learned to factorize x and y, but has not learned a consistent representation across all x’s (and likewise y’s). In Appendix C.3 Fig. 11, we quantitatively assess model’s ability to interpolate as a function of the held-out range width in the data manifold and found that model’s ability to interpolate gradually decrease as a function of the held-out region width.

To further verify our hypothesis that the model is capable of composition but not interpolation, we design an alternative dataset that includes a mixture of 2D Gaussian SOS data and 1D Gaussian stripe data (single stripes that are either horizontal or vertical). Specifically, we generate a set of 1D Gaussian stripe data across the entire range of $x$ and $y$ without any held-out regions, as shown in Fig. 4(b). Note that the 1D Gaussian stripe data respects the 2D structure of the conditional input embedding. The exact generation procedure for the 1D Gaussian stripe dataset is described in Appendix C.1. This 1D dataset is then combined with the previous 2D Gaussian SOS dataset with the same held-out test regions. The resulting 2D+1D dataset contains 1D data in all range of $x$ and $y$ but has not seen 2D data (composition of 1D) in the test regions. The model trained under this dataset, deemed the 2D+1D model, has high performance in generating both $\mu_{x}$ and $\mu_{y}$ across all of the test regions. This shows that given the necessary “ingredients" (1D Gaussian stripes of all range of $x$ and $y$), the model is able to compose them in a proper manner (2D Gaussian SOSs). These observations again confirm that the learned representation of the model is factorized but not consistent across a single feature, meaning that the model is able to compose independent features but not interpolate within a single feature dimension.

Finally, we test whether compositionality can be learned given no compositional examples in the training dataset. To test this, we train a 1D model on the 1D Gaussian stripe dataset as described above. Our results show that while model achieves respectable accuracy in generating the correct $\mu_{y}$ across the test regions, it does not simultaneously generate the correct $\mu_{x}$, indicating that the model did not learn to properly compose the 1D stripes. This is also evident from the generated image samples from the 1D model shown in Fig. 10(c). Without any compositional examples, the model fails to generate the intersection of two 1D Gaussian stripes and defaults back to generating one of the stripes at the correct location. While the 1D model has also learned to factorize $x$ and $y$ as independent concepts, it is not able to generalize here due to the lack of compositional examples. Hence, based on the performance of the 1D, 2D, and 2D+1D models that we have seen above, we conclude that complete compositional generalization requires data of all range of $x$’s and $y$’s and some compositional examples of the two to be present within the dataset.

Building on our previous results, we investigate whether combining the 2D Gaussian bump dataset with 1D stripe data can impart a different form of compositionality to the model – multiplicative rather than additive. To do this, we create a dataset that mirrors the 2D+1D Gaussian SOS set, substituting the 2D Gaussian SOS data with 2D Gaussian bump data while maintaining the same test regions. The latent representation of this dataset can be found in Fig. 9(h) in the Appendix. The resulting model trained under the 2D Gaussian bump + 1D Gaussian stripe dataset is able to generate 2D Gaussian bumps within the test regions up to 70% accuracy, suggesting that the model is adept at handling an alternative form of compositionality. We further investigated whether the number of 2D Gaussian bump examples in the dataset influences compositionality by subsampling the 2D Gaussian bumps in the in-distribution regions (Fig. 9(g)). The results show that models trained on datasets with varying percentages of subsampled 2D Gaussian bumps achieve similar performance across the test regions, even with as little as 20% of the original data. These findings suggest that training for compositional generalization can be efficiently achieved by using a full range of data for each individual factor of variation, along with a few compositional examples.

To validate the benefit of training with 2D + 1D datasets, we analyzed the data efficiency scaling of models trained on 2D Gaussian bumps + 1D Gaussian stripes, comparing them to models trained solely on 2D Gaussian bumps as a function of image size $N$. For both the baseline and augmented datasets, we excluded $20\%$ of the 2D Gaussian bumps from the center square region of the $N\times N$ canvas for out-of-distribution (OOD) evaluations. The accuracy of the models are plotted against the percentage of 2D Gaussian bumps subsampled in Fig. 5(a) for the baseline dataset and (b) for the augmented dataset. The results clearly demonstrate that the model trained on the augmented dataset performs better with a significantly lower percentage of 2D Gaussian bumps in the training dataset. We then plotted the sizes of the baseline and augmented datasets required to reach a $60\%$ accuracy threshold as a function of $N$ on a log-log scale in Fig. 5(c). We found that models trained on the augmented datasets only require linear number of data in $N$ while those trained on the baseline datasets require quadratic number of data in $N$. This result suggests that augmented datasets with samples of all independent factors plus a few compositional examples provide an efficient training method that has sample complexity linear in the number of factors. In Appendix C.2 and C.4, we further discuss the sample efficiency of learning compositionality from dataset of isolated factors of variation plus few compositional examples.

In summary, our findings indicate that the model excels at compositionality but struggles with interpolation. Consequently, effective generalization necessitates that the training set encompasses the full range of each individual factor of variation, alongside compositional examples. Furthermore, we discovered that models trained on datasets featuring explicitly factorized examples exhibit remarkable data efficiency and adaptability to various forms of compositionality. These results align with the observed factorization patterns in the model representation discussed in Sec. 3.1 and suggest promising strategies for developing more data-efficient training schemes that utilize tailored datasets.

We have established that the model learns factorized (but not fully continuous) and compositional representations. Next, we identify a potential connection between manifold learning in diffusion models and percolation theory in physics, which provides a normative explanation as to how factorization (or compositional generalization) emerges. To start, we note that our Gaussian bump datasets can be approximated as a simple Poisson Boolean (Gilbert’s disk) model Gilbert (1961); Hanisch (1980), a well-studied system in continuum percolation theory. Fig. 6(a) shows a schematic illustration of our dataset of Gaussian bumps approximated as disks of various widths on a 2D $32\times 32$ lattice of grid spacings $d_{x}$ and $d_{y}$. In percolation theory, the quantity of interest is the critical fraction of nodes that need to have non-zero overlaps in order for the entire system to be interconnected with high probability. For most systems (finite- and infinite-sized), there exist a phase transition that occurs at the critical fraction that can be either analytically derived or numerically estimated, with the transition becoming sharper as the system size scales. In our context, since the Gaussian bump images are pixelated, overlaps between neighboring Gaussian bumps are not smooth. In Fig. 6(b), we plot the normalized L2-norm of the neighboring distance-$d$-separated Gaussian bumps with various spread $\sigma$’s. We note that since the L2-norm measure of overlap is normalized to 1 between a given Gaussian bump and itself, the overlap between a given Gaussian bump and a slightly offset Gaussian bump temporarily goes beyond 1 due to the discrete nature of the data.

We hypothesize that below a threshold amount of training data, the diffusion model cannot construct a faithful representation of the training dataset. From the percolation perspective, if there does not exist a large enough interconnected component within the dataset, the model will fail to learn the relative spatial location of the data points, making it hard to learn a faithful 2D representation. To test this hypothesis, we first simulated the mass ratio of largest interconnected components as a function of the unit area data density $\lambda$ with our 2D Gaussian bump datasets. For simulation, we use a dataset of 1024 data points, corresponding to $d:=d_{x}=d_{y}=1.0$ on a $32\times 32$ lattice. We then randomly sample $\lambda\times 1024$ points on the lattice to compute the size of the largest interconnected cluster of Gaussian bumps. Here, we define a hyperparameter of threshold overlap beyond which we consider two Gaussian bumps as overlapping. In Fig. 6(c), we show the simulation results. averaged over 5 runs, with the chosen threshold overlap of 0.005 for datasets with various $\sigma$’s. Additional simulation results with various chosen overlap thresholds can be found in Appendix D.

Next, we quantify the connection between percolation theory and manifold formation in diffusion models. To account for the stochasticity in sampling the datasets and avoid significant overhead in model training, we choose a fix set of grid points on the $32\times 32$ lattice and generate datasets of varying $\sigma$’s from this fixed set of grid points. Moreover, we choose each sampled fraction of the dataset to be a strict superset of the smaller sampled fractions of the dataset. For example, dataset at $\lambda=0.2$ contains all of dataset at $\lambda=0.1$, and dataset at $\lambda=0.3$ contains all of dataset at $\lambda=0.2$ and hence all of $\lambda=0.1$, so on and so forth. For datasets of different $\sigma$’s, we chose the same grid points such that the only difference between datasets of different $\sigma$’s is just the $\sigma$ itself. This procedure eliminates all the stochasticity due to the probabilistic nature of percolation, and ensures that differences in training are only due to the differences in $\sigma$ and $\lambda$. To ensure fair comparison, we proportionally add training time to each model trained with fewer data such that the total amount of training batches are constant across all models.

In Fig. 6(e), we display the accuracy of models trained with different fractions of a training dataset of a total size of 1024 for various $\sigma$’s. We note that there are high variances in the model performance, especially for $\sigma=0.3,0.5$. This is due to stochasticity in model training itself as well as the gradual as opposed to a sharp phase transition, as shown in simulation Fig. 6(c). In addition, we notice that based on the simulation, $\sigma=0.1$ will not forge any connected cluster due to its minuscule spread. In Fig. 6(d), an sample image of Gaussian bumps of $\sigma=0.1$ is shown on the left, where it is a mere pixel centered at $(\mu_{x},\mu_{y})$. Nonetheless, in experiment Fig. 6(e), we see that the model is able to learn the $\sigma=0.1$ dataset at $\lambda=0.8,0.9$. This is partially due to the fact that the noising process in diffusion training smears out the Gaussian bump of $\sigma=0.1$ to give it a larger effective width, as shown in the generated bump image in Fig. 6(d). In general, our experimental results agree with the percolation theory prediction that smaller similarities between the samples makes it harder for the models to learn meaningful representations, despite the same quanitity of data. The portraits of the learned representations by the models generating Fig. 6(e) is shown in Fig. 15 of Appendix D.2.