Object centric architectures enable efficient causal representation learning
Authors: Amin Mansouri
Jason Hartford
Yan Zhang
Yoshua Bengio
Venue: ICLR 2024 poster (2024)
Keywords: object centric learning
representation learning
disentanglement
weakly supervised learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: Existing methods for causal representation learning fail on objects. Object-centric learning can fix it.

Abstract: Causal representation learning has showed a variety of settings in which we can disentangle latent variables with identifiability guarantees (up to some reasonable equivalence class). Common to all of these approaches is the assumption that (1) the latent variables are represented as $d$-dimensional vectors, and (2) that the observations are the output of some injective generative function of these latent variables. While these assumptions appear benign, we show that when the observations are of multiple objects, the generative function is no longer injective and disentanglement fails in practice. We can address this failure by combining recent developments in object-centric learning and causal representation learning. By modifying the Slot Attention architecture (Locatello et al., 2020), we develop an object-centric architecture that leverages weak supervision from sparse perturbations to disentangle each object's properties. This approach is more data-efficient in the sense that it requires significantly fewer perturbations than a comparable approach that encodes to a Euclidean space and we show that this approach successfully disentangles the properties of a set of objects in a series of simple image-based disentanglement experiments.

Published as a conference paper at ICLR 2024
OBJECT-CENTRIC ARCHITECTURES ENABLE EFFICIENT CAUSAL REPRESENTATION LEARNING
Amin Mansouri"
Jason Hartford' Valence Labs
Yan Zhang
Mila, Quehec Al Institute
Samsung- -AAT AI Lab, Montreal
Yoshua Bengio
Mila, Quebec Al Institute /
U. Montreal
/ CIFAR
ABSTRACT
Causal representation learning has showed variety of settings in which we can disentangle latent variables with identifiability guarantees (up to some reasonable equivalence class). Common to all of these approaches is the assumption that (1) the latent variables are represented as d-dimensional vectors, and (2) that the observations are the output of some injective generative function of these latent variables. While these assumptions appear benign, we show that when the observations are of multiple objects, the generative function is longer injective and disentanglement fails in practice. We can address this failure by combining recent developments in object-centric learning and causal representation learning. By modifying the Slot Attention architecture (Locatello et al., 2020b), we develop object-centric architecture that leverages weak supervision from sparse perturbations to disentangle each object's properties. This approach is more data-efficient the sense that requires significantly fewer perturbations than a comparable approach that encodes to Euclidean space and we show that this approach successfully disentangles the properties of a set of objects series of simple image-based disentanglement experiments.
INTRODUCTION
Consider the image in Figure (left). We can clearly see four different colored balls, each at different position. But asking. "Which the first shape? And which the second?" does not have a clear answer: the image just depicts an unordered set of objects. This observation seems trivial, but it implies that there exist permutations the objects which leave the image unchanged. For example, could swap the positions of the two blue balls without changing single pixel in the image.
In causal representation learning, the standard assumption that our observations x are "rendered" by some generative function g(-) that maps the latent properties of the image to pixel space (i.c. x = g(=));; ); the goal is to disentangle the image by finding an "inverse" map that recovers 2 from x up some irrelevant transformation. The only constraint g(-) that is assumed by all recent papers (for example Hyvarinen & Morioka, 2016; 2017: Locatello et al., 2020a; Khemakhem et al., 2020a;b; Lachapelle al., 2022; Ahuja al., 2022a;b; 2023), that g(-) is injective', such that g(1)) g(22) implies that 21 = 22. But notice that if we represent the latents 2 as some d-dimensional vectors in Euclidean space, then whenever we observe objects like those shown in Figure 1, this injectivity assumption fails: symmetries in the objects' pixel representation imply that there exist non con-trivial permutation matrices II, such that g(z) = (TTz).. This is not just a theoretical inconvenience: Figure (right) shows that when the identity of the balls is not distinguishable, the disentanglement performance of recent approach from Ahuja et al. (2022b) is upper-bounded by 1/k where & is the number of balls.
correspondence
n. mansouri @mi a quebecceecc...... or jason@va lence Tabs com
'Some papers place stronger constraints on g(-), such as linearity Hyvärinen & Oja, 2000; Squires et al., 2023, sparsity Moran et al., 2022; Zheng et al., 2022, or constraints g's Jacobian Gresele et al., 2021; Brady al., 2023 but injectivity is the weakest assumption common all approaches.

---

Published as a conference paper at ICLR 2024
1.0
MCC scoto
0.8
0.6
0.4
Injoctive
0.2
Non-injective ResNet
Disentangled Slot Attention
0.0
3
Number of balls
Figure 1: (Left) An example image of simple objects. (Right) Mean correlation coefficient (MCC) score which measures the correlation between inferred latent variables and their associated ground truth values. Ahuja et al. (2022b)'s 's approach achieves almost perfect MCC scores (i.e. score 1) when the ball color is used to make the generative function injective ("Injective ResNet"), but achieves an MCC score of at most where l is the number of objects when colors are selected randomly ("Non-injective ResNet"). We show that is possible to recover the injective performance by disentangling object-centric representations ("Disentangled Slot Attention").
In parallel to this line of work, there has been significant progress in the object-centric learning literature (c.g. van Steenkiste et al., 2018a; Goyal et al., 2019; Locatello et al., 2020b; Goyal et al., 2020; Lin et al., 2020; Zhang et al., 2023) that has developed a suite of architectures that allow us to separate observations into sets of object representations. Two recent papers (Brady et al., 2023; Lachapelle et al., 2023) showed that the additive decoders used in these architectures give rise to provable object-wise disentanglement, but they did not address the task of disentangling the objects' associated properties. In this paper we show that by leveraging object-centric architectures, we effectively reduce the multi-object problem to a set of single-object disentanglement problems which not only addresses injectivity failures, but also results perturbations we need to observe to disentangle properties using Ahuja et al. (2022b))ss 's approach. We illustrate these results by developing a property disentanglement algorithm that combines Zhang et al. (2023)'s SA-MESH object-centric architecture with Ahuja et al. (2022h)'s ss approach to disentangle- ment and show that our approach is very effective at disentangling the properties of objects on both 2D and 3D synthetic benchmarks.
significant reduction the number of
In summary, we make the following contributions:
We highlight two problems that arise from objects which violate standard assumptions used to identify latent variables (Section 3).
We show that these problems can be addressed by leveraging object-centte architectures, and that using object-centric architectures also enables to use factor of k fewer perturbations to disentangle properties, where & the number of objects (Section 4).
We implement the first object-centric disentanglement approach that disentangles object proper- ties with identifiability guarantees (Section 5).
We achieve strong empirical results² on both 2D and 3D synthetic benchmarks (Section 7),
BACKGROUND
Causal representation learning (Schölkopf et al., 2021) seeks reliably extract meaningful latent variables from unstructured observations such as images. This problem is impossible without additional structure because there are infinitely many latent distributions p(2) that are consistent with the observed distribution, p(z) f p(x) 2))p(z), only one of which corresponds to the ground truth distribution (Hyvärinen & Pajunen, 1999; Locatello et al., 2019). We therefore need to restrict the solution space either through distributional assumptions through assumptions on the functional form of the generative function g 2 - X' that maps from the latent space to the observed space (Xi & Bloem- Bloem-Reddy, 2023). A key assumption that (to the
the form of the latent distribution p(z).
"The code reproduce our results can be found http://ghhhaaccmmmmmmmsmrrrrooo
2

---

Published as a conference paper at ICLR 2024
best of our knowledge) is leveraged by all papers that provide identifiability guarantees, is that g(-) is injective such that if we see identical images, the latents are identical (i.e. if g(21) = (22) then 21 22).
Given these restrictions, we can analyze the identifiability of latent variables for a given inference algorithm by considering the set of optimal solutions that satisfy these assumptions. We say latent variables are identified if the procedure will recover the latents exactly in the infinite data limit. Typically, some irreducible indeterminacy will remain, so latent variables will be identified up to some equivalence class A. For example, if the true latent vector is and we have an algorithm for which all optimal solutions return linear transformation of such that, A {A : A2]. then we say the algorithm is linearly identifies latent variables. We will call latent variables disentangled if the learning algorithm recovers the true latent variables up to a permutation (corresponding to relabeling of the original variables), and element-wis transformation. That is, for all i, = h, (z)())) where t is a permutation, and h-(-) is an element-wise function; for the results we consider in this paper this function simply a scaling and offset, fi(2) = a; + d, corresponding to a change of units of measurement and intercept.
In this paper, we will build on recent line of work that leverages paired samples from sparse perturbations to identify latent variables (Locatello et al., 2020a; Brehmer et al., 2022; Ahuja et al., 2022b). Our approach generalizes Ahuja et al. (2022b) to address the non-injectivity induced by objects, so we will briefly review their main results. Ahuja et al. assume that they have access to paired samples, (z, x²) where x = g(z), x' = g(=),, and z; is perturbed by a set of sparse offsets A = (51..... .......... 5x }. such that 2/ 2, + 6, for all i (1....,k) k... They show that if g(-) is injective analytic function from Rd - X, every 8 E A is I-sparse, and at least d linearly independent offsets are observed, then an encoder, that minimizes the following objective recovers the true z up to permutations, scaling and an offset (Ahuja et al., 2022b, Theorem 1),
f E arg min p Baad' [1r((s) + ((()))
(())))
f(x)= = = II +c
where II is a permutation matrix, A is an invertible diagonal matrix and c is an offset.
3 OBJECTS RESULT IN NON-IDENTIHABILITY
We begin by formally characterizing the challenges that arise when images contain multiple objects.
Data generating process. We assume that set Z == { of & objects is drawn from some joint distribution, Z In order to compare set and vector representations, let vec, x (Z) denote a flattened vector representation of Z ordered according to some permutation group of permutations of objects; when π is omitted, vee(2) simply refers to an arbitrary default ordering (i.e. the identity element of the group). Each object is described by d-dimensional vector of properties* 2; E Rd and hence vec Z)) E Aad We say objects have shared properties if the coordinates of z; have consistent meaning across objects. For example, the objects in Figure I (left), each have x, y coordinates and a color which can be represented by its hue, so 2; = Ip pý. ht".. In general, the set of properties associated to an object can be different across objects, but for simplicity, our discussion will focus on properties that are fully shared between all objects.
Sym(k). the symmetric
The non-injectivity problem. We observe images x which are generated via a generative function g(-) that renders set of object properties into a scene in pixel space, such that x g(Z). While B(-) is a set function, we can define an equivalent vector generative function, g. which, by definition, produces the same output g (Z); i.e. for all Sym(k), g (Z)) = g(Z). This generative function g taking vectors as input is consistent with standard disentanglement assumptions except that it is not injective:
Proposition If gree (Z)) = (Z)) for all π E Sym(k), then g(-) is not injective.
'A natural extension of the perspective we take in this paper is to also treat properties as sets rather than ordered vectors; for example, see Singh et al. (2023). We leave understanding the identifiability of these approaches future work.
3

