Provable Compositional Generalization for Object-Centric Learning
Authors: Thaddäus Wiedemer
Jack Brady
Alexander Panfilov
Attila Juhos
Matthias Bethge
Wieland Brendel
Venue: ICLR 2024 oral (2024)
Keywords: compositional generalization
identifiability
object-centric learning
generalization
OOD generalization
unsupervised learning
slot attention
disentanglement
autoencoders
representation learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We show theoretical conditions under which compositional generalization is guaranteed for object-centric representation learning.

Abstract: Learning representations that generalize to novel compositions of known concepts is crucial for bridging the gap between human and machine perception. One prominent effort is learning object-centric representations, which are widely conjectured to enable compositional generalization. Yet, it remains unclear when this conjecture will be true, as a principled theoretical or empirical understanding of compositional generalization is lacking. In this work, we investigate when compositional generalization is guaranteed for object-centric representations through the lens of identifiability theory. We show that autoencoders that satisfy structural assumptions on the decoder and enforce encoder-decoder consistency will learn object-centric representations that provably generalize compositionally. We validate our theoretical result and highlight the practical relevance of our assumptions through experiments on synthetic image data.

PROVABLE COMPOSITIONAL GENERALIZATION FOR OBJECT-CENTRIC LEARNING
1,2.3~ Matthias Bethge's 13 Wieland Brendel 2.3,4
Thaddäus Wiedemer 1233+ 1223+ Jack Brady!.2.3-
Alexander Panfilov! Attila Juhos 1,2,3+
University Tübingen Max Planck Institute for Intelligent Systems
'Tübingen AI Center
*ELLIS Institute Tübingen
(thaddaeus. wiedemer, jack.brady, attila. juhos wtabbosjouebiingn..nng
ABSTRACT
Learning representations that generalize to novel compositions of known con- cepts crucial for bridging the gap between human and machine perception. One prominent effort is learning object-centric representations, which are widely con- jectured enable compositional generalization. Yet, remains unclear when this conjecture will be true, as principled theoretical or empirical understanding of compositional generalization lacking. this work, we investigate when com- positional generalization is guaranteed for object-centric representations through the lens of identifiability theory. We show that autoencoders that satisfy structural assumptions the decoder and enforce encoder-decoder consistency will learn object-centric representations that provably generalize compositionally. We vali- date theoretical result and highlight the practical relevance of our assumptions through experiments on synthetic image data.
INTRODUCTION
Despite tremendous advances in machine learning, large gap still exists between humans and machines terms of learning efficiency and generalization (Tenenbaum et al., 2011; Behrens et al., 2018; Schölkopf et al., 2021). A key reason for this is thought to be that machines lack the ability to generalize compositionaally which humans heavily rely on (Fodor and Pylyshyn, 1988; Lake et al., 2017; Battaglia al., 2018; Goyal and Bengio, 2022; Greff al., 2020). Namely, humans are able to recompose previously learned knowledge to generalize to never-before-seen situations.
Significant work has thus gone into the problem of learning representations that generalize compositionally. One prominent effort is object-centric representation learning (Burgess et al., 2019; Greff al., 2019; Locatello et al., 2020a; Lin al., 2020; Singh et al., 2022; Elsayed al., 2022; Seitzer et al., 2023), which aims to represent each object in an image via a distinct subset of the image's ' latent code. Due this modular structure, object-centric representations are widely conjectured to enable compositional generalization (Battaglia et al., 2018; Kipf et al., 2020; Greff al., 2020; Locatello al., 2020a). Yet, it remains unclear when this conjecture is actually true because a theoretical understanding of compositional generalization for unsupervised object-centric representations is lacking, and empirical methods are frequently not scrutinized for their ability to generalize compositionally. Consequently, it is uncertain to what extent advancements in object- centric learning promote compositional generalization and what obstacles still need to be overcome.
In this work, we take step towards addressing this point by investigating theoretically when compo- sitional generalization is possible in object-centric representation learning. To do this, we formulate compositional generalization as problem of identifiability under a latent variable model in which objects are described by subsets of latents called slots (see Fig. left). Identifiability provides a rig- orous framework study representation learning, but previous results have only considered identi- fiability of latents in-distribution (ID), i.e., latents that generate the training distribution (Hyvärinen
"Equal contribution, order decided by dice roll. Code git wtwtch...combbrrbccco.cc www.bbenneeeooooooooooooeeeeeoee -gwoo/groopppoooo grooppobjectt www/ottjecct..oo composi tional -genera 1 izztiion

---

DOD
train
add, decader e emoodar with consistency loss
addttive decoder
2dd. decoder + encoder
MSF
inferted
eoodeer
0,0
0
Salnin Seriibttoo
triining dish Bution
Kramiig ButrBbution
8.0
4.4
comp pommppaiiionall decader
slot
Figure 1: Compositional generalization in object-centric learning. We assume a latent variable model where objects in an image (here, triangle and circle) are described by latent slots. Our notion of compositional generalization requires a model to identify the ground-truie truth latent slots (slor identifiabilityy Def. 2) on the train distribution and to transfer this identifiability to out-of-distribute (OOD) combinations of slots (Def. 3). An autoencoder achieves slot identifiability on the train distri- bution if its decoder is compositional (Thm. 1). Further, we prove that decoders that are additive are able to generalize OOD as visualized in (A) via the isolated decoder reconstruction error over a 21) projection of the latent space (see App. B.3). However, this does not guarantee that the entire model generalizes OOD, as the encoder will generally not invert the decoder on OOD slot combinations, leading to a large overall reconstruction error (B). To address this, we introduce a compositional consistency regularizer (Def. 6), which allows the full autoencoder to generalize OOD (C, Thm.3).
et al., 2023). Compositional generalization, however, requires identifying the ground-truie latents not just ID, but also out-of-distribution (OOD) for unseen combinations of latent slots.
We pinpoint the core challenges in achieving this form of identifiability in Sec. 2 and show how they can be overcome theoretically by autoencoder models which satisfy two properties: additiv- ity (Def. 5) and compositional consistency (Def. 6). Informally, additivity states that the latents are decoded the sum of individual slot-wise decodings, while compositional consistency states that the encoder inverts the decoder ID as well as OOD. When coupled with previous identifiability re- sults from Brady et al. (2023), we prove that autoencoders that satisfy these assumptions will learn object-centric representations which provably generalize compositionally (Thm.3).
We discuss implementing additivity in practice and propose regularizer that enforces compositional consistency by ensuring that the encoder inverts the decoder on novel combinations of ID latent slots (Sec. 4). We use this to empirically verify our theoretical results in See. 6.1 and find that additive autoencoders that minimize our proposed regularizer compositionally. In Sec. 6.2, we study the importance of our theoretical assumptions for the popular object-centric model Slot Attention (Locatello et al., 2020a) on this dataset.
multi-object dataset are able to generalize
Notation Vectors or vector-valued functions are denoted by bold letters For vectors with fac- torized dimensionality (c.g.. z usually from I KM ) or functions with factorized output (usually mapping to KMM, indexing with l denotes the k-tth contiguous sub-vector (i.e., Z E RM or @x( RM). Additionally, for positive integer K we write the set f1...., K} [K].
PROBLEM SETUP
Informally, we say that a model generalizes compositionally if yields an object-centric represen- tation for images containing unseen combinations of known objects, i.c., objects observed during training (Zhao et al., 2022: Frady et al., 2023; Wiedemer et al., 2023). For example, a model trained on images containing red square and others containing blue triangle should generalize to images containing both objects simultaneously- even this combination has not previously been observed.
To formalize this idea, we first define scenes of multiple objects through a latent variable model. Specifically, we assume that observations æ of multi-object scenes are generated from latent vectors by a diffeomorphic generator f Z X' mapping from latent space Z to data space X C R.N, i.e., x = f(z) (see also App. 1) Each object in æ should be represented by a unique sub-vector z& in the latent vector z, which we refer to as a slot. Thus, we assume that the latent space Z factorizes into K slots Z₁ with M dimensions each:
2₁ C TR M
and
Z Z
Zx C KM
2

---

A slot 2₁ contains all possible configurations for the k-th object, while Z encompasses all possible combinations of objects. For our notion of compositional generalization, a model should observe all possible configurations of each object but not necessarily all combinations of objects. This corre- sponds to observing samples generated from a subset 2 of the latent space 2, where 2 contains all possible values for each slot. We formalize this subset below.
Definition 1 (Slot-supported subset). For 2 C 2 = 21 x ... x Zk. let 2 == {zu/z E 23).
g is said to be slot-supported subset of Z if 2 Zy for any is [K].
One extreme example of slot-supported subset g is the trivial case g 2; another is set containing the values for each slot exactly once such that 2 resembles a 1D manifold in R.M.
We assume observations x: from training space X are generated by slot-supported subset 2 i.c., X = f(2)). The following generative process describes this:
== ff)),, 2 Ps. supp(p=) = Z
(2)
Samples from such generative process are visualized in Fig. for a simple setting with two objects described only by their y-coordinate. We can see that the training space contains each possible configuration for the two objects but not all possible combinations of objects.
Now, assume we have inference model which only observes data on X generated according to
Eq. 2. In principle, this model could be any sufficiently expressive diffeomorphism: however, we will assume it to be autoencoder, is common in object-centric learning (Yuan et al., 2023). Namely, we assume the model consists of a pair of differentiable functions: an encoder 9 R + k and a decoder f : RKM + RN,, which induce the inferred latent space 2 = g(X) and the reconstructed data space & == f(2). The functions are optimized to invert each other on 2 by minimizing the reconstruction objective
Earpo Lree (X free (g: f, x=) Eu- ( [I((((=)) - =1[2]= supp(r=) X
(3)
We say that an autoencoder (g. f' produces an object-centric representation via é == ((x) each inferred latent slot 2, encodes all information from exactly one ground-truth latent slot zk, i.c., the model separates objects in its latent representation. We refer to this notion as slot identifiability, which we formalize below, building upon Brady al. (2023):
Definition 2 (Slot identifiability). Let f Z + X' be a diffeomorphismm. Let 2 be slot-supported subset of Z. An autoencoder (g. f) is said slot-identify z on z w.F.t via é = g(f()) minimizes Cre (X*) and there exists permutation π of [K] and a set of diffeomorphisms h4 (kk + 2n for any k E [K].
Intuitively, by assuming (g. f) minimizes (x')) we know that on the training space X 9 is diffeomorphism with f as its inverse. This ensures that é preserves all information from ground- truth latent z. Furthermore, requiring that the slots εn and (x(k) are related by a diffeomorphism ensures that this information factorizes in the sense that each inferred slot contains only and all information from a corresponding ground-truth slot." We can now formally define what it means for an autoencoder to generalize compositionally.
Definition 3 (Compositional generalization). Let f - X' be diffeomorphism and 2 be slot-supported subset of Z. An autoencoder (g. that slot-identifies on 2 w.r.t. is said to generalize compositionally w.E.t. z if also slot-identifies on Z w.r.t. f.
This definition divides training an autoencoder that generalizes compositionally into two challenges.
Challenge Identifiability Firstly, the model must slot-identify the ground-truth latents on the slot-supported subset 2 Identifiability is generally difficult and is known to be impossible without further assumptions on the generative model (Hyvärinen and Pajunen, 1999; Locatello et al., 2019). The majority of previous identifiability results have addressed this by placing some form of statistical independence assumptions on the latent distribution Pz (Hyvärinen et al., 2023). In our setting, however, P= is only supported on z which can lead extreme dependencies between latents (e.g.,
Note that when 2 Z2, we recover the definition of slot identifiability in Brady al. (2023).
3

---

Latent Space
undbsened:
Data Space
Inferred latents
Reconstructed data
Inferred latents
encoder
Coans
-
eigt 1ot miaaet
g
compesitiona & reducctble generstor
additire decoder
encoder
Bredd
Monnl
Figure 2: Overview of our theoretical contribution. (1) We assume access to data from a train- ing space X X. which generated from slot-supported subset 2 of the latent space Z (Def. 1), via compositional and irreducible generator f. (2) We show that an autoencoder with compositional decoder f trained via the reconstruction objective L on this data will slot-identify ground-truth latents z on g (Thm. 1). Since the inferred latents slot-identify z on 2 their slot-wise recombinations Z' slot-identify OOD on Z. However, the encoder is not guaranteed to infer OOD latents such that g(Xt å = Z'. (3) On the other hand, if the decoder is addi- tive, its reconstructions are guaranteed to generalize such that f(Z') X (Thm. 2). (4) Therefore, regularizing the encoder ğ to invert f using our proposed compositional consistency objective Coons (Def. 6) enforces 2 = Z', thus enabling the model to generalize compositionally (Thm. 3).
see Fig. 2 where slots are related almost linearly on 2 ). is thus much more natural to instead place assumptions the generator f sufficiently constrain the problem, in line with common practices in object-centric learning that typically assume a structured decoder (Yuan et al., 2023).
Challenge 2: Generalization Even if we can train an autoencoder (g. f) that slot-identifies
in-distribution (ID) on the slot-supported subset 2d, we still require it to also slot-identify x out-of- distribution (OOD) on all of Z. Empirically, multiple prior works have demonstrated in the context of disentanglement that this form of OOD generalization does not simply emerge for models that can identify the ground-truth latents (Montero al., 2021; 2022: Schott et al., 2022). From a theoretical perspective, OOD generalization of this form implies that the behavior of the generator f on the full latent space Z is completely determined by its behavior on 2 which could essentially be one-dimensional manifold. This will clearly not be the case if is arbitrary function, necessitating constraints on its function class to make any generalization guarantees.
COMPOSITIONAL GENERALIZATION IN THEORY
In this section, we show theoretically how the ground-truth generator f and autoencoder (g. can be constrained address both slot identifiability and generalization, thereby facilitating composi- tional generalization (complete proofs and further details are provided in App. A).
To address the problem of slot identifiability, Brady al. (2023) proposed to constrain the genera- tor ff via assumptions on its Jacobian, which they called compositionality and irreducibility. Infor- mally, compositionality states that each image pixel is locally a function of at most one latent slot, while irreducibility states that pixels belonging to the same object share information. We relegate a formal definition of irreducibility App. A.4 and only restate the definition for compositionality.
Definition 4 (Compositionnality) A differentiable f Z RN N is called compositional in z E Z if
0fn Ozk
05n (z) 0, for any k, [K], 1 + and any n E [N].
(z)
#
0z,
For a generator f satisfying these assumptions on 2, Brady et al. (2023) showed that an autoen- coder (g. ) with a compositional decoder (Def. 4) will slot-identify z on Z w.r.t. f. This result appealing for addressing Challenge I since does requires that the training space X is generated from the entire latent space Z. We here show an extension for cases when the training space X arises from convex, slot-supported subset 2
rely on assumptions on pzi however,
Theorem (Slot identifiability supt-sep subset). Let Z X be compositional and subset of Z. An autoencoder (9. f) irreducible diffeomorphism. Let 2 be convex, slot-supported subset
4

---

that minimizes (xgc(X)) for X identifies 2 on w.r.t. f in the sense of Def.
f2²) and whose decoder f is compositional g(x*). slot-
Thm. solves Challenge of slot identifiability on the slot-supported subset g but to generalize compositionally, still need to address Challenge 2 and extend this to all of Z. Because have slot identifiability on 2d we know each ground-truth slot and corresponding inferred slot are related by diffecomrphism hg. Since h& is defined for all configurations of slot ²z(k), the representation which slot-identifies (Z), ......... 2x) for any combination of slots (ID or OOD) in Z is given by
z' (h1 (2+(())...... hx (Z+(K))),
z' h1 (Zz()) x hx (Zx))))
(5)
Therefore, for autoencoder to generalize its slot identifiability from 2 to Z, it should match this representation such that for any = E Z.
(ff(2) = and f(z')= f(2).
We aim to satisfy these conditions by formulating properties of the decoder f such that it fulfills the second condition, which we then leverage to regularize the encoder ğ to fulfill the first condition.
3.1 DECODER GENERALIZATION VIA ADDITIVITY
We know from Thm. that f renders each inferred slot h& (Z+(k)) correctly to a corresponding object in æ for all possible values of (=(k). Furthermore, because the generator f satisfies compositionality (Def. 4), we know that these slot-wise renders should not be affected by changes to the value of any other slot 2,. This implies that for f to satisfy Eq. 6, we only need to ensure that its slot-wise renders remain invariant when constructing z' with an OOD combination slots z. We show below that additive decoders can achieve this invariance.
Definition 5 (Additive decoder). For autoencoder (g. the decoder f is said to be additive if
K
f(z) - Σ ($x), Am
where Pp IRM + R for any k E [K] and 2 E KMM
(7)
We can think of an additive decoder f as rendering each slot &n to an intermediate image via slor functions kp then summing these images to create the final output. These decoders are expressive enough to represent compositional generators (see App. A.7). Intuitively, they globally remove inter- actions between slots such that the correct renders learned on inferred latents of 2 are propagated to inferred latents of the entire 2. This is formalized with the following result.
Theorem 2 (Decoder generalization). Let f 2 - X be compositional diffeomorphism and z be slot-supported subset of Z. Let (g, f) be an auloencoder that slot-identifies on 2 w.r.t. f. If the decoder f is additive, then it generalizes in the following sense: f(z') = f()) any 2 E Z, where z' is defined according Eq. 5.
Consequently, f now injective on Z' and we get =(2')=
3.2 ENCODER GENERALIZATION VIA COMPOSITIONAL CONSISTENCY
Because the decoder f generalizes such that f(x') = f(x) (Thm. 2) and f(Z') = X, the condition on the encoder from Eq. 6 corresponds to enforcing that å inverts f on all of X. This is ensured ID on the training space X by minimizing the reconstruction objective Lree (Eq.3). However, there is nothing enforcing that ý also inverts f OOD outside of x (see Fig. 1B for a visualization of this problem). To address this, we propose the following regularizer.
Definition& 6 (Compositional consistency). Let qz' be distribution with supp(zz) = 2' (Eq. 5). An autoencoder (g. f) is said to be compositionally consistent if it minimizes the compositional consistency loss
Lcons (i,fz2) = -Ez'naz [|0(f(=))) - ='112] -
(8)
5

---

The loss can be understood first sampling an OOD combination of slots z' by composing inferred ID slots (k))) The decoder can then render z' to create an OOI) sample f(='). Re-encoding this sample such that (=))) z' then regularizes the encoder to invert the decoder OOD. We discuss how this regularization can be implemented in practice in Sec.
3.3 PUTTING IT ALL TOGETHER
Thm. showed how slot identifiability can be achieved ID on 2 if f satisfies compositionality, and Thm. 2, Def. showed how this identifiability can generalized to of 2 if the decoder is addi- tive and compositional consistency is minimized. Putting these results together, we can now prove conditions for which an autoencoder will generalize compositionally (see Fig. 2 for an overview).
Theorem (Compositionally generalizing autoencoder). Let f
X be a compositional and irreducible diffeomorphism. Let 2 be a convex, slot-supported subset of Z. Let (g. f) be an autoencoder with additive decoder f (Def. 5). f is compositional on ()X and 9. f solve
Live (gr x + Xo (g. f, 2') =0, for some X >>0 0,
then the autoencoder (g, f) generalizes compositionally w.r.t. 2 in the sense of Def. 3.
Moreover, 9 : X 2 inverts f: Z' X' and also -= z' h2 (2x(t) x ... x h (2x(2(())
COMPOSITIONAL GENERALIZATION IN PRACTICE
Compositionality Thm. 3 explicitly assumes that the decoder f satisfies compositionality on g(x*) but does not give a recipe to enforce this in practice. Brady et al. (2023) proposed reg- ularizer that enforces compositionality if minimized (see App. B.4), but their objective computa- tionally infeasible to optimize for larger models, thus limiting its practical use. At the same time, Brady et al. (2023) showed that explicitly optimizing this objective may not always be necessary, as the object-centric models used in their experiments seemed to minimize it implicitly, likely through the inductive biases in these models. We observe a similar phenomenon (see Fig. 4, right) and thus rely these inductive biases satisfy compositionality in experiments in Sec. 6.
Additivity trivial to implement an additive decoder by parameterizing the slot functions 4o from Eq. as, e.g., deconvolution neural networks. This resembles the decoders typically used in object-centric learning, with the key difference being the use of slot-wise masks m. Specifically, existing models commonly decoder of the form
f(z)= mi Dk
m (m)),, (mg, spx (2x),
(10)
where O is element-wise multiplication and (-) denotes the softmax function. Using masks my in this way facilitates modeling occluding objects but violates additivity as the softmax normalizes masks across slots, thus introducing interactions between slots during rendering. We empirically investigate how this interaction affects compositional generalization in Sec. 6.2.
Compositional Consistency The main challenge with implementing the proposed compositional consistency loss Lcons (Def. 6) sampling z' from dz' with support over 2. First, note that defined Z' in Eq. 5 through use of the functions h4 but can equivalently write
h (z*(())) = gv (f())) and z' ='= g'=)) x ... (x())).
(11)
The reformulation highlights that we can construct OOD samples in the consistency regulariza- tion from ID observations by randomly shuffling the slots of two inferred ID latents 21)) z(2) via Pk ~ ((1,2). Because z is slot-supported subset, constructing z' as
z' (z(01) (p:)
=(ex)), where for {1,2) 200 9 x (1). (1)) P
(12)
ensures that the samples z' cover the entire Z'. Practically, we sample 2)), 2(2) uniformly from the current batch. The compositional consistency objective with this sampling strategy is illustrated

---

dal 2
slot2
recombnation
stop-gradient
Looos
Compositional consistency regularization. In addition to the reconstruction objective, fomss is minimized on recombined latents z'. Recombining slots of the inferred latents & of two ID samples produces a latent z', which can be rendered to an OOD sample x' due to the decoder f generalizing OOD. The encoder is optimized re-encode this sample match z'.
Figure 3:
in Fig. 3. Note that the order of slots is generally preserved between (f(()) and z' so that we pair slots using the Hungarian algorithm (Kuhn, 1955) before calculating the loss. Furthermore, enforcing the consistency loss can be challenging in practice if the encoder contains stochastic op- erations such as the random re-initialization of slots in the Slot Attention module (Locatello et al., 2020a) during each forward pass. We explore the impact this Sec. 6.2.
RELATED WORK
Theoretical Analyses of Compositional Generalization Prior works have addressed identifiabil- ity and generalization theoretically in isolation. For example, several results show how identifiability can be achieved through assumptions on the latent distribution Hyvärinen and Morioka (2016; 2017) 1); Hyvärinen et al. (2019); Khemakhem et al. (2020a;b); Shu et al. (2020); Locatello et al. (2020b); Gresele et al. (2019); Lachapelle et al. (2021); Klindt al. (2021); Hälvä et al. (2021); ); von Kügelgen et al. (2021); Liang et al. (2023) or via structural assumptions on the generator function (Gresele et al., 2021; Horan et al., 2021; Moran et al., 2022; Buchholz et al., 2022: Zheng et al., 2022; Brady et al., 2023). However, none of these deal with generalization. On the other hand, frameworks for OOD generalization were proposed in the context of object-centric world models (Zhao et al., 2022) and regression problems (e.g., Netanyahu al., 2023), with latent variable formulations that closely resemble our work. In this context, OOD generalization was proven for additive inference mod- els (Dong and Ma, 2022) or slot-wise functions composed with a known nonlinearity (Wiedemer et al., 2023). Yet, these results are formulated in a regression setting, which assumes the problem of identifiability solved a priori. Concurrent work from Lachapelle al. (2023) also considers identi- fiability and generalization. Similar to us, they leverage additivity to achieve decoder generalization and show that additivity is sufficient for identifiability under additional assumptions on the decoder, while allowing more general supports. However, they only focus on decoder generalization, while we show theoretically and empirically how to enforce that the encoder also generalizes OOD.
Compositional Consistency Regularization Our compositional consistency loss (Def. 6), which generates and trains on novel data by composing previously learned concepts, resembles ideas in both natural and artificial intelligence. In natural intelligence, (Schwartenbeck et al., 2021; Kurth- Nelson al., 2022; Bakermans et al., 2023) propose that the hippocampal formation implements a form of compositional replay in which new knowledge is derived by composing previously learned abstractions. In machine learning, prior works Rezende and Viola (2018); Cemgil et al. (2020); Sinha and Dieng (2021); Leeb al. (2022) have shown that encoder can fail to correctly encode samples generated by decoder, though not the context of compositional generalization. For pro- gram synthesis, Ellis al. (2023) propose training recognition model on compositions of learned programs. In object-centric learning, Assouel et al. (2022) also train an encoder using on images from recomposed slots; however, the model tailored a specific visual reasoning task.
EXPERIMENTS
This section first verifies main theoretical result (Thm. 3) on synthetic multi-object data (Sec. We then ablate the impact each of our theoretical assumptions on compositional gen- eralization using the object-centric model Slot Attention (Locatello et al., 2020a) (Sec. 6.2).

---

1.0
with feo without Leans
Identifiability Score OOD
Compositional Contrast
0.8
1000
10.0
0.66
100
0.4
10
1.0
0.2
0.0
0.1
10.0
200
Ccons
Epoch
Figure 4: Experimental validation of Thm. 3. Left: Slot identifiability is measured through- out training as a function reconstruction loss (Cres, Eq.3 3) and compositional consistency (Leonss Def. 6). As predicted by Thm. 3, models which minimize Lree and Leons learn representations that are slot identifiable OOD. Right: Compositional contrast (see App. B.4) decreases throughout train- ing, indicating that the decoder is implicitly optimized to be compositional (Def. 4).
Data We generate multi-object data using the Spriteworld renderer (Watters et al., 2019). Images contain two objects on a black background (Fig. 6), each specified by four continuous latents (x/y position, size, color) and one discrete latent (shape). To ensure that the generator satisfies com- positionality (Def. 4), we exclude images with occluding objects. Irreducibility almost certainly satisfied due the high dimensionality of each image, argued in (Brady al., 2023). We sample latents on slot-supported subset by restricting support to a diagonal strip in 2 (see App. B.
Metrics To measure decoder's compositionality (Def. 4), rely the compositional contrast regularizer from Brady et al. (2023) (App. B.4), which was proven to be zero if function is com- positional. To measure slot identifiability, Brady et al. (2023) and fit nonlinear regressors to predict each ground-truth slot z, from an inferred slot & for every possible pair of slots. The regressor's 's fit measured by R² score quantifies how much information &, encodes about 2, We subsequently determine the best assignment between slots using the Hungarian algorithm (Kuhn, 1955) and report the R² averaged over all matched slots.
follow Locatello et al. (2020a); Dittadi et al. (2021);
6.1 VERIFYING THE THEORY
Experimental Setup We train an autoencoder with an additive decoder on the aforementioned multi-object dataset. The model uses two latent slots with 6 dimensions each. We train the model to minimize the reconstruction loss Lree (Eq. 3) for 100 epochs, then introduce the compositional consistency loss Lcons (Def. 6) and jointly optimize both objectives for an additional 200 epochs.
Fig. 4 (Right) shows that compositional contrast decreases over the course of training without additional regularization, thus justifying our choice not to optimize it explicitly. Fig. 4 (Left) visualizes slot identifiability of OOD latents as a function of Liec and Leoms OOD slot identifiability maximized exactly when free and Leons minimized, predicted by Thm. This is corrobo- rated by the heatmaps in Fig. IA-C, which illustrate that additivity enables the decoder to generalize predicted by Thm. 2 but minimizing Lcons required for the encoder also generalize OOD.
Results
6.2 ABLATING IMPACT OF THEORETICAL ASSUMPTIONS
Experimental Setup While Sec. 6.1 showed that our theoretical assumption can empirically en- able compositional generalization, these assumptions differ from typical practices object-centric models. We, therefore, ablate the relevance of each assumption for compositional generalization in the object-centric model Slot Attention. We investigate the effect of additivity by comparing a non- additive decoder that normalizes masks across slots using softmax with additive decoder that replaces the softmax with slot-wise sigmoid functions. We also train the model with and without op- timizing compositional consistency. Finally, we explore the impact of using deterministic encoder by replacing Slot Attention's random initialization of slots with a fixed initialization. All models use two slots with 16 dimensions each and are trained on the multi-object dataset from Sec. 6.1.
8

---

Decoder Error
Overall Error
noo adii-viiitee decoder (softmax) additive decoder (sigmoid)
withoout Coms
with Cons
MSE
9.8
Kraining Bltttiutionn
Brining distriution
Training distrbutton
Briiiing Bistibution
0.4
0.2
Bl1
slot
9101
Figure 5: Compositional generalization for Slot Attention, Visualizing the decoder reconstruc- tion error over a 2D projection of the latent space (see App.B.3 B.3 for details) reveals that the non- additive masked decoder Slot Attention does not generalize OOD on our dataset (A). Making the decoder additive by replacing softmax mask normalization with slot-wise sigmoid functions makes the decoder additive and enables OOD generalization (B, Thm. 2). The full model does not general- ize compositionally, however, since the encoder fails to invert the decoder OOD (C). Regularizing with the compositional consistency loss addresses this, enabling generalization (D, Thm. 3).
Compositional generalization for Slot Attention in terms of slot identifiability and reconstruction quality. Both metrics close optimal ID but degrade OOD with the standard assumptions in Slot Attention. Incorporating decoder additivity (Add.). compositional consistency (Coons) and deterministic inference (Det.) improves OOD performance.
Table 1:
Identifiability 12+
Reconstruction R²r
Add.
eonns
Det.
ID
OOD
ID
0OD
X
x
X
0.99
#1.7e-3
0.81 +9.0e-2
0.99
#1.0e=4
0.71 -1-9e-2
X
2-3-3
0.83
45.40-2
0.99
=5,80-4
0.72 =2.1o-2
/
X
0.99
=2.9c-2
0.92
#3.4o-2
0.99
48 4
0.79 =7.20-2
/
/
0.99
-1.9e3
0.94 22.2e-2
0.99
=1.9e-4
0.92 =2.1e-2
Results Fig. 5 illustrates that the non-additive decoder in Slot Attention does not generalize OOD on our multi-object dataset. Moreover, regularization via the consistency loss is required to make the encoder generalize. Tab. 1 ablates the effect of these assumptions for Slot Attention. We that satisfying additivity and compositional consistency and making inference deterministic (see Sec. 4) improves OOD identifiability and reconstruction performance.
DISCUSSION
Extensions of Experiments Our experiments in Sec. 6.2 provide evidence that compositional gen- eralization will not emerge naturally in object-centric models such as Slot Attention. However, to gain principled understanding of the limits of compositional generalization in these models, experiments with broader set of architectures on more datasets are required. Additionally, Loons, naive uniform manner, potentially
as implemented in Sec. samples novel slot combinations giving rise to implausible images in more complex settings, e.g., by sampling two background slots. Thus, a more principled sampling scheme should be employed to scale this loss.
Extensions of Theory The assumptions of compositionality and additivity on the decoder make progress towards theoretical understanding of compositional generalization, yet are inherently limiting. Namely, they do not allow slots to interact during rendering and thus cannot adequately model general multi-object scenes or latent abstractions outside of objects. Thus, key direction is to understand how slot interactions can be introduced while maintaining compositional generalization.
Conclusion Compositional generalization is crucial for robust machine perception; however, a principled understanding of how it can be realized has been lacking. We show in object-centric learning that compositional generalization is theoretically and empirically possible for autoencoders that possess structured decoder and regularize the encoder invert the decoder OOD. While results do not provide an immediate recipe for compositional generalization in real-world object- centric learning, they lay the foundation for future theoretical and empirical works.

---

REPRODUCIBILITY STATEMENT
Detailed versions of all theorems and definitions, as well as the full proofs for all results are included in App.A. A. We attach our codebase to facilitate the reproduction of our experiments. All hyperpa- rameters, model architectures, training regimes, datasets, and evaluation metrics are provided the codebase. Explanations for design choices given in Sec. 6 in the main text and App. B. The implementation of the compositional consistency loss is detailed Sec. 4, paragraph 3.
CONTRIBUTIONS
TW, JB, and WB initiated the project. JB and TW led the project. AP conducted all experiments with advising from TW and JB. AI, JB, and TW conceived the conceptual ideas behind the theorems. AJ developed the theoretical results and presented them in App. A. TW and JB led the writing of the manuscript with contributions to the theory presentation from AJ and additional contributions from AP, WB, and MB. TW created Figs. 1,2, and with insights from authors. AP and TW created 4 4 and 5 with insights from JB.
ACKNOWLEDGEMENTS
The authors would like to thank (in alphabetical order): Andrea Dittadi, Egor Krasheninnikov, Ev- genii Kortukov, Julius von Kügelgen, Prasanna Mayilvahannan, Roland Zimmermann, Sébastien Lachapelle, and Thomas Kipf for helpful insights and discussions.
This work supported by the German Federal Ministry of Education and Research (BMBF): Tübingen Al Center, FKZ: 011S 8039A, 011S 18039B. WB acknowledges financial support via an Emmy Noether Grant funded by the German Research Foundation (DFG) under grant BR 6382/1-1 and via the Open Philantropy Foundation funded by the Good Ventures Foundation. WB member of the Machine Learning Cluster of Excellence, EXC number 2064/1 Project num- ber 390727645. This research utilized compute resources at the Tübingen Machine Learning Cloud, DFG FKZ INST 37/1057-1 FUGG. The authors thank the International Max Planck Research School for Intelligent Systems (IMPRS-IS) ) for supporting TW and AJ.
REFERENCES
Joshua B. Tenenbaum, Charles Kemp, Thomas L. Griffiths, and Noah D. Goodman. How to grow mind: Statistics, structure, and abstraction. Science, 331:1279 1285, 2011.
Timothy E.J. Behrens, Timothy H. Muller, James C.R. Whittington, Shirley Mark, Alon B. Baram, Kimberly L. Stachenfeld, and Zeb Kurth-Nelson. What is cognitive map? organizing knowledge for flexible behavior. Neuron, 100(2):490-509 2018. ISSN 0896-6273. doi: https://doi.or//1)
1016/j.neuron.2018. 10.002..
Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, (0615))622-633) 2021.
Jerry A. Fodor and Zenon W. Pylyshyn. Connectionism and cognitive architecture: A critical anal- ysis. Cognition, 28:3-71, 1988.
Brenden M. Lake, Tomer D. Ullman, Joshua B. Tenenbaum, and Samuel J. Gershman. Building machines that learn and think like people. Behavioral and Brain Sciences, 40:e253, 2017. doi: 10.1017//01400555X1 160018877
Peter W. Battaglia, Jessica B. Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez,, Vinícius Flo- res Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, Çaglar Gülçehre, H. Francis Song, Andrew J. Ballard, Justin Gilmer, George E. Dahl, Ashish Vaswani, Kelsey R. Allen, Charlie Nash, Victoria Langston, Chris Dyer, Nicolas Man- fred Otto Heess. Daan Wierstra, Pushmeet Kohli, Matthew M. Botvinick, Oriol Vinyals, Yujia Li, and Razvan Pascanu. Relational inductive biases, deep learning, and graph networks. ArXiv, abs/1 Abs/1806.00 061, 2018.
10

---

Anirudh Goyal and Yoshua Bengio. Inductive biases for deep learning of higher-level cognition. Proceedings of the Royal Society A, 478(2266);202100668 2022.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.052088 2020.
Christopher P.
Burgess, Loic Mathee, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. MONet: Unsupervised Scene Decomposition and Repre- sentation, January 2019.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew M. Botvinick, and Alexander Lerchner. Multi-object representation learn- ing with iterative variational inference. ICML, volume 97 of Proceedings of Machine Learning Research, pages 2424 2433, 2019.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-Centric Learning with Slot At- tention. In Advances in Neural Information Processing Systems, volume 33, pages 11525-11538. Curran Associates, Inc., 2020a.
Zhixuan Lin, Yi-Fu Wu, Skand Vishwanath Peri, Weihao Sun, Gautam Singh, Fei Deng, Jindong Jiang. and Sungjin Ahn. Space: Unsupervised object-oriented scene representation via spatial attention and decomposition. In International Conference on Learning Representations, 2020.
Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate DALL-e learns to compose. In International Conference on Learning Representations, 2022.
Gamaleldin Elsayed, Aravindh Mahendran, Sjoerd van Steenkiste, Klaus Greff, Michael C Mozer, and Thomas Kipf. Savi++: Towards end-to-end object-centric learning from real-world videos. Advances in Neural Information Processing Systems, 35:28940-28955 2022.
Maximilian Seitzer, Max Horn, Andrii Zadaianchuk, Dominik Zietlow, Tianjun Xiao, Carl-Johann Gabriel, Tong He, Zheng Zhang. Bernhard Schölkopf, Thomas Brox, and Francesco Lo- catello. Bridging the gap to real-world object-centric learning. The Eleventh International Conference on Learning Representations, 2023.
Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models.
In International Conference on Learning Representations, 2020.
Aapo Hyvärinen, Ilyes Khemakhem, and Hiroshi Morioka. Nonlinear independent component anal- ysis for principled disentanglement in unsupervised deep learning. ArXiv, 1ets/2333.16555,, 2023.
Jack Brady, Roland S. Zimmermann, Yash Sharma, Bernhard Schölkopf, Julius Von Kügelgen, and Wieland Brendel. Provably learning object-centric representations. In Proceedings of the 40th International Conference Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 3038-3062. PMIR, 23-29 Jul 2023.
Linfeng Zhao, Lingzhi Kong, Robin Walters, and Lawson LS Wong. Toward compositional gener- alization in object-oriented world modeling. In International Conference on Machine Learning, PMLR, 2022.
pages
E Paxon Frady, Spencer Kent, Quinn Tran, Pentti Kanerva, Bruno A Olshausen, and Friedrich T Sommer. Learning and generalization of compositional representations of visual scenes. arXiv preprint (rXXi:2333.1369) 2023.
Thaddäus Wiedemer, Prasanna Mayilvahanan, Matthias Bethge, and Wieland Brendel. Composi- tional generalization from first principles. arXiv preprint arXiv:2307.05556, 2023.
Jinyang Yuan, Tonglin Chen, Bin Li, and Xiangyang Xue. Compositional scene representation
learning via reconstruction: A survey, IEEE Transactions on Pattern Analysis and Machine Intel- ligence, 2023.
Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural Networks, 12(3)422-433 439, 1999. ISSN 0893-6080.

---

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Rätsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions the unsupervised learn- ing of disentangled representations. In ICML, volume 97 of Proceedings of Machine Learning Research, pages 41 14-4124, 2019.
Milton Llera Montero, Casimir JH Ludwig, Rui Ponte Costa, Gaurav Malhotra, and Jeffrey Bow- ers. The role of disentanglement in generalisation. In International Conference on Learning Representations, 2021.
Milton Montero, Jeffrey Bowers, Rui Ponte Costa, Casimir Ludwig, and Gaurav Malhotra. Lost in latent space: Examining failures of disentangled models at combinatorial generalisation. In Advances in Neural Information Processing Systems, volume 35, pages 10136-10149 Curran Associates, Inc., 2022.
Lukas Schott, Julius Von Kügelgen, Frederik Träuble, Peter Vincent Gehler, Chris Russell, Matthias Bethge, Bernhard Schölkopf, Francesco Locatello, and Wieland Brendel. Visual representation learning does not generalize strongly within the same domain. In International Conference on Learning Representations, 2022.
H. W. Kuhn. The Hungarian method for the assignment problem. Naval Research Logistics Quar- terly, 2(1-2):83-97) March 1955. ISSN 00281441, 19319193. doi: 00.10022/a..3800020000
Aapo Hyvärinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ICA. In NIPS, pages 3765-3773, 2016.
Aapo Hyvärinen and Hiroshi Morioka. Nonlinear ICA of temporally dependent stationary sources. In AISTATS, volume 54 of Proceedings of Machine Learning Research, pages 460-469 2017.
Aapo Hyvärinen, Hiroaki Sasaki, and Richard E. Turner. Nonlinear ICA using auxiliary variables and generalized contrastive learning. AISTAS,, volume 89 of Proceedings of Machine Learning Research, pages 359-868. 2019.
Ilyes Khemakhem, Diederik P. Kingma, Ricardo Pio Monti, and Aapo Hyvärinen. Variational au- toencoders and nonlinear ICA: A unifying framework. In AISTATS, volume 108 of Proceedings of Machine Learning Research, pages 2207-2217, 2020a.
Ilyes Khemakhem, Ricardo Pio Monti, Diederik Kingma, and Aapo Hyvärinen, Ice-beem: Iden- tifiable conditional energy-based deep models based on nonlinear ICA. In Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Sys- tems 2020, NeurIPS 2020, December 6- 2020, virtual, 2020b.
Rui Shu, Yining Chen, Abhishek Kumar, Stefano Ermon, and Ben Poole. Weakly supervised disen- tanglement with guarantees. In ICLR, 2020.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael
Tschannen. Weakly-superviset disentanglement without compromises. In ICML, volume 1 119 of Proceedings of Machine Learning Research, pages
2020b.
Luigi Gresele, Paul K. Rubenstein, Arash Mchrjou, Francesco Locatello, and Bernhard Schölkopf. The incomplete rosetta stone problem: Identifiability results for multi-view nonlinear ICA. In Proceedings of the Thirty-Fifth Conference on Uncertainty in Artificial Intelligence, UAI 2019, Tel Aviv, Israel, July 22-25, 2019, volume 115 of Proceedings of Machine Learning Research, pages 217-227, 2019.
Sébastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica. First Conference on Causal Learning and Reasoning, 2021.
David A. Klindt, Lukas Schott, Yash Sharma, Ivan Ustyuzhaminov, Wieland Brendel, Matthias Bethge, and Dylan M. Paiton. Towards nonlinear disentanglement in natural data with tempo- ral sparse coding. In ICLR, 2021.
12

---

Hermanni Hälvä, Sylvain Le Corff, Luc Lehéricy, Jonathan So, Yongjie Zhu, Elisabeth Gassiat, and Aapo Hyvärinen. Disentangling identifiable features from noisy data with structured nonlinear ICA. In NeurlPS, pages 1624-1633, 2021.
Julius von Kügelgen. Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf. Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. In Advances in Neural Information Processing Systems, volume 34, pages 16451-16467, 2021.
Wendong Liang, Armin Keki'c, Julius von Kügelgen, Simon Buchholz, Michel Besserve, Luigi Gresele, and Bernhard Schölkopf. Causal component analysis. ArXiv, abs/2305511722. 2023.
Luigi Gresele, Julius von Kügelgen, Vincent Stimper, Bernhard Schölkopf, and Michel Besserve. Independent mechanism analysis, concept? Advances in Neural Information Processing Systems, 34:28233-28248, 2021.
Daniella Horan, Eitan Richardson, and Yair Weiss. When is unsupervised disentanglement ? In Advances in Neural Information Processing Systems, 2021.
Gemma Elyse Moran, Dhanya Sridhar, Yixin Wang, and David Blei. Identifiable deep generative models via sparse decoding. Transactions on Machine Learning Research, 2022. ISSN 2835- 8856.
Simon Buchholz, Michel Besserve, and Bernhard Schölkopf. Function classes for identifiable non- linear independent component analysis. In NeurlPS, 2022.
Yujia Zheng, Ignavier Ng, and Kun Zhang. On the identifiability of nonlinear ICA: sparsity and beyond. In NeurIPS, 2022.
Aviv Netanyahu, Abhishek Gupta, Max Simchowitz, Kaiqing Zhang, and Pulkit Agrawal. Learning to Extrapolate: Transductive Approach. In The Eleventh International Conference on Learning Representations, February 2023.
Kefan Dong and Tengyu Ma. First Steps Toward Understanding the Extrapolation of Nonlinear Models to Unseen Domains. In The Eleventh International Conference on Learning Representa- fions, September 2022.
Sébastien Lachapelle, Divyat Mahajan, Ioannis Mitliagkas, and Simon Lacoste-Julien. Additive decoders for latent variables identification and cartesian-pru -poducct extrapolation. arXiv preprint arXiv:2307.02598, 2023.
Philipp Schwartenbeck, Alon Baram, Yunzhe Liu, Shirley Mark, Timothy Muller, Raymond Dolan, Matthew Botvinick, Zeb Kurth-Nelson, and Timothy Behrens. Generative replay for com- positional visual understanding in the prefrontal-hippocamppa circuit. bioRxiv, 2021. doi:
1111001/2221.060044447220
Zeb Kurth-Nelson, Timothy Edward John Behrens, Greg Wayne, Kevin J. Miller, Lennart Luettgau, Raymond Dolan, Yunzhe Liu, and Philipp Schwartenbeck. Replay and compositional computa- tion. Neuron, 111:454-488 469, 2022.
Jacob J.W. Bakermans, Joseph Warren, James C.R. Whittington, and Timothy E.J. Behrens. Con- structing future behaviour in the hippocampal formation through composition and replay. bioRxiv, 2023. doi: 10. 10.1 1000/003..0..0..552222.
Danilo Jimenez Rezende and Fabio Viola. Taming vaes. arXiv preprint arXiv:1810.00597 2018.
Taylan Cemgil, Sumedh Ghaisas, Krishnamurthy Dvijotham, Sven Gowal, and Pushmeet Kohli. The autoencoding variational autoencoder. Advances in Neural Information Processing Systems, 33: 15077-15087, 2020.
Samarth Sinha and Adji Bousso Dieng.
Consistency regularization for variational auto-encoders.
Advances in Neural Information Processing Systems, 34: 12943494412244 2021.
13

---

Felix Leeb, Stefan Bauer, Michel Besserve, and Bernhard Schölkopf. Exploring the latent space of autoencoders with interventional assays. In Advances in Neural Information Processing Systems, 2022.
Kevin Ellis, Lionel Wong, Maxwell Nye, Mathias Sable-Meyer, Luc Cary, Lore Anaya Pozo, Luke Hewitt, Armando Solar-Lezama, and Joshua Tenenbaum. Dreamcoder: growing generalizable, interpretable knowledge with saeke-sleep bayesian program learning. Philosophical Transactions of the Royal Society A, 381(2251):2020055, 2023.
Rim Assouel, Pau Rodriguez, Perouz Taslakian, David Vazquez, and Yoshua Bengio. Object-centric compositional imagination for visual abstract reasoning. In ICLR2022 Workshop on the Elements of Reasoning: Objects. Structure and Causality, 2022.
Nicholas Watters, Loic Matthey, Sebastian Borgeaud, Rishabh Kabra, and Alexander Lerchner. Spriteworld: A flexible, configurable reinforcement learning environment. http://giiibbbcommmeemmmmdtttttt 2019.
Andrea Dittadi, Samuele Papa, Michele De Vita, Bernhard Schölkopf, Ole Winther, and Francesco Locatello. Generalization and robustness implications in object-centric leamning In International Conference on Machine Learning. 2021.
Christopher Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Desjardins, and Alexander Lerchner. Understanding disentangling in B-vae, 2018.
Iyy Loshchilov and Frank Hutter. Decoupled weight decay regularization. In International Confer- ence on Learning Representations, 2019.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury. Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Ed- ward 7.. Yang. Zachary De Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- performance deep learning library. In NeurIPS, pages 8024-8035, 2019..
Klaus Greff, Francois Belletti, Lucas Beyer, Carl Doersch, Yilun Du, Daniel Duckworth, David Fleet, Dan Gmanapragasam, Florian Golemo, Charles Herrmann, et al. Kubric: scalable dataset generator. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recog- nition, pages 3749-3761, 2022.
14

---

Table 2: General notation and nomenclature.
K
number of slots
M
dimensionality slots
N
dimensionality of observations
[n]
the set {1, 2, ...
f A
function f is defined on subset of A, i.. dom(f) S A
&E
function f defined on A with values in B is k-times continuously differen- tiable
E C* Diffeo (A, B)
function f is C* di-difffeomorhiis around A with values in B (Def. 7) (total) derivative of function f in point æ space of all linear operators between linear spaces V and W a generic subset of RRM (and RR whose role will usually be fulfilled by either 2 or (and 2 or X)
(x)
(V, W)
(and X)
projection of - onto the k-th slot space Zx
85n
fz)
Uzk
// (=)
{n E [N]] (z)m(() # #0}
fs(z)
subvector of f=)) corresponding to coordinates S [N]
A DEFINITIONS, THEOREMS, AND PROOFS
In this section, we detail the theoretical contributions of the paper, including all the proofs. Although there is change their contents, the formulation of some definitions and theorems are slightly altered here to be more precise and cover edge cases omitted in the main text. Hence, the numbering of the restated elements is reminiscent of that used the main text.
Throughout the following subsections, let the number of slots K, the slot dimensionality M, and the observed dimensionality N be arbitrary but fixed positive integers such that N KM.
A.1 INTRODUCTION AND DIFFEOMORPHISSS
This subsection aims provide accesible and completely formal definition of what would call throughout Sec. and a diffeomorphism.
Definition (C* Difffccoomom)hhimm Let Ž be a closed subset of RMM A function f F::RKA KM - R is said to be (e-diffeomorphism around ž, denoted by f E Diffeoo 22,
i) f is defined in an open neighbourhood W C KR of Ž,
ii) f is k-times continuously differentiable on W. i.e. f E C**(W)
ii is injective on Ż, i.e. bijective between 2 and f(ž) and
iv) for any E 2 the derivative Df(z) Lin(RKM ( injective (or equivalently, of full column-rank).
Remark. In the special case when a function h is a diffeomorphism around 2 C R KM but also maps to RKM for any 2 E å the derivative (2) is bijective, hence invertible linear transforma- tion of R KMM Besides that, because of iii), h. is injective. Based the inverse function theorem may conclude that h is injective in open neighbourhood W of Z and h (((w)) also k-times continuously differentiable (where, of course, h22 C h(W")).
Hence, arrive to the more familiar definition of diffeomorphism, i.e. h being bijective with both and h 1 being continuously differentiable. However, the latter definition cannot be applied the more general case when the dimensionality of the co-domain larger than that of the domain, since f 2() cannot be differentiable lower dimensional submanifold f f(2).
15

---

With some variations, the mathematical object with the properties listed in Def. 7 is often called an embedding and serves as a generator or a parametrization of a lower dimensional submanifold of R. This is closer to the interpretation we employ here.
Remark. The literature on diffeomorphisms often defines them as smooth, i.e. infinitely many times differentiable functions with smooth inverses. In our results and proofs, twice differentiability will suffice.
A.2 DEFINITION OF GENERATIVE PROCESS, AUTOENCODER
Here recall the definition of slot-supported subsets from Def. and provide definition of the latent variable model that encapsulates all of our assumptions on the generative process outlined in Sec. 2.
Definition (Projection onto slot space). For any 2 = 2, x 2k. let
2 { zk z E 2$}
(13)
be the projection of 8 onto the k-th slot space Zk. Definition 1b (Slot-supported subset). Let 2 = 21 x ... x Zx. A set 2 C 2 is said to be slot-supported subset of if
g Def.8
(Z E 2$} Zk for any & E [K].
(14)
Definition 9 (Partially observed generative process, POGP). A triplet 25,2) called a par- rially observed generative process (POGP), if
i) Z Z1 x ... ZK for convex, closed sets Z C RM i) g C 2 is closed, slot-supported subset of 2 (Def. 1b) and
i f E Difleo (2, R is a diffeomorphism around 2 (in the sense of Def. 7).
f), 2 For given POGP 2
we refer Z as the latent spuce and z as the training latent space. f is the generator, X' == f(Z) and XS == f(ZS) are the data space and training space respectively.
We also provide a definition for our object-centric model alongside the optimization objective. Definition 10 (Autoencoder, AE), A pair (9. f) is called an autoencoder (AE), if -R RKM and f K + N are continuously differentiable functions, i.e. 9, f E C². We refer to the function Ò as the encoder and f as the decoder.
In case a POGP (z, - ff) is given, we refer to = g(X) = ((f(Z) as the inferred latent space and X f(2) as the reconstructed data space.
for given set & N it holds that ğ E + (() bijective and its inverse f g({) + X (also invertible), then we say that the autoencoder (g. is consistent on X.
Definition 11 (Reconstruction Loss). Let (g. f) be autoencoder. Let x C R be arbitrary set and Pa be an arbitrary continuous distribution function with supp(r=) X. The following quantity is called the reconstruction loss of (g. f) with respect to p
Lee fP.)) Brrpp (1 [l(=())))) x 27
(15)
The following simple lemma tells us that for Cree (g. f. Pw.) vanish, the exact choice of Pæ is irrelevant, and that in this case the decoder (left) inverts the encoder on L. In other words, the autoencoder is consistent on X.
Lemma (Vanishing Ciee implies invertibility on X). Let (G. be an autoencoder and let L C N be arbitrary set. The following statements are equivalent:
i) there exists a cont. distribution function Pz with supp(pæ) = X such that Liee (g. f. Pa) = O:
ii for any cont, distribution function Pa with supp(p=) X we have Cree (g. f. Pw) 0:
ii) f o @gg22 id or, in other words, (g. f) is consistent on X.
16

---

Notation Since (g. Pw) vanishes either for all Pa none at the same time, henceforth in the context vanishing reconstruction loss, we are going to denote the reconstruction loss of (9.9991 f) W.r.l. Pm as Lie (0. f. X') or just (A))
Proof of Lem. 4. We prove the equivalence of these statements by proving ii) + i) = iii) = ii).
The implication ii) i) trivial.
For the implication = ait), let us suppose that there exists Pr with supp(p=) X such that Ciee (g. f. Pa) = O. Equivalently:
Ex-pa [1f(((())) = =.
(16)
This shows that [[(((()) - =II2 = 0 or f(g(z)) = x holds almost surely w.r.t. a ~ Pæ- Since both g. f are continuous functions, this implies that ff o gl id, which is exactly what was be proven.
Now, for iii + it), let us suppose that f f g1 = id and let Pa be any continuous distribution function with supp(p=) X. Since, (g(z)) æ holds for any non-random æ X, then with probability 1. f((())) also holds. From this, we conclude that:
Lree (g. f) Exrp -||
(17)
A.3 DEFINITION OF SLOT IDENTIFIABILITY AND COMPOSITIONAL GENERALIZATION
We are now ready to recall the definition of slot identifiability (Def. 2, originally in Brady et al. (2023)) in its most abstract form, followed by the definition of compositional generalization (Def. 3). Definition 2b (Slot identifiability). Let f E C³ Diffeo Z, )), where Z = Zi closed sets 2g C RM, and let 2 C Z be a closed, slot-supported subset of Z (e.g. Z or 8 from a POGP).
X ZK for
An autoencoder (g. (g. is said to slot-identify z 2 w.r.t via h(z) g (f)) if it minimizes free for X = f(2) and there exists a permutation æ of [K] and set of diffeomorphisms h4 E C Diffco (Zx(M). (2)) such that (z)(2) = h2 h*<==(()) for any and z, where (x) the projection of the set h(2) (c.f. Def.8).
Definition 3b (Compositional generalization). Let (2, 2,, be a POGP. An autoencoder (g, f) that slot-identifies z on 8 w.r.t. f is said to generalize compositionally w.r.. 28, if it also slot- identifies z on Z w.r.t. f.
A.4 COMPOSITIONALITY AND IRREDUCIBILITY ASSUMPTIONS
In this subsection we present the formal definitions of compositionality and irreducibility, already mentioned in Sec. 3 and originally introduced in Brady et al. (2023). These two properties represent sufficient conditions for POGPs (Def. 9) to slot-identifible on the training latent space (Def. 2b).
Let E C¹ (Z, Z C RRRM 0zz/ In this case, let us define
[K]. 2. For the sake of brevity let us denote (z)
(2) {n [M] (z)( + #0}
(18)
the set of coordinates locally influenced (i.e., in point 2) by slot 24 w.r.t. the generator f.
Definition 4b (Compositionality). We say that a function f (Z, S RRM, is compo-
sitional on 2 Z if for any z E 2 and E
{f)))
(19)
17

---

Let fs(z) denote the subvector of f() corresponding to coordinates S $ [N] and let (=)(=) Lin(RKM, (I be the corresponding derivative.
Definition 12 (Imeducibilityy). We say that a function E (2,RR) 2 S T KM is irreducible on ž C Z if for any z 2 and E [M] and any non-trivial partition 1z) S1 U S2 (i.e., S n S2 = 2 and Si, S2 + 2) we have:
rank (Dfs,(z)) + rank (Dfs,(z)) > rank (Dfif(e)
(20)
Remark. Given linear operators Lin(U,)) B e Lin(U, W), the following upper bound holds for the rank of (A, B) e Lin(U, V x W):
(A) + rank(B) 2 rank (A,B))
(21)
Therefore, it holds in general that:
rank (D)s + rank (Dfs,(z)) 2 rank (Df
(22)
hence, irreducibility only prohibits equality.
A.5 IDENTIFIABILITY ON THE TRAINING LATENT SPACE
Here recall Thm. 1, originally presented in slightly less general form in Brady al. (2023).
Theorem 1b (Slot identifiability slot-supported subset). Let 18, be POGP (Def. such that
2 is convex and
ii f is compositional and irreducible on 2
Let ().f) be an autoencoder (Def. 10) such that ii f) minimizes Cree (2 ) for x = f(z)) and
iv) f is compositional on 2 == ()
Then (). slot-identifies on 8 w.r.t in the sense of Def. 2b.
Remark. Due to Lem. 4, assumption iii) is equivalent to (g. f) being consistent on 2 i.e. 9lxx injective and its inverse fl
In its original framework, Brady al. (2023) assumed the training latent space be the entirety of KKM, In our case, however, 8 is a closed, convex, slot-supported subset of Z. Therefore, it is required reprove Thm. lb.
A.6 PROOF OF SLOT IDENTIFIABILITY ON SLOT-SUPPORTED SUBSET
In this subsection reprove Thm. lb via 3 steps. First, prove that the latent reconstruction function h ğ f is, under the consistency of the autoencoder, diffeomorphism. Secondly, restate Prop. 3 of Brady et al. (2023) using our notation. It is a result that locally describes the behaviour of h. irrespective of the shape of the latent training space. Hence, no proof is required. The third and final step is concluding Thm. l itself.
Lemma 5 (Latent reconstruction diffeomorphism). Let f E C* Diffeol (iffeo(2,R)) ) for 2 closed and (g. f) an autoencoder consistent on E = f(2). Then = f diffeomorphism around 2 (Def. 7).
CRM
RKM
In particular, for any z ž, Dh(z) invertible linear transformation of K continuously depending on z.
Proof. Since f is C in open neighbourhood of å and is C2 R- it follows that h 9 is also C in an open neighbourhood of 2.
18

---

The autoencoder (g. f) is consistent on X, hence 9 is injective and h 2 ((() is bijective.
Moreover, the inverse of ølæ is f restricted to (A)). Therefore, h - 9 o f implies that for any E 2 it holds that f(((z)) f(z). After differentiation we receive that for any E 2:
(f()))) Dh(z) =Df(z).
(23)
However, f is C¹ di-ddfffooooopismm consequently Df(=) E Lin(RR (IN) is injective. On the left-hand side, we then have an injective composition of linear functions. Therefore, Dh(z) e
Lin(RKM, I K)) is injective and, of course, bijective.
Consequently, h is c'-diffeomorphism.
Lemma 6 (Prop. 3 of Brady et al. (2023)). Let (2,25, f) POGP and (g. autoencoder satisfy 28,2 the assumptions Thm. 1b and let h 9 f (now C' ddiffoomorrpiss around 2 because of remark after Thm. Ib and Lem. 5).
Then for any z E 2 and any j E [K] there exists a unique k E [K] such that 0, (z) + O. For this particular k it holds that 0,hx(z) Lin(R* )) is invertible.
Remark. Since h is diffeomorphism, Dh(z) is invertible. Thus, the statement of Lem. 6 is equiv- alent to saying that for any z g and any k E [K] there exists a unique j E [K] such that 0, (5h2(2) #0.
Proof of Thm. lb. Let = 6o f. Based Lem, 6 and the latest remark, for any z and any k there exists unique j such that 0, hz() o O, and for this j. 12)4((2) Lin(R RKM is invertible. Firstly, we claim that, in this case, the mapping k + j is bijective and independent of z. More precisely, we show that there exists a permutation π of [K] such that
Step 1.
for any z, k and j: (zhx(2)+0 #0 = -==m( (k))
(24)
and, in the latter case, az(ue) (2) is invertible.
To prove this, conclude from the invertibility of Dx((2) that for any there exists such a. Now
suppose that there exist two distinct points z(1), (2) E z and indices k and ji * ja from [K] such that 02 (=)( + 0, 0, (2) * 0. Then, since 2 is path-connected (as being convex), provides us with a continuous function Φ [0, Let
2s such that (00)) (1)) and (01) = (2)
sup(t E [0,1] 0, hu (N((1) # 0).
(25)
On one hand, 0, he (o)1) 0z )))(2()(2)) * 0, hence 0g he (0(1)) 0 and for any t > 1": 0 hy (+(1) = 0. Therefore, because 0j h o Φ is continuous, it follows that 0, hu (+b(()) = 0.
On the other hand, 0g hz(((0)) = 0 (1)) + O implies that t* * 0 and there exists vergent sequence (t..) C [0, 1*) with lim in = t" such that 0,, (.((ttt)) + O. Therefore, for any # jı, 0, (db(t)) 0. From the continuity of 0. & conclude that for any jı. (+b(1)))) 0.
Subsequently, for any j E [K] (either j + ji or j ) holds that get Dhu (+((t) = O, which contradicts h being a diffcomorphism. Hence, there exists x such that for any 2, k: 0, (z) + + j = (k)..
0. Thus,
Step 2. Secondly, we now prove that h acts slot-wise with permutation m, i.e. for any & there exists h Ci-diffoomorii around (x(x) such that for any 2, he(z) = h=(=x())) By Def. 2b this would imply that (g. ff slot-identifies on 2 w.r.t. f.
19

---

To see this, let 2 =²(²] E 2 with (1) (2)
[x(k] =a(k)- To be proven: ha()(1)) ha(2(21). Since 2
is convex, the path / E [0, + z(1) = x(1) + +1. +z((2()) =20) is inside 2 Then
hx(2(21) h2(2(1)) h4 (=0000000 S'' dt
d
(1)(2(0)))dt
(26)
= -fo (22 2() = S. Σ (2 (2) -2 -
K
(27)
First using the fact that (z))(z(0)) #0 + =k(k and then substituting **()) = (2)
(kk)
**k) we
receive:
- h = [// (n 2 2 ) O.
(28)
A.7 ADDITIVITY AND CONNECTION TO COMPOSITIONALITY
This subsection presents a special subset of decoders that will allow autoencoders to generalize compositionally in the sense of Def. 3b. Definition 5b (Additivity). A function f 2 = 2, x x 2x RN is called additive on 2 if for any there exists Pp: Zx + RN such that
K
f(*) = Σ (x()) holds for any 2 E Z.
(29)
Lemma 7 (Compositionality implies additivity). Let Zn C RRM be convex sets and let f Z
21 x ... X 2x + N be C2 (i.e., twice continuously differentiable) and compositional on (in the sense of Def. 4b). Then f is additive on Z.
Proof of Lem. 7. The proof is broken down into two steps. First, we prove that the coordinate func- tions of compositional functions have a diagonal Hessian. Second, we prove that real-valued func- tions defined on convex set with diagonal Hessian are additive.
Step 1. Observe that is additive if and only if for any p E [N], fp is additive. Let p [N] be arbitrary but fixed and let q fp- To be proven: q is additive. Note that since is C2, q is C2 as well.
We first prove that q has diagonal Hessian, i.e., for any i,j KK] we have g)) 0 for any 2. Proving indirectly, let us assume there exist + j and E such that g() # 0.
The function q is C², thus 0g is continuous. Therefore, there exists a neighborhood V of such that 02 (z) * 0 for any E V. Hence, 0,g cannot be constant on V, for otherwise 0, (0.qg) 8 would be 0. Consequently, there exists =" E V such that 0, ([(²)) + O. Again, 0,g is continuous, hence there exists a neighborhood W V of =" such that 0,g( 2)) # for any z E W.
However, f is compositional, hence either 0zq(=) or 0zq(=) is O. Therefore 0,4(z) = 0 for any z E W. After taking the partial derivative with respect to slot z, we receive that 0º, (()) = (z)) = 0 for any E W C V. This contradicts the fact that 0z) * 0 for any 2 E W.
Step 2. Secondly, we prove that q defined on Z convex, having a diagonal Hessian, has to be additive. Observe that by slightly reformulating the property of diagonal Hessian, we receive:
for any z, i and j: 0,(0,9(22) + - j=
(30)
Comparing Eq. 30 to Eq. 24 from the proof of Thm. 1b, we realize that the derivative of Dq(=)) has a blockdiagonal structure, similar to h (except, the blocks may become 0). By repeating the same argument from Step 2. of the proof of Thm. 1b, we get that Dq is siot ambiguity with the identity permutation, iec for some C' difffccomoophissms Qs we have 0zq(z) = Qx(z)) for any 2, k.
20

---

However, then let 2, 20) 2 and let ((t) [0, Z be smooth path with Φ(0) =000)
x (Z being convex, such path exists). Then:
L."
d [a( = Dq((((t) ('(t)dt
g(z) - -(=
(31)
de
Σ/ Qx (+)(((t)) (t())dt.
0 ( (t())dt=
(32)
Functions Qs are continuous; hence, they can give rise to an integral function Φo.. Using the rule of integration by substitution, we receive:
K
k
a(=) - 00=2((=)))))))) = E Sx (b(1))
Σ (Gu(zz) - (x(
(33)
k=1
Denoting <px(zx) (2c)) (1((2²/))) + x (0), conclude that q is additive,
K
q(z) (zx))
(34) M
k=1
A.8 DECODER GENERALIZATION
In this subsection we recall a more precise format and prove Thm. 2. However, before that, we also precisely introduce the slot-wise recombination space (Z') and -function (h'), where Z' is an extension of Eq.5. 5. The latter was only defined for the case when our autoencoder slot-identified on the training latent space.
Definition 13 (Slot-wise recombination space and -function). Let (2, 22, be POGP and let (g. f) be an autoencoder. Let 28 =()) We call
2' 2
2
(35)
the slot-wise recombination space, where 2 is the projection of the set 2 (c.f. Def. 8).
In the case when (g. f) slot-identifies on g w.r.t. generator f. let the slot-wise recombination function be the concatenation all slot-functions h& (z(()())):
h hz) (Ea(2)))), ... hx hx(2x(k)))-
(36)
The space of all values taken by h'(z) is:
h' (2) h1 (2(()())) x ... x hx (x))))
(37)
Since in this case he (z)) -/ we have that 2' = -h'(2).
Theorem 2b (Decoder generalization.. Let (Z, ff be a POGP such that
is compositional on and
ii is (2-diffcomorphism around Z.
Let f). be an autoencoder such that
iii) f).f) slot-identifies z 2 w.r.t. f and
iv) f is additive (on R
Then & generalizes in the sense that (f)( = f(2) holds for any z E Z. What is more, f is injective Z' and we get f(z') (22 = X.
21

---

f(z and 25 -((x)). Condition fii) implies that (g. f) mini- Proof of Thm. 2b. Let X mizes ()) or equivalently, because of Lem. 4, 9 : X + 2 and f 2 + x invert each other. Also, the projection of ~ to the k-th slot is 2% = (2)))) Furthermore, from Def. 13 we know that g(f(z)) h=);,, for any z E 2 Hence, by applying on both sides, we receive:
f(z) f()(=)) for any 2 E z
(38)
Besides, f is additive. Due to ii) and Lem. 7, f is also additive 2. More precisely: for any E [K] there exist functions sp: 2x + RRN, Ox h2 (2(()) - RN such that:
f(z) = Σ (*x())) for any E 2 and
(39)
k=1
f(*) Σ Sp (xx())) for any & e Z'.
(40)
Substituting this into Eq. 38, we receive:
k Σ (x))) = Σ Su (1'x(z)) = Σ Sh (Hu() for any z E 2
K
k
(41)
is easily seen that functions spx; Q/: are C After differentiating Eq. 41 with respect to žk. we receive:
Dep(zE)) = D (+)+-((k) o h (x)) (zx) for any E 2 Since g is a slot-supported subset of Z, Eq. 42 holds for any 2s E Zn. Let us denote
(42)
.k- = - ( (R) o 1x))(-) E (Zx).
Since 2 is convex, let z Zn fixed and define the path t E [0, u(t) xx +t(zx 20
Zz. Then:
(**))) - (2 (0)) (x((t))) (u(tt))
= - Dipx(u(t)) d((t)
(43)
Due to Eq. 42, we may continue:
(2x())) ))(( = /'' (T)(((()) of 7h (u(1)) (x()) 7h (2x)) (2()) (0))
(44)
Consequently, there exist constants ck RRM such that for any k:
(*x) = Φ ((k) (())))))xx))) + + for any 2か E Zx.
(45)
After adding them up for all A and using Eq. 39:
K
f(z) = Σ Sk (2x)
Σ
(p 100 (hr (()((())(22))) + cx
(46)
k=1
k=1
K
sbx (bx (na()2(()))
Σ
Ck.
(47)
Now, based on Eq. 40, we receive:
K f(z) = f(h ( (x(()()))) +
K
c - f(h(x))))) +
Σ Ck
(48)
k=1
1=1
holds for any z E Z. In particular, Eq. 48 holds for z E z which together with 38 implies that
K
Σ
C₁ = 0.
k=1
22

---

Finally, arrive at:
f(z) = f(1(2z)) for any 2 E Z.
(49)
Note that previously Eq. only held for 2
Moreover, since h' 2 2 is diffeomorphism, we see that
=(z)= = =
(50)
Remark. the process the proof, we have also proven the identifiability of the slot-wise addi- tive components up constant translations: There exists permutation E S(K) and for any k constants C₁ such that
(*x(2)) = 1 (hx + Ca for any 2 E kx.
(51)
A.9 ENCODER GENERALIZATION
Definition 6b (Compositional consistency). Let (g. be an autoencoder. Let Z' C R KM be an arbitrary set and let !=' be an arbitrary continuous distribution with supp(1z) = 2'. The following quantity is called the compositional consistency loss of (g. f) with respect to fz'
Ccons (g. (g. f.f.gg 9x/) Egiigg [l0(f((2)) -= x'112]
(52)
We say that (g. f) compositionally consistent if the compositional consistency loss w.r.t. 9x vanishes.
We now state lemma that, similarly to Lem. 4, states that for Lcons (g. f, 42x) to vanish, the exact choice of q is irrelevant and that in this case the decoder (right) inverts the encoder on Z', meaning that the autoencoder is consistent on f(2') to Z'.
Lemma (Vanishing Lcons implies invertibility on Z'). For (g. f) and z' C RRK the following statements are equivalent:
there exists q2 with supp(yz) = Z' such that froms (am((g)) S f, 0: ii) for any qx' with supp(pæ) = Z' we have Lcons of (. 0;
iii) 9 flz- id or, in other words. (g. f) is consistent on f(Z') to Z'.
Remark. The proof is analogous to the one of Lem. 4. Henceforth, in the context of vanishing con- sistency loss, we are going to denote Ccom:(). f. 92-) simply either by f.oms (g. f 2') or Leons (Z').
Theorem 3b (Compositionally generalizing autoencoder). Let (Z, Z5, f) be a POGP (Def. 9) such that
i) 2 is convex,
ii f is compositional and irreducible on Z and
ûi) is C'-diffcomorphism around Z.
Let (g.f) be an autoencoder with S = f(Z5) (Def. 10) such that iv) (g.f) minimizes Eie (g. f. x A f. Z') for some where Z' the slot-wise recombination space (see definition of Z' in Def. 13),
v) f is compositional on 2 := 5) and
vi) f is additive (on R.KM K)
Then the autoencoder (G. f) generalizes compasitionally w.r.t. 2 in the sense of Def. 3b. More- over, 9 X 2 = = inverts f 2' X and 2 2' hn (2((1))) x... hg (x)()))
23

