SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Thaddäus Wiedemer<sup>1,2,3*</sup> Prasanna Mayilvahanan<sup>1,2,3*</sup>
Matthias Bethge $^{1,2\dagger}$  Wieland Brendel $^{2,3\dagger}$
$^{1}$ University of Tübingen  $^{2}$ Tübingen AI Center  $^{3}$ Max-Planck-Institute for Intelligent Systems, Tübingen
{thaddaeus.wiedemer, prasanna.mayilvahanan}@uni-tuebingen.de
Leveraging the compositional nature of our world to expedite learning and facilitate generalization is a hallmark of human perception. In machine learning, on the other hand, achieving compositional generalization has proven to be an elusive goal, even for models with explicit compositional priors. To get a better handle on compositional generalization, we here approach it from the bottom up: Inspired by identifiable representation learning, we investigate compositionality as a property of the data-generating process rather than the data itself. This reformulation enables us to derive mild conditions on only the support of the training distribution and the model architecture, which are sufficient for compositional generalization. We further demonstrate how our theoretical framework applies to real-world scenarios and validate our findings empirically. Our results set the stage for a principled theoretical study of compositional generalization.
Systematic compositionality [1] is the remarkable ability to utilize a finite set of known components to understand and generate a vast array of novel combinations. This ability, referred to by Chomsky [2] as the "infinite use of finite means", is a distinguishing feature of human cognition, enabling us to adapt to diverse situations and learn from varied experiences.
It's been a long-standing idea to leverage the compositional nature of the world for learning. In object-centric learning, models learn to isolate representations of individual objects as building blocks for complex scenes. In disentanglement, models aim to infer factors of variation that capture compositional and interpretable aspects of their inputs, for example hair color, skin color, and gender for facial data. So far, however, there is little evidence that these methods deliver substantially increased learning efficacy or generalization capabilities (Schott et al. [3], Montero et al. [4]). Across domains and modalities, machine learning models still largely fail to capture and utilize the compositional nature of the training data (Lake and Baroni [5], Loula et al. [6], Keysers et al. [7]).
To exemplify this failure, consider a model trained on a data set with images of two sprites with varying position, size, shape, and color overlaid on a black canvas. Given the latent factors, a simple multi-layer neural network can easily learn to reconstruct images containing compositions of these sprites that were covered by the training set. However, reconstruction fails for novel compositions—even if the individual components have been observed before (see Figure 1). Failure to generalize to
*Equal contribution † Equal supervision
Code available at https://github.com/brendel-group/compositional-ood-generalization
37th Conference on Neural Information Processing Systems (NeurIPS 2023), New Orleans.
Figure 1: Compositional generalization fails even in regression settings. Left: We train a model  $f$  to reconstruct images containing two sprites given their latent representation (x, y, shape, size, color). Center: In the training set (top row and left column), one sprite is fixed in its base configuration (orange triangle or red circle), while the other is varied randomly (in this example, sprite 1 varies in position, sprite 2 in shape and color). As a result, each sample in the test set (lower right block) can be expressed as a novel composition of known components. Right: While the model is able to fit the training data, it fails to generalize compositionally to the test data.
unseen data in even this simplistic regression setting demonstrates that compositional generalization does not automatically emerge simply because the data is of a compositional nature.
We therefore take a step back to formally study compositionality and understand what conditions need to be fulfilled for compositional generalization to occur. To this end, we take inspiration from identifiable representation learning and define a broad class of data generating processes that are compositional and for which we can prove that inference models can generalize to novel compositions that have not been part of the training set. More precisely, our contributions are as follows:
- We specify compositional data-generating processes both in terms of their function class and latent distributions (Sections 3.1 and 3.2) such that they cover a wide range of assumptions made by existing compositional methods.
- We prove a set of sufficient conditions under which models trained on the data are able to generalize compositionally (Section 3.3).
- We validate our theory in a range of synthetic experiments and perform several ablation studies that relate our findings to empirical methods (Section 4).
Representation learning Disentanglement and identifiable representation learning aim to learn succinct representations that both factorize the data space efficiently and are robust towards distributional changes [8-10]. However, the expectation that more compositional representations lead to better out-of-distribution (OOD) generalization has not been met, as demonstrated by Schott et al. [3] and Montero et al. [11]. Although our work does not directly address generalization issues in identifiable representation learning, our setup is directly inspired by it, and we examine data-generating processes similar to [12-14].
Empirical Approaches Many empirical methods use compositional priors and claim improved compositional generalization. The problem has been studied especially closely in language [15-17], but it remains far from being solved [5-7]. Object-centric learning is another domain in which compositionality plays a major role, and many approaches explicitly model the composition of scenes from object-"slots" [18-22]. The slot approach is also common in vector-symbolic architectures like [23] and [24]. For most of these works, however, compositional generalization
is not a focal point, and their actual generalization capability remains to be studied. There are also some architectures like transformers [25], graph neural networks [26], bilinear models [27], or complex-valued autoencoders [28] that have been claimed to exhibit some degree of compositional generalization, but again, principled analysis of their generalization ability is lacking. Our framework can guide the systematic evaluation of these methods. While we use the visual domain as an example throughout this work, our contributions are not tied to any specific data domain or modality.
Theoretical approaches to OOD generalization The OOD generalization problem for non-linear models where train and test distributions differ in their densities, but not their supports, has been studied extensively, most prominently by Ben-David and Urner [29] and Sugiyama et al. [30]. We refer the reader to Shen et al. [31] for a comprehensive overview. In contrast, compositional generalization requires generalizing to a distribution with different, possibly non-overlapping support. This problem is more challenging and remains unsolved. Ben-David et al. [32] were able to show that models can generalize between distributions with a very specific relation, but it is unclear what realistic distributions fit their constraints. Netyahu et al. [33] also study out-of-support problems theoretically but touch on compositional generalization only as a workaround for general extrapolation. Recently, Dong and Ma [34] took a first step towards a more applicable theory of compositional generalization to unseen domains, but their results still rely on specific distributions, and they do not consider functions with arbitrary (nonlinear) compositions or multi-variate outputs. In contrast, our framework is independent of the exact distributions used for training and testing, and our assumptions on the compositional nature of the data allow us to prove generalization in a much broader setting.
3 A framework for compositional generalization
Notation  $[N]$  denotes the set of natural numbers  $\{1,2,\dots,N\}$ . Vector-valued variables (e.g.,  $\mathbf{x}$ ) and functions (e.g.,  $\pmb{f}$ ) are written in bold. Id denotes the (vector-valued) identity function. We write the support of a distribution  $P$  as  $\operatorname{supp} P$ . To express that two functions  $\pmb{f}, \pmb{g}$  are equal for all points in the support of distribution  $P$ , i.e.,  $\pmb{f}(\pmb{x}) = \pmb{g}(\pmb{x}) \forall \pmb{x} \in \operatorname{supp} P$ , we write  $\pmb{f} \equiv_{P} \pmb{g}$ . Finally,  $\frac{\partial f}{\partial x}$  denotes the total derivative of a vector-valued function  $\pmb{f}$  by all its inputs  $\pmb{x}$ , corresponding to the Jacobian matrix with entries  $\frac{\partial f_i}{\partial x_j}$ .
3.1 Compositionality
Colloquially, the term "compositional data" implies that the data can be broken down into discrete, identifiable components that collectively form the whole. For instance, in natural images, these components might be objects, while in music, they might be individual instruments. As a running illustrative example, we will refer to a simple dataset similar to multi-dSprites [21], as shown in Figure 1. Each sample in this dataset is a composition of two basic sprites, each with a random position, shape, size, and color, size.
Drawing inspiration from identifiable representation learning, we define compositionality mathematically as a property of the data-generating process. In our example, the samples are generated by a simple rendering engine that initially renders each sprite individually on separate canvases. These canvases are then overlaid to produce a single image featuring two sprites. More specifically, the rendering engine uses the (latent) properties of sprite one,  $z_{1} = (z_{1,\mathrm{x}}, z_{1,\mathrm{y}}, z_{1,\mathrm{shape}}, z_{1,\mathrm{size}}, z_{1,\mathrm{color}})$ , to produce an image  $\tilde{x}_{1}$  of the first sprite. The same process is repeated with the properties of sprite two,  $z_{2} = (z_{2,\mathrm{x}}, z_{2,\mathrm{y}}, z_{2,\mathrm{shape}}, z_{2,\mathrm{size}}, z_{2,\mathrm{color}})$ , to create an image  $\tilde{x}_{2}$  of the second sprite. Lastly, the engine combines  $\tilde{x}_{1}$  and  $\tilde{x}_{2}$  to create the final overlaid rendering  $x$  of both sprites. Figure 2 demonstrates this process.
In this scenario, the individual sprite renderers carry out the bulk of the work. In contrast, the composition of the two intermediate sprite images  $\tilde{x}_1,\tilde{x}_2$  can be formulated as a simple pixelwise operation (see Appendix B.1 for more details). The rendering processes for each sprite are independent: adjusting the properties of one sprite will not influence the intermediate image of the other, and vice versa.
We posit that this two-step generative procedure—the (intricate) generation of individual components and their (simple) composition into a single output—is a key characteristic of a broad class of
Figure 2: Compositional representation of a function (Definition 1). The component functions  $\varphi_{k}$  map each component latent  $z_{k}$  to an intermediate component representation  $\tilde{x}_k$ . The composition function  $C$  composes all component representations into a final data point  $x$ .
Figure 3: Compositional support (Definition 2). A-D: Distribution  $P$  (blue) has compositional support w.r.t. to the entire latent space if it has full support over the margins. E: Gaps in the support require the model to interpolate/extrapolate rather than generalize compositionally. F: The support of the joint needs to be in an open set.
compositional problems. If we know the composition function, then understanding the basic elements (for example, the individual sprites) is enough to grasp all possible combinations of sprites in the dataset. We can thus represent any latent variable model  $f: \mathcal{Z} \to \mathcal{X}$ , which maps a latent vector  $z \in \mathcal{Z}$  to a sample  $x$  in the observation space  $\mathcal{X}$ , as a two-step generative process.
Definition 1 (Compositional representation).  $\{C, \varphi_1, \dots, \varphi_K, \mathcal{Z}_1, \dots, \mathcal{Z}_K, \tilde{\mathcal{X}}_1, \dots, \tilde{\mathcal{X}}_K\}$  is a compositional representation of function  $f$  if
$$
\forall \boldsymbol {z} \in \mathcal {Z} \quad \boldsymbol {f} (\boldsymbol {z}) = \boldsymbol {C} \left(\varphi_ {1} \left(\boldsymbol {z} _ {1}\right), \dots , \varphi_ {K} \left(\boldsymbol {z} _ {K}\right)\right) \quad \text {a n d} \quad \mathcal {Z} = \mathcal {Z} _ {1} \times \dots \times \mathcal {Z} _ {K}, \tag {1}
$$
where  $z_{i}$  denotes the canonical projection of  $z$  onto  $\mathcal{Z}_i$ . We refer to  $\varphi_k: \mathcal{Z}_k \to \tilde{\mathcal{X}}_k$  as the component functions, to  $\tilde{\mathcal{X}}_1, \ldots, \tilde{\mathcal{X}}_K$  as the (hidden) component spaces, and to  $C: \tilde{\mathcal{X}}_1 \times \dots \times \tilde{\mathcal{X}}_K \to \mathcal{X}$  as the composition function.
Note that in its most general form, we do not require the component functions to be identical or to map to the same component space. The compositional representation of a function  $f$  is also not unique. For instance, any  $f$  possesses a trivial compositional representation given by  $\{f, \mathbf{Id}, \dots, \mathbf{Id}\}$  (for the sake of clarity, we will omit the explicit mention of the latent factorization and component spaces henceforth). We will later establish conditions that must be met by at least one compositional representation of  $f$ .
Our definition of compositionality naturally aligns with various methods in the fields of identifiability, disentanglement, or object-centric learning. In the decoder of SlotAttention [18], for example, each component function is a spatial broadcast decoder followed by a CNN, and the composition function is implemented as alpha compositing. Frady et al. [24] model the component functions as element-wise multiplication of high-dimensional latent codes, which are then composed through a straightforward sum. A similar approach is chosen by Vankov and Bowers [23], except that interactions between components are modeled using matrix multiplication.
3.2 Compositional Generalization
The model in Figure 1 was trained supervisedly, i.e., it was trained to reconstruct samples  $\pmb{x}$  given the ground-truth latent factors  $(z_1, z_2)$  for each sprite (see Section 4 for more details). We denote this model as  $\hat{\pmb{f}}$ , indicating that it is meant to replicate the ground-truth generating process  $\pmb{f}$  of the data. The model  $\hat{\pmb{f}}$  indeed learned to fit  $\pmb{f}$  almost perfectly on the training distribution  $P$ , but failed to do so on the test distribution  $Q$ .
This failure is surprising because the test samples only contain sprites already encountered during training. The novelty lies solely in the combination of these sprites. We would expect any model that comprehends the compositional nature of the dataset to readily generalize to these test samples.
This compositional aspect of the generalization problem manifests itself in the structure of the training and test distribution. In our running example, the model was trained on samples from a distribution  $P$  that contained all possible sprites in each slot but only in combination with one base sprite in the other slot (illustrated in Figure 3A). More formally, the support of  $P$  can be written as
$$
\operatorname {s u p p} P = \left\{\left(\boldsymbol {z} _ {1} \in \mathcal {Z} _ {1}, \boldsymbol {z} _ {2} \in \mathcal {Z} _ {2}\right) | \boldsymbol {z} _ {1} = \boldsymbol {z} _ {1} ^ {0} \vee \boldsymbol {z} _ {2} = \boldsymbol {z} _ {2} ^ {0} \right\}, \tag {2}
$$
where  $z_{k}^{0}$  denotes the base configuration of a sprite (e.g., the orange triangle and red square in the samples shown in Figure 1).
The test distribution  $Q$  is a uniform distribution over the full product space  $\mathcal{Z}_1 \times \mathcal{Z}_2$ , i.e., it contains all possible sprite combinations. More generally, we say that a generalization problem is compositional if the test distribution contains only components that have been present in the training distribution, see Figure 3. This notion can be formalized as follows based on the support of the marginal distributions:
Definition 2 (Compositional support). Given two arbitrary distributions  $P, Q$  over latents  $z = (z_1, \ldots, z_K) \in \mathcal{Z} = \mathcal{Z}_1 \times \dots \times \mathcal{Z}_K$ ,  $P$  has compositional support w.r.t.  $Q$  if the support over all marginals  $P_{\mathbf{z}_k}, Q_{\mathbf{z}_k}$  is the same:
$$
\operatorname {s u p p} P _ {\boldsymbol {z} _ {k}} = \operatorname {s u p p} Q _ {\boldsymbol {z} _ {k}} \subseteq \mathcal {Z} _ {k} \quad \forall k \in [ K ]. \tag {3}
$$
Clearly, compositional generalization requires compositional support. If regions of the test latent space exist for which a component is not observed, as in Figure 3E, we can examine a model's generalization capability, but the problem is not compositional. Depending on whether the gap in the support is in the middle of a latent's domain or towards either end, the generalization problem becomes an interpolation or extrapolation problem instead, which are not the focus of this work.
3.3 Sufficient conditions for compositional generalization
With the above setup, we can now begin to examine under what conditions compositional generalization can be guaranteed to occur.
To make this question precise, let us assume for the moment that sprites don't occlude each other but that they are just summed up in pixel space. Then the compositional representation of the generative process is simply  $\{\sum (\cdot),\varphi_1,\varphi_2\}$ , i.e.
$$
\boldsymbol {f} (\boldsymbol {z}) = \varphi_ {1} \left(\boldsymbol {z} _ {1}\right) + \varphi_ {2} \left(\boldsymbol {z} _ {2}\right). \tag {4}
$$
The question becomes: Given training samples  $(z, x)$  from  $P$ , can we train a model  $\hat{f}$  that fitting this generative process  $\pmb{f}$  on  $P$  also guarantees fitting it on  $Q$ ? That is, we are looking for conditions such that the model generalizes from  $P$  to  $Q$ :
$$
\boldsymbol {f} \underset {P} {\equiv} \hat {\boldsymbol {f}} \Rightarrow \boldsymbol {f} \underset {Q} {\equiv} \hat {\boldsymbol {f}}. \tag {5}
$$
We assume that  $C$  is known, so in order to generalize, we must be able to reconstruct the individual component functions  $\varphi_{k}$ . For the simple case from equation 4, we can fully reconstruct the component functions as follows. First, we note that if  $\operatorname{supp} P$  is in an open set, we can locally reconstruct the hidden Jacobian of  $\varphi_{k}$  from the observable Jacobian of  $f$  as
$$
\frac {\partial \boldsymbol {f}}{\partial \boldsymbol {z} _ {k}} (\boldsymbol {z}) = \frac {\partial \varphi_ {k}}{\partial \boldsymbol {z} _ {k}} (\boldsymbol {z} _ {k}). \tag {6}
$$
Since the training distribution contains all possible component configurations  $z_{k}$ , we can reconstruct the Jacobian of  $\varphi_{k}$  in every point  $z_{k}$ . Then we know everything about  $\varphi_{k}$  up to a global offset (which can be removed if there exists a known initial point for integration).
Our goal is to extend this approach to a maximally large set of composition functions  $C$ . Our reasoning is straightforward if  $C$  is the identity, but what if we have occlusions or other nonlinear interactions between slots? What are general conditions on  $C$  and the support of the training distribution  $P$  such that we can still reconstruct the individual component functions and thus generalize compositionally?
Figure 4: Sufficient support condition (Definition 3). For a (compositional) diagonal support, all samples will contain sprites with similar positions, leading to heavy occlusions and making reconstruction of the background sprite impossible (left). Reconstruction of the background sprite is only possible if the support is chosen broad enough, such that the subset of points sharing the same background sprite  $P'$  contains samples with sufficient variance in the foreground sample. Specifically, each pixel of the background sprite must be observable at least once (right).
Let us now consider the sprites example with occlusions, where  $\varphi_{1}$  renders the background sprite that is occasionally occluded by the foreground sprite rendered by  $\varphi_{2}$ . Let us also assume that the support of  $P$  is basically a thin region around the diagonal; see Figure 4 (left). In this case, the two sprites are always relatively similar, leading to large overlaps for practically all samples of the training set. It is impossible to reconstruct the full Jacobian of the occluded sprite from a single sample. Instead, we need a set of samples for which the background sprite is the same while the foreground sprite is in different positions; see Figure 4 (right). With sufficient samples of this kind, we can observe all pixels of the background sprite at least once. Then reconstruction of the Jacobian of  $\varphi_{1}$  is possible again.
This line of thought brings us to a more general condition on the data-generating process: The composition function  $C$  and the support of the training set must be chosen such that the full Jacobian can be reconstructed for each component function and for all component latents. In other words, for each configuration of a given component,  $P$  must be sufficiently large so that it is possible to track how each dimension of the output depends on each dimension of the component representation. We formally define the concept of sufficient support below. Note that whether the support of  $P$  is sufficient or not depends on the choice of composition function  $C$ ; see Appendix C for examples.
Definition 3 (Sufficient support). A distribution  $P$  over latents  $\pmb{z} = (z_{1},\dots,z_{K})\in \mathcal{Z}$ , has sufficient support w.r.t. a compositional representation of a function  $\pmb{f}$ , if  $\operatorname{supp} P$  is in an open set and for any latent value  $\pmb{z}_{k}^{*}$ , there exists a finite set of points  $P^{\prime}(\pmb {z},k)\subseteq \{\pmb {p}\in \operatorname {supp}P|\pmb {p}_k = \pmb {z}_k\}$  for which the sum of total derivatives of  $\pmb{C}$  has full rank. That is,
$$
\operatorname {r a n k} \sum_ {\boldsymbol {p} \in P ^ {\prime} (\boldsymbol {z}, k)} \frac {\partial C}{\partial \varphi_ {k}} (\varphi (\boldsymbol {p})) = M, \tag {7}
$$
where  $M$  is the dimension of the component space  $\tilde{\mathcal{X}}_k\subseteq \mathbb{R}^M$
We are now ready to state our main theorem, namely that if  $f$ ,  $\hat{f}$  share the same composition function and if  $P$  has compositional and sufficient support, then the model  $\hat{f}$  generalizes to  $Q$  if it matches the ground-truth data-generating process  $f$  on  $P$ .
Theorem 1. Let  $P, Q$  be arbitrary distributions over latents  $\mathbf{z} = (\mathbf{z}_1, \dots, \mathbf{z}_K) \in \mathcal{Z}$ . Let  $\mathbf{f}, \hat{\mathbf{f}}$  be functions with compositional representations in the sense of definition 1 that share  $\{\mathbf{C}, \mathcal{Z}_1, \dots, \mathcal{Z}_K\}$ , but use arbitrary  $\{\varphi_1, \dots, \varphi_K, \tilde{\mathcal{X}}_1, \dots, \tilde{\mathcal{X}}_K\}$ ,  $\{\hat{\varphi}_1, \dots, \hat{\varphi}_K, \hat{\mathcal{X}}_1, \dots, \hat{\mathcal{X}}_K\}$ .
Assume the following assumptions hold:
(A1)  $C, \varphi_k, \hat{\varphi}_k$  are differentiable,  $C$  is Lipschitz in  $\varphi$ , and  $\varphi$  is continuous in  $z$ .
(A2)  $P$  has compositional support w.r.t.  $Q$  in the sense of definition 2.
(A3)  $P$  has sufficient support w.r.t.  $\pmb{f}$  in the sense of definition 3.
(A4) There exists an initial point  $\pmb{p}^0\in \operatorname {supp}P$  such that  $\varphi (\pmb{p}^{0}) = \hat{\varphi} (\pmb{p}^{0})$
Then  $\hat{f}$  generalizes to  $Q$ , i.e.  $f \equiv_{P} \hat{f} \Rightarrow f \equiv_{Q} \hat{f}$ .