Title: Identifiable Object-Centric Representation Learning via Probabilistic Slot Attention
Venue: NeurIPS 2024
Accepted: yes
Keywords: Object-centric learning

Probabilistic slot-attention

Identifiability

latent mixture models
Abstract: Learning modular object-centric representations is said to be crucial for systematic generalization. Existing methods show promising object-binding capabilities empirically, but theoretical identifiability guarantees remain relatively underdeveloped. Understanding when object-centric representations can theoretically be identified is important for scaling slot-based methods to high-dimensional images with correctness guarantees. To that end, we propose a probabilistic slot-attention algorithm that imposes an *aggregate* mixture prior over object-centric slot representations, thereby providing slot identifiability guarantees without supervision, up to an equivalence relation. We provide empirical verification of our theoretical identifiability result using both simple 2-dimensional data and high-resolution imaging datasets.

Identifiable Object-Centric Representation Learning via Probabilistic Slot Attention

Abstract

Learning modular object-centric representations is crucial for systematic generalization. Existing methods show promising object-binding capabilities empirically, but theoretical identifiability guarantees remain relatively underdeveloped. Understanding when object-centric representations can theoretically be identified is crucial for scaling slot-based methods to high-dimensional images with correctness guarantees. To that end, we propose a probabilistic slot-attention algorithm that imposes an aggregate mixture prior over object-centric slot representations, thereby providing slot identifiability guarantees without supervision, up to an equivalence relation. We provide empirical verification of our theoretical identifiability result using both simple 2-dimensional data and high-resolution imaging datasets.

1 Introduction

It has been hypothesized that developing machine learning (ML) systems capable of human-level understanding requires imbuing them with notions of objectness [48, 65]. Objectness notions can be characterised as physical, abstract, semantic, geometric, or via spaces and boundaries [83, 17]. Humans can generalise across environments with few examples to learn from [70], and this has been attributed to our ability to segregate percepts into object entities [64, 25, 45, 4].

Obtaining object-centric representations is deemed to be a key step for achieving true compositional generalization [5, 48, 3, 22], and uncovering causal influence between discrete concepts and their environment [57, 19, 20, 65, 4]. Significant progress in learning object-centric representations has been made [15, 16, 44, 66, 10, 67], particularly in unsupervised object discovery settings using an iterative attention mechanism known as Slot Attention (SA) [53]. However, most existing work approaches object-centric representation learning empirically, leaving theoretical understanding relatively underdeveloped. Establishing the identifiability [31, 29] of representations is important as it clarifies under which conditions object-centric representation learning is theoretically possible [6].

A well-known result shows that identifiability of latent variables is fundamentally impossible without assumptions about the data generating process [31, 51]. Therefore, understanding when object representations can theoretically be identified is important to scale object-centric methods to high-dimensional images. Recent works [6, 47] make important advances on this by explicitly

stating the set of assumptions necessary for providing theoretical identifiability of object-centric representations. However, they restrict their attention to properties of the mixing function , studying a class of models with additive decoders. Although there are merits to this approach, there are practical challenges with the so-called compositional contrast objective [6], as it involves computing Jacobians and requires second-order optimization via gradient descent. Consequently, satisfying the identifiability conditions explicitly (e.g. compositional contrast must be zero) is computationally restrictive for moderately high-dimensional data. In this work,

we present a probabilistic perspective that is not subject to the same scalability issues while still providing theoretical identifiability of object-centric representations without supervision. In Table 1, we list object-centric learning methods, their (sometimes implicit) modelling assumptions (see § 5 for additional information and Appendix A for a detailed breakdown and discussion), and their respective identifiability guarantees of object representations. Most methods do not guarantee identifiability, and make the B -disentanglement (1) and additive decoder (2) assumptions. Brady et. al. [6] do provide identifiability guarantees and additionally assume irreducibility (3) and compositionality (4). Our method provides identifiability guarantees by introducing latent structure (i.e. via a GMM prior) which generalizes to non-additive decoders. This is advantageous as the computational complexity of additive decoders scales linearly with the number of slots K - whereas our approach is invariant to K . Moreover, non-additive decoders have been found to significantly improve performance in practice [66, 67, 69], though the theoretical basis is underexplored. Finally, latent structure can reduce the complexity burden on the mixing function f (decoder), making it easier to learn in practice [18, 43].

Contributions. Our main contributions are the following: (i) We prove that object-centric representations (i.e. slots) are identifiable without supervision up to an equivalence relation (§ 5) under a latent mixture model specification. To that end, we propose a probabilistic slot-attention algorithm (§ 4) which imposes an aggregate mixture prior over slot representations. (ii) We show that our approach induces a non-degenerate ( global ) Gaussian Mixture Model (GMM) by aggregating per-datapoint ( local ) GMMs, providing a slot prior which: (a) is empirically stable across runs (i.e. identifiable up to affine transformations and slot permutations); (b) can be tractably sampled from. (iii) We provide conclusive empirical evidence of our theoretical object-centric identifiability result, including visual verification on synthetic 2-dimensional data as well as standard imaging benchmarks (§ 6).

2 Related Work

Identifiable Representation Learning. Identifiability of representations stems from early work in independent component analysis (ICA) [31, 29], and is making a resurgence recently [30, 32, 51, 37, 75, 46, 82]. Common strategies for tackling this identifiability problem are: (i) restricting the class of mixing functions; (ii) using non-i.i.d data, interventional data or counterfactuals; and (iii) imposing structure in the latent space via distributional assumptions. Regarding (i), restricting the class of the mixing functions to conformal maps [8] or volume-preserving transformations [81] has been found to produce identifiable models. For (ii), prior works [84, 52, 7, 2, 75] assume access to contrastive pairs of observations ( x , ˜ x ) obtained from either data augmentation, interventions, or approximate counterfactual inference. As for (iii), latent space structure is enforced via either: (a) using auxiliary variables to make latent variables conditionally independent [33, 37, 38]; or (b) distributional assumptions such as placing a mixture prior over the latent variables in a V AE [11, 80, 43]. In this work, we prove an identifiability result via strategy (iii) but within an object-centric learning context, where the latent variables are a set of object slots [53].

Object-Centric Learning. Much early work on unsupervised representation learning is based on the Variational Autoencoder (V AE) framework [41], and relies on independence assumptions between latent variables to learn so-called disentangled representations [5, 24, 40, 12, 59]. These methods are closely linked to object-centric representation learning [9, 15, 21], as they also leverage (iterative) variational inference procedures [58, 73, 50]. Alternatively, an iterative attention mechanism known as slot attention (SA) [53] has been the focus of much follow-up work recently [16, 67, 76, 68, 14].

Local GMMs

Encoder

Aggregate GMM

Sampled Slots

Generator

Environment

23

X1

X2

X3

文

Although slot attention-based methods show promising object binding [22] capabilities empirically on select datasets, they do not provide identifiability guarantees on the learned representations. Recently, [55, 56] assume the access to interventional data-generating process following [1] to demonstrate the identifiability of object-centric representations, while [6, 47] presented the identifiability results for object representations (i.e. slots), clarifying the necessary assumptions and properties of the mixing function (e.g. additive decoders). However, satisfying Brady et. al. [6]'s compositional contrast identifiability condition explicitly (must be zero) requires computationally restrictive secondorder optimization. In contrast, we shift the focus to learning structured object-centric latent spaces via Probabilistic Slot Attention (PSA), bridging the gap between generative model identifiability literature and object-centric representation learning. Notably, our PSA approach is also related to probabilistic capsule routing [27, 63, 62, 61] since slots are equivalent to universal capsules [26], but like slot attention, offers output permutation symmetry and does not face scalability issues.

3 Background

Notation. Let X ⊆ R H × W × C denote the input image space, where each image x is of size H × W pixels with C channels. Let f e : X → Z denote an encoder mapping image space to a latent space Z ⊆ R N × d , where each latent variable z consists of N , d -dimensional vectors. Lastly, let f d : S → X denote a decoder mapping from slot representation space S ⊆ R K × d to image space.

Slot Attention. Slot attention [53] receives a set of feature embeddings z ∈ R N × d per input x , and applies an iterative attention mechanism to produce K object-centric representations called slots s ∈ R K × d . Let W k , W v denote key and value transformation matrices acting on z , and W q the query transformation matrix acting on s . To simplify our exposition later on, let f s : Z × S → S be shorthand notation for the slot update function, defined as:

where q = W q s t ∈ R K × d , k = W k z ∈ R N × d , and v = W v z ∈ R N × d correspond to the query, key and value vectors respectively and A ∈ R K × N is the attention matrix. Unlike self-attention [74], the queries q in slot attention are a function of the slots s t , and are iteratively refined over T iterations. The initial slots s t =0 are randomly sampled from a standard Gaussian. The queries at iteration t are given by ˆ q t = W q s t , and the slot update process can be summarized as in Equation 1.

Compositionality. Compositionality as defined by Brady et al. [6] is a structure imposed on the slot decoder mapping f d which implies that each image pixel is a function of at most one slot representation, thereby enforcing a local sparsity structure on the Jacobian matrix of f d .

M

x

z

s

1

s

0

N

(0

,

I

)

s

2

s

T

. . .

(a) Slot Attention Encoder

z

u

µ

σ

π

x

K

N

M

(b) Probabilistic Slot Attention

z

u

̂

µ

̂

σ

̂

π

MK

N

(c) Aggregate Posterior Mixture

Definition 1 (Compositional Contrast) . For a differentiable mapping f d : Z → X , the compositional contrast of f d at z is given by:

Brady et al. [6]'s main result (Theorem 1) relies on compositionality and invertibility of f d to guarantee slot-identifiability when both the compositional contrast and the reconstruction loss equal zero. However, using C comp ( f d , z ) as a metric or as part of an objective function is computationally prohibitive 2 . Our method aims to minimize C comp ( f d , z ) implicitly [47], without additive decoders.

4 Probabilistic Slot Attention

In this section, we present a probabilistic slot attention framework which imposes a mixture prior structure over the slot latent space. This structure will prove to be instrumental in establishing our main identifiability result in Section 5. We begin by approaching standard slot attention [53] from a graphical modelling perspective. As shown in Figure 2 and explained in Section 3, applying slot attention to a deterministic encoding z = f e ( x ) ∈ R N × d yields a set of K object slot representations s 1: K := s 1 , . . . , s K . This process induces a stochastic encoder q ( s 1: K | x ) , where the stochasticity comes from the random initialization of the slots: s t =0 1: K ∼ N ( s 1: K ; 0 , I ) ∈ R K × d . Since each slot is a deterministic function of its previous state s t := f s ( z , s t -1 ) it is possible to randomly sample initial states s 0 and obtain stochastic estimates of the slots. 3 However, since each transition depends on z , which in turn depends on the input x , we do not get a generative model we can tractably sample from . This can conceivably be remedied by placing a tractable prior over z and using the VAE framework along the lines of [77], however, here we propose an entirely different approach which does not require making additional variational approximations (see Appendix G for further discussion).

Local Slot Mixtures. Probabilistic slot attention augments standard slot attention by introducing a per-datapoint (i.e. local) Gaussian Mixture Model (GMM) for learning slot distributions. Intuitively, a local GMM can be understood as a way to cluster features within a given image, encouraging the grouping of similar features into object representations. However, unlike regular clustering, here the clustered points are dynamically transformed representations of the actual data. Specifically,

we use an encoder function f e that maps each image x i ∈ R H × W × C in the dataset { x i } i M =1 , to a latent spatial representation z i ∈ R N × d . The latent variable z may be deterministic or stochastic, and we consider the case where N < HW to reflect a modest downscaling with respect to (w.r.t.) x . The goal is to dynamically map each of the N , d -dimensional vector representations in each z , to one-ofK object slot distributions within a mixture. A local GMMcan be fit to each posterior latent representation z ∼ q ( z | x i ) 4 on the fly by maximizing likelihood:

