SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Learning dynamics respect the structure of the concept graph (Fig. 5). We find the ability to compositionally generalize, i.e., produce samples from out-of-distribution concept classes, emerges at a rate which is inversely related to a class's concept distance with respect to classes seen in training. Specifically, in Fig. 5 (a) we show the learning dynamics of the model, where lightblue nodes denote concept classes within the training dataset, while pink and darkpink nodes respectively denote classes at concept distances of 1 and 2 from classes in the training dataset. As the model learns to fit its training data (lightblue nodes), it learns capabilities that can be composed to produce samples from out-of-distribution concept classes (pink / darkpink nodes). Fig. 5 (b) further shows that the learning dynamics of compositional generalization are influenced by the concept distance from the training set: the model first learns the concept classes in the training dataset (lightblue lines) and then generalizes to concept classes with a concept distance of 1 from the training dataset (pink lines). Thereafter, the model suddenly acquires the capability to compositionally generalize to a concept class with a concept distance of 2 from the training dataset (darkpink line). Fig. 5 (c) shows the images generated by the model over time. We observe that rough shapes and sizes are learned relatively early in training, by the 4th epoch, while the color is highly biased to be red, the majority color in the training dataset, up to the 10th epoch. Then, around the 20th epoch, the model learns to generate the underrepresented
Figure 5: Concept distance from the training set governs the order in which compositional capabilities emerge. (a) Concept graph (cube) depicting training data points (blue nodes) and concept distances for test data points, where pink nodes represent distance  $= 1$ , and darkpink nodes represent distance  $= 2$ . Each trajectory represents the learning dynamics of generated images given a tuple prompt. Each trajectory represents the learning dynamics of generated images based on each tuple prompt. During every training epoch, 50 images are generated and binary classification is performed to predict each concept, including color, shape, and size. (b) Compositional generalization happens in sequence, starting with concept distance  $= 1$  and progressing to concept distance  $= 2$ . The x-axis represents the number of epochs, and the y-axis represents the progress of compositional generalization. (c) Images generated as a function of time clearly show a sudden emergence of the capability to change the color of small, blue, triangles.
color (blue) for concept classes with a concept distance of 1. Finally, around the 40th epoch, the model learns to generate the underrepresented color (blue) for the class at distance 2, showing a sudden emergence of capability to generate samples from that class. The observations above generalize to higher-dimensional concept graphs as well (see Fig. 16). Specifically, we now introduce a fourth concept variable backgroundcolor  $= \{\text{white, black}\}$  and again find that compositional generalization occurs in the order governed by the distance from the training set: the accuracy begins to rise for concept distance of 1, continues to increase for distance 2, and peaks for 3.
Multiplicative influence of capabilities drives the sudden emergence of compositional generalization (Fig. 6). The interpretability of our experimental setup illustrates that the multiplicative reliance on underlying capabilities is the critical structure behind the sudden emergence of compositional generalization in our setup. For example, in Fig. 6 (a) we show the dynamics of an additive vs. multiplicative accuracy measure for generating concept class {111, triangle, small, blue}), which has a concept distance of 2 from the training data. We observe a sudden improvement of the multiplicative measure. To better understand the above, in Fig. 6 (b), we plot the accuracy of probes used for predicting each concept variable (shape, size, color). From the plot, we notice
that the model fails to acquire strong color transformation capabilities until the final stage of training, effectively bottlenecking compositional generalization of 111 class. This empirical observation leads us to the following hypothesis on emergence of compositional abilities in generative models.
Hypothesis. (Compositional Abilities Emerge Multiplicatively.) We hypothesize the nonlinear increase in capability observed in large neural networks as size and computational power scale up is partially driven by the task's compositionality. Models must learn all required concepts, but compositional generalization is hindered by the multiplicative, not additive, impact of learning progress on each concept. This results in a rather sudden emergence of capabilities to produce or reason about out-of-distribution data.
Figure 7: Toy Model of Compositional Emergence. (a) Probability of learning a compositional capability at a concept distance  $n$  w.r.t. atomic abilities as a function of time  $t$ . (b) Critical optimization time steps  $t^*$  required for the above probability to reach a threshold  $P^*$  as a function of concept distance  $n$ .
The core aspect worth noting in our hypothesis is that compositional tasks require the concurrent acquisition of all involved "atomic" capabilities, akin to an AND logical condition. Correspondingly, in generative models, the learning of individual atomic abilities leads to a manifestation of emergent capabilities<sup>1</sup>. We argue this requirement leads to a non-linear increase in a model's capabilities (see also concurrent work making similar claims [78, 117]). We analyze this further via a simple toy model below.
Toy analysis. We consider a toy compositional problem which demonstrates how complex capabilities can rapidly develop ("emerge") from a set of atomic abilities. Assume there are  $n$  atomic abilities, each with a probability  $p$  of being learned in a given time step, i.e., the dynamics of learning an ability can be modeled as a Bernoulli coin flip: once the coin turns from 0 to 1, the model learns the said ability. The probability that the ability will be learned in  $t$  steps is thus given by  $1 - (1 - p)^t$ . An implicit assumption in this model is that the learning dynamics of
Figure 6: Multiplicative influence of individual capabilities elicits an "emergence" of compositionality. (a) Accuracy of producing samples from the concept class \(\{111\), (triangle, blue, small)}\), which has a concept distance of 2 from the training data. A multiplicative measure (solid line) has a score of 1 when all concept variables of shape, color, and size are correctly predicted. Conversely, an additive measure (dashed line) independently relates each concept variable prediction accuracy to the compositional accuracy, deceptively suggesting smooth progress. (b) Learning dynamics of predicting each of the three concept variables: shape (blue), color (orange), and size (green).
This is related to the notion of "weak emergence" studied in psychology [115] and complex systems [116].
atomic abilities are independent (hence the name atomic) and that once an atomic ability is learned, it will not be forgotten. Now, our goal is to characterize the dynamics of learning a capability that involves composing these atomic abilities. The multiplicative emergence hypothesis argues that the learning such a capability by time  $t$  requires that the model have learned the  $n$  atomic abilities by that time. We then get the probability that the compositional capability has been learned by time  $t$  is  $P(n) = (1 - (1 - p)^t)^n$ . These dynamics are plotted in Fig. 7 (a) and look strikingly similar to the rapid learning (emergence) we found in our diffusion model setup (see Fig. 5 b): for increasing degree of compositionality, we see much sharper learning dynamics (i.e., emergence) that is controlled by the concept distance w.r.t. atomic abilities. Further, assuming a given threshold probability  $P^*$  at which one claims a capability has been learned, for a degree of composition  $n$ , we can solve for the critical time  $t^*$  at which a compositional capability is learned as follows:  $t^* = \left\lceil \frac{\log(1 - (P^*)^{1 / n})}{\log(1 - p)} \right\rceil$ . Plotting this in Fig. 7 (b), we see the progress on increasingly more compositional tasks is logarithmic w.r.t. the number of atomic concepts being composed (aka concept distance). Importantly, this implies if we allow the model to train infinitely and it learns several atomic abilities, it will have an explosion of capabilities due to the inherent compositionality of the data generating process—we further investigate this in a follow up work [118].
4.1 Challenges for Compositional Generalization
We have shown that, given a well-structured synthetic training dataset, a conditional diffusion model can learn to compose its capabilities to generate novel inputs not encountered in the training set. Next, we analyze an adversarial setups under which the model fails to learn relevant capabilities to compose concepts and if simple interventions can help mitigate these failures.
Critical frequency for learning capabilities (Fig. 8). We first probe the effect of changing the frequencies of samples from different concept classes in the training data and examine how this affects the model's ability to learn and compose capabilities involving that concept class. Results are shown in Fig. 8 and demonstrate how the frequency of color and size concepts in the training data impacts the generalization capabilities of the diffusion model. Specifically, we change the number of samples in the training data from 0 to 300 for concept class 001 (Fig. 8 (a)) and class 100 (Fig. 8 (b)). As can be seen, low frequencies of certain concept values degrade the accuracy of the model in both settings. Notably, training for out-of-distribution concept classes (pink lines) requires more samples than that for in-distribution ones (lightblue lines). This suggests that as the sample size grows, memorization occurs first, and generalization is achieved superlinearly as a function of data frequency. More importantly, we observe a critical number of samples are required before we can see the onset of capabilities to alter a concept. Specifically, in Fig. 8 (a), we can see that the model rapidly learns the color concept after being provided with 10 samples for a concept of large blue circle, 001. In contrast, in Fig. 8 (b), the model learns the shape concept only after reaching a certain threshold in the number of samples with a concept of large red triangle, 100.
We believe the results above are especially interesting because an often used strategy to prevent a generative model from learning harmful capabilities, such as the ability to generate images involving sensitive concepts, involves cleaning the dataset to remove images corresponding to such concepts, hopefully hindering the model's ability to generate samples containing the concept [33, 119, 120].
(b)
Figure 8: How does the frequency of data samples impact the learning of capabilities? We systematically control the frequency of a specific concept class in the training dataset and observe how it affects the model's learning of capabilities. (a) Capability to alter colors is quickly learned after introducing approx. 10 samples with a concept class of large blue circle, 001. (b) In contrast, a critical threshold (marked with dotted vertical lines) exists for learning a capability to alter the shape: as we gradually introduce samples, with a concept class of large red triangle, 100.
(b)
Figure 9: Fine-tuning is not a remedy for misgeneralization in an adversarial task. (a) The model struggles to generalize when the training setup is configured adversarially. As shown in the plot, the model incorrectly produces triangles (instead of circles) for both 011 and 010. (b) The model faces difficulty learning new concepts through fine-tuning. We add node 101 (triangle, large, blue) to the dataset and attempt to fine-tune the model. However, even with a large learning rate equal to the one used for training, the model fails to learn the capability to alter object size.
Such filtering is generally expensive and, arguably, statistically impossible to achieve perfectly in a large dataset, i.e., a few samples corresponding to the sensitive concept are likely to remain in the data. Our findings suggest that one might better consider the relationship between the probability of learning a concept and its frequency of occurrence in the data to determine the degree of filtering required: if their presence in the training data is below a critical threshold, that can be sufficient to deter the model from learning a capability to generate samples related to that concept, obviating the need for a perfect filtering of the dataset.
Fine-tuning is not a remedy for misgeneralization in an adversarial task. (Fig. 9). We next evaluate a setting where concept classes present in the training data are not neighbors, i.e., their concept distances are greater than 1, but nonetheless represent all concept variables to allow a model to learn relevant capabilities. Specifically, as shown in Fig. 9 (a), we use only the following four pairs of tuples and images for training:  $\{000, (\text{circle}, \text{red}, \text{large})\}$ ,  $\{100, (\text{triangle}, \text{red}, \text{large})\}$ ,  $\{001, (\text{circle}, \text{blue}, \text{small})\}$ ,  $\{111, (\text{triangle}, \text{blue}, \text{small})\}$ . Arguably, capabilities corresponding to shape change (000 to 100) and color change (000 to 001) can still be learned as the setup for these is similar to our prior experiments. However, our results show this yields an interesting failure mode; specifically, in Fig. 9 (a) we find the model fails to dissociate a second element value of 1 with the shape of a small triangle, and this bias causes the model to generate small triangles for both 011 and 010, when it should have produced small circles.
We further analyze if the failure above can be fixed via mere downstream fine-tuning (see Fig. 9 (b)). Specifically, we fine-tune the trained model on a dataset that includes the concept class of 101 (left), 010 (middle), and 011 (right). We find the model still generates small triangles for 011 and 010 even after fine-tuning. When the concept classes 010 (middle) and 011 (right) are added, the newly introduced concepts "overwrite" all existing concepts, causing the previously learned concepts (e.g., the color and shape concept for 101) to be forgotten. These results present a further challenge in addressing learned bias.
4.2 Additional Experimental Results with Real Data
To validate our hypotheses in more realistic settings, we conducted experiments using the CelebA dataset. We selected three attributes as the concepts: Gender (categorized as Female and "Male"), Smiling ("Smiling" and Not Smiling), and Hair Color ("Black Hair" and "Blonde Hair"). The concept graph corresponding to these attributes in the CelebA dataset is illustrated in Fig. 10 (a). For each individual concept, we trained a 3-layer CNN and adopted the product of accuracies from each concept as the evaluation metric. Although not all nodes achieved an accuracy of  $100\%$  due to limited training time, we find that our observations (partially) extend: (i) learning dynamics following the structure of the concept and (ii) delayed generalization of underrepresented attribute (gender), hold true, even at larger, more realistic scales. Fig. 10 (b) illustrates the learning dynamics using the real CelebA dataset. We again observe the sequential pattern in generalization. The accuracy begins to rise at concept distance 1 (denoted by pink lines), followed by concept distance 2 (red line). Notably, the node labeled 011—even though it is at concept distance 1—precedes the memorization stage (represented by blue lines). This exception can be attributed to the gender bias present in the training data. In Fig. 10 (c), we plotted the accuracy for the individual concept of gender (i.e., Female and Male). The Female concept class's training curve (red lines) reaches convergence faster
Figure 10: Concept distance from the training set governs the order in which compositional capabilities emerge. (a) Concept graph (cube) depicting training data points (blue nodes) and concept distances for test data points for CelebA dataset. (b) Compositional generalization happens in sequence, starting with concept distance 1 and progressing to concept distance 2. (c) Delayed emergence of abilities to generate underrepresented attribute (gender) for distant classes. Gender prediction accuracy based on generated samples at each step during training. We observe that the ability to generate underrepresented attributes occurs much later as concept distance increases.
compared to the Male concept class (blue lines). This disparity stems from the Female class being more predominant in the dataset, thereby making the Male class underrepresented. This observation offers useful insights for practitioners: to mitigate biased outcomes against underrepresented groups, we need to further train the diffusion model beyond its initial convergence on the training dataset.
In this work, we introduced a concept graph framework and used controlled interventions with a synthetic dataset to formulate and characterize four novel failure modes in the compositional generalization of diffusion models. We first hypothesized and verified that compositional generalization to a concept class far from the training data, based on concept distance, can fail, as it requires more optimization steps due to the sequential nature of generalization (Fig. 5). We show compositionality emerges in a sequence that respects the geometric structure of the concept graph, eliciting a multiplicative emergence effect that manifests as sudden increases in the model's performance to produce out-of-distribution data well after it has learned to produce samples from its training distribution. This behavior is reminiscent of the recently observed phenomenon of grokking in language modeling-like objectives [80, 81, 121]. Furthermore, when a particular concept value is underrepresented in the training dataset (e.g., an underrepresented color or gender), avoiding misgeneralization of the concept value, and thus achieving compositional generalization, requires far more additional optimization steps, even after achieving perfect performance on the training set. We also discovered a critical threshold for the inclusion of certain concepts in the training data, which is useful for deliberately making the model fail at learning harmful concepts (Fig. 8). Finally, we first analyze misgeneralizations in an adversarial task and find that fine-tuning does not provide a solution (Fig. 9). Overall, our synthetic data approach allowed us to apply controlled interventions that shed light on four different mechanisms behind the failure of diffusion models in compositional generalization.
Authors' Contributions
ESL initiated the research by defining the problem. ESL and HT formulated the theoretical framework, the concept graph, and collaborated with MO on the experimental design. MO conducted most experiments and, along with ESL and HT, produced the visualizations and figures. ESL and HT authored the introduction. ESL was responsible for the literature review and related work. MO, ESL, and HT jointly worked on the results and discussion sections. Throughout the study, MO, ESL, and HT collaborated on analyzing experimental results and refining the manuscript. HT and RPD provided supervisory guidance and feedback.
We thank Naomi Saphra, David Krueger and his group, Nikhil Vyas, Mohamed El Banani, and Anna Golubeva for fruitful discussions during the course of this project. ESL's time at University of Michigan was partially supported via NSF under award CNS-2008151.
A.1 Synthetic dataset
The dataset consists of 5,000 rendered images of 2D geometric shapes, along with corresponding concept classes. These synthetic images are generated using  $\mathrm{Blender}^2$  by creating a scene graph and rendering it. Each scene contains a single object placed on a blank background of size  $28 \times 28$ . The objects have three types of attributes: size, color, and shape. There are two shapes (circle and triangle), three colors (red, blue), and two sizes (large, small), resulting in up to eight different combinations of attributes. Each image is annotated with the corresponding object attributes, which can be utilized as conditional features for image generation. This enables us to directly evaluate the text-to-image generation capability of the diffusion model against ground truth images. Fig. 11 depicts example images with the corresponding concept classes.
Figure 11: Examples of data samples with pairs of images and corresponding concept classes.
A.2 Loss function
Diffusion models convert Gaussian noise into samples from a data distribution through an iterative denoising process. The sampling process starts with a noisy input  $\mathbf{x}_T$ , and denoised samples are generated through gradual iteration,  $\mathbf{x}_{T - 1}$ ,  $\mathbf{x}_{T - 2}$ , until the original input  $\mathbf{x}_0$  is obtained. Conditional diffusion models [10, 122] allow for a denoising process conditioned on texts or class labels. In all the experiments, we used the conditional diffusion model with the form  $p(\mathbf{x}|V)$ , where  $\mathbf{x}$  denotes an image and  $V = \{v_{1}, v_{2}, \dots, v_{n}\}$  denotes a set of  $n$  concept variables. To predict the noise  $\epsilon$  at each timestep  $t \in [0, T]$ , we follow the approach proposed in [3] by training a neural network. Specifically, we construct a neural network  $\epsilon_{\theta}(\mathbf{x}_t, t, \mathbf{a})$  and minimize the mean squared error (MSE) between the predicted Gaussian noise and the true noise:
$$
\mathcal {L} = \mathbb {E} _ {t \in [ 0, T ], \mathbf {x} _ {0} \sim q (\mathbf {x} _ {0}), \epsilon \sim \mathcal {N} (\mathbf {0}, \mathbf {I})} \left[ \| \epsilon - \epsilon_ {\theta} (\mathbf {x} _ {0}, t, V) \| ^ {2} \right], \tag {1}
$$
where  $q(\mathbf{x}_0)$  denotes the distribution of input image  $\mathbf{x}_0$ , and  $\mathcal{N}(\mathbf{0},\mathbf{I})$  denotes the standard Gaussian distribution.
A.3 Architecture
We use the conditional U-Net architecture [1], as in [3], for our neural network  $\epsilon_{\theta}(\cdot)$ . Our architecture comprises two down-sampling and up-sampling blocks, with each block consisting of  $3\times 3$  convolutional layers, GELU activation, the global attention, and pooling layers. The conditional information  $V$  are fed through an embedding layer and concatenated with the image feature maps at each stage of the up-sampling blocks. We illustrate our network architecture in Fig. 12.
$^{2}$ http://www.blender.org
Figure 12: The architecture of the conditional diffusion model. The architecture of the conditional diffusion model involves an iterative process comprising noise addition and denoising steps. The model leverages conditioning information, specifically concept classes, to guide the transformation of the input image towards a desired state. In our implementation, we utilize a U-Net to parameterize the denoising process. The U-Net architecture consists of three upsampling convolutional layers and three downsampling convolutional layers, which are connected through skip connections. Each layer within the U-Net includes a pooling layer, a global attention mechanism, and a GELU activation function.
A.4 Optimizer
We implemented the diffusion model using PyTorch and trained it on four Nvidia A100 GPUs. We performed a hyperparameter search based on a validation set. We tested batch sizes ranging from 32 to 256, the number of channels in each layer from 64 to 512, leaning rate from  $10^{-4}$  and  $10^{-3}$ , the number of steps in the diffusion process from 100 to 400. We employed the Adam optimizer [123] with  $\beta_{1} = 0.9$ ,  $\beta_{2} = 0.99$ , and weight decay of  $10^{-5}$ .
A.5 Evaluation metric