SlowGuess/NeurIPS_2023_OCR huggingface dataset:
For the evaluation, we follow a probing protocol popularly used in several prior works on learning disentangled representations [103, 109-114]. To evaluate the attributes of the images generated by our conditional diffusion model, we trained linear classifiers on these images for three specific attributes: shape, color, and size. For each attribute, we developed a dedicated classifier:  $f_{0}(\hat{x}_{0})$  for shape,  $f_{1}(\hat{x}_{0})$  for color, and  $f_{2}(\hat{x}_{0})$  for size. Here  $\hat{x}_0$  denotes the image generated by the conditional diffusion model. We utilized a cross-entropy loss function to train these classifiers. The output of each classifier fell into one of two categories: for shape, the categories were circle or triangle; for color, blue or red; and for size, large or small. We then calculated the accuracy for each attribute using the corresponding classifier outputs. To quantitatively assess the accuracy of predicted attributes aligning with their corresponding ground-truth concept classes, we utilize a multiplicative measure. This measure gaues the accuracy of all attributes, and defined by the product of individual accuracies for each attribute as follows:
$$
\text {A c c u r a c y} = \frac {1}{N _ {t}} \sum_ {n = 1} ^ {N _ {t}} \mathbb {1} \left(f _ {0} \left(x _ {0} ^ {(n)}\right), v _ {0} ^ {(n)}\right) \cdot \mathbb {1} \left(f _ {1} \left(x _ {0} ^ {(n)}\right), v _ {1} ^ {(n)}\right) \cdot \mathbb {1} \left(f _ {2} \left(x _ {0} ^ {(n)}\right), v _ {2} ^ {(n)}\right), \tag {2}
$$
where  $\mathbb{1}(\cdot)$  is the indicator function,  $n$  is the index of the test samples, and  $N_{t}$  is the total number of samples used for evaluation. For our experiments, we generated  $N_{t} = 50$  images for each input of concept classes.  $v_{0}^{(n)}$ ,  $v_{1}^{(n)}$ , and  $v_{2}^{(n)}$  denote the actual (ground truth) concepts classes for shape, color, and size, respectively. We trained them over 50 epochs using the training dataset comprising 5,000 pairs of concept classes and images. The trained linear classifiers achieved an accuracy rate of  $100\%$  on the test set drawn from the original synthetic dataset.
Fig. 13, 14, and 15 display the average pixel values of the generated images with the prediction outcomes from the linear probing classifiers for the shape, color, and size attributes of these images, respectively. The figures in each subplot represent the ratio of accurately classified images out of a total of 50 samples. The predictions made by the linear classifiers align with our intuitive expectations, thereby confirming their suitability for evaluating the generated images.
Figure 13: Average pixel values of the generated images and prediction outcomes of the linear probing classifier for shape
Figure 14: Average pixel values of the generated images and prediction outcomes of the linear probing classifier for color
Figure 15: Average pixel values of the generated images and prediction outcomes of the linear probing classifier for size.
A.6 Experiment with real data
We sourced our data from the CelebA dataset<sup>3</sup>. Our training dataset consists of 15,000 facial images, each labeled with its respective concept class. We delineated the concept classes based on these attributes: gender (classified as "Male" or otherwise), facial expression ("Smiling" or not), and hair color (either "Black Hair" or "Brown Hair"). We adjusted the images to a resolution of  $48 \times 48$  pixels.
For our evaluation, we designed three classifiers for each concept using a three-layer CNN with 64 units. The learning rate was established at  $10^{-3}$ , trained the model for 10 epochs. We employed the ReLU activation function for the hidden layers and the sigmoid function for the output layer. We divided the image dataset into training and test sets at a ratio of  $80\%$  to  $20\%$ . The trained classifier exhibited a  $96\%$  accuracy for gender,  $86\%$  for facial expression, and  $95\%$  for hair color. All additional experimental configurations conformed to our synthetic data experiment procedures.
In Fig. 16, we generalize the findings in Fig. 5 to a larger concept graph, i.e., one with four concept variables. Specifically, we now introduce a fourth concept variable backgroundcolor = {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 at a concept distance of 1, continues to increase at a distance of 2, and peaks at 3.
In Fig. 17, we investigate the bottlenecks that affect compositional generalization by examining the prediction accuracy of concept variables such as shape, color, and size, as well as the overall classification accuracy. To eval
uate this, we generated 50 samples in each epoch of training and classified them using a linear probe. Our findings reveal that various concept bottlenecks impede the learning process for different compositional objects. For instance, the generalization of blue-colored objects (011, 101, 111) is hindered by the color concept, indicating that the model struggles to correctly identify objects based on their color when it comes to these particular compositions. Similarly, the generalization to small red triangles (111) is bottlenecked by the shape concept, implying that the model faces difficulties in accurately recognizing the shape of these specific compositions. These observations highlight the challenges associated with compositional generalization and shed light on the specific concept bottlenecks that hinder the learning process. Understanding and addressing these bottlenecks can contribute to the development of more robust models capable of achieving better compositional generalization in various contexts.
In Fig. 18, 19, and 20 we present the experimental results for the diffusion model without global attention. Notably, we observe that these results exhibit a similar pattern to the ones obtained from the diffusion model with global attention, which is shown in Fig. 5 and 6. The experimental results consistently support our findings. Firstly, the order of emergence of compositional capabilities aligns with the compositional structure of the data-generating process, as evident in the results shown in Fig. 18. Secondly, we observe a notable delay in the emergence of the ability to generate underrepresented colors as the concept distance within a concept class increases. This emphasizes the importance of continued training beyond the point of achieving in-distribution generalization, as demonstrated in Fig. 19. Specifically, we observe that the model is capable of generating the majority color (red) much earlier than the underrepresented color (blue). Importantly, the training
Figure 16: Compositional generalization. (a) Consider a lattice representation of a concept graph corresponding to four concept variables. Blue nodes denote classes represented in the training data; a model can either memorize data from these classes or learn capabilities to transform samples from one class to another. If it learns capabilities, it can compose them with data-generating process of samples from a concept class in the training data and produce samples that are entirely out-of-distribution, denoted as pink and green nodes. (b) We find compositional generalization happens in a sequence, starting with concept distance of 1 and then progressing to concept distance of 2 and 3.
<sup>3</sup>https://mmlab.ie.cuhk.edu.hk/projects/CelebA.html
Figure 17: Analyzing the bottlenecks of compositional generalization. We examine the prediction accuracy of concept variables (shape, color, and size) and the overall classification accuracy. In each epoch of training, 50 samples were generated and classified using linear probe. We observe that various concept bottlenecks hinder the learning of different compositional objects. For example, compositional generalization to blue-colored objects (011, 101, 111) is bottlenecked by the color concept, while generalization to small red triangles (111) is bottlenecked by the shape concept.
(c)
Figure 18: Concept distance from the training set govern the order in which compositional capabilities emerge. (a) Concept graph (cube) depicting training data points (blue nodes) and concept distances for test data points. (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 capability to change color for small, red triangles.
time to achieve generalization based on a concept increases with concept distance. This observation provides important insights for training models with fairness in mind: even once generalization for in-distribution concept classes is achieved, stopping the training of a model will likely lead to a failure in generating underrepresented concepts, particularly for compositional generalization. Thirdly, our findings reveal a multiplicative impact of learning individual concepts on performance in compositional tasks, offering an explanation for the sudden emergence of capabilities, as depicted in Fig. 20.
In Fig. 21, we conducted experiments with different configurations by using varied sets of nodes in the training data. In Fig. 21(a) and Fig. 21(b), we inverted the axes of the concept graph from Fig. 9(a): (a) from bottom to top and (b) from left to right. In Fig. 21(c) and Fig. 21(d), we inverted the axes of the concept graph from Fig. 5(a): (c) from bottom to top and (d) from left to right. In Fig. 21(b) and Fig. 21(c), we observe that the model incorrectly generates triangles for both 011 and 010, even though circles are expected. In the other cases, the model outputs correct images for all the nodes. Our observations indicate that the model struggles to dissociate a second element value of 1 with the shape of a small triangle, exhibiting a bias to generate triangles rather than circles for small objects.
Figure 19: Delayed emergence of abilities to generate underrepresented colors for distant classes. Prediction accuracy of color based on generated samples at each epoch during training. We observe that the ability to generate underrepresented colors emerges significantly later as the concept distance of a concept class increases, highlighting the need for extended training beyond the point of achieving in-distribution generalization. This prolonged training enables effective composition of the underrepresented concept and leads to improved generalization.
Figure 20: Multiplicity underlies the sudden emergence of compositional capabilities. (a) Accuracy of producing samples from the concept class {111, (triangle, blue, small)}, which has a concept distance of 2 from the training data. (b) Learning dynamics of accuracies for predicting each of the three concept variables: shape (blue), color (orange), and size (green).
Figure 21: The additional experimental results with different configurations by using varied sets of nodes in the training data.