---

Proof of Thm. 3b. Observe that assumption iv) is equivalent Cos f. x)) + connss (g. f' 0,
(53)
i.. (g. minimizes both which may happen if and only if Lre (g. f. x3) Loom (g. 2') = i.. free (2r5) and Leoms(2) (2').
Firstly, i) f is compositional on Z, and hence on 2², 1i) f is compositional on == g({
and i (g. (g. minimizes (x))))) we may conclude based on Thm. lb that (g. f) slot-identifies on z Consequently, the slot-wise recombination function h' well-defined.
Secondly, i) f is compositional on 2; ii) f E C²Diffeo(Z); iii) f is additive and iv) (g. f) slot-identifies z 2 Therefore, Thm. 2b implies that f in injective on Z', f('(zz) f()) holds for any 2 E Z and f(Z') = f( XX
Finally, from Lem. 8 and the fact that (g. f) minimizes Leons (Z'), we deduce that (g. f) is consis- tent on f(z"), meaning that is injective f(z') X' and its inverse is f z' X. However, by definition, ((x) 2. Consequently, we proved that 2 Z'.
Furthermore, we recall that (h'(2)) = f()) holds for any # E Z. As f is invertible on 2 with inverse 9, may pre-apply 9 both sides and receive
g(f(z)) (z), for any E 2,
(54)
which proves that (g. f) also slot-identifies z on Z and concludes our proof.
B EXPERIMENT DETAILS
B. DATA GENERATION
The multi-object sprites dataset used in all experiments was generated using DeepMind's Sprite- world renderer (Watters et al., 2019). Each image consists of two sprites where the ground-truth latents for each sprite were sampled uniformly in the following intervals: x-position in J0.1, 0.9). y-position [0.2, 0.8), shape {0. corresponding to (triangle, square), scale [0.09, 0.22]. and color (HSV) in [0.05, 0.95] where saturation and value are fixed and only hue is sampled.
All latents were scaled such that their sampling intervals become equivalent to a hypercube [o, 1 (2 slots with 5 latents each) and then scaled back to their original values before rendering. The slot-supported subset 2 of the latent space was defined as a slot-wise band around the diagonal through this hypercube with width d 0.25 along each latent slot, i.e..
zZ
22) E [5] (z) =22) < V28
(55)
In this sampling region, sprites would almost entirely overlap for small 6. Therefore, we apply an offset to the x-position latent of slot ø2. Specifically, set it to (x + 0.5) mod 1, where the sampled x-position.
The training set and ID test set were then sampled uniformly from the resulting region, 2 while the OOD test set was sampled uniformly from Z \ Z Objects with Euclidean distance smaller than 0.2 in their position latents were filtered avoid overlaps. The resulting training consists of 100,000 samples, while the ID and OOD test set each consist of 5,000 samples. Each rendered image is of size 64 64 3.
B.2 MODEL ARCHITECTURE AND TRAINING SETTINGS
The encoder and decoder for the additive autoencoder used in Sec. 6.1 closely resemble the archi- tectures from Burgess et al. (2018). The encoder consists of four convolutional layers, followed by four linear layers, and outputs vector of dimensions 2 x h, where h represents the latent size of a single slot and is set to 16. The decoder consists of three linear layers, followed by four trans- posed convolutional layers, and is applied to each slot separately; the slot-wise reconstructions
24

---

Figure 6: Samples from dataset used in Sec. 6. Left: In-distribution samples with latents sampled from the diagonal region. Objects are highly correlated in all latents and use an offset in their x- position to avoid direct overlaps. Right: Out-of-distribution samples with latents sampled from the off-diagonal region.
subsequently summed produce the final output. The ReLU activations replaced with ELU in both the encoder and decoder. The Slot Attention model in Sec. 6.2 follows the implementation from the original paper Locatello et al. (2020a), where hyperparameters were chosen in accordance with the Object-Centric library (Dittadi et al., 2021) for the Multi-dSprites setting, but with the slot dimension set 16.
The additive autoencoder is trained for 300 epochs, while Slot Attention is trained for 400 epochs. Both models are trained using a batch size of 64. We optimize both models with AdamW (Loshchilov and Hutter, 2019) with warmup of eleven epochs. The initial learning rate set as 10-7 and doubles every epoch until it reaches the value of 0.0004. Subsequently, the learning rate is halved every 50 epochs until it reaches 1 x 10- Both models were trained using PyTorch (Paszke et al., 2019).
For the experiments verifying our theoretical results in Sec. 6.1, we only included models in our results which were able to adequately minimize their training objective. More specifically, we only selected models that achieved reconstruction loss the ID test set less than 2.0. For the experiments with Slot Attention in Sec. 6.2, we only reported results for models that were able to separate objects ID where seeds were selected by visual inspection. This was done since the primary purpose of these experiments was not to see the effect of our assumptions on slot identifiability ID but instead OOD. Thus, we aimed ensure that the effect of our theoretical assumptions on our OOD metrics were not influenced by the confounder of poor slot identifiability ID. Using these selection criteria gave us between and 10 seeds for both models, which were used to compute our results.
If used, the composition consistency loss is introduced with À = 1 from epoch 100 onwards for the additive autoencoder model and epoch 150 onwards for Slot Attention. The number of recombined samples z' in each forward pass of the consistency loss is equal to the batch size for all experiments. In our compositional consistency implementation, normalization of both the latents z' and the re- encoded latents (=(')) proved to be essential before matching them with the Hungarian algorithm and calculating the loss value. Without this normalization, encountered numerical instabilities, which resulted in exploding gradients.
B.3 MEASURING RECONSTRUCTION ERROR
For the heatmaps in Figs. and 5, we first calculate the normalized reconstruction MSE on both the ID and OOD test sets. We then project the ground-truth latents of each test point onto a 2D plane with the and y-axes corresponding to the color latent of each object. We report the average MSE in each bin. this projection, some OOD points would end up the ID region. Since we do not observe difference in MSE for OOD points that projected to the ID OOD region, we simply filter those OOD points to allow for a clear visual separation of the regions.
When reporting the isolated decoder reconstruction error in Fig. A and Fig. 5 A and B, we aim to visualize how much of the overall MSE can attributed to only the decoder. Since the models approximately slot-identify the ground-truth latents ID on the training distribution, the MSE of the full autoencoder serves as a tight upper bound for the isolated decoder reconstruction error. Thus, in the ID region, use this MSE when reporting the isolated decoder error. For the OOD region,
25

---

however, the reconstruction error could be attributed to failure of the encoder infer the correct latent representations or to a failure of the decoder in rendering the inferred latents. To remove the effect of the encoder's OOD generalization eπror, we do the following: For a given OOD test image, we find two II) test images that each contain one of the objects the OOD image in the same configuration. Because the encoder is approximately slot identifiable ID,, know that the correct representation OOD for each object in the image is given by the encoder's ID representation for the individual objects in both ID images. To get these ID representations, we must solve matching problem to find which ID latent slot corresponds to given object in each image. We do this by matching slot-wise renders of the decoder with the ground -tunh slot-wise renders for each object based on MSE using the Hungarian algorithm (Kuhn, 1955). Using this representation then allows us to obtain the correct representation for an OOD image without relying on the encoder to generalize OOD. The entire reconstruction error on this image can thus be attributed to failure of the decoder to generalize OOD.
B.4 COMPOSITIONAL. CONTRAST
The compositional contrast given by Brady et al. (2023) is defined as follows: Definition 14 (Compositional Contrast). Let f : Z + X be differentiable. The compositional
contrast of f at is
K
I5n Ozk
05n Ozj
Coomp (S,mp(112)
(5,z) n=1
(z)
(z)
(56)
This contrast function was proven to be zero and only if f is compositional according to Def. 4. The function can be understood as computing each pairwise product of the (L2) norms for each pixel's gradients w.r.t. any two distinet slots l * j and taking the sum. This quantity is non-negative and will only be zero if each pixel is affected by most one slot such that f satisfies Def. 4. We use this contrast function to measure compositionality of a decoder in our experiments in Sec. 6. More empirical and theoretical details on the function may be found in Brady et al. (2023).
B.5 COMPUTATIONAL COST OF THE PROPOSED CONSISTENCY LOSS
Computing the consistency loss illustrated in Fig. 3 poses some additional computational over- head. Namely, it requires additional passes through the encoder and decoder as well as computation of the encoder's gradients w.r.t. the loss. As outlined in Sec. 4, we computed the consistency loss on samples z' obtained by randomly shuffling the slots in the current batch, effectively doubling the batch size. We found this to increase training time by maximum of 28 % across runs.
For two slots, it would be possible to sample up to b(b - 1) novel combinations of the slots within the given batch, where b is the batch size. This number increases combinatorially to (b-nj! for n slots, which could pose severe computational overhead. While our experiments demonstrate that the loss works well with just samples, App. C.3 also shows that does not scale well more than two slots, and drawing more samples might be a way to remedy this.
On the other hand, there might ways to draw samples more effectively. One approach could be to sample slot combinations in proportion to their value w.r.t. the consistency loss to sample combinations according to their likelihood under prior over latents. Using a likelihood could avoid sampling implausible combinations; however, such a scheme is challenging it relies on the likelihood being valid for OOD combinations of slots. Another possibility would be to include heuristics to directly filter combinations based on a priori knowledge of the nata-gennrattinn process, e.g., to filter objects with similar coordinates which would intersect.
B.6 USING MORE REALISTIC DATASETS
Extending our experiments from the sprites dataset to more realistic data poses two main challenges. Firstly, in real-world datasets one generally does not have access to the ground-truth latents making our evaluation schemes inapplicable. Secondly, even if access to ground-trun latent information is available, experiments require being able to sample latents densely from a slot-supported subset.
26

---

without feons
with Conns
1.0
Reconstruction R OOD
Identifiabiln Score OOD
0.8
0.8
0.6
0.6
0.4
0.4
0.2
0.16
0.2
0.25
0.4
0.8
0.16
0.2
0.25
0.4
0.8
slot-wise band width 6
slot-wise band width 6
Figure 7: OOD Reconstruction quality and slot identifiability function of training region size. The width (0, 1) of the slot-wise band around the diagonal (see Eq. 55) if 2 is a line and 1 if 2 2; experiments in Sec. 6 used 8 0.25. In the absence of the consistency loss (Def. 6), a large is required for models to achieve high OOD performance for reconstruction and slot identifiability. In contrast, models trained with the consistency loss yield consistently high performance across all Results are averaged over at least four random seeds.
Specifically, our experiments rely sampling from a diagonal strip in the latent space with small width. If such region were sub-sampled from an existing dataset, this would leave only very small number of data points which are insufficient to train model. To this end, our experiments require access to the ground-truti renderer for dataset such that latents can be sampled densely. This is not available in most cases, however. interesting avenue address this would be to leverage recent rendering pipelines such as Kubric (Greff et al., 2022) to create more complex synthetic datasets. We leave this as an interesting direction for future work.
C ADDITIONAL EXPERIMENTS AND FIGURES
This section provides additional experimental results to the main experiments from Sec, 6.
C.I IMPACT OF TRAINING REGION SIZE
We ablate the impact of the size of slot-supported subset on OOD metrics in Fig. 7 by varying the width of the slot-wise band around the diagonal (see Eq. 55). All models use an additive decoder and differ in whether they optimize the consistency loss (Def. 6). We see that models which do not optimize the loss require an increasingly large order to achieve high OOD reconstruction and identifiability scores, while models which do optimize the loss, achieve consistently high scores on OOD metrics across all values of 8.
C.2 VIOLATING SLOT-SUUPPORTEDNESSS
We illustrate the effect of violating the assumption that g is slot-supported subset Z (recall Def. I) in Fig. on reconstruction Ioss for additive autoencoder trained with consistency loss. To do this, we create a gap in the slot-supported subset z by removing all occurrences of objects with hue-latent in the interval (0.5, 0.8). We can in Fig. 8 that this leads to poor reconstruction performance in the region containing the gap which propagates to the OOD regions as well.
C.3 IMPACT OF MORE THAN 2 OBJECTS ON CONSISTENCY LOSS
We also examine how the consistency loss scales as the number objects in the training data is increased from two three and four. We find that the number of objects grows, optimization of the consistency loss becomes more challenging (Fig. 9, bottom left) which makes sense considering that the number of possible slot combinations grows combinatorially with the number of slots. This,
27

---

Slot-supported
Support with gaps
MSE
1.8
8.8
training distribution
slot 2
8.6
0.4
8.2
slot
slot
Figure 8: Violating slot-supportedness prevents OOD generalization. We retrain the additive autoencoder from Sec. 6.1 on data arising from latent subset 2 in which all occurrences of objects with a hue-latent in the interval (0.5, 0.8) have been removed. This leads to a cross-shaped gap in the latent space (see App. B.3
the latent subset g which can be visualized via 2D projection for details) where the x- and y-axes correspond to the hue-latent of each object (right). Compared to a model trained without this gap (left), we can see that the model unable to reconstruct ID samples in this gap (i.e., samples where either object has this hue), and this error propagates outward to OOD samples.
1.0
without Coos with Care
Reconstruction re COD
Reconstruction 7P ID
0.6
1.0
Identifubility Score 00D
0.8
3-
0.6
0.4
0.2
0.0
Number
Number Slots
Figure 9: Impact of number of objects on consistency loss. We measure how the consistency loss scale as the number of objects in the training data is increased from two to three and four. We mea- sure this for additive autoencoders which explicitly optimize the consistency loss (red) and models which do not optimize (blue). Top and Bottom left: We high as the number of objects grow, but the consistency loss increases steeply across models. Top right: Consequently, the OOD reconstruction quality decreases the number of objects increases. Bottom right: This then prevents the encoder from slot-identifying the ground-truth latents OOD. However, training with the consistency loss still yields generally better results than training without it. All results are averaged over at least four random seeds. The consistency loss is normalized by number of slots, and Rt2 scores are clipped to zero.
that the ID reconstruction remains
in turn, leads to poor OOD reconstruction quality (Fig. 9, top right) which prevents the encoder from slot-identifying the ground-truth latents (Fig. 9, bottom right). As hypothesized in App. B.5, more principled schemes for sampling slot combinations could mitigate these scaling issues.
28


Reviews:
[{"review_id": "Cqyjwk0Hak", "reviewer": "Reviewer_KGuC", "summary": "The paper addresses the problem of compositional generalization in object-centric autoencoders.\nThe authors formalize this as requiring the model to identify the ground-truth object latents not just on the training distribution, but also on out-of-distribution combinations.\nThey make two key assumptions to achieve this: (1) The generative process satisfies compositionality, meaning each pixel depends on one object, and irreducibility, preventing objects from being decomposed.\n(2) The decoder is additive, decoding each object slot independently.\nUnder these assumptions, the authors prove autoencoders can identify objects in-distribution by minimizing reconstruction error.\nThe additive decoder then guarantees generalization out-of-distribution.\nHowever, the encoder may still fail to generalize.\nTo address this, the authors propose compositional consistency regularization.\nThis trains the encoder to invert the decoder on recombined object slots, enabling the full autoencoder to generalize.\nBy combining in-distribution identifiability and compositional consistency regularization, the authors prove autoencoders satisfying their assumptions will generalize compositionally.\nThrough synthetic experiments, they provide empirical evidence supporting their theoretical results.\nIn particular, they demonstrate the importance of additivity and compositional consistency for generalization.", "questions": "Please see above.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "This paper made contributions for \n- Formalizing compositional generalization as an identifiability problem\n- Theoretical guarantees for in-distribution identifiability\n- Showing an additive decoder enables out-of-distribution generalization\n- Introducing compositional consistency regularization\n- Providing overall theoretical guarantees for compositional generalization\n\nThe work makes theoretical progress on understanding compositional generalization in object-centric representation learning.", "weaknesses": "- The assumptions of compositionality and irreducibility are quite restrictive. Most real-world datasets likely violate these. \n- The additive decoder limits modeling of complex object interactions and relations.\n- The consistency regularization implementation requires sampling implausible object combinations. More principled schemes could improve results in complex environments.\n- Experiments only validate the theory on simple synthetic datasets. Testing on more diverse and realistic data would better demonstrate applicability, though the evaluation would also be more challenging.\n- The proposed methods, especially when ensuring encoder-decoder consistency and handling latent slots, might pose scalability issues for very large datasets or more complex models. A discussion on the scalability, computational costs, and potential solutions would make the paper more robust.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ClvWuLVjqS", "reviewer": "Reviewer_wPGp", "summary": "This paper presents conditions where compositional generalization is theoretically guaranteed for object-centric learning. \n\nSpecifically, they first extend the identifiability theory of object-centric representations to handle partial joint distribution supports, with an additional assumption/constraint on the decoder to be compositional. This ensures that slots are identifiable in the training distribution. They then ensure the generalizability of decoders (which, e.g., generate images given slot representations) with another assumption/constraint as the decoder being additive. The theoretical analysis is similar to those proving the compositional generalizability of any additive inference models.\n\nThe novel step is to enforce the compositional generalizability of encoders by learning with the synthesized data in new compositions of latent slots/symbols given the generalizable decoder. So, in order to learn an encoder that can generalize to unseen combinations of objects, they first build a dataset with new compositions of latent symbols/slots by permutating learned latent symbols/slots in the training distribution. They then generate the fake images using the \"supposedly generalizable\" decoders on new combinations. The encoder is trained to learn the inverse mapping of the decoder. This process is formulated as a compositional consistency regularization loss in practice.\n\nThe experimental results are aligned with the theories in a simple two-object synthetic image environment.", "questions": "Are there results in more complex environments? \n\nCan the theories be generalized to more general settings with weaker assumptions?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "This paper discusses an important problem: learning compositionally generalizable object-centric representations. The paper is well-written and easy to read. The connections with related works are also interesting and inspiring. \n\nThe reviewer especially appreciates the theoretical guarantees and analysis. Even though the assumptions are strong on both the functions to be approximated as well as the parameterization of learned functions, they are still aligned with the image object-centric representation learning setting, and the methods can be relaxed and realized using modern object discovery methods such as slot attentions. \n\nThe proposed regularization loss to enforce the compositional generalizability of encoders is interesting and seems easy to use. \n\nThe ablation study on the additive decoder (softmax v.s. sigmoid in slot attentions) is interesting and inspiring.", "weaknesses": "It would be great if the assumptions could be relaxed, e.g., to handle occluded objects or to handle general latent variable learning domains other than the image objects. \n\nThe \"contemporary\" work [1] discussed most parts of this paper except for the generalizable encoder. \n\nThe experimental environment is simple with two-object synthetic images. It would be more convincing to see results on multi-object real images. \n\n[1] S ́ ebastien Lachapelle, Divyat Mahajan, Ioannis Mitliagkas, and Simon Lacoste-Julien. Additive decoders for latent variables identification and cartesian-product extrapolation. arXiv preprint arXiv:2307.02598, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jbfwEpmJNf", "reviewer": "Reviewer_JA2G", "summary": "The authors theoretically and empirically show that compositional generalization can be achieved through:\n1. Structural constraints on the decoder (each data dimension is rendered as the sum of functions operating on slots separately), which ensures that the decoder compositionally generalizes.\n2. An encoder-decoder consistency loss (reconstruction loss for the representations) on slot-shuffled representations from the encoder output, which encourages the encoder to compositionally generalize with the additive decoder.\n\nThe paper provides a joint encoder-decoder framework for compositional generalization for autoencoders, where previous work has mostly focused on specific aspects of the setting.", "questions": "- How does the effect of the consistency loss scale with the number of slots?\n- What is the impact of how slot-supported the training data is? i.e. in Figure 2 (1), what is the impact of the width of the blue band on empirical effectivity?\n- How does the method hold up on non-synthetic data, especially if you slightly relax some constraints? For instance, what if you have expressive slot-wise decoding, but allow for a low-expressivity non-linear combination at the end for rendering?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well-written.\n- The theory is sound and significant for the community.\n- The joint encoder-decoder framework for compositional generalization in autoencoders is quite elegant.\n- The limitations of the framework and the additivity constraint on the decoder are adequately stated.", "weaknesses": "Although they support the theory, the experiments are quite limited. For instance, these are all with only two slots with 16 dimensions each. See the questions section for additional information that would be interesting to see from experimentation.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (oral)