---

Published as a conference paper at ICLR 2024
Proof. The contrapositive of the definition of injectivity states that 21 * 22 implies ((21) + g(22). but by definition g()), ), there exist 21 & 22 such that g(21) = g ). In particular, for any set Z and permutations T2 E Sym(k), the vectors vec., (Z) 21 * 22 vecr. (Z).
This proposition simply states that if images are composed of sets of objects, then if we model the generative function as a map from a Euclidean space, this map will not be injective by construction.
With the exception of Lachapelle et al. (2023), all of the causal representation learning papers cited in section 6 assume the generative function is injective. To see why injectivity necessary general, consider an image with two objects. If the two objects are identical, then there are two disentangled solutions corresponding to the two permutations, so it is not possible to identify a unique solution.
The object identity problem. When applying sparse perturbations on Z (see section 2), we are effectively perturbing one coordinate of one object. However, how multiple possible objects in Z we have perturbed? In the case of injective mappings. this is simple: since there is consistent ordering for them, we know that coordinate in vee (ee corresponds to the same object before and after the perturbation.
know which object of the
However, this is no longer the case in setting. Since the objects are actually part of set, cannot rely on their ordering: the perturbed object can, in principle, freely swap order with other objects; there is no guarantee that the ordering before and after the perturbation remains the same. In fact, we know that these ordering changes mus/ be present due to the responsibility problem:
Proposition (Zhang et al. (2020); Hayes et al. (2023)). If the data is generated according to the data generating process described above with ( == g(Z) and k >> thhn f(-) is discontinuous.
Proof Sketch. Consider Figure 2. notice that if we perform 90% rotation in pixel space of the image, the image is identical, but the latent space has been permuted, since each ball has swapped positions. Because the image on the left and the image on the right are identical in pixel space, any encoder, f: X + Rd will map them them to identical latents. There exists continuous pixel-space rotation from 09 to 90", but it must entail a discontinous swap in which latent is responsible for which part of pixel-space according the encoder.
A general proof can be found in Hayes et al. (2023). These discontinuities manifest themselves as changes in permutation from one vec,., (2) to another v (Z). In disentanglement approaches that leverage paired samples (e.g. Ahuja al., 2022b; Brehmer et al., 2022), continuity enables the learning algorithm to implicitly rely on the object identities to stay consistent. Without continuity, one cannot rely on the fact that vec(Z) (Z) and vec + Ō should be the same up the perturbation vector ó, because the perturbation may result in a discontinuous change of vec () + 8 when an observation encoded back to latent space. As a consequence, we lose track of which object we have perturbed in the first place, so naïve use of existing disentanglement methods fails.
Another challenge is that the encoder f (Equation 1) has to map observations to vec(Z) in discontinuoss way, which is traditionally difficult to model with standard machine learning techniques.
In summary, the unordered nature of objects in Z results in non-injectivity, losing track of object identities, and the need for learning discontinuous functions. These all contribute to the non- identifiability of traditional disentanglement methods in theory and practice.
4 OBJECT-CENTRIC CAUSAL REPRESENTATION LEARNING
A natural solution to this problem is to recognize that the latent representations of multi-object images are sets and should be treated as such by our encoders and decoders in order to enforce invariance among these permutations. Both Brady et al. (2023) and Lachapelle et al. (2023) showed that architectures that enforce an appropriate object-wise decomposition in their decoders provably disentangle images into object- blocks of latent variables. These results do not disentangle the properties of objects, but they solve an important precursor: the assumption that there exists an object-wise decomposition of the generative function is sufficient to partition the latents into objects.
4

---

Published as a conference paper at ICLR 2024
x = (((2) Roi Ray Rg))
X Za.
Driginal
Parmuledd:
Figure 2: An illustration of the object identity problem. Permuting the order of the latents [21, 22, 23, 23] is equivalent a 90 degree rotation in pixel-space.
Like these two papers, we will assume that natural images can be decomposed into objects,* each of which occupies a disjoint set of pixels. When this is the case, we say that an image is object-separable. To define object separability formally,
will need to consider partition P of an image into
disjoint subsets of pixels P = {x(1)) ..... (x)) indexed by an index set Ip = [I.....kk] k): further, denote an index set that indexes the set latent variables as Iz. We can then say,
image, I, object-separable if there exists an object-wise partition and bijection Ip + Iz that associates each subset of pixels in P with a particular element of the set of latents, 2,, such that each subset of pixels 100) P is the output of an injective map with respect to its ussociated latent Zo(s)- That is, for all i, (x Co)Z)), x0 C ((Z)), we have that x00 (8( implies 2ax(a) Za(i))
Definition
This definition says that an image can be separated into objects if it can be partitioned into parts such that each part is rendered via an injective map from some latent 2,. We can think of each 209) as a patch of pixels, with bijection a that relates each of the k patches of pixels in the partition (x() 1 to latent variable in Z {= Each patch "depends" on its associated latent via an injective map.
Brady et al. (2023) and Lachapelle et al. (2023) give two different formal characterizations of partitions P that are consistent with our object-wise definition. Brady et al. 's characterization requires that a differentiable generative function is compositional, in the sense that each x(() E P only functionally depends on a single 2j E Z, and irreducible in the sense no x:(4) E P can be further decomposed into non-trivial subsets that have functionally independent latents. Lachapelle et al. 's assumption is weaker than ours in that they only require that the generative function is defined as g(Z) <E,, (z())) where is an invertible function, and that is a diffcomorphism that is "sufficiently nonlinear" (see Assumption 2 Lachapelle et al., 2023); object-separabl images are special case with a as the identity function and each 9x(-) rendering a disjoint subset of a, and hence their results apply to our setting.
Disentangling properties with object-centric encoding. In section we showed that the assump- tions underlying sparse perturbation-banl based disentanglement approach violated in multi-object scenes. But, the results from Brady et al. (2023) and Lachapelle et al. (2023) show that the objects can be separated into disjoint (but entangled) sets of latent variables. This suggests natural approach to disentangling properties in multi-object scenes:
we can reduce the multi-object disentanglement problem to a single-object problem with an object-wise partition of the image. Within each patch of pixels z+() E P injectivity holds, and we no longer have multiple solutions patch level. This partition identifiable and
*This a pragmatic approximation that suffices for the purposes objects is far more subtle because what we interpret as an "object" often depends on a task or choice of hierarchy; for more nuanced treatment, Smith (2019)'s Chapter is an excellent introduction into the subtleties around demarcating "object".
paper. but careful treatment
Functional dependence is defined non-zero partial derivatives. i.e.
Oxd # 0.
8ay
5

---

Published as a conference paper at ICLR 2024
use object-centric architecture to learn the object-wise partition. We require that this object-centric architecture can handle the responsibility problem.
we leverage Ahuja et (2022b)'s approach to using weak supervision to disentangle the properties of each object individually. Since we assume that properties between objects are shared, this requires factor of k fewer perturbations in the perturbation set A, where the number of objects.
we address the object identity problem where we lose track of object identities after perturbations through an explicit matching procedure that re-identifies the object being perturbed.
See section for details of how we implement this. This approach not only addresses the challenges outlined in Section 3, it also significantly reduces the number of perturbations that we have to apply in order to disentangle shared properties.
Theorem (informal). If data generating process outputs observations with k objects that have shared properties, then an object-centric architecture of the form F(x) {f(z!0) Jate where P is an object-wise partition and :
X' - Rdd will disentangle in k times fewer perturbations than an encoder of the form X Rkdd
The proof is given in Appendix A. The main insight is that if we have object-centric architecture that learns an object-wise partition P and uses the same encoding function on every patch, then every perturbation provides weak supervision to every object, despite the fact that only one was perturbed. As a result, we do not need to disentangle each object's properties separately, and hence reduce the number of required interventions by a factor of k..
METHOD
Object-wise partitions There exist number of ways to decompose image into objects, but for our purposes, pixel segmentation-bass approaches (Greff et al., 2019; Locatello et al., 2020b; Zhang al., 2023) let us directly adapt existing disentanglement techniques to work with object-centric encoders. A pixel segmentation encoder f maps from images x to a set of slot vectors {sI..... } each of which depends on a subset of the pixels 2(() E P. Images are then reconstructed using slot decoder ' that maps from the set of slot representation back to pixel space. The dependence between slots and patches of pixels is typically controlled by soft-attention matrix, which will typically not result in a partition of the pixels. In our implementation, we Zhang et al.' 's SA-MESH modification of the original Locatello et al. slot attention architecture, which adds an entropy regularization term to learn sparse attention matrices that do approximately partition the input by encouraging the subsets of pixels (16 to be disjoint (for details on the architectures, see Appendix B). Importantly for us, Zhang et al. (2023) is exclusively multiset- -ulliset--quivariint (Zhang et al., 2022), which allows to model discontinuous functions, thus handling the responsibility problem.
Slot attention is usually trained with reconstruction loss from relatively high- high-dimensional per-object slot representations, dimensional latent description (in the simplest case, just the two dimensions representing the (æ, y) coordinates of each object). To disentangle these high- nigh-dimensional slot representations, we simply add projection head, p s; > 2,, that is trained by a latent space loss.
R,, but for the images that work with, we want relatively low
Disentanglement via weak supervision with matching Ahuja al. assume access to pairs of images (x.x²) that differ by sparse offset 5. They enforce this assumption via disentanglement loss that requires that the latent representations of this pair of images differ by 6, such that f(x) + 8 f(z'). When using slot attention architecture, we introduce matching step to the loss infer the object to which the offset 5 was applied. With 1-sparse ó vectors, the matching step reduces to simple minimization over a cost matrix that measures || (f) (2(x:()) -)1l for all pairs of slots j. In Appendix D, we provide a more general matching procedure that applies to settings with dense offsets 8. We jointly optimize the following reconstruction and disentanglement loss,
f. J.9.P E arg min Ea -000--gf(fo (( E. smmin (m - ()( (2)
f.9.P

---

Published as a conference paper at ICLR 2024
The first term in this loss enforces that the encoder / decoder pair f.9 capture enough information in the slot representations s; to reconstruct æ. The second term contains the matching term and ensures that the function that projects from slot representation to latents disentangles the slot representations into individual properties. The offset could be known or unknown to the model, and for the remainder of this paper, we focus on the more challenging and natural case of unknown offssess. See appendix C for more details.
RELATED WORK
Causal representation learning Our work builds on the nascent field of causal representation learning (Schölkopf et al., 2021). In particular, disentanglement approaches builds on ideas in Ahuja et al. (2022b) which uses the same assumptions as Locatello et al. (2020a) but relaxes the requirement that the latent variables are independently distributed. These approaches form part of larger body of recent work that shows the importance of sparsity and weak supervision from actions in disentanglement (Lachapelle et al., 2022; Lachapelle & Lacoste-Julien, 2022; Brehmer et al., 2022; Lippe et al., 2022: 2023b;a). In the appendix, we also show how known mechanisms from Ahuja al. (2022a) can be dealt with in our framework. closely related, but more general setting, the recent progress on disentanglement from interventional distributions which do not require paired samples (Ahuja et al., 2023; Buchholz et al., 2023; von Kügelgen al., 2023); we believe a useful extension of our approach would consider these settings. This literature builds on the foundational work from the nonlinear independent component analysis (ICA) literature (Hyvarinen & Morioka, 2016; 2017; Hyvarinen et al., 2019; Khemakhem et al., 2020a).
Object-centric learning. Natural data can often be decomposed into smaller that explain the data. The overarching goal of object-centric learning is to model such data in terms of these multiple objects. The reason for this is simple: set of relevant objects rather than, for example, a large grid of feature vectors. Representing data in this way has downstream benefits like better robustness (Huang al., 2020). An important line of research in this area is how to obtain such objects from data like images and video in the first place. Typically, a reconstruction setup is used: given an image input, the model learns the objects in the latent space, which are then decoded back into the original image with a standard reconstruction loss (Locatello et al., 2020b; van Steenkiste et al., 2018b). Nguyen et al. (2023) propose RSM, conceptually close idea to work. They jointly learn object-centric representations with modular dynamics model by minimizing a rolled out reconstruction loss. However, they do not obtain any disentanglement of object properties, and the form of our proposed weak-supervision provides insights to the effectiveness of their method for improving generalization.
usually easier to reason over small
We slot attention since it makes very few assumptions about the desired data. For instance, some methods model foreground differently from background. Additionally, DINOSAUR (Seitzer et al., 2022) shows recent success on more complex images, which demonstrates the versatility of the slot attention approach. While in general object-centric models operate on image inputs and thus identify visual objects, is in principle applicable to other domains like audio (Reddy et al., 2023) as well.
7 EMPIRICAL EVALUATION
Setup. We evaluated our method 2D and 3D synthetic image datasets that allowed us carefully control various aspects of the environment, such as the number of objects, their sizes, shapes, colors, relative position, and dynamics. Examples of our 2D and 3D datasets are shown in figures 1,3 respectively. The object-wise true latents in either dataset consist of 2 = (Pe, h, s, Φ), where Pz:Py denote the coordinates of the center of an object, followed by color hue h, shape s, size r, and rotation angle & about the z-axis. Therefore, we deal with both discrete and continuous properties. For further details on the dataset generation, see appendix F.6.
Disentanglement Metrics. We compared ?-the projections of non- non-background slots to the true latents of objects to measure the disentanglement of the properties in 2. We evaluated identifiability of the learned representations either up to affine transformations or up to permutation and scaling. These two metrics were computed by fitting linear regression between 2, and reporting the
7

---

Published as a conference paper at ICLR 2024
Figure 3: (Left) An example image before any perturbation. (Right) Possible perturbations in the synthetic 3D dataset, i.c. change in size, orientation, color, position, and shape.
coefficient of determination R?, and using the mean correlation coefficient (MCC) (Hyvarinen & Morioka, 2016; 2017).
Table 1: Permutation Disentanglement (MCC) scores on 2D shapes test set under unknown fully sparse perturbations. All results are averaged over 3 seeds except those requiring to train SA-MESH from scratch that were trained only once. SA-LR (supervised) achieves
of 1.0 in all settings. pos pos,, y. color, size, rotation
pos_ pos,
Model
2
= 4
3
Ours
1.00 #0.01
1.00 10.01
0.98
#0.01
0.95 10.01
0.93 10.00
0.94 10.01
SA-RP
0.80
0.90
0.82
0.58
0.52
0.50
SA-PC
1.00
1.00
1.00
0.86
0.85
0.84
CNNi
0.96 +0.02
0.99 =0.01
0.98
=0.02
0.91 =0.01
0.89 =0.01
0,90
#0.01
CNN
0.40 +0.01
0.25
+0.03
0.21 +0.01
0.58
0.00
0.42 +0000
0.27 +0.01
Table 2: LD scores on 3D shapes test set under unknown fully sparse perturbations, SA-LR achieves a score of 1.0 in all settings.
pos g yosy; color
pos posy. color, size, rotation
Model
n
n 3
n
1 2
n 4
Ours
0.99 +0.01
0.99 +0.00
1.00
+0.01
0.91
+0.03
0.95 +0.01
0.93
+0.01
SA-RP
0.67
0.58
0.58
0.51
0.56
0.60
SA-PC 0.64
0.62
0.64
0.56
0.76
0.76
Baselines. Our baselines were selected to use series of linear probes to evaluate how much of the property information is already present in the slot representations implementation which was only trained for reconstruction. We compare this with our approach which explicitly optimizes to disentangle the properties with the weakly supervised loss in Eq 2. For the basclines, we mapped from the slot representations to d-dimensional latent space with random projections (Ri)-which which preserve distances in the projected space obtain crude estimate of vanilla slot attention disentanglement-ten the first d principal components (PC), and linear regression (LR) which provides a supervised upper-bound on what is achievable with linear maps. Finally, we also included standard ResNet18 (He et al., 2016) (denoted CNN) trained with Ahuja et al. (2022b)'s 's procedure that does not address injectivity issues, and the same trained on a DGP which is modified to be injective (denoted CNNf).
of vanilla SA-MESH
2D Shapes. The results in Table (with additional results in appendix 1) confirmed that as long as the generative function is injective can empirically achieve identification (see CNNT). But
'We make g injective by using the properties that are not the target of disentanglement, i.c., if 22, y are the target properties, we will uniquely color each object based on its order in the default permutation. If 2,, y, are targets, we will use unique shapes for each object based on its order in the permutation. The same logic follows for other property combinations.
8

---

Published as a conference paper at ICLR 2024
Table 3: MCC scores on 3D shapes test set under unknown fully sparse perturbations. SA-LR achieves a score of 1.0 in all settings.
pos, yosy color
pos pos,- yy color, size, rotation
Model
71
2
7 3
4
TO = 2
T = 3
12 = 4
Ours
0.99 #0.01
0.99 #0.00
0.99
#0.01
0.89
#0.02
0.92 #0.03
0.92 #0.02
SA-RP
0.62
0.54
0.54
0.49 0.64
0.50
0.46
SA-PC 0.69
0.68
0.70
0.77
0.78
the moment we drop any ordering over the objects and render via non-injective function, then identification via Re which is suited only to injective generative functions, fails disastrously (see the corresponding to CNN in table 4. On the other hand, difficulty identifying object properties because it treats them as a set by leveraging slot attention and a matching procedure. Additionally, the shared structure of learned latents in our method significantly improves the sample efficiency for disentanglement (see appendix E.3). The strong performance of the principal components of vanilla SA-MESH on the position disentanglement task likely results from the positional encoding. On the more complex task that also involves color, size and rotation, MCC performance drops for SA-PC, though it is still surprisingly high given that the model is just trained for reconstruction. This is likely because these are very simple images with properties that were selected independently, uniformly at random so the slot principal components align with the ground-truth axes of variation in the data.
that our method has no
3D Shapes. Figure 3 shows examples of perturbations that the model observes and uses for disentangling object properties in tables 2 and 3. We present the disentanglement scores for various combinations of properties and environments with k = {2, 3, 4} number of objects in the scene. Since non-injective CNN failed consistently in the simpler 2D dataset, we do not evaluate it with 3D shapes. Results for our method are averaged over 3 seeds, but since the baselines require training SA-MESH from scratch, they were trained only once it is computationally expensive to obtain excellent reconstructions with slot attention. These results essentially confirm our findings in the simpler 2D dataset, and demonstrate how treating the scene as a set with our method results perfect disentanglement of object properties. For the results on other combinations of properties, please see appendix E.2.
CONCLUSION
This study establishes a connection between causal representation learning and object-centric learning. and (to the best of our knowledge) for the first time shows how achieve disentangled representations in environments with multiple interchangeable objects. The importance of recognizing this synergy two-fold. Firstly, causal representation learning has largely ignored the subtleties of objects in assuming injectivity and fixed Id representations. Conversely, object-centric learning has not dealt with the challenge of unsupervised disentanglement. Yet disentangled representations can significantly improve model's generalization capabilities under distribution shifts, and could also allow for learning parsimonious models of the dynamics when such proper representations are achieved, which we deem as important avenues for future research. In this study we provided empirical evidence showcasing the successful disentanglement of object-centric representations through the fusion of slot attention with recent advances in causal representation learning.
ACKNOWLEDGMENTS
The authors would like to thank Kartik Ahuja for the helpful discussions and feedback on the initial directions of this work. This research was enabled by compute resources and technical support provided by Mila, Quebec AI Institute (mila..uebbec) as well as the Digital Research Alliance of Canada, formerly known Compute Canada. Amin Mansouri was supported in part by the support from Canada CIFAR AI Chair Program, from the Canada Excellence Research Chairs Program, and Microsoft.

---

Published as a conference paper at ICLR 2024
REFERENCES
Kartik Ahuja, Jason Hartford, and Yoshua Bengio. Properties from mechanisms: equivariance perspective on identifiable representation learning. In International Conference Learning Representations, 2022a. URL https::/ https////ppeeeeeeeeeeeomee net htt/://oummmmmm//mee MMMM.
net.
jMq4M.
Kartik Ahuja, Jason Hartford, and Yoshua Bengio. Weakly supervised representation learning with sparse perturbations. In Neural Information Processing Systems, 2022b.
Kartik Ahuja, Divyat Mahajan, Yixin Wang, and Yoshua Bengio. Interventional causal representation learning, 2023.
Jack Brady, Roland S. Zimmermann, Yash Sharma, Bernhard Schölkopf, Julius von Kügelgen, and Wieland Brendel. Provably learning object-centric representations, arXiv preprint arXiv: 2305,14229, 2023.
Johann Brehmer, Pim De Haan, Phillip Lippe, and Taco Cohen. Weakly supervised causal representa- tion learning, arXiv preprint arXiv:2203. 2022.
Simon Buchholz, Goutham Rajendran, Elan Rosenfeld, Bryon Aragam, Bernhard Schölkopf, and Pradeep Ravikumar. Learning linear causal representations from interventions under general nonlinear mixing, 2023.
Michael Chang, Thomas L. Griffiths, and Sergey Levine. Object representations as fixed points: Training iterative inference algorithms with implicit differentiation. In ICLR Workshop on Deep Generative Models for Highly Structured Data, 2022. URL https: opw.reveee.e.eeee///e net/
Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using encoder-decoder for statistical machine translation. arXiv preprint arXiv: 1406. 1078, 2014.
Marco Cuturi. Sinkhorn distances: Lightspeed computation of optimal transportation distances. NEURIPS, 2013.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. arXiv preprint arXiv:1909. 1019999008999 10893, 2019.
Anirudh Goyal, Alex Lamb, Phanideep Gampa, Philippe Beaudoin, Sergey Levine, Charles Blundell, Yoshua Bengio, and Michael Mozer. Object files and schemata: Factorizing declarative and procedural knowledge dynamical systems. arXiv preprint 2xXi:20006162222 2020.
Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Christopher Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner, Multi-object representation learning with iterative variational inference. In International Conference on Machine Learning, pp.
2424 2433. PMLR, 2019.
Klaus Greff, Francois Belletti, Lucas Beyer, Carl Doersch, Yilun Du, Daniel Duckworth, David Fleet, Dan Gmanapragasam, Florian Golemo, Charles Herrmann, Thomas Kipf, Abhijit Kundu, Dmitry Lagun, Issam Laradji, Hsueh-Ti (Derek) Liu, Henning Meyer, Yishu Miao, Derek Nowrouzezahrai, Cengiz Oztireli, Etienne Pot, Noha Radwan, Daniel Rebain, Sara Sabour, Mehdi S. M. Sajjadi, Matan Sela, Vincent Sitzmann, Austin Stone, Deqing Sun, Suhani Vora, Ziyu Wang, Tianhao Wu, Kwang Moo Yi, Fangcheng Zhong, and Andrea Tagliasacchi. Kubric: a scalable dataset generator. Computer Vision And Pattern Recognition, 2022.
Luca Gresele, Julius V Kügelgen, Valentin Stimper, Bernhard Schölkopf, and Michel Besserve. Independent mechanism analysis, concept? In Advances in Neural Information Processing Systems, 2021.
Ben Hayes, Charalampos Saitis, and György Fazekas. The responsibility problem in neural networks with unordered targets, 2023, URL https:/ jd7Hy1 jdivv.
/ http://:////ppeeeee..... net wwtteforuttttteeeeeeeeeeeoeeetet forum? id=
10

---

Published as a conference paper at ICLR 2024
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
Qian Huang, Horace He, Abhay Singh, Yan Zhang, Ser-Nam Lim, and Austin R. Benson.
Better set representatiion for relational reasoning. In NeurIPS, URL httppa://proceelings
2020.
www.iinnnnnnneurreeeeeeeeereeree neurips htcpii.s.ccppppppeeeeee cc/pppex//022/ /hash/ 2 Abst:aattc...tttttttttt......... htmll.
09ccf31
Aapo Hyvarinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. Advances in Neural Information Processing Systems, 29, 2016.
Aapo Hyvarinen and Hiroshi Morioka. Nonlinear ica of temporally dependent stationary sources. In Artificial Intelligence and Statistics, pp. 460-469. PMLR, 2017.
Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural Networks, 12(3):429-445 1999. ISSN 0893-6080. doi: hitps: wowp.agggo...00000000.9366sss0e0 sci htw.eeneeerrttcceeeppe...eepcpce
URL https::////t.. //wp:://wwww./// sciencedi rect com/ 403.
Aapo Hyvarinen, Hiroaki Sasaki, and Richard Turner. Nonlinear ica using auxiliary variables and generalized contrastive learning. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 859-868. PMLR, 2019.
A. Hyvärinen and E. Oja. Independent component analysis: algorithms and applica- tions. Neural Networks, 13(4):411-430 2000. ISSN 0893-6080. doi: https://doi.or/// 1016/50899-608800000002226 URL https: /ttp://wwweeceeddd..// rect. rwwttccom/cccccce...oo..o....oct science////////cc atticlle/111/0888300000000
Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: A unifying framework. International Conference on Artificial Intelligence and Statistics, pp. 2207-22111 PMLR, 2020a.
Ilyes Khemakhem, Ricardo Monti, Diederik Kingma, and Aapo Hyvarinen. Ice-beem: Identifiable conditional energy-based deep models based on nonlinear ica. Advances in Neural Information Processing Systems, 33: 11112227777777 2020b.
Sebastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi LE PRIOL, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: new principle for nonlinear ICA. In First Conference on Causal Learning and Reasoning, 2022. URL http:: httttttt::///peeeeteee http::///peeeeeeeeeeeee net. wwt.:eet..orrrne forum? ww=cHhdFeedddd...
Sébastien Lachapelle and Simon Lacoste-Julien. Partial disentanglement via mechanism sparsity, 2022.
Sébastien Lachapelle, Divyat Mahajan, Ioannis Mitliagkas, and Simon Lacoste-Julien. Additive Decoders for Latent Variables Identification and Cartesian-Product Extrapolation, July 2023. URL pttt:://arxie orooooogbao.... abs 02598. arXiv:2307.02599 [cs, stat].
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. CoRR, abs/2001.02407, 2020. URL http: htttt::/arxiv..oot org/ abs/ 2001. 02407.
Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. Citris: Causal identifiability from temporal intervened sequences. arXiv preprint arXiv:2202.03169, 2022.
Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Efstratios Gavves. Biscuit: Causal representation learning from binary interactions. In The 39th Conference Uncertainty in Artificial Intelligence, 2023a. URL https: : /pwp/eeeeeee..e.//o///// net forum?
Id=Vs7Dn31xed xuB.
11

---

Published as a conference paper at ICLR 2024
Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M. Asano, Taco Cohen, and Efstratios Gavves. Causal representation learning for instantaneous and temporal effects in interactive systems, 2023b.
Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentan- gled representations. In international conference on machine learning, pp. 4114-4124. PMLR, 2019.
Francesco Locatello, Ben Poole, Gunnar Rätsch, Bernhard Schölkopf, Olivier Bachem, and Michael Tschannen. Weakly-supervised disentanglement without compromises. In International Conference on Machine Learning. pp. 6348 6359. PMLR, 2020a.
Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. CoRR, abs/2006. 15055, 2020b. URL https: /ttp///aavvvv/ee///ee/eeeeeaaevv org hwwgg//b////// 200/s////200.... 15055.
I. Loshchilov and Hutter. Decoupled weight decay regularization. International Conference on Learning Representations, 2017.
Gabriel E Moran, Dheevatsa Sridhar, Yining Wang, and David Blei. Identifiable deep generative models via sparse decoding. Transactions on Machine Learning Research, 2022.
Trang Nguyen, Amin Mansouri, Kanika Madan, Nguyen Duy Khuong, Kartik Ahuja, Dianbo Liu, and Yoshua Bengio. Reusable slotwise mechanisms. In Thirty-seventh Conference Neural Information Processing Systems, 2023. URL https: http///pppeeeeeeeeeeeeeeeeeeeeet net fwwrmmormm iid= CniUitfEY3.
Reddy, Scott Wisdom, Klaus GrefT, J. Hershey, and Thomas Kipf. Audioslots: slot-centric generative model for audio separation. IEEE International Conference on Acoustics, Speech, and Signal Processing Workshops (ICASSPW), 2023. doi: 00.110//AASSPP99222000.0.000
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, 109(5):612-634 2021.
Maximilian Seitzer, Max Horn, Andrii Zadaianchuk, Dominik Zietlow, Tianjun Xiao, Carl-Johann Simon-Gabricl Tong He, Zheng Zhang, B. Scholkopf. T. Brox, and Francesco Locatello. Bridging the gap to real-world object-centric learning. International Conference on Learning Representa- fions, 2022. doi: 0405500/
Pete Shinners. Pygame. htttp:/ppygame.cr. 2011.
Gautam Singh, Yeongbin Kim, and Sungjin Ahn. Neural systematic binder. In The Eleventh International Conference on Learning Representations, 2023. URL https: / h /ttp:///oppeeeeeeeeee net/
9t.
Richard Sinkhorn and Paul Knopp. Concerning nonnegative matrices and doubly stochastic matrices. Pacific Journal of Mathematics, 21:344-344, 1967. URL http:tts:://ee //ap:///pp... .ttap..smaanttt.. semant icscholarr orq/ CCrp/or 50329347.
Brian Cantwell Smith. The promise of artificial intelligence: reckoning and judgment. The MIT Press, 2019.
Christian Squires, Aryeh Seigal, Shreya S Bhate, and Caroline Uhler. Linear causal disentanglement via interventions. In International Conference on Machine Learning. PMLR, 2023.
Sjoerd van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schnidhuber. Relational neural expectation maxinization: Unsupervised discovery of objects and their interactions. CoRR, Abs/1800..10353 2018a. URL http: /tt:://axxxa /taaivivvvv.////.///.aaa//aaaavv org/ abs 110..101555
Sjoerd van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural expectation maximization: Unsupervised discovery of objects and their interactions. In Interna- tional Conference on Learning Representations, 2018b. URL https: openrevieww net/ forum? (rrum??ddryH220cbbb
12

---

Published as a conference paper at ICLR 2024
Julius von Kügelgen, Michel Besserve, Wendong Liang. Luigi Gresele, Armin Kekić, Elias Barein- boim, David M. Blei, and Bernhard Schölkopf. Nonparametric identifiability of causal representa- tions from unknown interventions, 2023.
Quanhan Xi and Benjamin Bloem-Reddy, Indeterminacy generative models: Characterization and strong identifiability. In Francisco Ruiz, Jennifer Dy, and Jan- Jan-Willem van de Meent (eds.), Proceedings of The 26th International Conference on Artificial Intelligence and Statistics, volume 206 of Proceedings of Machine Learning Research, pp. 6912-698 2-399. 1939. PMLR, 25-27 Apr 2023. URL https: : pww.eedinoceeeeesssooooo. mlr pwwss
/ v200 xi xi23a html.
Yan Zhang, Jonathon Hare, and Adam Prigel-Bennett. FSPool: Learning set representations with featurewise sort pooling. International Conference on Learning Representations, 2020. URL
https:///ggt /ttp::///ppeeeeeeeeeeee net/ [orum?id-HUUgA222 ott.ee/aewe.eee..aoooooaoaaaaaaa
Yan Zhang, David W Zhang, Simon Lacoste-Julien, Gertjan J Burghouts, and Cees GM Snoek. Multiset-cquivarann set prediction with approximate implicit differentiation. In International Conference on Learning Representations, 2022. URL https: http://pppeeeeeeeeeeeeeeeeeeeeee net. forum?
id=5K7RRqZE
Yan Zhang, David W Zhang, Simon Lacoste-Julien, Gertjan Burghouts, and Cees G. M. Snoek. Un- locking slot attention by changing optimal transport costs. In International Conference on Learning Representations, 2023, URL https: / www.revien.eeeee/////e net nwt forum? Fdd-FFMoWFFE 5d.
Yujia Zheng, Ignavier Ng, and Kun Zhang. On the identifiability of nonlinear ica: Sparsity and beyond. Advances in Neural Information Processing Systems, 35:16111-16422,, 2022.
A PROOF OF THEOREM 1
We want to compare the number of perturbations needed to disentangle shared properties with standard encoder to those needed by an object-centric encoder. Our strategy will be as follows,
1. Setup data generating process with multiple objects where injectivity holds by construction so that we can restate Theorem from Ahuja et al. (2022b) to show they need k x d perturbations.
2. Define an object-centric architecture in terms of the object-wise partitions that we defined in Definition 1.
3. Restate an analog of Theorem from Ahuja et al. based on the object-centric encoder.
4. Theorem 1 in the main text will follow as a collary of the difference between the number of perturbations used in the two theorems above.
We begin by defining data generating process such that g( ((vec,x()) is injective by construction. We can achieve this by appending an id, i, to each 2; in Z. such that Z = {2 + [i] where @ denotes concatenation, and then choosing g such that depends be rendered in a different color). Like Ahuja et al., we assume we have data that is perturbed by of I-sparse perturbations that perturbs each of the properties from each of the is objects. Taken together, we have the following data generating process (DGP),
(for example, each could
A :
Z = (+01)) Pz, =Z)) ža Z₁ 28 di.s E A, L g((2).
Zk
where each object has shared properties, 2, Rd. and 21, and 2 are of the same type e.g. position x, huc, etc.- for all j. As before, assume is injective, and define g" g(vec..(()) where n is the permutation that sorts Z by the index i, that g" is injective by construction.
Now, Ahuja et al. show that if the encoder, f f:x' X + Red is chosen minimize the following loss.
f E arg min En ((((x)+ + ((x'))
and the following assumptions hold,
13

---

Published as a conference paper at ICLR 2024
Assumption 1. The dimension of the span of the perturbations in equation 3 is kd, i.e. dim span A) kd.
Assumption 2. a(z) f is an analytic function. For each component E (1.... ... kd) of az) and each component j E {1,... ....,,dl kd] of z, define the set S {0 jjai(z b) V v)a((2) (0),(0)b, R&d), where fixed vector in Redd Each set S has a non-zero
Lebesgue measure in 12kd
Then we have,
Theorem 2 ((Ahuja et al., 2022b)). Assume we have data from the DGP equation 3 and assumption I and 2 hold and the number of perturbations per example equals the latent dimension, kd, then the encoder that solves equation 3 identifies true latents up to permutation and scaling, i.e. & = 1lAz + where A E Rkdxkd is an invertible diagonal matrix, I1 E R kd is permutation matrix and c an offset.
Proof. See Ahuja et al. for the proof.
Now, consider an object-centric architecture encoder of the form F(z) == {f(x"))='en where P is an object-wise partition and f: X + RR.. Let a E I denote a permutation of the latents from the set of k --permutations. Let:
Fxx) E arg min E [min [I((((=')) ((f(=)()) + 6)
o'EE
Note that since 8 non-zero for only pair of patches (0) ,,0 and otherwise, the minimizer over Σ is almost surely unique. Assumptions and 4 are analogs of and 2 above, but make reference the dimensionality of the co-domain of rather than
Assumption 3. The dimension of the span of the perturbations in equation 3 is d, iee,
dim (pann () ( dd
Assumption 4. a(=) fo (z) is an analytic function. For each component E {1, (1..... d} of a(2) and each component j E (1.....d) of z, define the set Sij = {0 Vaa,(z + b) = (3)()) ()a((2)+ (0), E Rd), where is a fixed vector in Rd. Each set Süj has a non-zero Lebesgue measure in Rd
With this setup, the following theorem follows directly from Theorem as a reduction from the multi-object single-object setting.
Theorem 3. Assume have data from the DGP in equation 3 and assumption 3 and 4 hold and the number of perturbations per example equals the latent dimension, m = d, then the encoder that solves equation for an object-wise partition P. identifies latents up permulation and scaling, Le. & = TT + c, where E Rkdxkdd is an invertible diagonal matrix, E Rledkkd is permutation matrix and an offset.
Proof. Because P is object-wise partition, the function that produces each (0) E P is injective with respect some 2a(() (i.c. one of the object's latents). Thus for each (()) the solution equation 5 is equivalent to the single object setting with & = 1, and thus theorem 2 applies, which implies that ((x:()) = & = IA + e for all Now let & vec,, ( ((&, (=)) Because each é; is identified up to permutation, scaling and offset, and for any x, there exists II such that IIA + which completes the result.
Corollary 1. If the assumptions for Theorem and 3 hold and a data generating process outputs observations containing k objects with shared properties, then object-centric architecture the form F(x) = (f(x(()) ((( Ja will disentangle in 1/k fewer perturbations than an encoder of the form f X > RRidd
Proof. This follows directly from comparing the the number of perturbations required in Theorems and 3.
14

---

Published as a conference paper at ICLR 2024
B BACKGROUND ON SLOT-AITENTION-BASE ARCHITECTURES
Slot attention (Locatello et al., 2020b) is a neural network component that, intuitively, summarizes the relevant information in the input set (most commonly, image features with position embeddings) into a smaller set of so-called "slots". Each slot is a feature vector that can be thought of as capturing information about one "object" in the input set, which usually comprises multiple elements of the input set. This is done by repeating cross-attention between the inputs and the slots to compute per-slot updates.
In the traditional set-up. these slots are then used to reconstruct the input with an auto-encoder objective: each slot is decoded into a separate image through shared image decoder, which is followed by merging these per-slot images into single reconstructed image. Ideally, slot attention is able to decompose the original image into distinct objects, each which is modeled by a single slot.
More concretely, slot attention takes input matrix X E IRax with the number of inputs and the dimensionality of each input. We also randomly initialize the slots ( E Rmxd We start by computing the query, key, and value matrices as part of cross-attention.
( z K = =
(6)
This is followed by the normalized cross-altention to determine the attention map IR then GRU (Cho et al., 2014) to apply this update to the slots.
A normalize (Q ((OKT
z (++)) GR) (V
(8)
The function normalize encourages slots to compete for inputs by applying softmax slots and normalizing the weights for each input to sum to one. After I steps, the algorithm outputs (TT), set of m embedding vectors {z] (T) m that can be used as input to shared image decoder.
SA-MESH. The specific version of slot attention we use in this paper is SA-MESH (Zhang et al., 2023). It makes regular slot attention more powerful by giving it the ability to break ties between slots more effectively. In practice, this improves the quality of the individual slot representations significantly due to less mixing of unrelated inputs into the same slot.
The key difference with regular slot attention is that it features entropy minimization procedure to approximate an optimal transport solution, which makes the attention map sparse. The connection to optimal transport made by the use of the standard Sinkhorn algorithm (Sinkhorn & Knopp, 1967; Cuturi, 2013).
(ESH(C) = arg min H(sinkhorn (C))
A (1) = sinkhorn(MESH( ((TK)))
(10)
The optimization problem is solved by unrolling gradient descent, with a noisy initialization to ensure that ties are broken.
C ALTERNATIVE PERTURBATION MECHANISMS
Dense vs. Sparse. We can have number of assumptions on the perturbation mechanisms and the nature model's knowledge about those mechanisms. In the most general case, suppose M M=(m-(( m²),...., (.....,,((.)) m²)-)) denotes the set of all possible perturbation mechanisms. To obtain x', the perturbed variant of x, we then select a subset of k & le objects as targets that undergo perturbations determined by a subset of K mechanisms M' C M from the set of all possible perturbation mechanisms. The correspondence between the k' mechanisms and k' perturbed objects decided by random permutation N. i.e. transition dynamics of object j to produce 2 (for objects that are not supposed to change from
7. MM]
[j] means that mechanism governs the
CConcretely, it makes the mapping from the initial slots to the final slots exclusively multiset-equivariant (Zhang al., 2022) rather than gemmuzzioonqquvrriinnsseeeaaaaaa
15

---

Published as a conference paper at ICLR 2024
t + a dummy mechanism with index -1 can be assumed which results in change). A mechanism m'(.) : R + Rd in M is a vector-valued function that operates on object-wise true latents 2f and outputs z/41 x + d; such that = _ Li]. Perturbation vectors ó; could be sparse or not. The subset M' can contain = k mechanisms to perturb all the & objects in the environment, and if none of the k' A resulting perturbations is sparse, we denote the set M' fully dense perlurbations, i.e., all of the properties of all objects will change from 1 + 1 + 1. M'' can also contain at least one object but not all of them (1 k' < <k with sparse or dense perturbations, may consist of only single object (k' = 1) that is perturbed by fully sparse mechanism, one that only alters single property and leaves the rest unchanged. We denote this scenario as fully sparse perturbation.
it
D MATCHING
Perturbations alter the properties of objects from + + 1 and the model has to infer which object's properties were perturbed to update its representations and minimize the latent loss equation 1. But recall that the model has no direct access to objects. It receives the observations at t, + and encodes each of them to a set of slots S,, S, These slots do not follow any fixed ordering, and moreover, there is no guarantee that each slot binds to exactly one unique object. Slots can also correspond to the background. Each perturbation d' changes the properties of some object z// so the model requires to find pair of slots (s s{+1) that are bound object 2 at and + L, respectively, Once the model figures out such matching, then the latent loss that results in disentanglement can be computed via the projections of these slots 2:
C- Σ
= p(s)), 2L41 pp(s))
(11)
The problem of finding correspondence between slot projections at 1, L +-1 1 and the perturbations is an instance of the 3-dimensional matching. However, recall that we use fully sparse perturbations for our mechanisms, thus, the problem is significantly simplified.
When the model is presented with fully sparse perturbations, the 3-dimensional matching reduces to finding the minimum element in 1-d cost array of size |S]? where [S] is the number of slots. The reason is that since only one property of single object 2 is being altered (and the model knows perturbations are fully sparse), the model can apply that 1-sparse perturbation 8 (or the transformed version in the unknown setting) all of slot projections $/ for j E (1...., [S]) at time f, and only find one pair of slots from the set of [S] possible pairs at t,t + 1, which correspond to the perturbed object:
(i.j) arg min 2,j'
|
(2) + +2
(12)
However, this minimization can be made simpler if we use the slots obtained at to initialize the slots at + 1. This way, in practice, the order of slots at t, 1 is very likely to be preserved, not only since the slots at ++1 are initialized with hints from t, but also the sparse perturbation makes it much easier for slots bind to the same object as only a small subset of the scene needs to be readjusted among the slots. Therefore the matching would reduce to a simple minimization over [S] elements:
arg min || | (2) + 6.)
(13)
Note however that we still would use all slot projections for evaluation, the only difference with this matching scheme is that gradient signals are only propagated from the perturbed object (as they also should, since there is no change in other slots, there is nothing there to be learned that could help disentanglement.)
E
FURTHER EXPERIMENTAL RESULTS
E
2D SHAPES
Table 4 gives the linear disentanglement scores for the same experiment as that shown in table in the main text.
16

---

Published as a conference paper at ICLR 2024
Table 4: Linear Disentanglement (LD) scores on 2D shapes test set under unknown fully sparse perturbations. All results are averaged over 3 seeds except those requiring to train SA-MESH from scratch that were trained only once. SA-LR, which is supervised by the ground truth latents, and is an upper bound on the disentanglement performance, achieves score of 1.0 in all settings.
pos=1 posy
pos pos,, y, color, size, rotation
Model
n = 2
3
71
= 4
2
n =
Ours
1.00 +0.00
1.00
+0.00
1.00
+0.00
0.95
+0001
0.93 +0.01
0.94
+0.02
SA-RP
0.92
0.96
0.94
0.75
0.70
0.68
SA-PC
1.00
1.00
1.00
0.93
0.88
0.86
CNN1
0.94 +0.05
0.99
#0.00
0.96
+0.03
0.87
=0.01
0.84
=0.01
0.86
=0.01
CNN
0.24 +0.01
0.13
+0.01
0.07
#0.01
0.35
=0.00
0.19 #0.00
0.08 #0.01
Tables 5,6 extend our results under unknown fully sparse perturbations on the 2D shapes dataset to more combinations of disentanglement target properties. We can observe that our results stay very close to the upper bound on the achievable performance which uses a supervised linear regression from slot projections& å to ground truth latents z. These tables highlight once again the pivotal role of the injectivity assumption in achieving identification with conventional encoders that ignore the object-centricity of the environment (see the performance drop from CNVt to CNN, where the latter drops the unrealistic injectivity assumption).
Table 5: Linear Disentanglement (LD) scores on 2D shapes test set under unknown fully sparse perturbations. All results are averaged over 3 seeds except those requiring to train SA-MESH from scratch that were trained only once. SA-LR achieves a score of 1.0 all settings.
pos pos,, ,os,,.oooo
pos pos,. Shape
Model
n1 =
n 3
11 =
n1
n =
n1
Ours
1.00 +00.01
+0.01 0.98
0.99
+0.00
1.00
=0.01
0.98 +0.01
0.99 +0.01
SA-RP
0.77
0.61
0.60
0.71
0.68
0.70
SA-PC
0.97
0.98
0.99
0.80
0.66
0.87
CNN
+:000
0.99
+0.01
0.98
+0000
1.00
+0.00
1.00 +0.00
0.99 =:0.00
CNN
0.35 +0.00
0.15 #0.00
0.07
#0.01
0.32 #0.01
0.15 #0.01
0.11 #0.01
pos
pos color, shape
Model
n
n 3
n
Ours
0.99
#0.00
0.98 #0.01
0.99 =0.00
SA-RP
0.69
0.73
0.60
SA-PC
0.74
0.75
0.52
CNN
1.00
+0.00
0.99
#0.01
1.00 +0.00
CNN
0.40 10.00
0.21 10.00
0.11 10.00
E.2
3D SHAPES
Quantitative Results. Tables 7, 8 extend our results under unknown fully sparse perturbations the 3D shapes dataset to more combinations of disentanglement target properties. Again, we can observe the applicability of our method to this more complex 3D dataset that contains artifacts related to depth, occlusion, and lighting, to name few. Again our results stay very close to the upper bound on the achievable performance which uses a supervised linear regression from slot projections to ground Iruth latents 2.
Qualitative Results. Figures 4-6 illustrate the learned disentangled (object-centric) representations. Each figure shows sequence of 3D samples evolving over 5 steps (shown the left), and how the learned representations respond to the perturbations (shown on the right). Perturbations include changing the object's pos,,, po color, á (rotation). The model used here is trained with 3 slots, and the learned representations are the result of projection layer learned via our weakly-supervise method applied to 64 14-ding dimensional object-centric representations. The projection maps each slot
17

---

Published as a conference paper at ICLR 2024
Table 6: Permutation Disentanglement (MCC) scores 2D shapes test under unknown fully sparse perturbations. All results are averaged over 3 seeds except those requiring to train SA-MESH from scratch that were trained only once. SA-LR achieves a score of 1.0 in all settings.
2: posy color
2os posy shape
Model
n
n
n
n
Ours
1.00 =0.01
0.95
+0.05
0.97 =0.02
0.99
=0.01
0.99
=0.01
0.99 +0.01
SA-RP
0.74
0.60
0.60
0.66
0.63
0.59
SA-PC
0.87
0.89
0.90
0.83
0.81
0.89
CNN
1.00 =0.00
0.99
#0.01
0,99
=0.01
1.00
=0.00
1.00 =0.00
0.99 +0.00
CNN
0.55 #0.01
0.35
#0.01
0.24
#0.01
0.52 #0.02
0.33 #0.02
0.28 #0.02
pos pos color, shape
Model
2
3
4
Ours
0.99 =00.01
0.98 10.01
0.99 =00011
SA-RP
0.54
0.68
0.55
SA-PC
0.64
0.63
0.57
CNN
1.00 10.00
0.99
+0.01
1.00 =00.00
CNN
0.61 +0.00
0.43 +0.00
0.30 +0.00
Table 7: Linear Disentanglement (LD) scores on 3D shapes set under unknown fully sparse perturbations. All results are averaged over 3 seeds except those requiring to train SA-MESH from scratch that were trained only once. SA-LR achieves score of 1.0 in all settings.
pos yosy. size
pos posy, color, rotation
Model
n
2
3
n1
4
n 3
Ours
0.98 +0.01
0.98 +0.01
0.98
+:000
0.98
+0.00
0.97 +0.01
0.98
#0.01
SA-RP
0.61
0.62
0.53
0.59
0.54
0.55
SA-PC
0.78
0.84
0.78
0.70
0.72
0.69
Table 8: Permutation Disentanglement (MCC) scores 3D shapes test under unknown fully sparse perturbations. All results are averaged over 3 seeds except those requiring to train SA-MESH from scratch that were trained only once. SA-LR achieves a score of 1.0 in all settings.
pos_: posy, size
pos,; pos,:: color, rotation
Model
n 2
n1
4
Ours
0.96 #0.02
0.96 +0.05
0.96
#0.03
0.98 00.0
0.98 =0.0
0.97 #0.01
SA-RP
0.60
0.57
0.52
0.55
0.51
0.49
SA-PC
0.87
0.90
0.86
0.73
0.76
0.74
18

---

Published as a conference paper at ICLR 2024
from R84 + R+, i.e., the disentanglement target space. In figures 4-6, the 4 dimensions of such projections for object slots are presented over 5 steps, i.c., each set of colored lines shows the evolution of the projection of slot corresponding to the object with the same color. Please refer to the figures for details of the perturbations. Lastly, we have kept the number of objects these scenes to two for clarity of the presentation, however tables 2,3, 7,8 show that achieve similar performances with other sets of properties and number of objects in the scene.
619
030
1
1.34
138
412-
4.44
410
Figure 4: (Left) From top to bottom, each steps perturbs the pos, coordinate of the red object by 0.2 in the ground-truth latent space, (Right) which is reflected (through affine transform) linear increase in pos,, while the rest of the properties for both objects remain the same, demonstrating that the learned representations some non-zero angle w.r.t. the origin, therefore, perturbations in the x direction appear as vertical displacements.
indeed disentangled. In generating these samples, camera has
E.3 COMPARISON OF SAMPLE EFFICIENCY
Figure demonstrates the sample efficiency of our object-centric model compared to ResNet that achieves disentanglement with an injective DGP. Both models are trained with varying number of training samples that contain & = 4 objects for which pos,, pos,, color, shape are the disentanglement target properties. Since sample 1-sparse perturbations uniformly. the training dataset size could be thought of as proxy for the number of different perturbations given configuration of objects would encounter. Although according to theory, the injective ResNet should require at least & times more perturbations to identify the latents up to affine transformations, we observe that the advantage of our object-centric model in terms of sample efficiency is much more pronounced in practice. Our method can achieve close to perfect disentanglement with as few as 100 training samples, while an injective ResNet takes 100 times more samples raise comparable performance. This highlights the practical importance of exploiting the inherent set structure of objects in a scene for representation learning.
F IMPLEMENTATION AND EXPERIMENTAL DETAILS
F.1 SA-MESH ARCHITECTURE
For the slot attention architecture, we closely follow Locatello et al. (2020b), in particular, we use the same CNN encoder and decoder they for CLEVR, except for the initial resolution of the spatial
19

---

Published as a conference paper at ICLR 2024
6.8
6.00
6.90
0.00
Figure (Left) From top to bottom, the objects are perturbed as follows; Red: Downward perturbation so it sits at the same coordinate as the Blue cube (1-2), Rotates counterclockwwss along the 2 z-axis to at the same orientation as the Blue cube (2-4), Further downward perturbation (2 times the displacement from step to 2). Blue: Rotates counterclockwise along the z-axis (4-5). (Right) Note how the learned representations mapping correctly reflects the similar positions and rotations in the ground-truth. truth, i.e., both by having properties of objects coincide at the same value, and by preserving the ratio of perturbations.
broadcast decoder with 3D shapes where we use 4 x since we are dealing with 64 x 64 images. We slot size of 64 and always scene. We use 3 iterations for the recurrent updates in SA-MESH. For details concerning SA-MESH also truncate the backpropagation through slot
+ 1 number of slots, where is the number of objects in the
follow Zhang al. (2023). Additionally, updates as suggested by Chang et al. (2022) to improve training stability.
F.2 DISENTANGLEMENT HEADS
SA-MESH outputs n + slots that are of size 64, yet we need to project each these slots to a d- dimenensional space so we can leverage the disentanglement method from Ahuja al. (2022a:b). We can simply achieve this projection by a single MIP, however, we decided to allocate more parameters for this projection and use separate projection heads mapping 64-dimensional vectors to d separate scalars. This way identification of different properties will not affect one another due to model capacity constraints. We stack the layers shown table obtain a projection head per each property. The same set of d projections will be shared among all slots.
Table Layers projection head for disentanglement.
Layer
Input Size
Output Size
Bias
Activation
Linear (1)
64
32
True
ReLU
Linear (2)
32
32
Truc
ReLU
Linear (3)
32
16
False
ReLU
Linear (4)
16
1
False
ReLU
20

---

Published as a conference paper at ICLR 2024
8.58
2.08
412
8.20
6.00
3.0
Figure 6: (Left) From top bottom, the objects are perturbed as follows; Since the objects change color, let us call the red cube in the top frame to be object 1, and the blue one to be object 2. Object 1: The color is perturbed to equal to object 2 (1-2), Moves to the right as its pos, is perturbed by 0.2 (2-3), The color hue is once again perturbed become purple (3-4). Object 2: Moves toward right (perturbation in pos, by 0.2) to be at the same coordinate as object (1-2), The color hue is decreased now the colors of objects 1,2 swapped (2-3). Moves by 0.2 in the y direction to align with the other object once again (3-4), Change its color hue twice the previous color perturbation with the opposite sign to match the color of the other object (4-5). (Right) Red curves correspond object 1, and the blue curves correspond to object 2. Again, notice the sections where the curves coincide, well as the ratio of jumps in the properties, showing consistency of the learned representations with the ground-truth causal representation that gives rise to these observations.
F.3 CONVNET BASELINE
As baseline for injective scenarios, we use a ResNet18 (He et al., 2016) with an output width of 128 that is passed through LeakyRe activation, which is then followed by d linear projection heads (for the same reason we use separate disentanglement heads) that map the 128-dimensional output of the CNN encoder to d separate 1-dimensional scalars that should correspond to the target dimensional space.
F.4 TRAINING
For each and for each set of disentanglement target properties, we first train SA-MESH for 2000 epochs with a batch size of 64 for 2D shapes (as the images 128 for 3D shapes (since images are 64 x 64) on a single A 1100 GPU with 40GB of memory. We used a fixed schedule for the learning rate at 2 10 4, and we used Adam W (Loshchilov & Hutter, 2017) with weight decay of 0,01 along with e = 10 10-8, , = 0.9, 32 = 0.999. SA-MESH was first solely trained by minimizing reconstruction error on the training set, then disentanglement performance was reported on the test set for projection-baste -based baselines (RP, PC, LR). Due to the high number of combinations target disentanglement properties and k. we just trained SA-MESH for each configuration only once.
128 128), and batch size of
Unsupervised disentanglement with our method has an additional stage which takes the aforemen- tioned pre- treitriined SA-MESH models and jointly minimizes the loss in equation 2. Note that at this stage, the SA-MESH model is not frozen, so the gradients flow through its network as well and help adjust the slot representations with the signal from the latent loss. Under known perturbations,
21

---

Published as a conference paper at ICLR 2024
10
Linear Discnlanglement
0.4
0,2
Injective Reshel
Disentangled Slot Attention
2000
4000
6000
8000
10000
Number of Samples
Figure 7: Comparing the disentanglement performance of an injective ResNet vs. our object-centric method based on the number of training samples. The dataset contains four 2D objects in which pos,, pos, color, shape can vary.
use the actual perturbations from the DGP as o to guide the model via equation 2, however, under unknown perturbations setting, replace all perturbations 6, by a hyperparameter C (see section 5). CNN baselines were trained similar to SA-MESH but for much shorter, i.e., 200 epochs, and usually converge very fast in less than 50 epochs.
F.5 HYPERPARAMETER OPTIMIZATION
We started around the hyperparameters used by Locatello (2020b) and Zhang et al. (2023) where applicable, and tuned on small subsets of the 2D shapes training data based linear and permutation disentanglement metrics. We considered 5 values for the learning rate [2 x 10-3 2 x 10-1, 10-1 6 x 10 10-5,2 x 10 10-51. Larger batch sizes were always better and we were only constrained by memory in the case of 128 x 128 images of the 2D shapes dataset. We considered values [0.1, 0.5] for [C], the fixed value representing all unknown perturbations, and found |C| 0.1 to perform better. We also considered slot sizes (64, 128] on a small subset of the 2D shapes training dataset. Lastly we considered 9 combinations for the relative importance of latent loss and reconstruction loss when training the disentanglement heads, i.e., we considered all combinations of Whatent E {1, 10, 100), Wrocons {1, 10, 100). and found the combination of Wrecons 100, Whatent 10 strike the optimal balance between maintaining good reconstructions and allowing the slot representations to give rise disentangled projections.
F.6 DATASETS
2D Shapes. We use pygame engine (Shinners, 2011) for generating multi-object 2D scenes. Object properties in both datasets include Pz: Py, color, shape, size, and rotation angle. Based on object properties, they are each rendered and placed on a white background (for the 2D dataset) or placed floor that is illuminated by souree lights and is being visited from somewhere above the floor (for the 3D dataset), and then aggregated to produce Ty. We discretized the range of color hues in order test the model's ability to obtain disentangled representations in the simultaneous presence of both continuous (position, size, and rotation angle) and discrete properties (color and shape).
In the 2D dataset:
22

---

Published as a conference paper at ICLR 2024
Pu. Py are generated randomly and uniformly in the [0, 1] range, i..., the boundaries of the scene, such that no two objects overlap and no object falls even partially outside the boundaries. Positional coordinates can be perturbed by any value in 1-0.2.0.2.
For color, we use HSV color representations and fix saturation (S) and value (V) at 0.6 and choose hue (H) from a set of values predefined before training (for instance [0.0, 0.25, 0.5, 0.75).. We adopted this 1-d representation to be consistent and have each property be represented by a scalar. Additionally we wanted to test the model's capacity when dealing with mixed discrete properties. Also, training Slot Attention or SA-MESH with discrete colors is computationally advantageous since the model will not have to deal with reconstructing all colors. However it should be noted that HSV is a cylindrical geometry with color hues being the angular dimension which results in values that have a distance of 1.0 being exactly the same color (given a fixed saturation and value). That is why a list of color hues such as [0.0, 0.33, 0.66, 1.00] would not work since 0.0 and 1.0 are the same color, yet our model interprets the difference as a perturbation with the amount of 1.0, which is clearly wrong. A change of color from color to j where i. index the list of color hues H would be provided to the model as perturbation in the amount of (H j] H where H denotes the number of colors in H.
Shape is also clearly discrete and is selected at random uniformly from the following set of shapes S {circle, square, triangle, heart, diamond). Note however, that the effects of perturbations need to be visible in the pixel space, and we should be wary of the dis- entanglement target properties; for instance if the rotation angle d, is a property we aim to disentangle with perturbations, then we should exclude circle from the set of possible shapes since a circle does not reflect in the pixel space the angle perturbations. A shape transformation from shape i to j where j.j index S, would be provided to the model as a perturbation with the amount of (j i)
Size is a continuous property in the range (0.12, 0.24] of the height width of the image which is 1. It can be perturbed by any amount in [-0.02, 0.02].
Rotation angle is also another continuous property in [0, x/4]. Similar to color hues, since this property is also angular, we have limited the range not to encounter situations that appear the same in the pixel space but have very different rotation angles (a square that is rotated 2 clockwise seems unaltered, /4 and 3T/ /4 rotations both look the same for square.). Angular perturbations can vary in the [-0 0.2,, 0.2) range.
We generate samples in pairs corresponding to 1 + 1. For fully dense perturbations, we generate k vectors of dimension d, where k is the number of objects. We repeat the generation until the conditions of non-overlapping objects and non-identifiability are met, i.e., no two objects at either t or t + 1 should overlap (before and after the perturbations), no object should fall in whole or partially out of the scene, and no two objects should be perturbed by d-dimensional offsets that are closer than some The last condition is necessary for fully dense perturbations as otherwise the matching has no way of distinguishing which perturbation to assign to which object since the matching solely relies on the difference between 1, + 1. For fully sparse perturbations, we are not constrained by the latter, and we only need to choose perturbations that do not push the chosen object out of boundaries, or make it overlap with another object. For any experiment we can have subset of {pos,,, pos,, color, shape, size, rotation angle) as the properties we wish to disentangle by observing perturbations in the pixel space, and we call them disentanglement target properties. In the generation process, any non-target property will be fixed for all objects in the whole dataset to avoid introducing unwanted variance to the disentanglement of target properties, i.e., if we choose {posz. pos,, y. color, shape, rotation angle} as target properties, then all the objects in all samples would have the same fixed size. Lastly, we can choose to make DGP injective or not. If we choose make a DGP injective, we would index the objects and choose a property to be set for objects according to the indices, i.e., we can choose to make the DGP injective by color; Suppose k = and the list of our color hues is [0.0. 0.25, 0.5. 0.75]. We would color the objects, which are now ordered according to some index set I. according to Z. Non-target properties (excluding the injectivity imposing property) will again be kept fixed for the whole dataset, and target properties are generated according to fully dense or fully sparse perturbation schemes. The perturbations to all properties are signed, and this is especially crucial for discrete properties such as shape. The reason is that disentanglement is achieved through observing relative distances in the pixel space, and having only positive or only negative perturbations deprives the model of having a reference for each property.
23

---

Published as a conference paper at ICLR 2024
For training generate 1000 pair per target property such that the model on average sees at least 500 samples for either positive or negative perturbations to each property, i... if we choose {poss. y color) target properties, we will generate 3000 samples for training. The validation and test sets always have 1000 samples. For the 2D dataset, we generate 128 x 128 images for better visual quality that is not distorted due to artifacts caused by perturbations. We then normalize and clip the image features (RGB values) to be in [-1, 1] range.
3D Shapes. For generating the 3D datasets we leverage kubric library (Greff et al., 2022) to obtain realistic scenes which we can highly customize. Objects sit on a floor, a perspective camera situated at (2.5, 0, 3.0) and looks at (0.0, 0.0, 0.0). Directional light illuminates the scene from (1.0, 0.0, 1.0) towards the center. The set of possible target properties are similar to 2D shapes, and the range of properties which each object is spawned is as follows:
pos, are generated randomly and uniformly in the [-1.5, 1.5) and (-1.0, 1.0] ranges respectively, such that no two objects overlap and no object falls even partially outside the boundaries. Note however, by overlap we mean that objects are spawned such that they do not mutually fill a volume in the 3D space, and we only prevent such occurrences, but we do allow occlusions from the perspective of the camera, which adds to the complexity of this synthetic dataset. pos, is never a disentanglement target property and is always set such that objects sit on the floor (except when rotated). The reason for fixing the coordinate is that any possible perturbation to the 3I) coordinates is always going to be interpreted on a 2D scene that is observed by a camera that placed somewhere above the floor. Therefore. introducing a third coordinate in the DGP and target properties has no point. Positional coordinates can be perturbed by offsets [-0.3, 0.3] range.
color is similarly parameterized by scalar in HSV format in 2D.
Shape can be any of (sphere, cube, cylinder, cone]. Again, since the effects of perturbations need be visible in the pixel space, we will not use spheres if the rotation angle ø is a property we aim to disentangle with perturbations, as sphere rotations do not reflect in the pixel space. For rotation in the 3D space we choose the z-axis as the axis of rotation so that angular perturbations are maximally visible (w.rt. the perspective camera's 's location).
Size is a continuous property in the range [0.3, 0.7] and can be perturbed by offsets in [-0.15, 0.15] range.
Rotation angle follows the convention of 2D shapes DGP, except that the rotations are around the 2-axis for better visual quality.
Since kubric c already generates high fidelity images, we use 64 x 64 images to lower the computa- tional burden of SA-MESH autoencoder. The number of samples is always fixed 20, 000 regardless of the target properties since the 3D dataset is more complex. We use similar transformation as in 2D, and normalize and clip the image features (RGB values) to be [-1, 1] range.
F.7 LIMITATIONS
Our study focuses showing when disentanglement is possible when treating object-centric envi- ronments as a set of representations instead of fixed-size vectors. We have analyzed the performance of our model comprehensively complexities of real-world scenarios. Yet, we believe and showed such analysis is a necessary first step identify the intricacies involved in making our algorithm work. Our analysis has been limited in a number of directions. First, while we do consider a wide range of continuous and discrete properties to be disentangled, the number of objects we use is rather low, which ideally should be scaled to real-world scenes containing more objects. Second, although our experiments include artifacts related to occlusion, depth, and lighting, in all of our experiments we simplify the problem by having the objects situated on homogeneous backgrounds, whereas real-world scenes would comprise more complex backgrounds. Such decisions were mainly due to (1) generating datasets of size more than 5k for each combination of properties being a computationally heavy task on its own, (2) training SA-MESH from seratch for each combination of properties and number of objects would quickly add up each training takes 12 hours single A100 GPU to achieve nice reconstructions, (3) details related to the background and the number of objects are tangential to the focus of this study, which is to demonstrate how to disentangle the causal factors in object-centric environment.
two synthetic datasets that are relatively limited in capturing the
24


Reviews:
[{"review_id": "fst19RVefa", "reviewer": "Reviewer_9Arp", "summary": "The authors focus their attention on disentanglement in an object centric environment. Authors formally introduce the challenges that may arise when multiple objects are present in the observations, especially when objects can be identical. Authors also show the effect experimentally by plotting the achieved MCC score. Authors combine the ideas of object wise partition and Ahuja et al.'s weak supervision disentanglement to solve the aforementioned problem with a claim that fewer perturbations are needed in such procedure due to shared latent spaces in the objects. Experiments are conducted with a single set of latent variables shared across multiple objects in the scene, with 2D and 3D synthetic dataset. Authors deal with both discrete and continuous latents. Experiments are conducted against a strong baseline.", "questions": "Minor Edits - \n - Double \"them\" in Proposition 2 Proof Sketch", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "- Clear communication of proposition and idea\n- Well thought through related work sections and explanations\n- Always felt I understood what authors are trying to say\n- Simple yet elegant idea to solve the multiple object disentanglement problem\n- Strong baseline for experimentation", "weaknesses": "- I feel the experimental setup was quite limiting for the proposed claim\n- How do we handle multiple set of latents? What happens when there are shared and not-shared latents across this set how does the disentanglement look like?\n- Authors mentions limitations on number of objects and background although I would like to see the effect on some real life dataset (maybe can be constructed with multiple colored mnist images in the scene?)", "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": "LFqO4WRRxW", "reviewer": "Reviewer_T5fz", "summary": "The paper studies causal representation learning in the context of multi-object scenes, challenging the assumptions commonly held in previous work on single-object images (such as injectivity). The paper proposes a new approach for disentangling object properties via weak supervision under sparse perturbations and assesses its effectiveness on synthetic datasets.", "questions": "- It's stated in Sec4 that the object-centric model requires k times fewer perturbations. Is there any empirical evidence to support this claim?\n- For slot matching, is it necessary to solve an expensive optimization problem (Eq13)? How does it perform compared to a simple initialization trick (e.g, using the slot of the 1st image to initialize that of the 2nd)?\n- Could the authors provide a more detailed explanation on `simply add a projection head, that is trained by a latent space loss`? Why is the projection head needed on top of the object-centric encoder?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "-\tThe study effectively bridges recent advances in causal representation learning with object-centric learning, both from theoretical and methodological perspectives. The theoretical analysis identifies and discusses the issue of non-identifiability (e.g., non-injectivity) in multi-object scenarios. The method part nicely adapts the weakly-supervised approach from disentangled/causal representations to multi-object contexts.\n-\tThe paper is very well written and enjoyable to read.", "weaknesses": "- Insight: While technical sound, the proposed method itself appears not as interesting. If I understand correctly, the method is essentially nothing but adding a slot matching step to the existing weak-supervised disentangled representation learning.\n- Positioning: The claim of being the first to `achieve disentangled representations in environments with multiple interchangeable objects` is overstated. Similar objectives have been pursued in recent studies, e.g., Block-Slot Representation [1]. Slot matching has also been considered in recent work [2]. A more thorough positioning in the literature about multi-object scenes might be needed.\n- Theory: the proof sketch for Proposition 2 appears less than compelling in its current form. Part of my doubt is attributed to the use of identical objects. It would strengthen the paper if the author could prove Proposition 2 in more common scenarios.\n\n[1] Neural Systematic Binder, ICLR'23 \\\n[2] Causal Triplet, CLeaR'23\n\n------------------------------------------------   \n\nPost-rebuttal Review:\n\nThe response from the authors has addressed the primary concerns I had previously. I thus raise the rating from 5 to 6. \\\nThe theoretical part is definitely a solid contribution, whereas the insight and positioning of the current manuscript remain somewhat weak. \\\nI would recommend the authors to consider incorporating their rebuttal into the final version of their paper.", "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": "nxbdxdo1Ut", "reviewer": "Reviewer_HmRq", "summary": "This work addresses the issue of non-identifiability in causal representation learning models using object-centric approaches. The authors show that using sparse transitions between pairs of images and models that exploit such regularities can allow to solve the responsibility problem when attempting to invert a data-generation process in the image domain. Specifically, the exploit the fact that object centric models partition images into objects to reduce the number of interventions needed in order to identify the properties that correspond to each object thus allowing them to be identified. To encourage identifiability the authors introduce a penalty that seeks to keep all slots in the object-centric representation constant for both images, which can them only vary for the object that was perturbed. The authors then show that this leads to an increase in the ability of downstream models to predict the different properties of the objects in a scene.", "questions": "My only question is related to how good the prediction of each individual object property is. My understanding is that the prediction scores are averaged over all properties, but previous research has show that some properties are easier to predict than others (eg. rotation is harder than position). Thus, I would like to see the scores for these different properties in isolation.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The paper is well motivated based on the limitations of previous work in causal representations learning and object-centric approaches.\n2. The authors correctly justify their approach with easy to follow but formal statements (or propositions) that theoretically ground their approach.\n3. While maybe a bit light on experiments, I believe that for their particular purposes the results substantiate their results, with one caveat that I will address later.", "weaknesses": "1. The authors do however miss the opportunity to connect their research with previous work on disentanglement learning in a more systematic way. Indeed their approach of encouraging sparse transitions has already been explored in disentangled models (albeit not object-centric ones) in Klindt et al., (2021) and Montero et al., (2022). and thus these are both relevant work that should be discussed (though I don't think a comparison with these methods is required).\n2. While I don't have an issue with the evaluation per se, I do think that the scores of the models should be separated by properties as well as averaged over all of them (as is currently presented according to my understanding).\n\n\n\nRefs:\n[1] Klindt, D., Schott, L., Sharma, Y., Ustyuzhaninov, I., Brendel, W., Bethge, M., & Paiton, D. (2020). Towards nonlinear disentanglement in natural data with temporal sparse coding. arXiv preprint arXiv:2007.10930.\n[2] Montero, M. L., Bowers, J. S., Costa, R. P., Ludwig, C. J., & Malhotra, G. (2022). Lost in latent space: Disentangled models and the challenge of combinatorial generalisation. arXiv preprint arXiv:2204.02283.", "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 (poster)