SlowGuess/NeurIPS_2023_OCR huggingface dataset:
The proof follows roughly the intuition we developed above in that we show that the Jacobians of the component functions can be reconstructed everywhere. Bear in mind that this is simply a construction for the proof: The theorem holds whenever  $\hat{\pmb{f}}$  fits the output of  $\pmb{f}$  on the training distribution  $P$ , which we can achieve with standard supervised training and without access to the ground-truth Jacobians. It should also be emphasized that since the compositional representation is not unique, the theorem holds if there exists at least one for which the assumptions are fulfilled. Note also that the initial point condition (A4) is needed in the proof, but in all practical experiments (see below), we can generalize compositionally without explicit knowledge of that point. We relegate further details to Appendix A.
We validate our theoretical framework on the multi-sprite data. All models were trained for 2000 epochs on training sets of 100k samples using an NVIDIA RTX 2080 Ti; all test sets contain 10k samples. Table 1 summarizes the reconstruction quality achieved on the in-domain (ID) test set  $(P)$  and the entire latent space  $(Q)$  for all experiments.
Motivating experiment We implement the setup from Figure 1 to demonstrate that a compositional model does indeed generalize if the conditions from Theorem 1 are met. We model the component functions as four fully-connected layers followed by four upsampling-convolution stages, mapping the 5d component latent to  $64 \times 64$  RGB images. For training stability, the composition function is implemented as a soft pixel-wise addition using the sigmoid function  $\sigma(\cdot)$  as
$$
\boldsymbol {x} = \sigma \left(\tilde {\boldsymbol {x}} _ {1}\right) \cdot \tilde {\boldsymbol {x}} _ {1} + \sigma \left(- \tilde {\boldsymbol {x}} _ {1}\right) \cdot \tilde {\boldsymbol {x}} _ {2}, \tag {8}
$$
which allows component 1 to occlude component 2. We contrast this to a non-compositional monolithic model, which has the same architecture as a single component function (with adjusted layer sizes to match the overall parameter count of the compositional model). Both models are trained on samples  $(z, x)$  from the training set using an MSE reconstruction loss. We show that both models have the capacity to fit the data by training on random samples covering the entire latent space (Table 1, #1,2). We then train on a distribution with orthogonal support as in equation 2, albeit with two planes for the foreground component to satisfy the sufficient support condition (Definition 3) as explained in Figure 4. Both models can reconstruct ID samples, but only the compositional model generalizes to the entire latent space (Table 1, #3,4).
Flexible compositional support Next, we demonstrate the variety of settings that fulfill the compositional support assumption as illustrated in Figure 3B and C. To this end, we repeat the experiment on training sets  $P$  sampled from (i) a normal distribution with orthogonal support (Table 1, #5) and (ii) a uniform distribution over a diagonal support chosen broad enough to satisfy the sufficient support condition (Table 1, #6; see also Appendix C for details on how the support was chosen). The model generalizes to the entire latent space in both settings. Since the generalization performance is already close to the performance ceiling, broadening the support of both distributions (Table 1, #7,8) does not further increase performance.
Violating Conditions Finally, we look at the effect of violating some conditions.
- Gaps in support (Table 1, #9) Gaps in the support of the training set such that some component configurations are never observed (Figure 3E) violate the compositional support condition (Definition 2). While the overall reconstruction performance only drops slightly, visualizing the reconstruction error over a 2d-slice of the latent space in Figure 5 illustrates clearly that generalization fails exactly where the condition is violated.
- Insufficient training variability (Table 1, #10) Reducing the width of the diagonal support violates the sufficient support condition (Definition 3) as soon as some parts of the background component are always occluded and can not be observed in the output anymore (Compare Appendix C for details). We can clearly see that reconstruction performance on the entire latent space drops significantly as a result.
- Collapsed Composition Function (Table 1, #11) Changing the output of each component function from RGB to RGBa and implementing the composition as alpha compositing yields a model that is still compositional, but for which no support can satisfy the sufficient support condition since
#	Train Set	Model	R2ID↑	R2all ↑	ΔR2↓
1	Random	Monolithic	0.931±5.8e-4	0.931±5.8e-4	0.000±8.2e-4
2	Random	Compositional	0.957±1.0e-3	0.957±1.0e-3	0.000±1.4e-3
3	Orthogonal	Monolithic	0.948±1.7e-3	-0.500±6.7e-2	1.448±6.7e-2
4	Orthogonal	Compositional	0.957±6.4e-4	0.951±1.4e-3	0.006±1.5e-3
5	Ortho. ~ N	Compositional	0.957±5.5e-4	0.951±1.0e-3	0.006±1.1e-3
6	Diagonal	Compositional	0.954±5.4e-3	0.945±1.6e-2	0.009±1.7e-2
7	Ortho. (broad)	Compositional	0.959±9.7e-4	0.954±1.3e-3	0.005±1.6e-3
8	Diag. (broad)	Compositional	0.957±1.2e-3	0.955±1.2e-3	0.002±1.7e-3
9	Ortho. (gap)	Compositional	0.954±1.7e-3	0.895±7.4e-3	0.059±7.6e-3
10	Diag. (narrow)	Compositional	0.867±7.5e-2	0.589±2.8e-1	0.278±2.9e-1
11	Orthogonal	Comp. (RGBa)	0.984±1.8e-4	0.979±1.4e-4	0.005±2.3e-4
Table 1: We report the reconstruction quality measured as variance-weighted  $R^2$  score (closer to 1 is better) on the in-domain (ID) test set and the entire latent space. As the ID region occupies a tiny fraction of the entire latent space, the difference in performance  $(\Delta R^2)$  indicates how well a model generalizes OOD. All results are averaged over 5 random seeds.#1-4 The results demonstrate that a monolithic model fails to generalize in the setup from Figure 1, but a compositional model performs well on the entire latent space. #5-8 Generalization can occur in a variety of settings that fulfill the sufficient conditions from Theorem 1. #9,10 Violating the compositional and sufficient support condition prohibits generalization while choosing a more complex function class still works (#11). Table 2 in the appendix additionally reports the MSE for all experiments.
Figure 5: Heatmap of the reconstruction error over a  $z_{1,\mathrm{x}} - z_{2,\mathrm{x}}$ -projection of the latent space with overlaid training support (red). Generalization can occur when the support is compositional (left) but fails exactly where the support is incomplete at  $z_{1,\mathrm{x}} \in [0.14, 0.46]$  (right).
the derivative of transparent pixels will always be zero and the Jacobian matrix can therefore never have full rank (see Appendix B.1 for more details). However, we observe that the model still generalizes to the entire latent space and achieves even lower reconstruction error than the original model. This emphasizes that what we present are merely sufficient conditions. We include this experiment to motivate future work to find weaker conditions.
We presented a first step and a framework to study compositional generalization in a more principled way. Clearly, there remain many open questions and limitations that we leave for future work.
Supervised setting We only studied a supervised regression setting in which the model has access to the ground-truth latents of each training sample. Or own findings underlined the results of previous works, e.g., Schott et al. [3] that compositional generalization is not trivial even in this setting. Ultimately, we are of course interested in the unsupervised setting akin to what is typically studied in identifiable representation learning. The unsupervised setting comes with inherent ambiguities as the
relationship between ground-truth latent space and inferred representations is unknown. No prior works exist that address this identifiability problem when training on a subset  $P$  of the latent space, which makes generalizations guarantees as presented in this work harder to derive. Still, the results in this paper build an important foundation for future studies as sufficient conditions in the supervised setting can be considered necessary conditions in the unsupervised setting.
Jacobian and initial point The proof of Theorem 1 utilizes the Jacobian of the ground-truth model. We emphasize again that this construction is necessary only for the proof and does not mean that we require access to the data-generating processes' full Jacobian for training. Similarly, the existence of an initial point  $p^0$  is a technicality of the proof that is not reflected in the experiments. While it is not yet clear whether it is possible to complete the proof without the initial point condition, we believe there is a self-consistency condition that might alleviate the need for this condition. The experiments thus hint at the existence of alternative proof strategies with relaxed assumptions.
Known composition function We also assume the composition function to be known which is approximately true in many interesting scenarios, such as object composition in scenes or the composition of instruments in music. In fact, many structured representation learning approaches like e.g. SlotAttention [18] incorporate structural components that are meant to mimic the compositional nature of the ground-truth-generating process. In other interesting cases like language, however, the composition function is unknown a priori and needs to be learned. This might be possible by observing how the gradients of  $C$  change with respect to a fixed slot, at least if certain regularity conditions are fulfilled.
Inductive biases Some of the conditions we derived can be relaxed in the presence of certain inductive biases. For example, models with an inductive bias towards shift invariance might be able to cope with certain gaps in the training support (e.g., if sprites are not visible in every position). Similarly, assuming all component functions  $\varphi$  to be identical would substantially simplify the problem and allow for much smaller sufficient supports  $P$ . The conditions we derived do not assume any inductive bias but are meant to formally guarantee compositional generalization. We expect that our conditions generalize to more realistic conditions as long as the core aspects are fulfilled.
Error bounds Our generalization results hold only if the learned model perfectly matches the ground-truth model on the training distribution. This is similar to identifiable representation learning, where a model must find the global minimum of a certain loss or reconstruction error for the theory to hold. Nonetheless, extending our results towards generalization errors that are bounded by the error on the training distribution is an important avenue for future work.
Broader impact Compositional generalization, once achieved, has the potential to be benefit many downstream applications. By increasing sample and training efficiency, it could help to democratize the development and research of large-scale models. Better generalization capabilities could also increase the reliability and robustness of models but may amplify existing biases and inequalities in the data by generalizing them and hinder our ability to interpret and certify a model's decisions.
Machine learning, despite all recent breakthroughs, still struggles with generalization. Taking advantage of the basic building blocks that compose our visual world and our languages remains unique to human cognition. We believe that progress towards more generalizable machine learning is hampered by a lack of a formal understanding of how generalization can occur. This paper focuses on compositional generalization and provides a precise mathematical framework to study it. We derive a set of sufficient conditions under which compositional generalization can occur and which cover a wide range of existing approaches. We see this work as a stepping stone towards identifiable representation learning techniques that can provably infer and leverage the compositional structure of the data. It is certainly still a long road toward scalable empirical learning techniques that can fully leverage the compositional nature of our world. However, once achieved, there is an opportunity for drastically more sample-efficient, robust, and human-aligned machine learning models.
We would like to thank (in alphabetical order): Jack Brady, Simon Buchholz, Attila Juhos, and Roland Zimmermann for helpful discussions and feedback.
This work was supported by the German Federal Ministry of Education and Research (BMBF): Tübingen AI Center, FKZ: 01IS18039A. WB acknowledges financial support via an Emmy Noether Grant funded by the German Research Foundation (DFG) under grant no. BR 6382/1-1 and via the Open Philanthropy Foundation funded by the Good Ventures Foundation. WB is a member of the Machine Learning Cluster of Excellence, EXC number 2064/1 - Project number 390727645. This research utilized compute resources at the Tübingen Machine Learning Cloud, DFG FKZ INST 37/1057-1 FUGG. We thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) for supporting TW and PM.
Author contributions
The project was led and coordinated by TW. TW and PM jointly developed the theory with insights from WB. TW implemented and conducted the experiments with input from PM and WB. TW led the writing of the manuscript with help from WB, PM, and MB. TW created all figures with comments from PM and WB.
[1] Jerry A. Fodor and Zenon W. Pylyshyn. Connectionism and cognitive architecture: A critical analysis. Cognition, 28(1):3-71, March 1988. ISSN 0010-0277. doi: 10.1016/0010-0277(88) 90031-5.
[2] Noam Chomsky. Aspects of the Theory of Syntax, volume 11. MIT press, 2014.
[3] Lukas Schott, Julius von Kugelgen, Frederik Träuble, Peter Gehler, Chris Russell, Matthias Bethge, Bernhard Schölkopf, Francesco Locatello, and Wieland Brendel. Visual Representation Learning Does Not Generalize Strongly Within the Same Domain. arXiv:2107.08221 [cs], February 2022.
[4] Milton Llera Montero, Casimir JH Ludwig, Rui Ponte Costa, Gaurav Malhotra, and Jeffrey Bowers. The role of Disentanglement in Generalisation. In International Conference on Learning Representations, February 2022.
[5] Brenden Lake and Marco Baroni. Generalization without systematicity: On the compositional skills of sequence-to-sequence recurrent networks. In International Conference on Machine Learning, pages 2873-2882. PMLR, 2018.
[6] Joao Loula, Marco Baroni, and Brenden M Lake. Rearranging the familiar: Testing compositional generalization in recurrent networks. arXiv preprint arXiv:1807.07545, 2018.
[7] Daniel Keysers, Nathanael Scharli, Nathan Scales, Hylke Buisman, Daniel Furrer, Sergii Kashubin, Nikola Momchev, Danila Sinopalnikov, Lukasz Stafiniak, Tibor Tihon, et al. Measuring compositional generalization: A comprehensive method on realistic data. arXiv preprint arXiv:1912.09713, 2019.
[8] Francesco Locatello, Gabriele Abbati, Thomas Rainforth, Stefan Bauer, Bernhard Schölkopf, and Olivier Bachem. On the fairness of disentangled representations. Advances in neural information processing systems, 32, 2019.
[9] Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE Transactions on Pattern Analysis and Machine Intelligence, 35(8): 1798-1828, 2013. doi: 10.1109/TPAMI.2013.50.
[10] Sjoerd van Steenkiste, Francesco Locatello, Jürgen Schmidhuber, and Olivier Bachem. Are disentangled representations helpful for abstract visual reasoning? In H. Wallach, H. Larochelle, A. Beygelzimer, F. dAlché-Buc, E. Fox, and R. Garnett, editors, Advances in Neural Information Processing Systems, volume 32. Curran Associates, Inc., 2019.
[27] Zhang-Wei Hong, Ge Yang, and Pulkit Agrawal. Bi-linear value networks for multi-goal reinforcement learning. In International Conference on Learning Representations, 2021.
[28] Sindy Löwe, Phillip Lippe, Maja Rudolph, and Max Welling. Complex-Valued Autoencoders for Object Discovery, November 2022.
[29] Shai Ben-David and Ruth Urner. Domain adaptation-can quantity compensate for quality? Annals of Mathematics and Artificial Intelligence, 70:185-202, 2014.
[30] Masashi Sugiyama, Matthias Krauledat, and Klaus-Robert Müller. Covariate shift adaptation by importance weighted cross validation. Journal of Machine Learning Research, 8(5), 2007.
[31] Zheyan Shen, Jiashuo Liu, Yue He, Xingxuan Zhang, Renzhe Xu, Han Yu, and Peng Cui. Towards Out-Of-Distribution Generalization: A Survey. arXiv:2108.13624 [cs], August 2021. doi: 10.48550/arXiv.2108.13624.
[32] Shai Ben-David, John Blitzer, Koby Crammer, Alex Kulesza, Fernando Pereira, and Jennifer Wortman Vaughan. A theory of learning from different domains. Machine learning, 79: 151-175, 2010.
[33] Aviv Netanyahu, Abhishek Gupta, Max Simchowitz, Kaiqing Zhang, and Pulkit Agrawal. Learning to Extrapolate: A Transductive Approach. In The Eleventh International Conference on Learning Representations, February 2023.
[34] Kefan Dong and Tengyu Ma. First Steps Toward Understanding the Extrapolation of Nonlinear Models to Unseen Domains. In NeurIPS 2022 Workshop on Distribution Shifts: Connecting Methods and Applications, October 2022.
[35] Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Li Fei-Fei, C. Lawrence Zitnick, and Ross Girshick. CLEVR: A Diagnostic Dataset for Compositional Language and Elementary Visual Reasoning. arXiv:1612.06890 [cs], December 2016.
We reiterate the setup and notation introduced in the paper here for ease of reference.
Notation  $[N]$  denotes the set of natural numbers  $\{1,2,\dots,N\}$ . Vector-valued variables (e.g.,  $\pmb{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}$ .
Setup We are given two arbitrary distributions  $P, Q$  over latents  $\mathbf{z} = (z_1, \dots, z_K) \in \mathcal{Z}$ . Each latent  $\mathbf{z}_k$  describes one of the  $K$  components of the final data point  $\mathbf{x}$  produced by the ground-truth data-generating process  $\mathbf{f}$ . A model  $\hat{\mathbf{f}}$  is trained to fit the data-generating process on samples of  $P$ ; the aim is to derive conditions on  $P$  and  $\hat{\mathbf{f}}$  that are sufficient for  $\hat{\mathbf{f}}$  to then also fit  $\mathbf{f}$  on  $Q$ .
We assume that  $f, \hat{f}$  are chosen such that we can find at least one compositional representation (Definition 1) for either function that shares a common composition function  $C$  and factorization of the latent space  $\mathcal{Z}_1 \times \dots \times \mathcal{Z}_K = \mathcal{Z}$ .
Proof of Theorem 1. For  $\hat{\pmb{f}}$  to generalize to  $Q$ , we need to show fitting  $\pmb{f}$  on  $P$  implies also fitting it on  $Q$ , in other words
$$
f \equiv_ {P} \hat {f} \Rightarrow f \equiv_ {Q} \hat {f} \tag {9}
$$
Step 1. Since  $C$  is known and fixed, we immediately get
$$
\varphi \equiv_ {Q} \hat {\varphi} \Rightarrow f \equiv_ {Q} \hat {f}, \tag {10}
$$
i.e., it suffices to show that the component functions generalize. Note, however, that since  $C$  is not generally assumed to be invertible, we do not directly get that agreement of  $f, \hat{f}$  on  $P$  also implies agreement of their component functions  $\varphi, \hat{\varphi}$  on  $P$ .
Step 2. We require  $P$  to have compositional support w.r.t.  $Q$  (Definition 2 and Assumption (A2)). The consequence of this assumption is that any point  $\pmb{q} = (\pmb{q}_1, \dots, \pmb{q}_K) \in Q$  can be constructed from components of the  $K$  support points  $\pmb{p}^k = (\pmb{p}_1^k, \dots, \pmb{p}_K^k) \in P$  subject to  $\pmb{p}_k^k = \pmb{q}_k$  as
$$
\boldsymbol {q} = \left(\boldsymbol {p} _ {1} ^ {1}, \dots , \boldsymbol {p} _ {K} ^ {K}\right). \tag {11}
$$
A trivial consequence, then, is that points  $\tilde{\pmb{x}}\in \tilde{\mathcal{X}}$  in component space corresponding to points in  $Q$  in latent space can always be mapped back to latents in  $P$
$$
\varphi (\boldsymbol {q}) = \left(\varphi_ {1} \left(\boldsymbol {q} _ {1}\right), \dots , \varphi_ {K} \left(\boldsymbol {q} _ {K}\right)\right) = \left(\varphi_ {1} \left(\boldsymbol {p} _ {1} ^ {(1)}\right), \dots , \varphi_ {K} \left(\boldsymbol {p} _ {K} ^ {(K)}\right)\right) \tag {12}
$$
because each component function  $\varphi_{k}$  only depends on the latents  $z_{k}$  of a single component. This is also the case for the component functions  $\hat{\varphi}$  of  $\hat{f}$  so that we get
$$
\varphi \underset {P} {\equiv} \hat {\varphi} \Rightarrow \varphi \underset {Q} {\equiv} \hat {\varphi}. \tag {13}
$$
Step 3. We now only need to show that  $\varphi \equiv \hat{\varphi}$  follows from  $f \equiv \hat{f}$ . As noted above, this is not guaranteed to be the case, as  $C$  is not generally invertible (e.g., in the presence of occlusions). We, therefore, need to consider when a unique reconstruction of the component functions  $\varphi$  (and correspondingly  $\hat{\varphi}$ ) is possible, based on only the observations  $x = f(z)$  on  $Q$ .
As explained in the main paper, we can reason about how a change in the latents  $z_{k}$  of some slot affects the final output, which we can express through the chain rule as
$$
\underbrace {\frac {\partial \boldsymbol {f}}{\partial \boldsymbol {z} _ {k}} (\boldsymbol {z})} _ {N \times D} = \underbrace {\frac {\partial \boldsymbol {C}}{\partial \varphi_ {k}} (\boldsymbol {\varphi} (\boldsymbol {z}))} _ {N \times M} \underbrace {\frac {\partial \varphi_ {k}}{\partial \boldsymbol {z} _ {k}} (\boldsymbol {z} _ {k})} _ {M \times D} \quad \forall k \in [ K ]. \tag {14}
$$
Here,  $N$  is the dimension of the final output (e.g.,  $64 \times 64 \times 3$  for RGB images),  $M$  is the dimension of a component's representation  $\tilde{\pmb{x}}_k$  (e.g., also  $64 \times 64 \times 3$  for RGB images if composition happens in image space), and  $D$  is the dimension of a component's latent description  $\pmb{z}_k$  (e.g., 5: x-position, y-position, shape, size, hue for sprites). Note that we can look at the derivative component-wise because each component function  $\varphi_k$  only depends on the latents  $\pmb{z}_k$  of its component. However, the combination function still depends on the (hidden) representation of all components, and therefore  $\frac{\partial C}{\partial \varphi_k}$  is a function of the entire  $\varphi$  and  $z$ .
In Equation 14, the left-hand side (LHS)  $\frac{\partial f}{\partial z_k}$  can be computed from the training, as long as  $\operatorname{supp} P$  is an open set. On the right-hand side (RHS), the functional form of  $\frac{\partial C}{\partial \varphi_k}$  is known since  $C$  is given, but since  $\varphi(z)$  is still unknown, the exact entries of this Jacobian matrix are unknown. As such, Equation 14 defines a system of partial differential equations (PDEs) for the set of component functions  $\varphi$  with independent variables  $z$ .
Before we can attempt to solve this system of PDEs, we simplify it by isolating  $\frac{\partial\varphi_k}{\partial z_k}$ . Since all terms are matrices, this is equivalent to solving a system of linear equations. For  $N = M$ ,  $\frac{\partial C}{\partial\varphi_k}$  is square, and we can solve by taking its inverse as long as the determinant is not zero. In the general case of  $N\geq M$ , however, we have to resort to the pseudoinverse to write