where µ i = ( µ i 1 , . . . , µ iK ) , σ 2 i = ( σ 2 i 1 , . . . , σ 2 iK ) and π i = ( π i 1 , . . . , π iK ) are the respective means, diagonal covariances and mixing coefficients of the i th K -component mixture. Figure 2b illustrates the resulting probabilistic graphical model (PGM) in more detail.

To maximize the likelihood in Equation 2 per datapoint x i , we present a bespoke expectation-maximisation (EM) algorithm for slot attention, yielding closed-form update equations for the parameters as shown in Algorithm 1, and explained next.

Probabilistic Projections. A powerful property of slot attention and cross-attention more broadly [74], is its ability to decouple the agreement mechanism from the representational content. That is, the dot-product is used to measure agreement between each query (slot) vector and all the key vectors, to dictate how much of each value vector (content from z i ) should be represented in each slot's revised representation. To retain this flexibility and decouple the attention computation from the content, we incorporate key-value projections into our probabilistic approach. For brevity, the i subscript is im-

Algorithm 1 Probabilistic Slot Attention

plicit in the following, keeping in mind that these are local quantities (per-datapoint x i ). The parameters of the K Gaussian slot distributions are initialized (at attention iteration t = 0 ) as follows:

The respective queries q , keys k , and values v are then given by:

where W q , W k , W v ∈ R d × d , whereas q ( t ) denotes the queries at attention iteration t . To measure agreement between each input feature (key) and slot (query), we evaluate the normalized probability density of each key under a Gaussian model defined by each slot:

where A nk corresponds to the posterior probability that slot (query) k is responsible for input feature (key) n . This process yields the slot attention matrix A ∈ R N × K . As shown in Algorithm 1, the mixture parameters π ( t ) , µ ( t ) , σ ( t ) 2 are then updated using the attention matrix and the values v . If the values are chosen to be equal to the keys v := k , then the procedure is more in line with standard EM, but the agreement mechanism and the content become entangled. After T probabilistic slot attention iterations, the resulting K Gaussians serve as slot posterior distributions:

where µ ( T ) and σ ( T ) 2 are the parameters of all the Gaussians in the mixture given a particular datapoint x . The slots s ( T ) 1: K are then used for input reconstruction, e.g. by maximizing a (possibly Gaussian) likelihood p ( x | s ( T ) 1: K ) parameterized by a (possibly additive) decoder f d .

-

10

-

5

0

5

10

z

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

p

(

z

)

GMM #1

GMM #2

GMM #3

Agg

.

GMM

-

10

-

5

0

5

10

z

0

.

0

0

.

1

0

.

2

0

.

3

GMM #1

GMM #2

GMM #3

Agg

.

GMM

-

10

-

5

0

5

10

z

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

GMM #1

GMM #2

GMM #3

Agg

.

GMM

-

10

-

5

0

5

10

z

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

GMM #1

GMM #2

GMM #3

Agg

.

GMM

Computational Complexity. Probabilistic slot attention (PSA) retains the O ( TNKd ) computational complexity of slot attention. The additional operations we introduce for calculating slot mixing coefficients and slot variances (under diagonal slot covariance structure) have complexities of O ( NK ) and O ( NKd ) respectively, which do not alter the dominant term. When PSA is combined with an additive decoder, it can lower computational complexity by eliminating the need to decode inactive slots. In the following, we outline a principled approach for pruning inactive slots.

Automatic Relevance Determination of Slots. An open problem in slot-based modelling is the dynamic estimation of the number of slots K needed for each input [51, 44]. Probabilistic slot attention offers an elegant solution to this problem using the concept of Automatic Relevance Determination (ARD) [60]. ARD is a statistical framework that prunes irrelevant features by imposing data-driven, sparsity-inducing priors on model parameters to regularize the solution space. Since the output mixing coefficients π ( T ) ∈ R K are input dependent (i.e. local), irrelevant components (slots) will naturally be pruned after T attention iterations, i.e. π ( T ) k → 0 for any unused slot k . We can either use a probability threshold τ ∈ [0 , 1) to prune unused slots or place a Dirichlet prior over the mixing coefficients to explicitly induce sparsity. For simplicity, we take the former approach:

where s τ denotes the set of active slots with mixing coefficient greater than τ , and each slot is (optionally) sampled from its Gaussian distribution: s ( T ) k ∼ N ( µ ( T ) k , σ ( T ) 2 k ) .

Aggregate Posterior Gaussian Mixture. As previously explained, probabilistic slot attention goes beyond standard slot attention by introducing a per-datapoint (i.e. local) GMM to learn distributions over slot representations. This imposes structure over the latent space and gives us access to posterior slot distributions after the attention iterations. Rather than constraining slot posteriors to be close to a tractable prior - e.g. via the V AE framework [41, 76] which requires further variational approximations - we leverage our probabilistic setup to compute the optimal (global) prior over slots.

Wepropose to compute the aggregate slot posterior by marginalizing out the data: q ( z ) = ∑ M i =1 q ( z | x i ) /M , given a pre-trained probabilistic slot attention model (Fig. 3). In § 5, we prove that the aggregate posterior is a tractable, non-degenerate Gaussian mixture distribution which: (i) Serves as the theoretically optimal prior over slots; (ii) Is empirically stable across runs (i.e. identifiable up to an affine transformation and slot permutation, § 5); (iii) Can be tractably sampled from and (optionally) used for slot-based scene composition tasks. Since GMMs are universal density approximators given enough components (even GMMs with diagonal covariances), the resulting aggregate posterior q ( z ) is highly flexible and multimodal. It often suffices to approximate it using a sufficiently large subset of the dataset, if marginalizing out the entire dataset becomes computationally restrictive, although we did not observe this to be the case in practice in our set of experiments.

5 Theory: Slot Identifiability Result

In this section, we leverage the properties of the probabilistic slot attention method proposed in Section 4 to prove a new object-centric identifiability result. We show that object-centric representations (i.e. slots) are identifiable without supervision (up to an equivalence relation) under mixture model-like assumptions about the latent space. This contrasts with existing work, which provides identifiability guarantees within a specific class of mixing functions, i.e. additive decoders [47]. Our result unifies generative model identifiablity [32, 37, 43] and object-centric learning.

Definition 2 (Identifiability.) . Given an observation space X , a probabilistic model p with parameters θ ∈ Θ is said to be identifiable if the mapping θ ∈ Θ ↦→ p θ ( x ) is injective:

Remark 1 . Definition 2 says that if any two choices of model parameters lead to the same marginal density, they are equal. This is often referred to as strong identifiability [32, 37], and it can be too restrictive, as guaranteeing identifiability up to a simple transformation (e.g. affine) is acceptable in practice. To reflect weaker notions of identifiability, we let ∼ denote an equivalence relation on Θ , such that a model can be said to be identifiable up to ∼ , or ∼ -identifiable.

Definition 3 ( ∼ s -equivalence) . Let f θ : S → X denote a mapping from slot representation space S to image space X (satisfying Assumption 8), the equivalence relation ∼ s w.r.t. to parameters θ ∈ Θ is defined bellow, where P ∈ P ⊆ { 0 , 1 } K × K is a slot permutation matrix, H ∈ R d × d is an affine transformation matrix, and c ∈ R d :

Lemma 1 (Aggregate Posterior Mixture) . Given that probabilistic slot attention induces a local (per-datapoint x ∈ { x i } i M =1 ) GMM with K components, the aggregate posterior q ( z ) obtained by marginalizing out x is a non-degenerate global Gaussian mixture with MK components:

Proof Sketch. The full proof is given in Appendix B. The result is obtained by integrating the product of the latent variable posterior density q ( z | x ) and the (local) GMM density given z , w.r.t. x . We then proceed by verifying that the mixing coefficients sum to one over all the components in the new mixture (Corollary 4), proving aggregated posterior to be a well-defined probability distribution, this can be empirically confirmed in Figure 3. Next, we use q ( z ) in our identifiablity result.

Theorem 1 (Mixture Distribution of Concatenated Slots) . Let f s denote a permutation equivariant PSA function such that f s ( z , P s t ) = Pf s ( z , s t ) , where P ∈ { 0 , 1 } K × K is an arbitrary permutation matrix. Let s = ( s 1 , . . . , s K ) ∈ R Kd be a random variable defined as the concatenation of K individual slots, where each slot is Gaussian distributed within a K -component mixture: s k ∼ N ( µ k , Σ k ) ∈ R d , ∀ k ∈ { 1 , . . . K } . Then, s is also GMM distributed with K ! mixture components:

Proof Sketch. The proof is in Appendix B. We observe that the permutation equivariance of the PSA function f s induces K ! ways of concatenating sampled slots s k , where each permutation maps to a different mode with block diagonal covariance structure in a GMM living in R Kd (e.g. Fig. 6, 7).

Theorem 2 ( ∼ s -Identifiable Slot Representations) . Given that the aggregate posterior q ( z ) is an optimal, non-degenerate mixture prior over slot space (Lemma 1), f : S → X is a piecewise affine weakly injective mixing function (Assumption 8), and the slot representation, s = ( s 1 , . . . , s K ) can be observed as a sample from a GMM (Theorem 1), then p ( s ) is identifiable as per Definition 3.

Proof Sketch. The proof is given in Appendix D. Lemma 1 and Corollary 4 state that the optimal latent variable prior in our case is GMM distributed, non-degenerate and equates to the aggregate posterior q ( z ) . This permits us to extend [43]'s result to show that if q ( z ) is distributed according to a non-degenerate GMM and the mixing function f d is piecewise affine and weakly injective, then the slot distribution representation, p ( s ) which is also a GMM (Theorem 1) is identifiable up to an affine transformation and arbitrary slot permutation.

Corollary 3 (Individual Slot Identifiability) . If the distribution over concatenated slots p ( s ) , where s = ( s 1 , . . . , s K ) ∈ R Kd , is ∼ s -identifiable (Theorem 2) then this implies q ( z ) is identifiable up to an affine transformation and permutation of the slots s k . Therefore, each slot distribution s k ∼ N ( µ k , Σ k ) ∈ R d , ∀ k ∈ { 1 , . . . K } is also identifiable up to an affine transformation.

z

2

Run #1

z

1

z

2

Run #2

z

1

z

2

Run #3

z

1

z

2

Run #4

z

1

z

2

Run #5

z

1

AE

SA

PSA-Proj (Ours)

PSA (Ours)

PSA-NoA (Ours)

SA-NoA

MONET

AE

SA

PSA-Proj (Ours)

PSA

PSA-NoA (Ours)

SA-NoA

MONET

0.7

SlotMeanCorrelationCoefficient

0.6

0.5

0.4

0.3

0.4

0.3

Slot IdentifiabilityScore

0.2

0.1

0.0

6 Experiments

Given that the focus of this work is theoretical, the primary goal of our experiments is to provide strong empirical evidence of our main identifiability result (ref. Figures 4, 5). With that said, we also extend our experimental study to popular imaging benchmarks to demonstrate that our method scales to higher-dimensional settings (ref. Tables 2, 4).

Datasets & Evaluation Metrics. Our experimental analysis involves standard benchmark datasets from object-centric learning literature including SPRITEWORLD [6], CLEVR [34], and OBJECTSROOM [35]. We report the foreground-adjusted rand index (FG-ARI) and FID [23] to quantify both object-level binding capabilities and image quality. Our main goal is to measure slot-identifiability, so we use the slot identifiability score [6] and the mean correlation coefficient (MCC) across slot representations - we call the latter slot-MCC (SMCC). For two sets of slots { s i } i M =1 , and { ˜ s i } i M =1 , where s i ∈ R K × d , ˜ s i ∈ R K × d , extracted from M images { x i } i M =1 , the SMCC between any s and ˜ s is obtained by matching the slot representations and their order. The order is matched by mapping slots in ˜ s w.r.t s assigned by π , followed by a learned affine mapping A between aligned ˜ s π ( i ) and s :

