SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Yi Ren
University of British Columbia*
renyi.joshua@gmail.com
Samuel Lavoie
Université de Montréal & Mila
samuel.lavoie.m@gmail.com
Mikhail Galkin
Intel AI Lab
mikhail.galkin@intel.com
Danica J. Sutherland
University of British Columbia & Amii
dsuth@cs.ubc.ca
aaron.courville@gmail.com
Compositional generalization, the ability of an agent to generalize to unseen combinations of latent factors, is easy for humans but hard for deep neural networks. A line of research in cognitive science has hypothesized a process, "iterated learning," to help explain how human language developed this ability; the theory rests on simultaneous pressures towards compressibility (when an ignorant agent learns from an informed one) and expressivity (when it uses the representation for downstream tasks). Inspired by this process, we propose to improve the compositional generalization of deep networks by using iterated learning on models with simplicial embeddings, which can approximately discretize representations. This approach is further motivated by an analysis of compositionality based on Kolmogorov complexity. We show that this combination of changes improves compositional generalization over other approaches, demonstrating these improvements both on vision tasks with well-understood latent factors and on real molecular graph prediction tasks where the latent structure is unknown.
Deep neural networks have shown an amazing ability to generalize to new samples on domains where they have been extensively trained, approaching or surpassing human performance on tasks including image classification [62], Go [70], reading comprehension [13], and more. A growing body of literature, however, demonstrates that some tasks that can be easily solved by a human can be hard for deep models. One important such problem is compositional generalization ([18], comp-gen for short). For example, Schott et al. [65] study manually-created vision datasets where the true generating factors are known, and demonstrate that a wide variety of current representation learning methods struggle to learn the underlying mechanism. To achieve true "artificially intelligent" methods that can succeed at a variety of difficult tasks, it seems necessary to demonstrate compositional generalization. One contribution of this paper is to lay out a framework towards understanding and improving compositional generalization, and argue that most currently-common training methods fall short.
*Work done in part during an internship at Mila.
In wondering how deep networks can learn to compositionally generalize, we might naturally ask: how did humans achieve such generalization? Or, as a particular case, how did human languages evolve components (typically, words) that can systematically combine to form new concepts? This has been a long-standing question in cognitive science and evolutionary linguistics. One promising hypothesis is known as iterated learning (IL), a procedure simulating cultural language evolution [41]. Aspects of this proposal are supported by lab experiments [42], a Bayesian model [7], the behavior of neural networks in a simple emergent communication task [60], and real tasks like machine translation [50] and visual question answering [76].
To link the study in cognitive science and deep learning, we first analyze the necessary properties of representations in order to generalize well compositionally. By linking the compositionality and the Kolmogorov complexity, we find iteratively resetting and relearning the representations can introduce compressibility pressure to the representations, which is also the key to the success of iterated learning. To apply iterated learning in a general representation learning problem, we propose to split the network into a backbone and a task head, and discretize the representation at the end of the backbone using simplicial embeddings (SEM, [45]). This scheme is more practical than LSTM [34] encoders previously used for neural iterated learning [60]. We observe in various controlled vision domains that SEM-IL can enhance compositional generalization by aligning learned representations to ground-truth generating factors. The proposed method also enhances downstream performance on molecular graph property prediction tasks, where the generating process is less clear-cut.
2 Compositional Generalization
Generalization is a long-standing topic in machine learning. The traditional notion of (in-distribution) generalization assumes that training and test samples come from the same distribution, but this is insufficient for many tasks: we expect a well-trained model to generalize to some novel scenarios that are unseen during training. One version of this is compositional generalization (comp-gen) [17], which requires the model to perform well on novel combinations of semantic concepts.
2.1 Data-generating assumption and problem definition
Any type of generalization requires some "shared rules" between training and test distributions. We hence assume a simple data-generating process that both training and test data samples obey. In Figure 1, the semantic generating factors, also known as latent variables, are divided into two groups: the task-relevant factors (or semantic generating factors)  $\mathbf{G} = [G_1,\dots,G_m]$ , and task-irrelevant (or noise) factors  $\mathbf{O}$ . This division depends on our understanding of the task; for example, if we only want to predict the digit identity of an image in the color-MNIST dataset [3], then  $m = 1$  and  $G_{1}$  represents the digit identity. All the other generating factors such as color, stroke, angle, and possible noise are merged into  $\mathbf{O}$ . If we want to predict a function that depends on both identity and color, e.g. identifying blue even numbers, we could have  $\mathbf{G} = [G_1,G_2]$  with  $G_{1}$  the identity and  $G_{2}$  the color.
Each input sample  $\mathbf{x} \in \mathcal{X}$  is determined by a deterministic function  $\operatorname{GenX}(\mathbf{G}, \mathbf{O})$ . The task label(s)  $\mathbf{y} \in \mathcal{Y}$  only depend on the factors  $\mathbf{G}$  and possible independent noise  $\epsilon$ , according to the deterministic function  $\operatorname{GenY}(\mathbf{G}, \epsilon)$ . Note  $(\mathbf{x}, \mathbf{O}) \perp (\mathbf{y}, \epsilon) \mid \mathbf{G}$ , and that  $\mathbf{O}, \mathbf{G}$ , and  $\epsilon$  are independent. The data-generating distribution  $P(\mathbf{x}, \mathbf{y})$  is determined by the latent distributions  $P(\mathbf{G})$  and  $P(\mathbf{O})$ , along with the GenX and GenY. We assume GenX and GenY are fixed across environments (the "rules of production" are consistent), while  $P(\mathbf{G})$  and  $P(\mathbf{O})$  might change between training and test.
For compositional generalization, we wish to model the problem of generalizing to new combinations of previously seen attributes: understanding "red circle" based on having seen "red square" and "blue circle." Thus, we may assume that the supports of  $P(\mathbf{G})$  are non-overlapping between train and test. (If this assumption is not true, it only makes the problem easier.) In summary, our goal is to find an algorithm  $\mathcal{A}$  such that, when trained on a dataset  $\mathcal{D}_{train} \sim P_{train}^n$ ,  $\mathcal{A}$  achieves small test risk  $\mathcal{R}_{P_{test}}(\mathcal{A}(\mathcal{D}_{train}))$ . Here  $P_{train}$  and  $P_{test}$  should satisfy these conditions:
-  $P_{train}$  and  $P_{test}$  have  $\mathbf{G},\mathbf{O},\epsilon$  jointly independent, and  $\mathbf{x} = \mathrm{GenX}(\mathbf{G},\mathbf{O}),\mathbf{y} = \mathrm{GenY}(\mathbf{G},\epsilon)$ .
- GenX and GenY are the same deterministic functions for  $P_{train}$  and  $P_{test}$ .
- In challenging cases, we may have  $\operatorname{supp}[P_{train}(\mathbf{G})] \cap \operatorname{supp}[P_{test}(\mathbf{G})] = \emptyset$ .
2This differs from the classical setting of covariate shift:  $P(\mathbf{y}\mid \mathbf{x})$  might change due to the shift in  $P(\mathbf{G})$
(a). Data generating assumption
(b). Ladder of compositionality
Figure 1: Left: the data-generating assumption and a typical representation learning method (in red). We use the model  $(g\circ h)(\mathbf{x})$  for downstream predictions. Right: the ladder of compositionality stating the requirements of  $\mathbf{z}$  using the entropy-related measurements; see Appendix A for more.
2.2 Representation Learning and Ladder of Compositionality
For compositional generalization, we expect that the model must extract atomic semantic features from the training data, and systematically recombine them in a procedure akin to how the data is generated [41]. We thus consider a typical representation learning framework, which resembles the inverse of the data generation process (Figure 1(a), bottom). We use a backbone  $h: \mathcal{X} \to \mathcal{Z}$  to convert the input signal  $\mathbf{x}$  into a representation  $\mathbf{z}$ , and a task head  $g: \mathcal{Z} \to \mathcal{Y}$  to solve the given task based on that representation  $\mathbf{z}$ . The prediction of the model is  $\hat{\mathbf{y}} = (g \circ h)(\mathbf{x})$ .
Intuitively, we would like our learned  $\mathbf{z}$  to uncover the hidden  $\mathbf{G}$ , and  $g(\mathbf{z})$  to recover  $\mathrm{GenY}(\mathbf{G},\epsilon)$ . We thus analyze how the relationship between  $\mathbf{z}$  and  $\mathbf{G}$  influences the model's generalization capability, building off principles such as information bottleneck [74]. Inspired by the "ladder of causation" [55], we propose a "ladder of compositionality" in Figure 1(b), which outlining a series of conditions on  $\mathbf{z}$  and  $\mathbf{G}$ . We hypothesize that comp-gen roughly requires reaching the highest rung of that ladder:
Hypothesis 1. To generalize compositionally, the learned  $\mathbf{z}$  should capture exactly the information in  $\mathbf{G}$  and nothing more (  $\mathbf{G}$  to  $\mathbf{z}$  should be a bijection), and moreover it should preserve the "structure" of  $\mathbf{G}$  (i.e. the mapping from  $\mathbf{G}$  to  $\mathbf{z}$  should be an isomorphism).
More on this hypothesis, the ladder, and relationship to models of disentanglement [32] are discussed in Appendix A. In short, we find that a model trained using common learning methods relying on mutual information between input  $\mathbf{x}$  and supervision  $\mathbf{y}$  cannot reliably reach the final stage of the ladder – it is necessary to seek other inductive biases in order to generalize compositionally.
3 Compressibility pressure and Compositional mapping
From the analysis above, we need to find other inductive biases to obtain compositional mappings. Inspired by how compositionality emerges in human language, $^3$  we speculate that the compressibility pressure is the key. Note that this pressure does not refer to compressing information from  $\mathbf{x}$  to  $\mathbf{z}$  (as in Stage III does), but whether a mapping can be expressed in a compact way by reusing common rules. In this section, we will first link compressibility pressure to Kolmogorov complexity by defining different mappings using group theory. As the Kolmogorov complexity is hard to compute, making explicit regularization difficult, we propose to implicitly regularize via iterated learning, a procedure in cognitive science proposed to increase compositionality in human-like language.
3.1 Compositional mappings have lower Kolmogorov complexity
From Occam's razor, we know efficient and effective mappings are more likely to capture the ground truth generating mechanism of the data, and hence generalize better. The efficiency is determined by how compressed the mapping is, which can also be measured by Kolmogorov complexity [47, 71].
<sup>3</sup>Human languages are examples of compositional mapping [35]: words are composed of combinations of reusable morphemes, and those words in turn are combined to form complex sentences following specific stable rules. These properties make our language unique among natural communication systems and enable humans to convey an open-ended set of messages in a compositional way [42]. Researchers in cognitive science and evolutionary linguistics have proposed many explanations for the origin of this property; one persuasive method for simulating it is iterated learning [41].
To build a link between compositionality and Kolmogorov complexity, we can first describe different bijections between  $\mathbf{z}$  and  $\mathbf{G}$  using group theory, and then use the description length to compare the complexity of a typical element. Specifically, assuming  $\mathbf{z} \in \mathbb{Z}, \mathbf{G} \in \mathcal{G}$  and  $|\mathcal{Z}| = |\mathcal{G}|$ , the space of all bijections between  $\mathbf{z}$  and  $\mathbf{G}$  is an isomorphism of a symmetric group  $S_{|\mathcal{G}|}$ . If  $\mathbf{G} = [G_1, \dots, G_m]$  and each  $G_m$  has  $v$  different possible values,  $|\mathcal{G}| = v^m$ . For clarity in the analysis, we assume  $\mathbf{z}$  also has the same shape. Then, any bijection between  $\mathbf{z}$  and  $\mathbf{G}$  can be represented by an element in  $S_{v^m}$ .
The space of compositional mapping, which is a subset of all bijections, has more constraints. Recall how a compositional mapping is generated (see Appendix A.4 for more details): we first select  $z_{i}$  for each  $G_{j}$  in a non-overlapping way. Such a process can be represented by an element in  $S_{m}$ . After that, we will assign different "words" for each  $z_{i}$ , which can be represented by an element in  $S_{v}$ . As we have  $m$  different  $z_{i}$ , this procedure will be repeated  $m$  times. In summary, any compositional mapping can be represented by an element in the group  $S_{v}^{m} \rtimes S_{v} \in S_{v^{m}}$ , where  $\rtimes$  is the semidirect product in group theory. The cardinality of  $S_{v^{m}}$  is significantly larger than  $S_{v}^{m} \rtimes S_{v}$ , and so a randomly selected bijection is unlikely to be compositional. Thus
Proposition 1 (Informal). For  $m, v \geq 2$ , among all bijections, any compositional mapping has much lower Kolmogorov complexity than a typical non-compositional mapping.
We prove this by constructing descriptive protocols for each bijection. As a compositional mapping has more reused rules, its description length can be smaller (see Appendix B.1 for more details).
3.2 Compressibility pressure is amplified in iterated learning
Now, our target is finding bijections with higher compositionality and lower Kolmogorov complexity, which are both non-trivial. Because the ground truth  $\mathbf{G}$  is usually inaccessible and the Kolmogorov complexity is hard to calculate. Fortunately, researchers find that human language also evolved to become more compositional without knowing  $\mathbf{G}$ . Authors of [42] hypothesize that the compressibility pressure, which exists when an innocent agent (e.g., a child) learns from an informed agent (e.g., an adult), plays an important role. Such pressure is reinforced and amplified when the human community repeats this learning fashion for multiple generations.
However, the aforementioned hypothesis assumes that simplicity bias is inborn in the human cognition system. Will deep neural agents also have similar preferences during training? The answer is yes. By analyzing an overparameterized model on a simple supervised learning problem, we can strictly prove that repeatedly introducing new agents to learn from the old agent (then this informed agent becomes the old agent for the next generation) can exert a non-trivial regularizing effect on the number of "active bases" of the learned mapping. Restricting the number of active bases encourages the model to reuse the learned rules. In other words, this regularization effect favors mappings with lower Kolmogorov complexity, which is exactly what we expect for compositional generalization. Due to the space limits, we left the formulation and proof of this problem in Appendix B.2.
3.3 Complete the proposed solution
We thus expect that iteratively resetting and relearning can amplify the compressibility pressure, which helps us to reach the final rung of the ladder from the third. Before that, we need another pressure to reach third rung (i.e., ensure a bijection between  $\mathbf{z}$  and  $\mathbf{G}$ ). Expressivity pressure, constraining the learned mapping to be capable enough to accomplish the downstream tasks, is what we need.
The complete iterated learning hypothesis of Kirby et al. [42] claims that the compositional mapping emerges under the interaction between the compressibility pressure (i.e., efficiency) and the expressivity pressure (i.e., effectiveness). Inspired by this, we propose to train a model in generations consisting of two phases. At the  $t$ -th generation, we first train the backbone  $h$  in an imitation phase, where a student  $h_t^S$  learns to imitate  $\mathbf{z}$  sampled from a teacher  $h_t^T$ . As analyzed above, iteratively doing so will amplify the compressibility pressure. Then, in the following interaction phase, the model  $g_t \circ h_t^S$  follows standard downstream training to predict  $\mathbf{y}$ . The task head  $g_t$  is randomly initialized and fine-tuned together with the backbone in this phase. By accomplishing this phase, the expressivity pressure is introduced. The fine-tuned backbone  $h_t^S$  then becomes the teacher  $h_{t+1}^T$  for the next generation, and we repeat, as illustrated in Figure 2 and Algorithm 1.
Another problem with applying iterated learning to deep neural networks is how to create the discrete message, i.e.,  $\mathbf{z}$ . Discretization is not necessary: for example, the imitation phase could use  $L_{2}$  loss to
Figure 2: An illustration of iterated learning and SEM layer design.
match a student's continuous representations to the teacher's. We find greatly improved performance with our discretization scheme, however, due to much-increased compressibility pressure. It is also possible [60] to use e.g. an LSTM encoder at the end of  $h(\mathbf{x})$  to produce discrete  $\mathbf{z}$ , and an LSTM decoder at the start of  $g(\mathbf{z})$ . The interaction phase is then not directly differentiable; though many estimator options exist [6, 39, 78], training tends to be difficult due to high bias and/or variance.
Instead, we consider a simplicial embedding layer (SEM, [45]), which has proven effective on many self-supervised learning tasks. As illustrated in Figure 2(c), a dense representation  $\mathbf{h}$  (the output of the original backbone) is linearly transformed into  $m$  vectors  $z_{i} \in \mathbb{R}^{v}$ . Then we apply a separate softmax with temperature  $\tau$  to each  $z_{i}$ , yielding  $\bar{z}_{i}$  which are, if the temperature is not too high, approximately sparse; the  $\bar{z}_{i}$  are then concatenated to a long vector  $\mathbf{z}$ . The overall process is
$$
\bar {z} _ {i} = \operatorname {S o f t m a x} _ {\tau} \left(z _ {i}\right) = \left[ \frac {e ^ {z _ {i j} / \tau}}{\sum_ {k = 1} ^ {V} e ^ {z _ {i k} / \tau}} \right] _ {j} \in \mathbb {R} ^ {v} \quad \mathbf {z} = \left[ \bar {z} _ {1} ^ {\top} \quad \dots \quad \bar {z} _ {m} ^ {\top} \right] ^ {\top} \in \mathbb {R} ^ {m v}. \tag {1}
$$
By using an encoder with a final SEM layer, we obtain an approximately-sparse  $\mathbf{z}$ . In the imitation phase, we generate discrete pseudo-labels by sampling from the categorical distribution defined by each  $\bar{z}_i$ , then use cross-entropy loss so that the student is effectively doing multi-label classification to reconstruct the teacher's representations. In the imitation phase, the task head  $g$  operates directly on the long vector  $\mathbf{z}$ . The full model  $g \circ h$  is differentiable, so we can use any standard task loss. Pseudocode for the proposed method, SEM-IL, is in the appendix (Algorithm 1).
4 Analysis on Controlled Vision Datasets
We will first verify the effectiveness of the proposed SEM-IL method on controlled vision datasets, where the ground truth  $\mathbf{G}$  is accessible. Thus, we can directly observe how  $\mathbf{z}$  gradually becomes more similar to  $\mathbf{G}$ , and how the compressibility and expressivity pressures affect the training process. In this section, we consider a regression task on 3dShapes [9], where recovering and recombining the generating factors is necessary for systematic generalization. The detailed experimental settings and results on additional similar datasets, dSprites [52] and MPI3D-real [23], are given in Appendix C.
4.1 The Effectiveness of SEM-IL
Better comp-gen performance We first show the effectiveness of the proposed method using results on 3dShapes, containing images of objects with various colors, sizes, and orientations against various backgrounds. Here  $\mathbf{G}$  numerically encodes floor hue, wall hue, object hue, and object scale into discrete values, and the goal is to recover a particular linear function of that  $\mathbf{G}$ . (Results for a simple nonlinear function were comparable.)
We compare five algorithms:
- Baseline: directly train a ResNet18 [31] on the downstream task.
- SEM-only: insert an SEM layer to the baseline model.
- IL-only: train a baseline model with Algorithm 1, using MSE loss during imitation.
- SEM-IL: train an SEM model with Algorithm 1.
- Given-G: train an SEM model to reproduce the true  $\mathbf{G}$  (which would not be known in practice), then fine-tune on the downstream task.
Figure 3: Left: compositional generalization performance on a regression task. Right: topological similarity for IL and non-IL methods. Note the values of  $\rho$  in the two panels are not comparable, as the structure of  $\mathbf{z}$  in the two settings (with or without SEM) is different.
In the first panel of Figure 3, we see that the baseline and SEM-only models perform similarly on the training set; IL-based methods periodically increase in error at the beginning of each generation, but are eventually only slightly worse than the baselines on training data. On the test set, however, evaluating compositional generalization by using values of  $\mathbf{G}$  which did not appear in training, SEM-IL brings significant improvement compared with other methods. Using only SEM or only IL gives no improvement over the baseline, however; it is only their combination which helps, as we will discuss further shortly. The (unrealistic) oracle method Given-G is unsurprisingly the best, since having  $\mathbf{z}$  similar to  $\mathbf{G}$  is indeed helpful for this task.
How  $\mathbf{z}$  evolves during learning To see if better generalization ability is indeed achieved by finding  $\mathbf{z}$  that resembles the structure of  $\mathbf{G}$ , we check their topological similarity<sup>4</sup>
$$
\rho (\mathbf {z}, \mathbf {G}) \triangleq \operatorname {C o r r} \left(d _ {z} \left(\mathbf {z} ^ {(i)}, \mathbf {z} ^ {(j)}\right), d _ {G} \left(\mathbf {G} ^ {(i)}, \mathbf {G} ^ {(j)}\right)\right) \tag {2}
$$
where  $d_{z}$  and  $d_{G}$  are distance metrics,  $\mathbf{z}^{(i)}$  is the predicted representation of  $\mathbf{x}^{(i)}$ , and  $\mathbf{G}^{(i)}$  is the corresponding ground-truth generating factors. This measurement is widely applied to evaluate the compositionality of the mappings in cognitive science [8] and emergent communication [60]. Following existing works, we use the Hamming distance for  $\mathbf{G}$  and discretized  $\mathbf{z}$  in SEM-based methods, and cosine distance for continuous  $\mathbf{z}$  in non-SEM methods. We expect  $h(\mathbf{x})$  to map  $\mathbf{x}$  with similar  $\mathbf{G}$  to close  $\mathbf{z}$ , and dissimilar  $\mathbf{x}$  to distant  $\mathbf{z}$ , so that  $\rho (\mathbf{z},\mathbf{G})$  will be high.
The third panel of Figure 3 shows that the SEM-only model quickly reaches a plateau after 200 epochs and then slowly decreases, while SEM-IL, after briefly stalling at the same point, continues to increase to a notably higher topological similarity. In the last panel, however, the IL-only method doesn't improve  $\rho$  over the baseline: it seems both parts are needed.
4.2 Discretized Representation is Beneficial for the Imitation Phase of IL