For more details on the metrics please refer to Appendix F.

Models & Baselines. We consider three ablations on our proposed probabilistic slot attention (PSA) method: (i) PSA base model (Algorithm 2); (ii) PSA-PROJ model (Algorithm 1); and (iii) PSA-NOA model, which is equivalent to PSA-PROJ but without an additive decoder. We experiment with two types of decoders: (i) an additive decoder similar to [79]'s spatial broadcasting model; and (ii) standard convolutional decoder. In all cases, we use LeakyReLU activations to satisfy the weak

-

15

-

10

-

5

0

5

10

15

0

.

00

0

.

05

0

.

10

0

.

15

0

.

20

0

.

25

0

.

30

N

(

µ

1

, σ

2

1

)

N

(

µ

2

, σ

2

2

)

N

(

µ

3

, σ

2

3

)

GMM

-

10

-

5

0

5

10

15

-

10

-

5

0

5

10

15

-

10

-

5

0

5

10

15

(

s

1

,

s

2

,

s

3

)

(

s

1

,

s

3

,

s

2

)

(

s

2

,

s

1

,

s

3

)

(

s

2

,

s

3

,

s

1

)

(

s

3

,

s

1

,

s

2

)

(

s

3

,

s

2

,

s

1

)

injectivity conditions (Assumption 8). In terms of object-centric learning baselines, we compare with standard additive autoencoder setups following [6], slot-attention (SA) [53], and MONET [9].

Verifying Slot Identifiability: Gaussian Mixture of Objects. To provide conclusive empirical evidence of our identifiability claim (Theorem 2), we set up a synthetic modelling scenario under which it is possible to visualize the aggregate posterior across different runs. The goal is to show that PSA is ∼ s -identifiable (Theorem 2) in the sense that it can recover the same latent space distribution up to an affine transformation and slot order permutation. For the data generating process, we defined a K =5 component GMM, with differing mean parameters { µ 1 , . . . , µ 5 } , and shared isotropic covariances. The 5 components emulate 5 different object types in a given environment. To create a single data point, we randomly chose 3 of the 5 components and sampled 128 points uniformly at random from each mode. Figure 9 shows some data samples, where different colours correspond to different objects. We used 1000 data points in total for training our PSA model. As shown in Figure 4, the aggregate posterior is either rotated, translated, skewed, or flipped across different runs as predicted by our theory - this contrasts with all baselines wherein the aggregate posterior is intractable. We observed an SMCC of 0 . 93 ± 0 . 04 , and R2-score of 0 . 50 ± 0 . 08 .

Case Study: Imaging Data. Although our focus is primarily theoretical, we now demonstrate that our method generalizes/scales to higher-dimensional imaging modalities. To that end, we first use the SPRITEWORLD [78] dataset to evaluate the SMCC and SIS w.r.t. ground truth latent variables. Figure 5 presents our identifiability results against the baselines. Similar to [6], we observe higher SIS when compositional contrast and reconstruction error decreases. However, when the mixing function is not additive, the compositional contrast does not decrease drastically while maintaining higher SMCC and SIS - this verifies our identifiability claim using only piecewise affine decoders. As shown in Figure 5, we also observe that PSA routing with additive decoder models achieves lower compositional contrast and reconstruction errors when compared with other methods. This indicates that stronger identifiability of slot representations is achievable when combining both slot latent structure and inductive biases in the mixing function. Unlike for the SPRITEWORLD dataset, all the ground truth generative factors are unobserved for the CLEVR and OBJECTSROOM datasets we use. Therefore, for evaluation in these cases, we train multiple models with different seeds and compute the SMCC and SIS measures across different runs. This is similar to our earlier synthetic experiment and standard practice in identifiability literature. Table 2 presents our main identifiability results on CLEVR and OBJECTSROOM. We observe similar trends in favour of our proposed PSA method as measured by both SMCC and (slot averaged) R2 score relative to the baselines.

Case Study: Complex Decoder Structure. To empirically test slot identifiability using more complex non-additive decoders, we used transformer decoders following SLATE [67], and simply replaced the slot attention module with probabilistic slot attention. On the CLEVR dataset, we observed a significantly improved SMCC of 0.73 ± . 01 and R2 of 0.55 ± . 06 relative to Table 2.

To demonstrate that probabilistic slot attention can scale to large-scale real-world data we ran additional experiments on the Pascal VOC2012 [49] dataset, following the exact DINOSAUR strategies and setups described in [66, 36] for fairness, then simply swapping out the slot attention module with probabilistic slot attention. As shown in Table 3, we find that probabilistic slot attention is competitive with standard slot attention on real-world data. We also tested more complex, non-additive decoders based on autoregressive transformers, following the DINOSAUR [66] setup. For our PSA TRANSFORMER (W/ DINO) model, we observed an MBO i of 0.447 , and MBO c of 0.521 which is competitive with the DINOSAUR TRANSFORMER baseline [66] of 0.44 and 0.512 respectively. In this case, we found that a lower maximum learning rate of 10 -4 was beneficial for stabilizing

PSA training. In summary, our experiments corroborate our theoretical results and suggest why non-additive decoder structures can still work well given the appropriate latent structure and inference procedures are in place. With that said, there is a trade-off between identifiability and expressivity induced by the choice of decoder structure [47], so depending on the use case, it may indeed be advantageous to combine both latent and additive decoder structures in practice.

7 Discussion

Understanding when object-centric representations can theoretically be identified is important for scaling slot-based methods to high-dimensional images with correctness guarantees. In contrast with existing work, which focuses primarily on properties of the slot mixing function , we leverage distributional assumptions about the slot latent space to prove a new slot-identifiability result. Specifically, we prove that object-centric representations are identifiable without supervision (up to an equivalence relation) under mixture model-like distributional assumptions on the latent slots. To that end, we proposed a probabilistic slot-attention algorithm that imposes an aggregate mixture prior over slot representations which is both demonstrably stable across runs and tractable to sample from. Our empirical study primarily verifies our theoretical identifiability claim and demonstrates that our framework achieves the lowest compositional contrast without being explicitly trained towards that objective, which is computationally infeasible beyond toy datasets. In summary, we show how slot identifiability can be achieved via probabilistic constraints on the latent space and piecewise decoders. These piecewise decoders manifest as e.g. standard MLPs with LeakyReLU activations and are generally less restrictive than additive decoders. When coupling probabilistic and additive decoder structures, we observe further performance improvements relative to either one in isolation.

Limitations & Future Work. We recognize that our assumptions, particularly the weak injectivity of the mixing function, may not always hold in practice for different types of architectures (see Appendix C for sufficiency conditions). Although generally applicable, the piecewise-affine functions we use may not always accurately reflect valid assumptions about real-world problems, e.g. when the model is misspecified. Like all object-centric learning methods, we also assume that the mixing function is invariant to permutations of the slots in practice, which technically makes it non-invertible. We deem this aspect to be an interesting area for future work, as an extension to accommodate permutation invariance would strengthen and generalize the identifiability guarantees we provide. Additionally, we do not study cases where objects are occluded, i.e. when are shared or bordering other objects. This limitation is not unique to our work [53, 6, 14, 15, 44] and overcoming it requires further investigation by the community. Nonetheless, our theoretical results capture the important concepts in object-centric learning and represent a valuable extension to the nascent theoretical foundations of the area. In future work, it would be valuable to further relax slot identifiability requirements/assumptions and study slot compositional properties of probabilistic slot attention.

References

Kartik Ahuja, Jason S Hartford, and Yoshua Bengio. Weakly supervised representation learning with sparse perturbations. Advances in Neural Information Processing Systems , 35:1551615528, 2022. [Cited on page 3.]

Kartik Ahuja, Yixin Wang, Divyat Mahajan, and Yoshua Bengio. Interventional causal representation learning. arXiv preprint arXiv:2209.11924 , 2022. [Cited on page 2.]

Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv:1806.01261 , 2018. [Cited on page 1.]

Timothy EJ Behrens, Timothy H Muller, James CR Whittington, Shirley Mark, Alon B Baram, Kimberly L Stachenfeld, and Zeb Kurth-Nelson. What is a cognitive map? organizing knowledge for flexible behavior. Neuron , 100(2):490-509, 2018. [Cited on page 1.]

Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions on pattern analysis and machine intelligence , 35(8):17981828, 2013. [Cited on pages 1 and 2.]

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. [Cited on pages 1, 2, 3, 4, 8, 9, 10, 17, 21, and 24.]

Johann Brehmer, Pim De Haan, Phillip Lippe, and Taco S Cohen. Weakly supervised causal representation learning. Advances in Neural Information Processing Systems , 35:38319-38331, 2022. [Cited on page 2.]

Simon Buchholz, Michel Besserve, and Bernhard Schölkopf. Function classes for identifiable nonlinear independent component analysis. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems , 2022. [Cited on page 2.]

Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. arXiv preprint arXiv:1901.11390 , 2019. [Cited on pages 2 and 9.]

Michael Chang, Thomas L Griffiths, and Sergey Levine. Object representations as fixed points: Training iterative refinement algorithms with implicit differentiation. arXiv preprint arXiv:2207.00787 , 2022. [Cited on page 1.]

Nat Dilokthanakul, Pedro AM Mediano, Marta Garnelo, Matthew CH Lee, Hugh Salimbeni, Kai Arulkumaran, and Murray Shanahan. Deep unsupervised clustering with gaussian mixture variational autoencoders. arXiv preprint arXiv:1611.02648 , 2016. [Cited on pages 2 and 25.]

Cian Eastwood and Christopher K. I. Williams. A framework for the quantitative evaluation of disentangled representations. In International Conference on Learning Representations , 2018. [Cited on page 2.]

Gamaleldin Elsayed, Aravindh Mahendran, Sjoerd van Steenkiste, Klaus Greff, Michael C Mozer, and Thomas Kipf. Savi++: Towards end-to-end object-centric learning from real-world videos. Advances in Neural Information Processing Systems , 35:28940-28954, 2022. [Cited on page 2.]

Patrick Emami, Pan He, Sanjay Ranka, and Anand Rangarajan. Slot order matters for compositional scene understanding. arXiv preprint arXiv:2206.01370 , 2022. [Cited on pages 2 and 10.]

Martin Engelcke, Adam R Kosiorek, Oiwi Parker Jones, and Ingmar Posner. Genesis: Generative scene inference and sampling with object-centric latent representations. arXiv preprint arXiv:1907.13052 , 2019. [Cited on pages 1, 2, and 10.]

Martin Engelcke, Oiwi Parker Jones, and Ingmar Posner. Genesis-v2: Inferring unordered object representations without iterative refinement. Advances in Neural Information Processing Systems , 34:8085-8094, 2021. [Cited on pages 1 and 2.]

Russell A Epstein, Eva Zita Patai, Joshua B Julian, and Hugo J Spiers. The cognitive map in humans: spatial navigation and beyond. Nature neuroscience , 20(11):1504-1513, 2017. [Cited on page 1.]

Fabian Falck, Haoting Zhang, Matthew Willetts, George Nicholson, Christopher Yau, and Chris C Holmes. Multi-facet clustering variational autoencoders. Advances in Neural Information Processing Systems , 34:8676-8690, 2021. [Cited on page 2.]

Tobias Gerstenberg, Noah D Goodman, David A Lagnado, and Joshua B Tenenbaum. A counterfactual simulation model of causal judgments for physical events. Psychological review , 128(5):936, 2021. [Cited on page 1.]

Alison Gopnik, Clark Glymour, David M Sobel, Laura E Schulz, Tamar Kushnir, and David Danks. A theory of causal learning in children: causal maps and bayes nets. Psychological review , 111(1):3, 2004. [Cited on page 1.]

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 , pages 2424-2433. PMLR, 2019. [Cited on page 2.]

Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. arXiv preprint arXiv:2012.05208 , 2020. [Cited on pages 1 and 3.]

Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to a local nash equilibrium. Advances in neural information processing systems , 30, 2017. [Cited on page 8.]

Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-VAE: Learning basic visual concepts with a constrained variational framework. In International Conference on Learning Representations , 2017. [Cited on page 2.]

Geoffrey Hinton. Some demonstrations of the effects of structural descriptions in mental imagery. Cognitive Science , 3(3):231-250, 1979. [Cited on page 1.]

Geoffrey Hinton. How to represent part-whole hierarchies in a neural network. Neural Computation , pages 1-40, 2022. [Cited on page 3.]

Geoffrey E Hinton, Sara Sabour, and Nicholas Frosst. Matrix capsules with em routing. In International conference on learning representations , 2018. [Cited on page 3.]

Matthew D Hoffman and Matthew J Johnson. Elbo surgery: yet another way to carve up the variational evidence lower bound. In Workshop in Advances in Approximate Bayesian Inference, NIPS , volume 1, 2016. [Cited on pages 18 and 25.]

A Hyvarinen and E Oja. Independent component analysis: algorithms and applications. Neural Networks , 13(4-5):411-430, 2000. [Cited on pages 1 and 2.]

Aapo Hyvarinen and Hiroshi Morioka. Unsupervised feature extraction by time-contrastive learning and nonlinear ica. Advances in neural information processing systems , 29, 2016. [Cited on page 2.]

Aapo Hyvärinen and Petteri Pajunen. Nonlinear independent component analysis: Existence and uniqueness results. Neural networks , 12(3):429-439, 1999. [Cited on pages 1 and 2.]

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 , pages 859-868. PMLR, 2019. [Cited on pages 2, 6, and 7.]

Aapo Hyvarinen, Hiroaki Sasaki, and Richard Turner. Nonlinear ica using auxiliary variables and generalized contrastive learning. In Proceedings of the Twenty-Second International Conference on Artificial Intelligence and Statistics , volume 89, pages 859-868. PMLR, 2019. [Cited on page 2.]

Justin Johnson, Bharath Hariharan, Laurens Van Der Maaten, Li Fei-Fei, C Lawrence Zitnick, and Ross Girshick. Clevr: A diagnostic dataset for compositional language and elementary visual reasoning. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 2901-2910, 2017. [Cited on page 8.]

Rishabh Kabra, Chris Burgess, Loic Matthey, Raphael Lopez Kaufman, Klaus Greff, Malcolm Reynolds, and Alexander Lerchner. Multi-object datasets. https://github.com/deepmind/multiobject-datasets/, 2019. [Cited on page 8.]

Ioannis Kakogeorgiou, Spyros Gidaris, Konstantinos Karantzalos, and Nikos Komodakis. Spot: Self-training with patch-order permutation for object-centric learning with autoregressive transformers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 22776-22786, 2024. [Cited on page 10.]

Ilyes Khemakhem, Diederik Kingma, Ricardo Monti, and Aapo Hyvarinen. Variational autoencoders and nonlinear ica: A unifying framework. In International Conference on Artificial Intelligence and Statistics , pages 2207-2217. PMLR, 2020. [Cited on pages 2, 6, 7, 17, and 24.]

Ilyes Khemakhem, Ricardo Monti, Diederik Kingma, and Aapo Hyvarinen. Ice-beem: Identifiable conditional energy-based deep models based on nonlinear ica. In Advances in Neural Information Processing Systems , volume 33, 2020. [Cited on pages 2, 17, and 24.]

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:12768-12778, 2020. [Cited on page 24.]

Hyunjik Kim and Andriy Mnih. Disentangling by factorising. In Proceedings of the 35th International Conference on Machine Learning , Proceedings of Machine Learning Research, 2018. [Cited on page 2.]

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013. [Cited on pages 2 and 6.]

Thomas Kipf, Gamaleldin F Elsayed, Aravindh Mahendran, Austin Stone, Sara Sabour, Georg Heigold, Rico Jonschkowski, Alexey Dosovitskiy, and Klaus Greff. Conditional object-centric learning from video. arXiv preprint arXiv:2111.12594 , 2021. [Cited on page 2.]

Bohdan Kivva, Goutham Rajendran, Pradeep Ravikumar, and Bryon Aragam. Identifiability of deep generative models without auxiliary information. Advances in Neural Information Processing Systems , 35:15687-15701, 2022. [Cited on pages 2, 6, 7, 17, 21, and 22.]

Avinash Kori, Francesco Locatello, Fabio De Sousa Ribeiro, Francesca Toni, and Ben Glocker. Grounded object centric learning. arXiv preprint arXiv:2307.09437 , 2023. [Cited on pages 1, 2, 6, and 10.]

Tejas D Kulkarni, William F Whitney, Pushmeet Kohli, and Josh Tenenbaum. Deep convolutional inverse graphics network. Advances in neural information processing systems , 28, 2015. [Cited on page 1.]

Sébastien Lachapelle, Pau Rodríguez López, Yash Sharma, Katie Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Nonparametric partial disentanglement via mechanism sparsity: Sparse actions, interventions and sparse temporal dependencies. arXiv preprint arXiv:2401.04890 , 2024. [Cited on page 2.]

Sébastien Lachapelle, Divyat Mahajan, Ioannis Mitliagkas, and Simon Lacoste-Julien. Additive decoders for latent variables identification and cartesian-product extrapolation. Advances in Neural Information Processing Systems , 36, 2024. [Cited on pages 1, 3, 4, 6, 10, and 17.]

Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel J Gershman. Building machines that learn and think like people. Behavioral and brain sciences , 40:e253, 2017. [Cited on page 1.]

Kunpeng Li, Ziyan Wu, Kuan-Chuan Peng, Jan Ernst, and Yun Fu. Tell me where to look: Guided attention inference network. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 9215-9223, 2018. [Cited on page 10.]

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. arXiv preprint arXiv:2001.02407 , 2020. [Cited on page 2.]

Francesco Locatello, Stefan Bauer, Mario Lucic, Gunnar Raetsch, Sylvain Gelly, Bernhard Schölkopf, and Olivier Bachem. Challenging common assumptions in the unsupervised learning of disentangled representations. In international conference on machine learning , pages 41144124. PMLR, 2019. [Cited on pages 1, 2, and 6.]

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 , pages 6348-6359. PMLR, 2020. [Cited on page 2.]

Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. Advances in Neural Information Processing Systems , 33:11525-11538, 2020. [Cited on pages 1, 2, 3, 4, 9, and 10.]

Sindy Löwe, Phillip Lippe, Francesco Locatello, and Max Welling. Rotating features for object discovery. Advances in Neural Information Processing Systems , 36, 2024. [Cited on page 10.]

Amin Mansouri, Jason Hartford, Kartik Ahuja, and Yoshua Bengio. Object-centric causal representation learning. In NeurIPS 2022 Workshop on Symmetry and Geometry in Neural Representations , 2022. [Cited on page 3.]

Amin Mansouri, Jason Hartford, Yan Zhang, and Yoshua Bengio. Object-centric architectures enable efficient causal representation learning. arXiv preprint arXiv:2310.19054 , 2023. [Cited on page 3.]

Gary F Marcus. The algebraic mind: Integrating connectionism and cognitive science . MIT press, 2003. [Cited on page 1.]

Joe Marino, Yisong Yue, and Stephan Mandt. Iterative amortized inference. In International Conference on Machine Learning , pages 3403-3412. PMLR, 2018. [Cited on page 2.]

Emile Mathieu, Tom Rainforth, N Siddharth, and Yee Whye Teh. Disentangling disentanglement in variational autoencoders. In Proceedings of the 36th International Conference on Machine Learning , 2019. [Cited on page 2.]

Radford M. Neal. Bayesian Learning for Neural Networks . Springer-Verlag, Berlin, Heidelberg, 1996. [Cited on page 6.]

Fabio De Sousa Ribeiro, Kevin Duarte, Miles Everett, Georgios Leontidis, and Mubarak Shah. Learning with capsules: A survey. arXiv preprint arXiv:2206.02664 , 2022. [Cited on page 3.]

Fabio De Sousa Ribeiro, Georgios Leontidis, and Stefanos Kollias. Capsule routing via variational bayes. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 34, pages 3749-3756, 2020. [Cited on page 3.]

Fabio De Sousa Ribeiro, Georgios Leontidis, and Stefanos Kollias. Introducing routing uncertainty in capsule networks. In Advances in Neural Information Processing Systems , volume 33, pages 6490-6502, 2020. [Cited on page 3.]

Irvin Rock. Orientation and form. 1973. [Cited on page 1.]

Bernhard Schölkopf and Julius von Kügelgen. From statistical to causal learning. Proceedings of the International Congress of Mathematicians , 2022. [Cited on page 1.]

Maximilian Seitzer, Max Horn, Andrii Zadaianchuk, Dominik Zietlow, Tianjun Xiao, CarlJohann Simon-Gabriel, Tong He, Zheng Zhang, Bernhard Schölkopf, Thomas Brox, et al. Bridging the gap to real-world object-centric learning. International Conference on Learning Representations , 2023. [Cited on pages 1, 2, and 10.]

Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate dall-e learns to compose. International Conference on Learning Representations , 2022. [Cited on pages 1, 2, and 9.]

Gautam Singh, Yeongbin Kim, and Sungjin Ahn. Neural block-slot representations. arXiv preprint arXiv:2211.01177 , 2022. [Cited on page 2.]

Gautam Singh, Yi-Fu Wu, and Sungjin Ahn. Simple unsupervised object-centric learning for complex and naturalistic videos. Advances in Neural Information Processing Systems , 35:18181-18196, 2022. [Cited on page 2.]

Joshua B Tenenbaum, Charles Kemp, Thomas L Griffiths, and Noah D Goodman. How to grow a mind: Statistics, structure, and abstraction. science , 331(6022):1279-1285, 2011. [Cited on page 1.]

Jakub Tomczak and Max Welling. Vae with a vampprior. In International Conference on Artificial Intelligence and Statistics , pages 1214-1223. PMLR, 2018. [Cited on page 25.]

Aaron Van Den Oord, Oriol Vinyals, et al. Neural discrete representation learning. Advances in neural information processing systems , 30, 2017. [Cited on page 25.]

Sjoerd Van Steenkiste, Karol Kurach, Jürgen Schmidhuber, and Sylvain Gelly. Investigating object compositionality in generative adversarial networks. Neural Networks , 130:309-325, 2020. [Cited on page 2.]

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017. [Cited on pages 3 and 5.]

Julius Von Kügelgen, Yash Sharma, Luigi Gresele, Wieland Brendel, Bernhard Schölkopf, Michel Besserve, and Francesco Locatello. Self-supervised learning with data augmentations provably isolates content from style. Advances in neural information processing systems , 34:16451-16467, 2021. [Cited on page 2.]

Yanbo Wang, Letao Liu, and Justin Dauwels. Slot-vae: Object-centric scene generation with slot attention. arXiv preprint arXiv:2306.06997 , 2023. [Cited on pages 2 and 6.]

Yanbo Wang, Letao Liu, and Justin Dauwels. Slot-VAE: Object-centric scene generation with slot attention. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett, editors, Proceedings of the 40th International Conference on Machine Learning , volume 202 of Proceedings of Machine Learning Research , pages 36020-36035. PMLR, 23-29 Jul 2023. [Cited on pages 4 and 25.]

Nicholas Watters, Loic Matthey, Sebastian Borgeaud, Rishabh Kabra, and Alexander Lerchner. Spriteworld: A flexible, configurable reinforcement learning environment. https://github.com/deepmind/spriteworld/, 2019. [Cited on page 9.]

Nicholas Watters, Loic Matthey, Christopher P Burgess, and Alexander Lerchner. Spatial broadcast decoder: A simple architecture for learning disentangled representations in vaes. arXiv preprint arXiv:1901.07017 , 2019. [Cited on page 8.]

Matthew Willetts and Brooks Paige. I don't need u: Identifiable non-linear ica without side information. arXiv preprint arXiv:2106.05238 , 2021. [Cited on page 2.]

Xiaojiang Yang, Yi Wang, Jiacheng Sun, Xing Zhang, Shifeng Zhang, Zhenguo Li, and Junchi Yan. Nonlinear ICA using volume-preserving transformations. In International Conference on Learning Representations , 2022. [Cited on page 2.]

Dingling Yao, Danru Xu, Sébastien Lachapelle, Sara Magliacane, Perouz Taslakian, Georg Martius, Julius von Kügelgen, and Francesco Locatello. Multi-view causal representation learning with partial observability. 2024. [Cited on page 2.]

Alan Yuille and Daniel Kersten. Vision as bayesian inference: analysis by synthesis? Trends in cognitive sciences , 10(7):301-308, 2006. [Cited on page 1.]

Roland S Zimmermann, Yash Sharma, Steffen Schneider, Matthias Bethge, and Wieland Brendel. Contrastive learning inverts the data generating process. In International Conference on Machine Learning , pages 12979-12990. PMLR, 2021. [Cited on page 2.]

-

2

-

1

0

1

2

0

.

0

0

.

5

1

.

0

1

.

5

2

.

0

2

.

5

N

(

µ

1

, σ

2

1

)

N

(

µ

2

, σ

2

2

)

N

(

µ

3

, σ

2

3

)

GMM

-

2

-

1

0

1

2

-

2

-

1

0

1

2

-

2

-

1

0

1

2

(

s

1

,

s

2

,

s

3

)

(

s

1

,

s

3

,

s

2

)

(

s

2

,

s

1

,

s

3

)

(

s

2

,

s

3

,

s

1

)

(

s

3

,

s

1

,

s

2

)

(

s

3

,

s

2

,

s

1

)

-

6

-

4

-

2

0

2

4

6

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

1

.

0

1

.

2

N

(

µ

1

, σ

2

1

)

N

(

µ

2

, σ

2

2

)

N

(

µ

3

, σ

2

3

)

GMM

-

7

.

5

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

7

.

5

-

7

.

5

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

7

.

5

-

7

.

5

-

5

.

0

-

2

.

5

0

.

0

2

.

5

5

.

0

7

.

5

(

s

1

,

s

2

,

s

3

)

(

s

1

,

s

3

,

s

2

)

(

s

2

,

s

1

,

s

3

)

(

s

2

,

s

3

,

s

1

)

(

s

3

,

s

1

,

s

2

)

(

s

3

,

s

2

,

s

1

)

A List of Assumptions

Assumption 1 ( B - Disentanglement, [47]) . Let s = { s B , ∀ B ∈ B} be a set of features wrt partition set B . The learned mixing function f d is said to be B disentangled wrt true decoder ˜ f d if there exists a permutation respecting diffeomorphism v B = ˜ f -1 d ◦ f d ∀ B ∈ B which for a given feature s can be expressed as v B ( s ) = v B ( s B ) .

Assumption 2 (Additive Mixing Function) . A mixing function f d : S → X , is said to be additive if there exist a partition set B and functions f B d : s → R |X| such that: f d ( s ) = ∑ B ∈B f B d ( s B ) .

Assumption 4 (Compositionality) . Compositionality as defined by is a structure imposed on the slot decoder f d which implies that each image pixel is a function of at most one slot representation, thereby enforcing a local sparsity structure on the Jacobian matrix of f d [6].

Assumption 3 (Irreducibility) . Given an an object x k ∈ x , a model is considered as irreducible if any subset of an object, y ⊆ x k is not functionally independent of the complement of the subset contained within the object, y c ∩ x as expressed by the Jacobin rank inequality in equation 5 in [6].

Remark 2 . The compositionality is guaranteed by explicitly minimising the compositionality contrast but was empirically observed to be implicitly satisfied in the case of additive decoders [6]. Later, [47] showed this as the property of additive decoder models. However, the additive decoders studied by [47] are not expressive enough to represent the 'masked decoders' typically used in object-centric representation learning, which stems from the normalization of the alpha masks. This means some care must be taken in extrapolating the results in [47] to the models we use in practice. Additionally, additive decoders scale linearly in the number of slots K , so some less significant scalability issues remain relative to state-of-the-art non-additive decoders (e.g. using Transformers).

Assumption 5 ( u task) . Conditioning latent variables on an observed variable to yield identifiable models. The main assumption is that conditioning on a (potentially observed) variable u renders the latent variables independent of each other [37].

Assumption 6 (Object Sufficiency) . A model is said to be object sufficient iff there are no additional objects in the original data distributions other than the ones expressed in training data.

Assumption 7 (Decoder Injectivity) . The function f d : S → X mapping from slot space to image space is a non-linear piecewise affine injective function. That is, it specifies a unique one-to-one mapping between slots and images.

Remark 3 . In practice, we use a monotonically increasing decoder with leakyReLU activation which should encourage injectivity behaviour [38, 37].

Assumption 8 ( Weak Injectivity [43]) . Let f : Z → X be a mapping between latent space and image space, where dim( Z ) ≤ dim( X ) . The mapping f d is weakly injective if there exists x 0 ∈ X and δ > 0 such that | f -1 ( { x } ) | = 1 , ∀ x ∈ B ( x 0 , δ ) ∩ f ( Z ) , and { x ∈ X : | f -1 ( { x } ) | = ∞} ⊆ f ( Z ) has measure zero w.r.t. to the Lebesgue measure on f ( Z ) .

Remark 4 . In words, Assumption 8 says that a mapping f d is weakly injective if: (i) in a small neighbourhood around a specific point x 0 ∈ X the mapping is injective - meaning each point in this neighbourhood maps to exactly one point in the latent space Z ; and (ii) while f d may not be globally injective, the set of points in X that map back to an infinite number of points in Z (non-injective points) is almost non-existent in terms of the Lebesgue measure on the image of Z under f d .

B Aggregate Posterior Proofs

Lemma 1 (Aggregate Posterior Mixture) Given that probabilistic slot attention induces a local (per-datapoint x ∈ { x i } i M =1 ) GMM with K components, the aggregate posterior q ( z ) obtained by marginalizing out x is a non-degenerate global Gaussian mixture with MK components given by:

Proof. We begin by noting that the aggregate posterior q ( z ) is the optimal prior p ( z ) so long as our posterior approximation q ( z | x ) is close enough to the true posterior p ( z | x ) , since for a dataset { x i } i M =1 we have that:

where we approximated p ( x ) using the empirical distribution, then substituted in the approximate posterior and marginalized out x . This observation was first made by [28] and we use it to motivate our setup.

In our case, probabilistic slot attention (Algorithm 1) fits a (local) GMM to each latent variable sampled from the approximate posterior: z ∼ q ( z | x i ) , for i = 1 , . . . , M . Let f ( z ) denote the (local) GMM density, its expectation is given by:

where we again used the empirical distribution approximation of p ( x ) , and the following basic identity of the Dirac delta to simplify: ∫ δ ( x -x ′ ) f ( x ) d x = f ( x ′ ) .

For the general case, however, we must instead compute the product of q ( z | x ) and f ( z ) rather than use a Dirac delta approximation as in Equation 23. To that end we may proceed as follows w.r.t. to

each datapoint x i :

where the posterior parameters of the resulting mixture are given in closed-form by:

which are the standard distributional parameters obtained from a product of two Gaussians.

For the updated mixture coefficients ̂ π ik , we propose a principled way to include a posterior-weighted contribution of each mode to the new mixture coefficients. First, note that ( π i 1 , π i 2 , . . . , π iK ) are parameters of a multinomial distribution as ∑ K k =1 π ik = 1 , for each datapoint x i . Since the Dirichlet distribution is the conjugate prior of the multinomial distribution, we can place a Dirichlet prior over the mixing coefficients for each datapoint, then update it to a posterior using the data. Concretely, we place a symmetric Dirichlet prior over the mixing coefficients as follows:

where α i ∈ R K are the concentration parameters of the i th Dirichlet distribution, and ∀ i, k : α ik = 1 , indicating uniformity over the open standard ( K -1) -simplex. To compute the posterior Dirichlet distribution we calculate 'pseudo-counts' by integrating the product of the posterior density q ( z | x i ) with each one of the K modes of the Gaussian mixture, thereby measuring a posterior-weighted contribution of each mode k to the new aggregate mixture:

which we can then use as pseudo-counts to compute the Dirichlet posterior:

for i = 1 , 2 , . . . , M . Each posterior probability is then readily given by the mean estimate

Putting everything together, the aggregated posterior is therefore given by:

which concludes the proof.

Corollary 4. The aggregate posterior q ( z ) is a non-degenerate Gaussian mixture, in the sense that it is a well-defined probability distribution, as the updated mixture coefficients sum to 1 over the number of components M × K .

Proof. Recall from Lemma 1 that the aggregate posterior q ( z ) - obtained by marginizaling out x from a probabilistic slot attention model - is a mixture distribution of M × K components with the following parameters:

To verify that q ( z ) is a non-degenerate mixture, we observe the following implication:

due to the Dirichlet posterior update in Equation 33, and therefore

which says that the scaled sum of the mixing proportions of all K components in all M GMMs must equal 1, proving that the associated aggregate posterior mixture q ( z ) is a well-defined probability distribution.

Theorem 1 (Mixture Distribution of Concatenated Slots) . Let f s denote a permutation equivariant probabilistic slot attention function such that f s ( z , P s t ) = Pf s ( z , s t ) , where P ∈ { 0 , 1 } K × K is an arbitrary permutation matrix. Let s = ( s 1 , s 2 , . . . , s K ) ∈ R Kd be a random variable defined as the concatenation of K individual sampled slots, where each slot is Gaussian distributed within a K -component mixture: s k ∼ N ( s k ; µ k , Σ k ) ∈ R d , ∀ k ∈ { 1 , . . . K } . Then, it holds that s is also Gaussian mixture distributed comprising K ! mixture components:

̸

Proof. Each slot s k ∼ N ( µ k , Σ k ) is sampled independently from s j ∼ N ( µ j , Σ j ) , for any j = k , meaning they are conditionally independent given the latent mixture component assignment. Thus, the concatenated slots variable s = ( s 1 , s 2 , . . . , s K ) , can be described by a Kd -dimensional multivariate Gaussian distribution with a block diagonal covariance structure as follows:

where π : [ K ] → [ K ] is a permutation function of the set [ K ] := { 1 , 2 , . . . , K } . Since the slot attention function f s is permutation equivariant, there exist K ! possible ways to concatenate K slots, and each permutation induces a mode within a Gaussian mixture living in R Kd space. Since each permutation of the slots is equally likely, the mixture coefficients are given by:

which concludes the proof.

Remark 5 . Based on the above result, it is evident that concatenating K ≥ 2 unique slots can be viewed as a sample from a GMM with K ! components. Constructing a scene requires at least two unique slots, one for the background and one for an object, thus supporting our theory regarding slot composition.

C Injective Decoders: Sufficient Conditions

In this section, we provide a theoretical overview of the decoder architecture we use and offer sufficient conditions for a leaky-ReLU decoder to be weakly injective in the sense of Assumption 8 as shown and adapted from [43].

Definition 4 (Piece-wise Decoders, [43]) . Let s = { s 1 , s 2 , . . . , s K } denote a given set of sampled z i in each mixture component ( K slots) in the GMM, P ( Z ) . Let σ : R → R denote the leaky-ReLU activation function, and let m = n 0 , n 1 , n 2 , . . . , n t = n and H ( n 1 , n 2 ) denote the set of full-rank affine functions h i : R n i → R n j . We consider piece-wise functions mapping each slot representation s ∈ S ∈ R m × K to an image x ∈ X ∈ R n in the output space, F mK → n σ : S → X , of the form below:

The following Corollary, corollary and proofs are adapted from [43].

Corollary 5. Given f d ∈ F m↪ → n σ where f d = h t ◦ σ ◦ h t -1 ◦ σ ◦ . . . σ ◦ h 1 , f d is injective.

Proof. Each affine function h i has full column rank and is therefore both injective and invertible. Since the activation function σ is also injective, we get that f d is injective and invertible.

Corollary 6. Let f d = h t ◦ σ ◦ h t -1 ◦ σ ◦ . . . σ ◦ h 1 ∈ F m → n σ where m = n 0 ≤ n 1 ≤ · · · ≤ n k = n . Given h i is affine and invertible, then for almost all x ∈ f d ( R m ) there exists δ such that f -1 d is a well-defined affine function, h i on B ( x, δ ) ∩ f d ( R m ) .

Proof. We know f d is a piecewise affine function which is invertible. Therefore, it simply follows, ∀ y ∈ B ( x, δ ) there exists δ such that f -1 d is an affine function in the domain B ( x, δ ) . One can therefore deduce there exists some affine function h i where f -1 d = h -1 i in the domain B ( x, δ ) .

D Slot Identifiability Proof

Definition 5 (Slot Identifiability [6]) . Given a diffeomorphic ground truth function f : S → X , and inference model ˆ g : X → S , ˆ g correctly slot identifies every object x j ∈ X with the ground slot s k ∈ S via ˆ s k = ˆ g ( f ( s k )) if there exists a unique slot k ∈ [ K ] for all x j ∈ X , and there exists an invertible diffeomorphism, or in our case an affine transformation, such that s k = h (ˆ s k ) for all s k ∈ S .

Remark 6 . [6] established that compositionality (Assumption 4) and irreducibility (Assumption 3) of the decoder f d is required for slot identifiability in the sense of Definition 5.

Summary & Intuition. The following theorems and proof extend the identifiability results of [43] to slot-attention models, we include all the proofs and details for the sake of completion. In this work, we do not consider the irreducibility criteria in [6] and define slot identifiability only by an injective mapping of each slot to subspaces representing objects in a scene to satisfy compositionality without the use of computationally heavy methods such as additive decoders and compositional contrast. We show identifiability of each slot representation up to affine transformation by passing a concatenation of samples from each mixture component which represents slots through a piecewise function in Definition 4 representing the decoder, f d . The trick in this proof lies in our observation of the fact that a concatenation of samples from each slot mixture component is a sample from a high dimensional GMM p ( s ) , as highlighted in Theorem 1. We then use the identifiability results of [43] to show p ( s ) is identifiable up to affine transformation. This then implies identifiability up to affine transformation of the aggregate posterior q ( z ) , a non-degenerate GMM by Lemma 1 where a sample from a mixture component in q ( z ) represents an individual slot representation. This contrasts with [43] which does not consider identifiability of the aggregate posterior and its mixture components in the context of slot representation learning.

In order to proceed, we begin by stating three key theorems defined and proven in the work of [43] which are essential for our slot identifiability proof. First, we restate the definition of a generic point as outlined by [43] below.

Definition 6. Apoint x ∈ f d ( R m ) ⊆ R n is generic if there exists δ > 0 , such that f d : B ( s , δ ) → R n is affine for every s ∈ f -1 d ( { x } )

Theorem 7 (Kivva et al. [43]) . Given f d : R m → R n is a piecewise affine function such that { x ∈ R n : | f -1 d ( { x } ) | = ∞} ⊆ f d ( R m ) has measure zero with respect to the Lebesgue measure on f d ( R m ) , this implies dim( f d ( R m )) = m and almost every point in f d ( R m ) (with respect to the Lebesgue measure on f d ( R m ) ) is generic with respect to f d .

Theorem 8 (Kivva et al. [43]) . Consider a pair of finite GMMs in R m :

Assume that there exists a ball B ( x , δ ) such that y and y ′ induce the same measure on B ( x , δ ) . Then y ≡ y ′ , and for some permutation τ we have that π i = π ′ τ ( i ) and ( µ i , Σ i ) = ( µ ′ τ ( i ) , Σ ′ τ ( i ) ) .

Theorem 9 (Kivva et al. [43]) . Given z ∼ ∑ J i =1 π i N ( z ; µ i , Σ i ) and z ′ ∼ ∑ J ′ j =1 π ′ j N ( z ′ ; µ ′ j , Σ ′ j ) and f d ( z ) and ˜ f d ( z ′ ) are equally distributed. We can assume for x ∈ R n and δ > 0 , f d is invertible on B ( x , 2 δ ) ∩ f d ( R m ) . This implies that there exists x 1 ∈ B ( x , δ ) and δ 1 > 0 such that both f d and ˜ f d are invertible on B ( x 1 , δ 1 ) ∩ f d ( R m ) .

We next propose our slot identifiability result below.

Theorem 2 ( ∼ s -Identifiable Slot Representations) . Given that the aggregate posterior q ( z ) is an optimal, non-degenerate mixture prior over slot space (Lemma 1), f d : S → X is a piecewise affine weakly injective mixing function (Assumption 8), and the slot representation, s = ( s 1 , . . . , s K ) can be observed as a sample from a GMM (Theorem 1), then p ( s ) is identifiable as per Definition 3.

Proof. The proof extends from [43] to slot-based models. Given two piece-wise affine functions f d , ˜ f d : S → X , ∀ k ∈ [ K ] , let s = ( s 1 , . . . , s K ) , ∋ s k ∼ N ( s k ; µ k , Σ k ) and s ′ = ( s ′ 1 , . . . , s ′ K ) , ∋ s ′ k ∼ N ( s ′ k ; µ ′ k , Σ ′ k ) be a pair of slot representations constructed by sampling and concatenating each mixture component (i.e. slots) from two distinct GMMs. As proven in Theorem 1, given individual sampled slots s k are conditionally independent given the mixture component k , then a concatenated sample is from a higher dimensional GMM in R Kd . Now, suppose for the sake of argument that f d ( S ) and ˜ f d ( S ′ ) are equally distributed. We assume that there exists x ∈ X and δ > 0 such that f d and ˜ f d are invertible and piecewise affine on B ( x , δ ) ∩ f d ( S ) , which implies dim f d ( S ) = |S| .

We now restrict the space B ( x , δ ) to a subspace B ( x ′ , δ ′ ) where x ∈ B ( x ′ , δ ′ ) such that f d and ˜ f d are now invertible and affine on B ( x ′ , δ ′ ) ∩ f d ( S ) . Next, we let L ⊆ X be an |S| -dimensional affine subspace (assuming |X| ≥ |S| ), such that B ( x ′ , δ ′ ) ∩ f d ( S ) = B ( x ′ , δ ′ ) ∩ L . We also define h f , h ˜ f : S → L to be a pair of invertible affine functions where h -1 f ( B ( x ′ , δ ′ ) ∩ L ) = f -1 d ( B ( x ′ , δ ′ ) ∩ L ) and h -1 ˜ f ( B ( x ′ , δ ′ ) ∩ L ) = ˜ f -1 d ( B ( x ′ , δ ′ ) ∩ L ) .

Therefore, this implies h f ( s ) and h ˜ f ( s ′ ) are finite GMMs which coincide on B ( x ′ , δ ′ ) ∩ L and h f ( s ) ≡ h ˜ f ( s ′ ) based on Theorem 8. Given, h = h -1 ˜ f ◦ h f and h f ( s ) and h ˜ f ( s ′ ) then h is an affine transformation such that h ( s ) = s ′ .

Given Theorems 7 and 9, there exists a point x ∈ f d ( S ) that is generic with respect f d and ˜ f d and invertible on B ( x , δ ) ∩ f d ( S ) . Having established that there is an affine transformation h ( s ) = s ′ and two invertible piecewise affine functions f d and ˜ f d on B ( x , δ ) ∩ f d ( S ) , this implies that p ( s ) is identifiable up to an affine transformation and permutation of s k ∈ s , which concludes the proof.

Corollary 3 (Individual Slot Identifiability) . If the concatenated slot distribution p ( s ) is ∼ s -identifiable (Theorem 2) then this implies q ( z ) is identifiable up to affine transformation and permutation of the slots, s k and therefore each slot distribution s k ∼ N ( s k ; µ k , Σ k ) ∈ R d , ∀ k ∈ { 1 , . . . K } is also identifiable up to an affine transformation.

Proof. Given Theorem 8, we know that each higher dimensional mixture component in p ( s ) induces the same measure on B ( x , δ ) and hence for some permutation τ we have that ( µ π ( i ) , Σ π ( i ) ) =

Algorithm 2 Probabilistic Slot Attention (no k or v )

Input:

k

z

(0)

,

π

for

=

k

t

e

f

(

x

1

)

/K

= 0

∀

A

, . . . , T

nk

ˆ

nk

A

µ

σ

(

t

(

t

(

π

K

j

∑

=1

A

nk

+1)

k

+1)

π

(

t

2

k

+1)

return

▷ input representation

▷

normalize attention

▷

update slot mean update mixing coeff.

▷ K

slot distributions

,

R

µ

∈

-

N

×

(0)

1

t

k

)

π

/

(

N

t

(

)

N

l

=1

∑

µ

(

k

T

)

,

σ

N

n

=1

∑

N

n

=1

∑

N

n

∑

(

T

=1

)

2

( µ ′ τ ( π ( i )) , Σ ′ τ ( π ( i )) ) . Therefore, each mixture component s π ( i ) is identifiable up to affine transformation, and permutation of slots representations in s . Now, given sampling s k is equivalent to obtaining K samples from the GMM, q ( z ) and concatenating, this makes q ( z ) identifiable up to affine transformation, h and permutation of slot representations in s .

It now trivially follows that each slot representation s k ∼ N ( s k ; µ k , Σ k ) ∈ R d , ∀ k ∈ { 1 , . . . , K } is identifiable up to affine transformation, h based on the following observed property of GMMs:

which concludes the proof.

E Algorithm

As a result of projection with separate query matrix is considered, the mean and variance updates are coupled with some non-linear affine transformation; for EM to have an exact solution, this transformation needs to be the identity (this can be seen when we derive the update equations for mean and variance). The resulting algorithms for these two cases are illustrated in algorithms 2 and 3.

d

k

∼ N

z

n

j

(0

;

N

q

W

µ

(

A

lk

ˆ

nk

A

ˆ

A

nk

A

nk

z

n

;

z

n

(

n

z

/N

,

(

t

W

q

I

d

)

k

)

,

µ

-

(

,

σ

σ

t

µ

(

(0)

(

t

)

2

k

)

)

2

k

j

t

,

(

σ

+1)

k

t

)

)

2

)

j

2

1

d

▷

F Metrics

SIS: Slot identifiability score [6], mainly focus on R2 score between ground-truth and the estimated slot representations wrt to maximum R2 score from the models fit between each and every inferred slots. By design SIS requires a model fitting at every validation step to compute this relativistic measure, due to which the metric seems to vary quite a bit across runs, as observed in Figure 8.

SMCC: Mean correlation coefficient is a well studied metric in disentangled representational learning [39, 37], we extend this with additional permutational invariance on slot dimensions. SMCC measures the correlation between estimated and ground truth representations (or estimated representations across runs, in the case when ground truth representations are unavailable) once the slots are matched using Hungarian matching. To compute identifiability up to affine transformation along the representational axis and permutation in slots (check definition 3), similar to weak identifiability as per the definition in the paper and in [38], we use the MCC up to some affine mapping A , which we learned by matching slot representations across runs. In summary, SMCC can be computed with the following three steps:

Matching the order of slots using Hungarian matching;

Affine transformation of slot representation;

Followed by computing mean correlation.

Apart from variations described in Figure 8, we further analyse both the metrics by fixing the model and data; and by computing both the metrics for 10 times. We then considered the mean and variance in the performance, which is reflected as follows: SIS: 35.26 ± 6.46 , SMCC: 77.83 ± 0.36 . Here, the resulting variation is only the reflection of the metric, which clearly indicates the stability of SMCC over SIS.

Slot Mean Correlation Coefficient

Slot Identifiability Score

Epoch

Epoch

15

15

15

15

15

15

15

10

10

10

10

10

10

10

5

5

5

5

5

5

5

0

0

0

0

0

0

0

5

5

5

5

5

5

5

10

10

10

10

10

10

10

15

15

15

15

15

15

15

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

15

15

15

15

15

15

10

10

10

10

10

10

10

5

5

5

5

5

5

5

0

0

0

0

0

0

0

5

5

5

5

5

5

5

10

10

10

10

10

15

15

10

15

15

15

15

10

10

10

10

5

5

5

5

15

15

10

5

0

5

10

15

0

0

5

5

10

10

15

15

15

15

15

15

10

5

0

5

10

15

10

10

5

5

0

0

5

5

10

10

15

15

15

15

15

15

10

5

0

5

10

15

10

10

5

5

0

0

5

5

10

10

15

15

15

15

10

10

5

5

0

0

15

15

10

5

0

5

10

15

5

5

10

10

15

15

15

15

10

5

0

5

10

15

15

15

10

10

5

5

0

0

5

5

10

10

15

15

15

15

10

5

0

5

10

15

15

15

10

10

5

5

0

0

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

15

15

15

15

15

15

10

15

10

15

10

5

5

0

0

5

5

10

10

15

10

15

15

15

10

15

15

10

5

0

5

10

15

10

15

15

10

5

0

5

10

15

15

15 15 15 15 15 G Comparison with Autoencoding Variational Bayes

10

10

10

10

10

10

15

15

10

10 15 15 10 5 0 5 10 15 15 10 5 0 5 15 10 5 0 5 10 15 15 10 5 0 5 15 10 5 0 5 10 15 15 10 5 0 5 15 10 5 0 5 10 15 15 10 5 0 5 15 10 5 0 5 10 15 15 10 5 0 5 15 10 5 0 5 10 15 15 10 5 0 5 5 0 5 10 15 5 0 5 10 15 5 0 5 10 15 5 0 5 10 15 5 0 5 10 15 5 0 5 10 15 As explained in Section 3, applying slot attention to a deterministic encoding z = f e ( x ) ∈ R N × d yields a set of K object slot representations s 1: K := s 1 , . . . , s K . In combination, this process induces a stochastic encoder q ( s 1: K | x ) , where the stochasticity comes from the random initialization of the slots in the first iteration: s t =0 1: K ∼ N ( s 1: K ; 0 , I ) ∈ R K × d . Since each slot is a deterministic function of its previous state s t := f s ( z , s t -1 ) it is possible to randomly sample initial states s 0 and obtain stochastic estimates of the slots. However, since each transition depends on z , which in turn depends on the input x , we do not get a generative model we can tractably sample from .

10 15 15 10 5 0 5 10 15 15 10 15 10 5 0 5 10 15 15 10 15 10 5 0 5 10 15 15 10 15 10 5 0 5 10 15 15 10 15 10 5 0 5 10 15 15 10 15 10 5 0 5 10 15 15 10 This can be remedied by placing a tractable prior over z and using the VAE framework along the lines of [77]. Specifically, Wang et al. [77] propose the Slot-VAE, which is a generative model that integrates slot attention and the V AE framework under a two-layer hierarchical latent model. However, under their formulation, there is a key challenge in calculating the KL term as the slot attention function is permutation equivariant meaning the slots have no fixed order across the posterior and prior. To compensate for this, the authors introduce a heuristic auxiliary loss and a parallel image processing path with an additional slot attention operation which is computationally costly.

In contrast, our approach does not suffer from such drawbacks. This was achieved by designing the slot attention operation itself as probabilistic, and proving that having a local (per-datapoint) GMM results in the aggregate slot posterior and the concatenated slots being GMM distributed (Section 4). We then proved a new slot identifiability result using this insight (Section 5). Our use of the aggregate posterior is inspired by but differs substantially in both method and application from previous works on VAEs [28, 71, 72]. Our primary goal is not to learn a better generative prior but to obtain slot identifiable representations. Lastly, since the concatenated slots are provably GMM distributed, our model is reminiscent of a GMM-VAE [11], but with a unique slot-based structure.

H Automatic Relevance Determination of Slots

For evaluation, we calculate the MAE between the estimated and ground truth numbers of slots and measure the reduction in FLOPs achieved using the proposed ARD method. We observe MAE values of 1.03 and 0.58, and significant savings of up to 41% and 62% in FLOPs on the CLEVR and Objects-Room datasets respectively, when compared to using a fixed number of slots K .

15

15

15

15

10

5

0

5

10

15

15

10

5

0

5

10

15

Input

Reconstruction

Slot 2; pi=0.02

Slot 3; pi=0.62

Slot 4; pi=0.33

Slot 6; pi=0.02

Slot 7; pi=0.02

Input

Reconstruction

Slot 1; pi=0.02

Slot 2; pi=0.31

Slot 3; pi=0.02

Slot 4; pi=0.61

Slot 5; pi=0.02

Slot 6; pi=0.02

Slot 7; pi=0.02

Input

Reconstruction

Slot 1; pi=0.02

Slot 2; pi=0.03

Slot 3; pi=0.70

Slot 4; pi=0.22

Slot 5; pi=0.02

Slot 7; pi=0.02

Input

Reconstruction

Slot 1; pi=0.56

Slot 2; pi=0.28

Slot 3; pi=0.11

Slot 6; pi=0.02

Slot 7; pi=0.03

Input

Reconstruction

Slot 1; pi=0.05

Slot 2; pi=0.02

Slot 3; pi=0.65

Slot 4; pi=0.02

Slot 5; pi=0.03

Slot 6; pi=0.05

Slot 7; pi=0.19

Input

Reconstruction

Slot 1; pi=0.06

Slot 2; pi=0.73

Slot 3; pi=0.02

Slot 7; pi=0.19

Input

Reconstruction

Slot 1; pi=0.47

Slot 2; pi=0.02

Slot 3; pi=0.03

Slot 5; pi=0.12

Slot 6; pi=0.35

Slot 7; pi=0.02

Input

Reconstruction

Slot 2; pi=0.01

Slot 3; pi=0.73

Slot 4; pi=0.02

Slot 5; pi=0.17

Slot 6; pi=0.06

Input

Reconstruction

Slot 1; pi=0.02

Slot 2; pi=0.02

Slot 3; pi=0.20

Slot 4; pi=0.02

Slot 5; pi=0.73

Slot 7; pi=0.02

I Compositionality Results

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

Slot 1

Slot 2

Slot 3

Slot 4

Slot 5

Slot 6

Slot 7

Composed

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: We clearly state the limitations in existing literature and our proposed analysis to overcome these limitation, which we further back with theoretical guarantees and experimental observations.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]

Justification: We provide a detailed discussion on limitations and explicitly list all the assumptions made, which serves nicely for future works.

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide rigorous mathematical proof all our theoretical claims.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We clearly list all the set of hyperparameters used in the analysis.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [Yes]

Justification: We provide codebase with the described set of hyper-parameters for reproducability.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We discuss hyper-parameters in the paper along with the codebase for earier reproducability.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: We report all our analysis across 5 random runs.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: Detailed in appendix.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: -

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: This paper proposes a Probabilistic Slot Attention algorithm, whose goal is to achieve identifiable object-centric representations. The work primarily makes theoretical advancements in the field of object-centric learning, and as such it has little immediate societal or ethical consequences. Our method might be a step towards interpretable and aligned models which are desired properties of trustworthy AI.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: -

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: -

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: -

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: -

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: -

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper studies the theoretical identifiability of object-centric representations (slots) in object-centric learning (OCL). Prior works study the same problem, but are limited to OCL models with an additive decoder. This work relaxes the constraint to non-additive decoders, which have shown important to scale up to more complex data in recent OCL works (e.g., TransformerDecoder, Diffusion Model). To tackle this generalized setting, authors propose Probabilistic Slot Attention (PSA), which applies the Gaussian Mixture Model (GMM) to produce slots from each data sample (e.g. an image). The authors prove the effectiveness of PSA both theoretically and empirically, using results on both low-dimensional data and 2D images.\n\nSTRENGTHS:\n- This paper tackles an important problem -- what assumptions are required to learn identifiable object slots?\n- The paper is generally clearly written and well presented.\n- I appreciate the efforts in experimenting with common OCL image datasets such as CLEVR and ObjectsRoom. They are missed in previous identifiable OCL works.\n- The experimental results with PSA and PSA-Proj (both using additive decoders) are solid.\n\nWEAKNESSES:\nI have only one big concern. However, this concern is closely related to the main contribution of this paper (correct me if I am wrong), and I cannot accept the paper if it is not well addressed.\n- This paper claims to be the first work learning identifiable slots with a non-additive decoder (NoA). This is great, as recent OCL works show impressive results using a Transformer-based or a diffusion-based decoder.\n- However, the paper does not really have experimental results supporting this claim. PSA-NoA seems to underperform PSA-Proj consistently on all datasets & metrics.\n- What's worse, even compared to vanilla SA, PSA-NoA still underperforms in FG-ARI (Table 3), and even on some identifiability-related metrics (Table 2). These results are concerning to me, as we are not sure if the proposed algorithm can really scale to recent OCL models with better decoders.\n\nMinor:\n- Please unify the citation format -- for papers that are published at conferences, please use their conference version, e.g., [5] (ICML), [45] (NeurIPS).\n- In Sec.4, the paper claims that PSA can dynamically determine the number of object slots. While Fig.10 in the Appendix shows a few results, I don't think that's enough to claim PSA \"offers an elegant solution to this problem\". [1] is a recent work that studies this problem and provides in-depth analysis. The authors can conduct experiments following [1]'s setting. But I understand that the dynamic slot number is not the main contribution of this work.\n- In line 51, the authors claim that the computation cost of non-additive decoders is invariant to the number of slots K. Is this true? In recent Transformer-based and diffusion-based decoders, they use cross-attention to condition the reconstruction on slots, and the computation cost of attention is quadratic to the token size, i.e., K in this case.\n\n[1] Fan, Ke, et al. \"Adaptive Slot Attention: Object Discovery with Dynamic Slot Number.\" CVPR. 2024.\n\nQUESTIONS:\n- Can the authors apply PSA to a better non-additive decoder? For example, the Transformer decoder proposed in [61]. Otherwise, it is hard to assess the contribution of this work.\n- What is the implementation detail of the non-additive decoder \"standard convolutional decoder\"? How to decode an image from K 1D slots using a CNN decoder?\n\nMinor:\n- How can we apply the theory proposed in this paper to more complex datasets, e.g., real-world COCO images? This is not required for a theory paper, but I'm curious about the authors' thoughts.\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations in Sec.7.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2024-06-15"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper addresses the problem of identifiability of object-centric representations. In contrast to prior works which achieve identifiability via assumptions on the generator, this paper explores identifiability via assumptions on the latent distribution. To do this, the authors introduce a probabilistic variant of the popular Slot Attention algorithm and prove theoretically that this method identifies the ground-truth object representations. The authors verify their theory on toy data as well as test their method on high dimensional image data, showing improved performance over baseline methods.\n\nSTRENGTHS:\n* The authors address an important problem in representation learning. Namely, understanding when learning object-centric representations is theoretically possible.\n\n&NewLine;\n\n* The paper is well written, well structured, and, in general, easy to understand.\n\n&NewLine;\n\n* The authors position their contribution well within the broader representation learning literature and provide a good review of prior work.\n\n&NewLine;\n\n* Section 4 is well written. I particularly appreciated being able to look at Algorithm 1 while reading the section to guide my intuition.\n\n&NewLine;\n\n* Exploring probabilistic constraints for identifiability in object-centric learning is an important problem to understand, thus I think the author’s contribution is of interest for the representation learning community.\n\n&NewLine;\n\n* The authors achieve identifiability by proposing an adaptation to a widely used method, making their method potentially easy to adopt in practice.\n\n&NewLine;\n\n* The authors conduct a solid empirical study and achieve promising results, in particular when coupling their method with structured decoders.\n\n&NewLine;\n\n* The figures in the manuscript are conceptually helpful and aesthetically well done.\n\nWEAKNESSES:\n### __Paper Positioning/Storyline__\n\nOne of the main issues I have with this work is that I do not think that the paper’s storyline i.e. how the authors motivate and position their contribution, accurately reflects the actual contributions of this work. As I understand, the current pitch of the paper is the following: \n\n&NewLine;\n\n*Previous works from [1, 2], prove identifiability of object-centric representations by making assumptions on the decoder. Enforcing the assumptions in [1] is not tractable in practice, however, due to scalability issues with the compositional contrast in [1]. In this work, we remedy this by exploring probabilistic constraints for identifiability which yield identifiability theoretically and empirically but do not suffer from the same empirical scalability issues as prior works.*\n\nI think this pitch is problematic for the following reasons:  \n\n&NewLine;\n\n__Firstly__, it is important to note that the compositional decoders explored in [1] were proven to be a subset of the additive decoders explored in [2]. Consequently, if the ground-truth decoder is compositional and one uses an additive decoder for the inference model, then assuming the assumptions from [2] are met, the inference model is slot identifiable i.e. it will implicitly minimize compositional contrast. \n\n&NewLine;\n\nThis does not completely dismiss the scalability issues noted by the authors, since as mentioned in Lines 51-52, additive decoders also suffer from some scalability issues. I think, however, that the author’s current discussion of scalability, in particular wrt [1], misses the key nuance discussed above. I would suggest the authors incorporate this discussion into the paper by altering their writing and positioning of their contribution accordingly. \n\n&NewLine;\n\n__Secondly__, the storyline and writing in the manuscript give the impression that one could dispense with decoder structure all together in favor of probabilistic structure on the latent space. As the authors show empirically in Section 6, however, this is not exactly the case. While probabilistic structure gives identifiability gains relative to baselines, without incorporating decoder structure, identifiability drops non-trivially across all metrics.  \n\n&NewLine;\n\nMoreover, one of the core motivations for object-centric learning is learning representations which generalize compositionally [2, 3]. Such compositional generalization is only possible through decoder structure where additivity is one such structure (see [3] Section 2.). If one only focuses on enforcing structure on the latent distribution, it is not clear to me how such compositional generalization can be achieved.\n\n&NewLine;\n\nWith all of this being said, I think a more accurate and superior pitch for the contributions in this work should focus on the advantage of using probabilistic and decoder structure in tandem opposed to suggesting that probabilistic structure is somehow superior. Something like: \n\n&NewLine;\n\n*In this work, we show how identifiability can be achieved via probabilistic constraints on the latent space. We show how such constraints can be naturally and tractably incorporated into Slot Attention. We verify our theory and method on toy data. We then show on image data that our probabilistic structure leads to improved identifiability over unstructured baselines. Furthermore, when coupled with structured decoders, our method yields performance which outperforms both probabilistic structure and decoder structure in isolation.* \n\n&NewLine;\n\nFor the reasons stated above, I would encourage the authors to alter their writing in the manuscript to adhere closer to a storyline along the lines of the one presented above. As things currently stand, the messaging in the paper feels a bit misleading and over-claiming. \n\n&NewLine;\n\n$\\newline$\n\n&NewLine;\n\n### __Theory Explanation__\nI would have appreciated a short paragraph giving some intuition on how exactly the probabilistic constraints imposed by the authors sufficiently restrict the problem such that slot identifiability is possible. Specifically, unsupervised identifiability via probabilistic structure is challenging and, in most cases, impossible [4., 5.]. Thus, I think it would be helpful if the authors could explain which structure in their method is key towards overcoming these well known unidentifiability issues. For example, is it the permutation invariance, the complex aggregate posterior etc.?\n\n&NewLine;\n\n$\\newline$\n\n&NewLine;\n\n### __Metrics__\nThe authors use two main metrics to validate identifiability: the slot identifiability score (SIS) from [1] and a new metric, slot MCC (SMCC). I found the authors explanations of the differences in these metrics unclear, which made it difficult to interpret some of the empirical results. My understanding is that SMCC is just a linear/affine version of SIS in terms of the predictor fit between latents. Thus, it is a bit unclear to me why the values for SIS should be so much lower than SMCC in the experiments on toy data in Section 6.\n\n&NewLine;\n\n$\\newline$\n\n&NewLine;\n\n### __Experiments__\nThe authors experiments on image data focus primarily on simple decoders opposed to e.g. Transformers and visually complex datasets. To assess the scalability of the author’s method it would be important to test the method on more complex models/datasets. I view this as a more minor weakness of this work, however, given that the contribution is largely theoretical.\n\n&NewLine;\n\n$\\newline$\n\n&NewLine;\n\n### __Appendix__\nIn section A. of the Appendix, the authors review some definitions from prior works e.g. in [1]. These definitions, however, are presented in an informal, imprecise way. I would encourage the authors to be precise in this section if they are going to include definitions from prior works.\n\n&NewLine;\n\n$\\newline$\n\n&NewLine;\n\n## __Bibliography__\n\n1. Provably Learning Object-Centric Representations (https://arxiv.org/abs/2305.142290)\n\n&NewLine;\n\n2. Additive Decoders for Latent Variables Identification and Cartesian-Product Extrapolation (https://arxiv.org/abs/2307.02598)\n\n&NewLine;\n\n3. Provable Compositional Generalization for Object-Centric Learning (https://arxiv.org/abs/2310.05327)\n\n&NewLine;\n\n4. Challenging Common Assumptions in the Unsupervised Learning of Disentangled Representations (https://arxiv.org/abs/1811.12359)\n\n&NewLine;\n\n5. Variational Autoencoders and Nonlinear ICA: A Unifying Framework (https://arxiv.org/abs/1907.04809)\n\nQUESTIONS:\n* Do the authors view their method as being a replacement for structured decoders in object-centric learning or do they view the method as being better used in tandem with structured decoders?\n\n* What is the key theoretical assumption/constraint which allows for identifiability to be possible? \n\n\n* What are the key differences between SMCC and SIS?\n\n\n* Do the authors have an explanation for the low SIS score on the toy data experiments?\n\n\n* Does “R2” refer to SIS when used or is this a different score? \n\n\n* Do the authors have any intuition about how well their method would perform for more complex models/datasets?\n\nLIMITATIONS:\nThe authors include a limitation section which discusses some limitations of their theory. I would encourage the authors to discuss some limitations in their experiments, as discussed above in the “Weaknesses” section. I would also include some discussion of the limitations of purely probabilistic structure for object-centric learning as it pertains to compositional generalization, as discussed above.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-10"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSolving the problem of identifiability is necessary to find consistent and interpretable factors of variation. There are two approaches to do so: a) place restrictions on the decoders, and b) impose distributional constraints on the latent space. This work takes the second approach and aims to impose a GMM on the latent space. The paper does this by proposing a modification to the vanilla slot attention framework that they call probabilistic slot attention. Under this framework the papers shows theoretical and empirical identifiability of slots.\n\nSTRENGTHS:\n+ The proposed probabilistic slot attention framework is an intuitive extension of the vanilla slot attention with the updates in each iteration resembling the familiar EM algorithm.\n+ The proposed framework offers a possible solution to the problem of dynamically determining the required number of slots.\n+ This is the first work to experiment with imposing a distribution on the latent space where prior works either focus on the generator or the decoder.\n+ Recovering the latent space upto an affine transformation is shown in the synthetic modelling scenario with theoretical guarantees provided.\n+ The paper studies an important premise – the theoretical understanding of object-centric representations. The paper is written well, with clear motivations for the proposed contributions.\n\nWEAKNESSES:\n* The paper states that the framework allows for tractably sampling from the aggregate posterior distribution and using it for scene composition tasks; however, this is not empirically qualified anywhere.\n* Experiments only validate the theory on simple synthetic datasets. Testing on more diverse and realistic data would better demonstrate applicability, though the evaluation would also be more challenging. Generally speaking, I would be concerned about the scalability of such an approach leveraging GMMs. I understand however that the objective of this work is theoretically study the identifiability of object-centric representations under less-constrained settings as compared to previous work. Maybe, having a short discussion on how such an approach can be scaled would be nice to see.\n* Table 1 lists $\\beta$-disentanglement and weak injectivity as core assumptions. While the former is common to all other related methods, weak injectivity is newly introduced. The implications of this assumption is hence important, but is missing in the paper. \n* In fig. 3, it is not clear what experiment the latents (x-axis) correspond to.\n* It might be helpful to visit slot-identifiability in the related work section, considering its literature is the most closely related to this work.\n* Adequate details about the encoder and decoder in the synthetic modelling scenario have not been provided.\n\nMinor/editorial\n* Typo in 769, “Obsereved” -> observed; Typo in 234, “emphirically” -> empirically\n* Some references are repeated. Eg. 36-37. Please check carefully across the full list.\n\nQUESTIONS:\n* Could you elaborate on assumption 6 and why it is not necessary in the current work?\n* In Algorithm 1, line 6, where attention $A_{nk}$ is calculated, mean is calculated as $W_q \\mu(t)_k$ but variance is not calculated as $W_q \\sigma_k^2$ ?\n* Should we not return $\\pi (t)$ at the end of the algorithm?  \n* The work explicitly mentions that having additive decoders is not a necessity of the current work, but any experimental results are shown only with additive decoders. Perhaps it may be enlightening to see experiments akin to [45] for non-additive decoders (transformer based auto-regressive decoders). There are no results to go along with the choice of a convolutional decoder as the non-additive variant which is mentioned in L277. Did I miss something here?\n* R2 score is reported in Sec 6 without definition. I assume this is the correlation? This is defined as MCC in the paper.\n\nLIMITATIONS:\nThe statistical distributional assumption on the latent space excludes for identification of any causal dependencies between objects which could be made explicit in the “Limitations and Future work” section.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-13"}

comments: SUMMARY:
This paper propose a probabilistic slot attention method that can learn identifiable object-centric representation. Compared with former work on identifiable object centric representation methods, the proposed model can scaling slot-based methods to high-dimensional images. Both theory analysis and experiments verify the effectiveness of proposed model.

STRENGTHS:
1. This paper propose a novel probabilistic slot attention model, which can learn Identifiable Object-Centric Representation.

2. the proposed model can scale to high-dimension image dataset. 

3. the proposed model is seem to be solid.

4. the paper is well written and read well.

WEAKNESSES:
1. the experiments are only take on two toy datasets.

QUESTIONS:
1. How much computational complexity has increased?

LIMITATIONS:
N/a
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2024-07-18