Object-Centric Learning with Slot Mixture Module
Authors: Daniil Kirilenko
Vitaliy Vorobyov
Alexey Kovalev
Aleksandr Panov
Venue: ICLR 2024 poster (2024)
Keywords: Object-centric representations
Gaussian Mixture Model
Slot Attention
Set Prediction Task
Primary area: representation learning for computer vision, audio, language, and other modalities
TL;DR: We proposed a generalization of a slot-based approach for object-centric representations as a Slot Mixture Model that allows state-of-the-art performance in the set property prediction and object discovery tasks.

Abstract: Object-centric architectures usually apply a differentiable module to the entire feature map to decompose it into sets of entity representations called slots. Some of these methods structurally resemble clustering algorithms, where the cluster's center in latent space serves as a slot representation. Slot Attention is an example of such a method, acting as a learnable analog of the soft k-means algorithm. Our work employs a learnable clustering method based on the Gaussian Mixture Model. Unlike other approaches, we represent slots not only as centers of clusters but also incorporate information about the distance between clusters and assigned vectors, leading to more expressive slot representations. Our experiments demonstrate that using this approach instead of Slot Attention improves performance in object-centric scenarios, achieving state-of-the-art results in the set property prediction task.

Published as a conference paper at ICLR 2024
OBJECT-CENTRIC LEARNING WITH SLOT MIXTURE MODULE
Daniil Kirilenko Vitaliy Vorobyov', Alexey K. Kovalev Aleksandr I. Panoy 1.34
1 FRC CSC RAS, Moscow, Russia
2 Università della Svizzera italiana, Lugano, Switzerland
3 AIRI, Moscow, Russia
4 MIPT, Dolgoprudny, Russia
daniil
ch, kovalev@airi.net
ABSTRACT
Object-centric architectures usually apply a differentiable module to the entire feature map decompose it into sets entity representations called slots. Some of these methods structurally resemble clustering algorithms, where the cluster's s center in latent space serves as a slot representation. Slot Attention an exam- ple of such a method, acting as a learnable analog of the soft k-means algorithm. Our work employs a learnable clustering method based on the Gaussian Mixture Model. Unlike other approaches, we represent slots not only centers of clusters but also incorporate information about the distance between clusters and assigned vectors, leading more expressive slot representations. Our experiments demon- strate that using this approach instead of Slot Attention improves performance object-centric scenarios, achieving state-of-the-art results in the set property pre- diction task.
INTRODUCTION
In recent years, interest in object-centric representations has significantly increased (Greff et al.
2019 BBurgess et al. 2019 Li et al. 2020 Engelcke et al.] 2019; 2021 Locatello et al. 2020
Seitzer al. 2022] Chakravarthy et al. [2023] Kirilenko et al. (023).. Such representations have
the potential to improve the generalization ability of machine learning methods in many domains.
such as reinforcement learning (Keramati et al., 2018 Watters et al. 2019a Kulkarni et al.
2019
Berner et al. 2019: Sun et al. (019). scene representation and generation (El-Nouby et al.
2019
Matsumori et al. 2021 Kulkarni et al. 2019: Kirilenko et al. (2022). reasoning (Yang et al.
2020
Kovalev et al. (022). object-centric visual tasks (Groth et al. 2018 Yi et al. 2020 Singh et al.
2022). and planning (Migimatsu & Bohg 2020).
Automatic segmentation of objects in the scene and the formation of structured latent space can be achieved in various ways (Greff et al.) 2020):: augmenting features with grouping information, using tensor product representation, incorporating ideas from attractor dynamics, etc. However, the most common and effective method for learning object-centric representations is Slot Attention (Locatello) al. 2020). Slot Attention maps the input feature vector received from the convolutional encoder to a fixed number of output feature vectors, referred to as slots. Through training, each object is assigned corresponding slot. If the number of slots exceeds the number of objects, then some slots remain empty (i.e., they do not contain objects). This approach has shown significant results in object-centric tasks like set property prediction and object discovery. The limitations of this model are that significant object discovery results are only achieved on simple synthetic data, such as CLEVR (Johnson et al. (0177). Tetrominos, and Multi-dSprites (Kabra al. 2019). On complex synthetic datasets such ClevrTex (Karazija et al. 2021) real data such as COCO (Lin et al.) 2014). Slot Attention shows underperforming results. Also, on CLEVR, the standard dataset for the property prediction task, Slot Attention, shows significant results only with weak thresholds, i.e., in cases where the model can be very mistaken regarding the position of objects. Another drawback of Slot Attention is its training instability, making it difficult to reproduce the results.
is worth noting that most modern work evaluates object-centric methods mainly only on the object discovery task, i.e., recovering object masks, which be unsupervised segmentation

---

Published as a conference paper at ICLR 2024
problem. We believe that the results the set property prediction task are also good indicators for object-centric methods, since this task requires distinguishing objects from each other and assigning them to given set of objects.
In our work, we propose a generalization of the Slot Attention approach, replacing the k-means algorithm with Gaussian Mixture Model (GMM) (Bauckhage 2015), which refer to the Slot Mixture Moduld" This leads to more expressive slot representations by incorporating information about the distance between clusters and their assigned vectors. Considering not only the weighted mean of a group of vectors as the representation gives an advantage in distinguishing similar but different groups. For example, two sets of vectors, sampled from different distributions but with the same expectation, would be indistinguishable if their mean were the only representation. This method enhances object-centric croblem-solving across various tasks and datasets, surpassing even highly specialized models in the set property prediction task.
We conduct thorough experimental studies to substantiate the results obtained. Through extensive experiments, we show that the proposed Slot Mixture Module achieves the state-of-the-art perfor- mance in the set property prediction task on the CLEVR dataset (Johnson et al., 2017), outperform- ing even highly specialized models (Zhang et al.) (019). We provide experimental results for the image reconstruction task on five datasets: four with synthetic images (CLEVR- (Singh et al,, 2022). ShapeStacks (Groth et al. 2018), ClevrTex (Karazija al. 202I), Biumoji (Graux 2021 and one with real-life images (COCO-2017 (Lin et al. 2014)). The proposed Slot Mixture Module consistently improves the performance of the GPT decoder. We show that the Slot Mixture Module outperforms the original Slot Attention on the object discovery task on such a complex dataset as ClevrTex. We also give an example of the use of advances in our approach, extending the idea of concept library by estimating concept distributions. To justify the replacement of k-means by GMM clustering, we compare k-means and GMM clustering approaches on the set property prediction task and show that GMM clustering is a better choice for object-centric learning.
The main contributions of our paper are as follows: 1) We propose generalization of the slot- based approach for object-centric representations, referred to as the Slot Mixture Module (Sec- tion 3ion3):: 2) Through extensive experiments, we demonstrate that our approach consistently improves performance in both unsupervised (image reconstruction in Section 4.1 and object discovery in Sections 4.3) and supervised (set property prediction in Section 4.2) object-oriented tasks across various achieving
tftte-of-theaate results in the set property prediction task (Section 4.2) with strin- gent thresholds on the CLEVR dataset; 3) We demonstrate how additionally estimated slot mixture weights might aid in constructing clear concept library by distinguishing between empty and filled slots. We also extend this approach to concept sampling (Section 4.4p.
2 BACKGROUND
2.1 SL.OT ATTENTION
Slot Attention (SA) module (Locatello et al. 2020) is an iterative attention mechanism designed to map distributed feature map x RN to set of K slots K Slots initialized randomly, and a trainable matrix E DxD is used to get query projections of the slots, while k. DxD matricies are used to get key and value vectors from the feature map x. Dot-product attention (Luong et al. (015)) between g and l projections with SoftMax across the g dimension implies competition between slots for explaining parts of the input. Attention coefficients A I N K are used to assign projections to slots via a weighted mean.
1
Ma.
M VD {(x)((s) ERAXK
Aij
Σ/= Mi.s
Wo Ara
s" = WT ((x) ERKXD
(2)
Σ Ar.
The Gated Recurrent Unit (GRU) (Coo et al. 2014 is employed for slot refinement, taking pre- updated slot representations as hidden states and the updated slots s* as inputs. A key characteristic of SA is its permutation invariance to input vectors and permutation equivariance for slots. This makes SA apt for handling sets and object-centric representations.
'The code is available at /ettpa://ithub.... com RI Institute/smm
2

---

Published as a conference paper at ICLR 2024
Technically, Slot Attention is learnable analog of the k-means clustering algorithm with ad- ditional trainable GRU update step and dot product (with trainable q. k. v) projections instead of Euclidean distance the measure of similarity between the imput vectors and cluster centroids. At the same time, k-means clustering can be considered as special case of the Gaussian Mixture Model.
2.2 MIXTURE MODELS
Mixture Models (MM) is a class of parametric probabilistic models, in which it is assumed that each z; from some observations X = [#1,.... xy E x is sampled from the mixture distribution with K mixture components and prior mixture weights RR:
zi ~ 0(x
Σ (xp((;(00))
P( II 0).
Σ Tk ===1
(3)
These models can seen as the models with latent variables 24,6 E {.1...., 2 that indicate which component æ; came from. The problem is find such K groups of component parameters 02 and component assignments each sample I; that will maximize the likelihood of the model P(X The Expectation Maximization (EM) algorithm is iterative algorithm that addresses this problem and includes two general steps. The Expectation (E) step: evaluate the expected value of the complete likelihood P(X, Z 0*) with respect to the conditional distribution of P 0):
N K II <.p(o, (0)) ((2
10 π = -E Ep log P Z(0"), P(X, 21x)
II
where I(+) is an indicator function.
The Maximization (M) step: find 0" n that maximize Q x 1= TT):
(0, argmax(or..-) - 0, IT).
(5)
One of the most widely used models of this kind is Gaussian Mixture Model (GMM), where each mixture component is modeled as Gaussian distribution parameterized with its mean values and covariance matrix, which is diagonal in the simplest case: P(x,(0,) = Na / Ex), Σg
diag(o In this case, the EM algorithm is reduced to the following calculations. E step:
p ((z))p(x;
TxN pek. Ex)
p(zz
(6)
p( Zz (
Σ
K
NNN Mes Ex)
M step:
1=
Thka
Eis The EE Thii
E
ΣN
7h,( Mi) M4)
ΣN
Tk.i
The key difference between the Gaussian Mixture Model and k-means clustering is that GMM con- siders not only the centers clusters but also the distance between clusters and assigned vectors with the prior probabilities of each cluster.
SLOT MIXTURE MODULE
For object-centric learning, we propose modified Gaussian Mixture Model approach called the Slot Mixture Module (SMM). This module uses GMM and M steps (Section 2.2) to map feature maps from the convolutional neural network (CNN) encoder to the set of slot representations, where slots are a concatenation of mean values and diagonal of the covariance matrix. Like Slot Atten- tion. SMM uses GRU, which takes current and previous mean values input and hidden states. In SA, slot representations are cluster centers, so SA is limited by the information contained and represented in these cluster centers. Considering only the weighted mean of group of vectors
3

---

Published as a conference paper at ICLR 2024
Algorithm The Slot Mixture Module pseudocode. л is initialized as a uniform categorical distri- bution, μr and Dding are initialized from Gaussian distributions with trainable parameters.
Input: x E RNxD
flattened CNN feature map with added positional embeddings; μ. Ediag E R TERK SMM initialization parameters.
Output: slots Nx2D a MLP(LayerNorm(zz) for t = 0...7' do
= footss- μt, Ediag)
7 SoftMax(logits log æ, dim=k)
π = >>mean(dim=:N)
4 WeightedMeuntweighhsyy values=z)
H = GRU(input=4t hidden=pu)
4 += MLLP(LayeNNrr(())
E dieg WeightedMean values=(xx 11)3))
end for
slots cоnсаt([μ,, Edigg,, dim=D)
return slots
as the representation gives an advantage in distinguishing similar but different groups. For example, two sets of vectors, sampled from different distributions but with the same expectation, would be indistinguishable if their mean were the only representation. This set of slots is further used in the downstream task. We also use the same additional neural network update step for the mean values before updating the covariance values:
U RNN(input=u hidden=44), 14 MLP(LayerNorm((u)) + Mk-
(8)
These two steps serve the needs of the downstream task by linking the external and internal mod- els. The internal model (E and M steps in SMM) tries to update its parameters fs, Σ so that the input vectors
assigned slots with the maximum likelihood. In contrast, the external model takes these parameters as input. The Algorithm [ presents the full pseudocode. A function fa(æ, ,t, Eding) : R-NxD x R R. stands for log-Gaussian density function which results in a matrix of log-probabilities, where each element reflects the estimated bog-probability of a vector from being generated by corresponding Gaussian distribution.
It should be noted that SMM does not perform EM updates explicitly because the algorithm uses neural network updates that are not trained to maximize expectations as traditional EM. Instead, SMM combines key update steps from the classical EM algorithm with trainable updates specific to neural networks. By incorporating trainable updates into the algorithmic cycle, leverages the principles and insights of the EM algorithm while taking advantage of the flexibility and adaptability of neural networks. Combining classic EM steps with trainable updates allows for a more dynamic and data-driven approach, potentially leading to improved performance and adaptability to different scenarios.
The Slot Mixture Module can be seen extension of the Slot Attention Module (Figure I with the following key differences: (1) SMM updates not only the mean values but also the covariance values and prior probabilities, (2) the Gaussian density function is used instead of the dot-product attention, and (3) slots are considered not only as mean values of the cluster but as the concatenation of mean and covariance values.
EXPERIMENTS
Since our model can be seen as an extension of the Slot Attention, we use it as our main competi- tor. Our approach predicts the mean μ and covariance Ediag. which doubles the intermediate slot representation, so we apply an additional matrix to effectively reduce the dimensionality of the in- termediate slot representations by a factor of two (the matrix size is 2D D). Thus, the final size of the slots produced by SMM is the same as that of the slots produced by SA. is also important to note that both SMM and SA have the same number of trainable parameters because SMM does not use and e projections (the size of the and v matrices D x D), which ensures a fair comparison between the two approaches.
4

---

Published as a conference paper at ICLR 2024
Randors Int
Mothell Scole
Random
Spttilee
GRU
GRU
Comre
CNN
Mestbal
Temes
Timmss
Slot Mixture Module
Slot Attention Module
Figure I: Visualized architectures of Slot Mixture Module (ours) and Slot Attention Module. Green color is used for steps involved in both modules, SMM involves the estimation of cluster centers (μ)), the distance between cluster centers and assigned vectors (a, orange steps), and prior mixture weights (I, red steps). The concatenation of Il and a serves as slot representations, and π can be used identify empty slots that do not contain information about any objects, f(z, pt, or) the log of the Gaussian PDF SA module estimates only centers of clusters.
In experiments, we trained two versions of the same model: the one with the Slot Attention module and another with SMM, keeping the same training conditions. We use the technique of detaching slots from the gradient computational graph at the last iteration of the algorithm from Chang et al. (2022) to achieve stable training and avoid exploding gradients. This work also shows that training with a small number of iterations (1 iteration) can lead to significantly worse performance in specific tasks than training with larger number of iterations iterations). We use five iterations default choice in the trade-off between best performance and excessive computation. Since the op- timal number of steps may depend on the datasets and tasks, we did not perform hyperparameter search, as this would lead to difficulties in comparing different tasks and datasets. We choose the av- erage number of based on the most related works: Locatello al. (2020) uses 3.[Singh et al. (2022) uses 7 for some datasets and 3 for others, Chang al. (2022) shows the difference in performance between and iterations. Additional architectural details are presented in Appendix
good
4.1 IMAGE RECONSTRUCTION USING TRANSFORMER
We use the SLATE model (Singh al. 2022), substituting SA with our SMM for an unsupervised image-to-imes task comparison. SLATE employs an Image GPT (Chen et al. 2020 decoder con- ditioned on slot representations autoregressively reconstruct discrete visual tokens from a discrete VAE (dVAE) (Im et al. 2017). It treats pre-computed slot representations query vectors and image latent code vectors as key/value vectors. Although SLATE excels at capturing complex slot-pixel interactions in synthetic images, its performance wanes with real-world data. The training signals for the dVVAE encoder, decoder, and latent discrete tokens are derived from the MSE between the input and reconstructed images. SA/SMM modules and Image GPT are trained with cross-entropy using the compressed image into dVAE tokens as the target distribution. This process is isolated from the rest of the model (i.e., dVAE), but both systems are trained simultaneously.
Setup. We consider the following datasets: CLEVR-Mirror (Singh et al.) (2022), ClevrTex Karazija et al. 2021), StapeStacks (Groth al., 2018), and COCO-2017 (Lin et al.) 2014). CLEVR-Mimor is an extension of the standard CLEVR dataset, which requires capturing global relations between local components due the presence of mirror. ShapeStacks tests the ability of the model to describe complex local interactions (multiple objects stacked on each other), and ClevrTex examines the model's capabilities in textural-rich scenes. For ShapeStacks, ClevrTex, and COCO, we used images rescaled to the resolution of 96 x 96, and CLEVR- Training conditions with hyperparameters corresponding et al. (2022), except that we use batch size equal to 64 and 2.5 x 10 training iterations for all experiments.
images are rescaled to 64 x 64. certain dataset are taken from [Singh
Results. shows the metrics of the image reconstruction performance for the test parts of dif- ferent datasets. Frechet Inception Distance (FID) (Heusel al. (017)) computed with the PyTorch- Ignite library (Fomin al. 2020 as a measure of the quality of the generated images. We also evaluated cross-entropy between tokens from dVAE (as ooe-hot distributions) and predicted distri- butions by Image GPT that conditioned by certain slot representations, as FID can be limited by the discrete VAE. SLATE has two parts that are trained separately: the dVAE and the autore- gressive transformer. The transformer takes slot representations input and aims reconstruct
5

---

Published as a conference paper at ICLR 2024
Table 1: Reconstruction performance (mean + std for 4 seeds). SLATE(SA) the original SLATE
model, SLATE(SMM) - modified SLATE with SMM instead of SA. For ceiling performance, FID results for dVAE are provided.
FID
PROSS-ENTRONT
DATA
SLATE(SA)
SLATE(SMM)
DVAE
SLATE(SA)
SLATE(SMM)
CLEVR-MIRRON
35.4 1
34.8 0.9
32.2 1 0.9
0.82 0.05
0.20 0.02
SHAPESTACKS
2.2
50.4 + 1.84
40.4 1 1.3
88.3 2.2
62 1 2.1
CLEVRTEX
116
113 4.3
82.8 1 3.9
566 = 18.1
507 21.7
21.7
COCO
129 5.2
122 4.1
92.0
540 1
461 + 12.1
CELEBA
39,2 1.5
41.3
35.7 =
=3.99199
655 = 28.6
BITMOJI
30.8 + 0.52
27.0
77.00..3 0.57
2.4
0.16
13.3 0.21
tokens from the dVAE encoder by minimizing the cross-entropy between ground truth tokens and predicted tokens. Therefore, we consider lower cross-entropy evidence of higher expressiveness of SMM representations compared to SA representations since it helps to reconstruct original to- kens accurately. The better prediction of dVAE-encoded tokens leads to better image reconstruction performance, even though was not directly trained for this since only the dVAE was trained to minimize image reconstruction loss.
In Table we report results for the SLATE model and its modification with SMM. SLATE does not output masks like object discovery model from the SA paper (Locatello et al. 2020) because SLATE uses attention maps as a heuristic to visualize each slot's information. would be incorrect object discovery.
treat them as masks for computing the FG-ARI since the task
In Appendix B we provide validation cross-entropy curves during training. Figures 2 and dis- play reconstructed images from ShapeStacks, ClevrTTxx and CLEVR-Mirror datasets, highlighting SMM's ability to recreate the correct object order from original images more accurately.
ShapeStacks
ClevrTex
CLEVR-Miroo
Figure 2: Examples of image generation with Image GPT conditioned different slot representa- tions. Images in the blue borders are from the model with the Slot Attention module, and images in green borders are generated using slots from the Slot Mixture Module. Red color stands for input images.
4.2 SET PROPERTY PREDICTION
Neural networks for sets are involved various applications across many data modalities (Carion et al. 2020 Achlioptas et al. 2017: Simonovsky & Komodakis 2018 Fujita al. 2019). Set Prop- erty Prediction is supervised task that requires the model to predict an unordered set of vectors representing the properties of objects from the input image. Sets of predicted and target vectors matched using a Hungarian algorithm (Kuhn 1955) and the learning signal provided by Huber Loss (Zhang al. 2019) between matched vectors. The Slot Mixture Module is suitable for oper- ating with sets as it preserves permutation equivariance regarding mixture components (slots) and initializes them randomly.
Setup. We use the CLEVR dataset because it is standard for comparing object-centric models on the set property prediction task. We rescaled images resolution of 128 128 as the data source.

---

Published as a conference paper at ICLR 2024
Source image
SMM
reconstruction
SA reconstruction
Figure 3: Examples of all the qualitatively incorrectly generated images from the random batch of 64 samples. In 6 cases, reconstruction using Slot Attention gave wrong order of objects (blue circle) or lost one object (red circle). In the remaining two samples, both modules gave incorrect reconstruction (green circle).
modules except SA/SMM are the same as the (Locatello et al., 2020). Each model trained with Adam optimizer (Kingma & Ba 2015) for 1.5 10 iterations with OneCycleLR (Smith & Topin 2019) learning rate scheduler at a maximum learning rate of x 10-4, We use a batch size of 512. SA/SMM number of iterations set to 5 during training and 7 during evaluation. The number of slots equals 10 since CLEVR images contain 10 or fewer objects.
Results. ShbleeSsowccass the quantitative results of experiments, where compute Average Precision (AP) using varied distance thresholds for object property and coordinate prediction. Our findings indicate that modifications the original Slot Attention such detaching slots before the final refinement iteration and rescaling coordinates - can enhance its performance. However, still trails behind the state-of- oftte-of-the-ar iDSPN model. With non-strict thresholds, other SA modifi- cations (SA-MESH (Zhang et al.) (2023) also produce near-ideal results. In contrast, SMM consis- tently outperforms a highly specialized model such as iDSPN and achieves state-of-the-an results with stringent thresholds (starting at 0.25).
Table 2: Set prediction performance on the CLEVR dataset (AP in %, mean + std for 4 seeds in our experiments). Slot Attention (Locatello et al. 2020) and iDSPN (Zhang et al. (019)) results are from the original papers. SA* is the Slot Attention model trained with the same conditions our SMM: detached slots at the last iteration rescaled coordinates to the range of [-1, 1]. SA-MESH results are taken from the original paper (Zhang et al. 2023).
MODEL
AP TD
AP1
A.o.s APo.25
APo.
APo.0625
SA
94.3 + 1.1
86.7 + 1.4
56.0 +
10.8 1 1.7
0.9 + 0.2
SA"
97.1 0.7
4..5
88.3
62.5 1 5.4
23.6 - 1.4
4.6 - 0.3
SA-MESH
99.4
99.2
98.9 0.2
91.1
47.6 0.8
12.5
IDSPN
98.8 0
0.6
9.6
95.8
76.9 = 2.5
32.3
9 3.9
SMM (OURS)
0.2
99.3
98.8
0.7
92.1 = 1.2
47.3 + 2.5
4.3
OВJЕCT DISCOVERY
Another unsupervised object-centric image-to-image task Object Discovery. We use model from Locatello et al. (2020) for this task. This approach decodes each slot representation into the 4-channel image using a Spatial Broadcast decoder (Watters et al., 2019b). The resulting recon- struction in the pixel space is estimated as mixture of decoded slots, where the first three channels are responsible for the reconstructed RGB image, and the fourth channel is for the weights of the mixture component that serve masks.
Setup. Since SA shows almost perfect metrics on such simple datasets as CLEVR, we conducted experiments complex dataset ClevrTex. We consider the same training setup from the original work Locatello et al. (2020) as for CLEVR, and the only difference is the ClevrTex image resolution is reduced to 64x64.

---

Published as a conference paper at ICLR 2024
S
Figure 4: Example of object-discovery on ClevrTex dataset. The first column represents ground- truth images. The second one is broadcast-decoder reconstructions. The next columns per slots of attention masks.
Results. BableBBsowe the similarity between ground truth segmentation masks of objects (excluded background) and mixture coefficients estimated via the Foreground Adjusted Rand Index (FG-ARI) score on ClevrTex (mean + std for 3 seeds are reported). SMM outperforms SA on such complex dataset. The examples of SMM object-discovery is shown Figure Additional results for SMM and SA are shown in Appendix C Figure
We have also conducted experiments on such simple datasets as CLEVRIU 10, Multi-dSprites, and Tetromi- noes. Since SA shows almost perfect results with long enough training these datascts, reduced the training steps to 300k (instead of 500k) to com- pare the convergence of the approaches. The FG-ARI results are: i) CLEVRIO SA 85.6 + 1.2,
Table 3: FG-ARI score for ClevrTex.
Dataset
SA
SMM
ClevrTex 62.40 2.33 74.7 1.3
SMM 91.3 + 0.6; ii) Multi-dSprites - SA 81.0 + 0.7, SMM 89.6 + 0.6; iii) Tetrominoes - SA 75.4 1.0, SMM 83.2 + 0.7. Our method shows faster convergence, resulting better object segmentation.
4.4 CONCEPT SAMPLING
Building on the concept library approach introduced in Singh et al. (2022), we propose an extension that enhances slot collection by discarding empty slots. This refinement is possible thanks to the SMM's evaluation of prior mixture weights т. We then apply GMM clustering on this streamlined set, which retains the ability assign slots clusters, identify underlying concepts, and enable concept sampling due to the estimation of Gaussian distribution parameters.
Our approach to assessing image editing quality involves extracting and sampling new slots for each image from the validation dataset using the Gaussian parameters of their assigned clusters. These sampled slots are then input to the decoder to generate the resulting image. Experiments on Bitmoji, ShapeStacks, and CLEVR- datasets yielded FID scores of 32.1, 59, and 41.4, respectively, indicating a high quality edited images comparable to regular autocncoding (Table TT. Figure illustrates this concept sampling applied to a single slot.
4.5 COMPARING VANILLA CLUSTERING
Slot Attention and Slot Mixture modules can be reduced to k-means and Gaussian Mixture Model clustering approaches by removing GRU/MLP updates, trainable q, k, v projections, and LayerNorm layers (Ba et al., 2016). Table 4 shows the training set property prediction model results for the CLEVR dataset using these vamilla clustering methods. Our experiments demonstrate that GMM clustering is better for object-centric learning, even without trainable layers. The results of the ablation studies are presented the Appendix[]
Table 4: Average Precision (mean + std for four seeds) with different distance thresholds for set property prediction task on the CLEVR dataset after 100k training iterations.
MODEL
APa
AP;
APo.5
APo.25
APo.125
K-MEANS
81.7 + 1.6
49.1 + 2.9
7.2 + 1.7
1.4 + 0.3
0.2 + 0.1
GMM
88.6 + 2.0
53.3 2.2
9.2 + 1.6
2.3 + 0.3
0.5 + 0.1
8

---

Published as a conference paper at ICLR 2024
5 RELATED WORKS
Object-centric representation. Object-centric learning involves the unsupervised identification of objects in a scene, a critical element of representation learning. The domain has evolved, starting with earlier methods that capitalized sequential attention mechanisms such as AIR (Eslami al. 2016 and its successors SQAIR (Kosiorek et al. 2018) and R-SQAIR (Stanié & Schmidhuber, 2019). The field also witnessed the advent of spatial mixture model-baset methods like Tagger (Gr- eff et al. (016) and NEM (Greff et al. 2017). which are aimed to optimize the complete likelihood of the mixture model at the pixel level, with the latter focusing unsupervised learning and per- ceptual grouping. R-NEM (Van Steenkiste et al. 2018) improved upon NEM by incorporating the discovery of object interactions. Recent advancements have integrated these principles, resulting in more sophisticated approaches like IODINE (Greff et al.) 2019), MONET (Burgess et al., 22119) and GENESIS (Engelcke et al. 2019 2021) that leverage the VAE framework (Kingma Welling 2014 Rezende et al. 20144). While these methods employ multi-step encoding and decoding, hybrid models like SPACE (Lin al. 2020) utilize spatial attention for scene decomposition. Despite the diversity of approaches, a common characteristic is the move towards more efficient and effective representation of complex scenes. Sequential extensions to SA (Kipf et al.. [02T] Elsayed et al.) 2022) work with video data, while Biza et al. (2023) introduced Invariant Slot Attention, enhanc- ing object representation invariance to geometric transformations, and Chakravarthy et al. (2023) integrated spatial- prior into object-centric models. (Seizzer al. (2022) introduces DI- NOSAUR that leverages self-supervised learning to further enhance SA efficacy in real-world data scenarios. Recent advancements such as Hénaff et al. (2022) and Wang al. (2023), while not di- rectly related to learning object representations, have made significant contributions to the related field of unsupervised object segmentation.
Source Image
Edited through concept sampling
Set prediction. Set-predicting neural network models are applied in diverse machine learn- et al. 2020 Fujita et al.J 2019 Simonovsky & Ko- modakis 2018). Despite their practicality, tra- ditional models often struggle with set repre- sentation. However, approaches like Deep Set Prediction Network (DSPN) and its improved version iDSPN (Zhang et al. 2019) account for sets' unordered nature through an inner gradi- ent descent loop and approximate implicit dif- ferentiation, respectively. Set-equivariant self- attention layers are also utilized in models like TSPN for set structure representation (Kosiorek al. 2020).
ing tasks (Achlioptas et al.
Example of editing images from ShapeStacks and Bitmoji using concept sampling. This method allows precise editing of individual concepts: the first row demonstrates the capabil- ity of changing shape, color of single objects or manipulating background. The second row gives an example of achieving similar yet slightly dif- ferent hairstyles with concept sampling.
Figure 5:
CONCLUSION AND DISCUSSION
This paper presents the Slot Mixture Module, a novel slot-hased approach to object-centric repre- sentations that generalizes the Gaussian Mixture Model (Bauckhage) 2013). Our approach uniquely considers cluster centers and inter-cluster distances enhance slot representations. Empirical ev- idence from the CLEVR dataset (Johnson et al., 2017) confirms our model's state-of-the-art per- formance in the set property prediction task with stringent thresholds and superior performance in the image reconstruction task. Furthermore, we demonstrate improved reconstruction performance on synthetic datasets like CLEVR-Mirror (Singh et al.) 2022). ShapeStacks (Groth al.) 2018), and ClevrTex Karazija et al. (202T). SMM also outperfoems SA on suach complex dataset as ClevrTex on th object discovery task. However, the limited effectiveness of modern models (Seitzer al. 2222)) on real-life images, like COCO-17 17 (Lin et al., 2014), underscores the need for improved generalization to real data. We discuss limitations of the proposed approach in Appendix F
Reproducibility Statement To improve the reproducibility of the results, we provide detailed
description of the method with the algorithm in pseudocode Section 3. we provide all nec- essary technical details in Section 4 and Appendix A we use only open datasets in the experi- ments, and we make the source code of our model publicly available at
com/
ATRT- Inst tutee e

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
This work was supported by Russian Science Foundation, grant No. 20-71-10116, https:// rscf. wwwpen//projeect/00-11110000
REFERENCES
Panos Achlioptas, Olga Diamanti, Ioannis Mitliagkas, and Leonidas Guibas. Representation learn- ing and adversarial generation of 3d point clouds. ArXiv, abs/1707.02392. 2017.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint arXiv: 160.006500 2016.
Christian Bauckhage. Lecture notes data science: k-means clustering is gaussian mixture mod- eling. 2015.
Christopher Berner, Greg Brockman, Brooke Chan, Vicki Cheung, Przemyslaw Debiak, Christy Dennison, David Farhi, Quirin Fischer, Shariq Hashme, Christopher Hesse, Rafal Józefowicz, Scott Gray, Catherine Olsson, Jakub W. Pachocki, Michael Petrov, Henrique Pondé de Oliveira Pinto, Jonathan Raiman, Tim Salimans, Jeremy Schlatter, Jonas Schneider, Szymon Sidor, Ilya Sutskever, Jie Tang, Filip Wolski, and Susan Zhang. Dota with large scale deep reinforcement learning. ArXiv, abs/1912.06t
2019.
Ondrej Biza, Sjoerd van Steenkiste, Mehdi SM Sajaadi, Gamaleldin F Elsayed, Aravindh Mahen- dran, and Thomas Kipf. Invariant slot attention: Object discovery with slot-centric reference frames. arXiv preprint arXiv:2302.04993. 202.3
Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and represen- tation. arXiv preprint 1911v:190..11390 2019.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. ArXiv, abs/2005. 2020.
Ayush Chakravarthy, Trang Nguyen, Anirudh Goyal, Yoshua Bengio, and Michael C. Mozer. Spot- light attention: Robust object-centric learning with spatial locality prior. arXiv preprint 1r555 2023.
Michael Chang. Thomas L. Griffiths, and Sergey Levine. Object representations as fixed points: Training iterative inference algorithms with implicit differentiation. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022. URL https:// openreview wwt.:eee//oooo hent//frummiieee (oorumnnd=SSsxx0O0ogg
Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever. Generative pretraining from pixels. In Hal Daumé III and Aarti Singh (eds.), Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 1691-1703. PMLR, 13-18 Jul 2020. URL https: http::::///roceeings..
mlr. plr.sress/vi wtwzess/v119/chen22// html
Kyunghyun Cho, Bart van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder- decoder for statistical machine translation. In Proceedings of the 2014 Conference on Ëm- pirical Methods in Natural Language Processing (EMNLP), pp. 1724-1734, Doha, Qatar, Oc- tober 2014. Association for Computational Linguisties. doi: 10.3115/v1/D14-1177, URL
https:
org/D14 1179
Alaaeldin El-Nouby, Shikhar Sharma, Hannes Schulz, Devon Hjelm, Layla El Asri, Samira Ebrahimi Kahou, Yoshua Bengio, and Graham W.Taylor. Tell, draw. and repeat: Gen- erating and modifying images based on continual linguistic instruction. 2019 JEEE/CVF Interna- fional Conference on Computer Vision (ICCV), pp. 10303-10311, 2019.
10

---

Published as a conference paper at ICLR 2024
Gamaleldin Elsayed, Aravindh Mahendran, Sjoerd 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-2888 2022.
Martin Engelcke, Adam R Kosiorek, Oiwi Parker Jones, and Ingmar Posner. Genesis: Gener- ative scene inference and sampling with object-centric latent representations. arXiv preprint arXiv: 2019.
Martin Engelcke, Oiwi Parker Jones, and Ingmar Posner. Genesis-v2: Inferring unordered object representations without iterative refinement. In NeurlPS, 2021.
SM Eslami, Nicolas Heess, Theophane Weber, Yuval Tassa, David Szepesvari, Geoffrey E Hinton, et al. Attend, infer, repeat: Fast scene understanding with generative models. Advances in neural information processing systems, 29, 2016.
V. Fomin, J. Anmol, S. Desroziers, J. Kriss, and A. Tejani. High-level library to help with training http://www..ommo/tttrr/////g/gna 2020.
neural networks in pytorch.
Yusuke Fujita, Naoyuki Kanda, Shota Horiguchi, Kenji Nagamatsu, and Shinji Watanabe. End-to- end neural speaker diarization with permutation-free objectives. In INTERPEECH,,00 2019.
Romain Graux. Bitmoji dataset. http:///tpppp/../o......../////p metadata/ 2021.
Klaus Greff, Antti Rasmus, Mathias Berglund, Tele Hao, Harri Valpola, and Jürgen Schmidhuber. Tagger: Deep unsupervised perceptual grouping. Advances in Neural Information Processing Systems, 29, 2016.
Klaus Greff, Sjoerd Van Steenkiste, and Jürgen Schmidhuber. Neural expectation maximization. Advances in Neural Information Processing Systems, 30, 2017.
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. International Conference on Machine Learning, pp. 2424 2433. PMLR, 2019.
Klaus Greff, Sjoerd van Steenkiste, and Jürgen Schmidhuber. On the binding problem in artificial neural networks. ArXiv, abs/2002.05200 2020.
Oliver Groth, Fabian Fuchs, Ingmar Posner, and Andrea Vedaldi. Shapestacks: Learning vision- based physical intuition for generalised object stacking. ArXiv, abs/ 0/s/1804.08018 88 2018.
Olivier J Hénaff, Skanda Koppula, Evan Shelhamer, Daniel Zoran, Andrew Jaegle, Andrew Zisser- man, João Carreira, and Relja Arandjelović. Object discovery and representation networks. In European Conference on Computer Vision, pp. 123-143. Springer, 2022.
Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. In Proceedings of the 31st International Conference Neural Information Processing Systems, NIPS' 17, pp. 6629-66400 Red Hook, NY, USA, 2017. Curran Associates Inc. ISBN 978151086096
Daniel Jiwoong Im, Sungjin Ahn, Roland Memisevic, and Yoshua Bengio. Denoising criterion for variational auto-encoding framework. Proceedings of the Thirty-Firsi AAAI Conference on Artificial Intelligence, AAAI 17, pp. 2059-2000 2065. AAAI Press, 2017.
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. CVPR, 2017.
Rishabh Kabra, Chris Burgess, Loie Matthey, Raphael Lopez Kaufman, Klaus Greff, Mal- colm Reynolds, and Alexander Lerchner. Multi-object datasets. https://gihubb.oo//oollll deepmind/multi object datasets/, 2019.
11

---

Published as a conference paper at ICLR 2024
Laurynas Karazija, Iro Laina, and Christian Rupprecht. Clevrtex: A texture-rich benchmark for unsupervised mbltt- segmentation. CoRR, abs/2111.10265, 2021. URL https:
www.//rggg//aaaeeee 10265
Ramtin Keramati, Jay Whang, Patrick Cho, and Emma Brunskill. Fast exploration with simplified models and approximately optimistic planning in model based reinforcement learning. arXiv: Artificial Intelligence. 2018.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Yoshua Bengio and Yann LeCun (eds.), 3rd International Conference on Learning Representations, ICLR 2015, San Diego CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http: //arxiv. org. abs/ 11111.1111.00 6980
Diederik P. Kingma and Max Welling. Auto-encoding variational bayes. CoRR, abs/1312.6114, 2014.
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, abs/2111.12594 2021.
Daniil Kirilenko, Alexey K. Kovalev, Yaroslav Solomentsev, Alexander Melekhin, Dmitry A. Yudin, and Aleksandr I. Panov. Vector symbolic scene representation for semantic place recognition. In 2022 International Joint Conference on Neural Networks (IJCNN). pp. 1-8, 2022. doi: 10.1 109/
Daniil Kirilenko, Alexandr Korchemnyi, Konstantin Smirnov, Alexey K Kovalev, and Aleksandr I Panov. Quantized disentangled representations for object-centric visual tasks. In International Conference on Pattern Recognition and Machine Intelligence, pp. 514-522. Springer, 2023.
Adam Kosiorek, Hyunjik Kim, Yee Whye Teh, and Ingmar Posner. Sequential attend, infer, repeat: Generative modelling of moving objects. Advances in Neural Information Processing Systems, 31, 2018.
Adam R. Kosiorek, Hyunjik Kim, and Danilo Jimenez Rezende. Conditional set generation with transformers. ArXiv, 11s/20006.18441. 2020.
Alexey K. Kovalev, Makhmud Shaban. Evgeny Osipov, and Aleksandr I. Panov. Vector semi- otic model for visual question answering. Cognitive Systems Research, 71:52-63, 2022. ISSN 1389-0417. doi:
0001 URL https://www:
sciencedi rect. com/ sc www.eeeeeeeeeeeeeeeeeeeeeeeeeeee http:///ce..eeee.....aa///aa////
H. W. Kuhn. The hungarian method for the assignment problem. problem.. Naval Research Logistics Quar- terly, 2(1-2):83-97, 1955. doi: htts://doi.oog http://.oo.oo/11....1/. wr.//o....0////1//./..... 100000002/a
Tejas D. Kulkarni, Ankush Gupta, Catalin Ionescu, Sebastian Borgeaud, Malcolm Reynolds, An- drew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for perception and control. ArXiv, abs/1906. 1883, 2019.
Nanbo Li, Cian Eastwood, and Robert Fisher. Learning object-centric representations of multi- object scenes from multiple views. Advances in Neural Information Processing Systems, 33: 5656-5666, 2020.
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft coco: Common objects in context. In David Fleet, Tomas Pajdla, Bernt Schiele, and Tinne Tuytelaars (eds.), Computer Vision - ECCV 2014, pp. 750-755, Cham, 2014. Springer International Publishing. ISBN 978-3-319-10602-11
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.
Francesco Locatello, Dirk Weissenborn. Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot atten- tion. arXiv preprint arXiv:2006. 050.0000655555 /5055, 2020.
12

---

Published as a conference paper at ICLR 2024
Thang Luong, Hieu Pham, and Christopher D. Manning. Effective approaches to attention-based neural machine translation. In Proceedings of the 2015 Conference on Empirical Methods in Nat- ural Language Processing, pp. 4412-1421, Lisbon, Portugal, September 2015. Association for / ogy.
Computational Linguistics. doi: 10. 10.18855//1/D55-1111
org/D15-1166
Shoya Matsumori, Kosuke Shingyouchi, Yukikoko Abe, Yosuke Fukuchi, Komei Sugiura, and Mi- chita Imai. Unified questioner transformer for descriptive question generation in goal-oriented visual dialogue. 2021 IEEE/CVF International Conference on Computer Vision (ICCV), pp.
1878-1887, 2021.
Toki Migimatsu and Jeannette Bohg. Object-centric task and motion planning in dynamic environ- ments. IEEE Robotics and Automation Letfers, 5:844-851 2020.
Danilo Jimenez Rezende, Shakir Mohamed, and Daan Wierstra. Stochastic back pook-propggttion and variational inference in deep latent gaussian models. ArXiv, abs/ 1401.4084, 2014.
Maximilian Seitzer, Max Horn, Andrii Zadaianchuk, Dominik Zietlow, Tianjun Xiao, Carl-Johann Simon-Gabriel, Tong He, Zheng Zhang, Bernhard Schölkopf, Thomas Brox, et al. Bridging the gap to real-world object-centric learning. arXiv preprint arXiv:2209. 14860, 2022.
Martin Simonovsky and Nikos Komodakis. Graphvae: Towards generation small graphs using variational autoencoders. In ICANN, 2018.
Gautam Singh, Fei Deng, and Sungjin Ahn. Illiterate DALL-e learns to compose. In International Conference on Learning Reprexentations, 2022. URL /ttR/ttts///opeeeviieeee http://wwwwwweeffrrooeeeee
hdd-hoorvowe30n
Leslie N Smith and Nicholay Topin. Super-convergence: Very fast training of neural networks using large learning rates. In Artificial intelligence and machine learning for multi-domane operations applications, volume 11006, pp. 369-386. SPIE, 2019.
Aleksandar Stanić and Jürgen Schmidhuber. R-sqair: relational sequential attend, infer, repeat. arXiv preprint arXiv:1910.052. 2019.
Chen Sun, Per Karlsson, Jiajun Wu, Joshua B. Tenenbaum, and Kevin Murphy. Stochastic pre- diction of multi-agent interactions from partial observations. ArXiv, 10./9992.,9941, 2019.
Sjoerd Van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural ex- pectation maximization: Unsupervised discovery of objects and their interactions. arXiv preprint arXiv:1802.10335. 2018.
Xudong Wang, Rohit Girdhar, Stella X Yu, and Ishan Misra. Cut and learn for unsupervised object detection and instance segmentation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 3124 3134, 2023.
Nicholas Watters, Loic Matthey, Matko Bosnjak, Christopher P. Burgess, and Alexander Lerchner. Cobra: Data-efficient model-based rl through unsupervised object discovery and curiosity-driven exploration. ArXiv, 0abs//995500999 2019a.
Nicholas Watters, Loie Matthey, Christopher P. Burgess, and Alexander Lerchner. Spatial broad- cast decoder: A simple architecture for learning disentangled representations in vaes. ArXiv, abs/1901.07011,, 2019b.
Jianwei Yang, Jiayuan Mao, Jiajun Wu, Devi Parikh, David Cox, Joshua B. Tenenbaum, and Chuang Gan. Object-centric diagnosis of visual reasoning. ArXiv, abs/2012.15587 2020.
Kexin Yi, Chuang Gan, Yunzhu Li, Pushmeet Kohli, Jiajun Wu, Antonio Torralba, and Joshua B. Tenenbaum. Clevrer: Collision events for video representation and reasoning. ArXiv, abs/1910.01442, 2020.
13

---

Published as a conference paper at ICLR 2024
Yan Zhang, Jonathon Hare, and Adam Prugel-Bennett. Deep set prediction networks. In H. Wallach, H. Larochelle, A. Beygekzimer, F. d'Alché-Buc, E. Fox, and R. Gar- nett (eds.), Advances in Neural Information Processing Systems, volume 32. Curran Asso- ciates, Inc., 2019. URL http:sss/pprrrccci..g......pssss http:///ps..ccpppp//////////////
[e 7eeTao00ee777e2555aa//33330--pp pdf
Yan Zhang, David W Zhang. Simon Lacoste-Julien, Gertjan Burghouts, and Cees GM Snoek. Unlocking slot attention by changing optimal transport costs. arXiv preprint arXiv:2301.13192) 2023.
14

---

Published as a conference paper at ICLR 2024
A ARCHITECTURE DETAILS
Tables 5 G describe hyperparameters for our experiments. the case of using SMM instead of SA, additional dimensionality reduction for slots via trainable matrix multiplication.
Table 5: Architecture of the CNN encoder for property prediction and object discovery tasks experiments. The set prediction model uses a stride of in the layers with while the object discovery model uses stride of in these layers.
Layer
Channels
Activation
Params
Conv2D x
64
ReLU
stride:
Conv2D 5
64
ReLU
stride: 1/2*
Conv2D 5 x
64
ReLU
stride: 1/2*
Conv2D 5
64
ReLU
stride:
Position Embedding
absolute
Flatten
dims: w,
LayerNorm
Linear
64
ReLU
Linear
64
Table 6: Spatial broadcast decoder for object discovery task on the CLEVR and ClevrTex datasets.
Layer
Channels/Size
Activation
Params
Spatial Broadcast
8 x
Position Embedding
absolute
Cony ToonTraasppse2D
5
64
ReLU
stride: 2
Conv TonnTraaspSse2D x 5
64
ReLU
stride: 2
Conv SounTraasppse2D
5
64
ReLU
stride: 2
Conv Transpsspp
x 5
64
ReLU
stride: 2
Conv Transpsanppse21 5
64
ReLU
stride:
Conv Ton/Transpos22> x 3
4
stride: 1
Split Channels
RGB (3), mask (1)
Softmax on masks (slots dim)
Combine components
Table 7: Spatial broadcast decoder for object discovery experiments on Tetrominoes and Multi- dSprites datasets.
Layer
Channels/Size
Activation
Params
Spatial Broadcast
64 x 64
Position Embedding
absolute
Conv TranTrasspoSDD
x 5 32
ReLU
stride: 1
Conv SonnTraaspose2D 5
32
ReLU
stride: 1
Conv Transpose2D x 5
32
ReLU
stride: 1
TrnnTraaspose2D> 3
stride:
Split Channels
RGB (3), mask (1)
Softmax on masks (slots dim)
Combine components
15

---

Published as a conference paper at ICLR 2024
Table 8: Hyperparameters used for experiments with the SLATE architecture.
Module
Parameter
Value
Image Size
96
Encoded Tokens
576
dVAE
Vocab size
4096
dVAE
Temp. Cooldown
1.0 to 0.1
dVAE
Temp. Cooldown Steps
30000
dVAE
LR (no warmup)
0.0003
Transformer
Layers
8
Transformer
Heads
8
Transformer
Hidden Dim.
192
SA/SMM
Num. slots
12
SA/SMM
Iterations
7
SA/SMM
Slot dim.
192
B ADDITIONAL RESULTS FOR IMAGE RECONSTRUCTION
Figure shows validation cross-entropy curves during training.
CLEVR-Mirror
CLEVR-Tex
ShapeStacks
COCO
SUB1
9MM
102
020
53k
2004
1008
2004
Figure 6: Validation cross-entropy during training for different datasets. Our experiments show that using the SMM module instead of SA consistently improves the validation performance of the autoregressive transformer by about 10 percent during training. The result maintained for all the datasets that we use.
16

---

Published as a conference paper at ICLR 2024
C
ADDITIONAL RESULTS FOR OBJECT DISCOVERY
Figure shows examples of object-discovery ClevrTTex dataset for SMM and SA.
Figure 7: Examples of object-discovery on ClevrTex dataset for SMM and SA. The first column represents ground-truth images. The second one is broadcast-decoder reconstructions. The next columns are per slots attention masks. Examples in the red borders are for SMM and examples in the blue borders are for SA.
D ABLATION STUDY
We also conducted the following ablation experiments. We trained Slot Attention (SA) and SMM object discovery models on the CLEVR6 dataset (scenes with six or fewer objects) with seven slots and three update iterations. Then we evaluated the trained models using different setups.
Evaluating the models on the CLEVR6 test split with 1, 2, .... and 7 update iterations gives the following FG-ARI scores (see T
Table 9; FG-ARI score for different number of update iterations (from I to 7)) at test time on the CLEVR6 dataset.
MODEL
2
3
4
5
7
SA
68.0
90.1
99.3
99.5
99.5
99.5 99.4
99.6
SMM
67.0
78.9
99.8
99.7
99.0
97.1
Evaluating the models on the CLEVR test images with only 7, 8, 9, and 10 objects (and a corre- sponding number of slots) gives the following results (see Table 10
17

---

Published as a conference paper at ICLR 2024
Table 10: FG-ARI score for different number of objects (from 7 to 10) at test time.
MODEL
7
9
10
SA
97.0
95.1
94.0
93.3
SMM
97.1
96.6
94.5
94.5
The results show that SMM is more robust the number of out-of-distribution objects and less robust to different numbers of update iterations at test time.
E HIGHER RESOLUTION OBJECT DISCOVERY
On Figure g and Figure Mare shown examples of object-discovery and image reconstruction on ClevrTex dataset with 128x 128 resolution for SMM and SA. Model architecture presented on Table]land T
Figure 8: Examples of object-discovery on ClevrTTex dataset with 128 x 128 image resolution for SMM and SA. The first column represents ground-truth images. The second one broadcast- decoder reconstructions. The next columns are per slots of attention masks. Examples in the red borders are for SMM and examples in the blue borders are for SA.
Figure 9: Image reconstructions on ClevrTex dataset with 128 x 128 resolution for SMM and SA.
The first row represents ground-truth images. The second one broadcast-decoder reconstructions. Examples in the red borders are for SMM and examples in the blue borders are for SA.
18

---

Published as a conference paper at ICLR 2024
Table 11: Spatial broadcast decoder for object discovery task on the ClevrTex datasets with 128x 128
image resolution.
Layer
Channels/Size
Activation
Params
Spatial Broadcast
8 x
Position Embedding
absolute
Conv TrnnTssasDDSe2 5
64
ReLU
stride: 2
Cony Connsraaspose2D> x 5
64
ReLU
stride: 2
Conv TrnnTssaspsse2D 5
64
ReLU
stride: 2
Conv Transpose2D X 5
64
ReLU
stride: 2
Conv SoanTraaspOSe2D 5
64
ReLU
stride: I
Conv TranTTosssp2 x 3
4
stride: 1
Split Channels
RGB (3), mask (1)
Softmax on masks (slots dim)
Combine components
Table 12: Architecture of the CNN encoder for ClevrTex with image size 128 x 128.
Layer
Channels
Activation
Params
Conv2D 5
64
ReLU
stride: 2
Conv2D x 5
64
ReLU
stride: I
Conv2D 5
64
ReLU
stride: I
Conv2D x 5
64
ReLU
stride: 1
Position Embedding
absolute
Flatten
dims: w, h
LayerNorm
Linear
64
ReLU
Linear
64
F LIMITATIONS
The main limitations SMM are due to the fact that the model belongs to the class of slot models, since it is based on prominent representative of this class Slot Attention. These limitations are: the need to specify the number of slots advance, thereby setting the maximum possible number of objects in the image; over/under segmentation based on the specified number of slots; as a result, poor quality with a large number of slots; poor results and not aligned with human perception segmentation of real data.
While we have enhanced real-world object-centric image generation, the overall quality remains subpar and the attention maps scarcely resemble human object-focused vision. Future research should address the challenge of scaling and refining these visual models for complex real-world images. Also, as shown in the ablation study, SMM is less robust different numbers of update iterations at test time.
We also want draw attention the potential negative societal impact of object-centric models, since modification and replacement of objects represented by slot in an image can be used for malicious purposes.
19


Reviews:
[{"review_id": "9K4YfxtH8G", "reviewer": "Reviewer_spVv", "summary": "This paper proposes a generalization of the Slot Attention approach by Locatello et al. 2020, replacing the k-means algorithm with a Gaussian Mixture Model to improve the expressiveness of the slot representations. In Slot Attention (SA), slot representations are cluster centers, which means that SA is limited by the information contained and represented in these cluster centers, whereas SMM represents slots not only as centers of clusters but also incorporate information about the distance between clusters and assigned vectors. Experiment results on standard benchmark datasets showed improved performance over SA.", "questions": "My main concerns mainly pertain to the quality of the reconstructed images. Any improvements in that aspect would strengthen the manuscript.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The idea presented in this manuscript is quite sound and intuitive. Overall, the manuscript is well-written and easy to digest. Given the growing interest in Slot Attention, this paper comes timely. It proposes a new direction into how object-centric learning via slot attention could be approached without drastically departing from the main concept while achieving better performance.", "weaknesses": "1. There is a typo in Equation 7. The covariance matrix $\\Sigma^{\\*}$ should be computed based on the updated mean $\\mu^{\\*}$\n\n2. The qualitative results could be improved, in my opinion. The images depicted in Figure 2 are quite blurry. This makes it quite difficult to assess whether SMM brings actually any substantial improvements over SA or not. I'd recommend the authors to provide higher quality images to strengthen their manuscript. \n\n3. In Figure 2, the images produced from SMM seem quite distorted. These distortions seem more pronounced for the ClevrTex dataset. I am under the impression that given how expressive SMM (when compared to SA) is the reconstructions from the learned slots would be more naturally-looking. It is not clear to me whether the distortions result from the considered Image GPT model, or whether it's because the learned slot representations are not that informative.\n\n4. Comparing *quantitatively* the attention maps learned by SMM against those learned by SA would have been quite helpful. I am under the impression that for more complex image scenes, the attention maps learned by SA would be more accurate than those learned by SMM due to potentially the high-variance involved.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "rZDRGRwjgM", "reviewer": "Reviewer_HCjc", "summary": "This paper introduces an object-centric architecture that is able to decompose the scene into a set of slots, useful for several downstream tasks ranging from image reconstruction, object discovery and property prediction. While the traditional Slot Attention model generates the slots using a learnable k-means clustering, the proposed method clusters the pixels using a learnable version of the gaussian mixture modeling. Concretely, they are using an iterative approach to estimate not only the centroids of the clusters, but also the covariance associated with each component. The resulting model proves to be beneficial compared to the basic Slot Attention architecture, especially in more difficult scenarios such as low resolution or harder datasets.", "questions": "Please see the Weaknesses section", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- The idea of replacing the k-means algorithm with GMM in the slot attention architecture represent an interesting and novel idea. Given that k-means clustering is a particular case of GMM, the resulting method has potential to mimic, and go beyond the capabilities of the slot attention models.\n- The method shows improvement on various tasks", "weaknesses": "- Since k-means is a particular case of the GMM framework (with the GMM allowing for a learnable variance in clusters), a discussion regarding the advantages of going for the more general model should be included. What are some real world scenarios where a gaussian-based method perform better?\n- From the results in Appendix, Table 9 the SMM model seems to be more sensitive to then number of iterations compared to the traditional SA. Is this a consequence of the EM algorithm not converging or are there other optimisation issues that causes this phenomena?\n- As mentioning in Section 3, the SMM differs from SA in 3 aspects: the dot product is replaced by the Gaussian density function; both covariance and mean values are updated in the iterative process and the slot representation is a combination of mean and covariance. It would be insightful to see an ablation study showing the contribution of each one of the changes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "gE3AViKxSJ", "reviewer": "Reviewer_7AcU", "summary": "The authors propose a modification of the popular Slot-Attention module in which the slots parameterize a gaussian mixture model, rather than a k-means model (which is how the original SA is typically conceptualized). This allows slots to model both the means and variances (diagonalized covariance) of a latent distribution leading to improved performance on various benchmarks, whilst also allowing \"empty\" slots to be identified and discarded based on their learned prior mixture weights. The results on various benchmarks are compelling, and the experiments are well-thought out - offering fair comparisons against a vanilla (with implicit gradients) SA baseline, and with some ablations.", "questions": "Some minor questions which likely reflect ignorance on the part of the reviewer:\n* _Redundant Slots_: Do the authors every observe multiple slots sharing the encoding the same object in a scene, or \"competition\" between the spatial attention over the CNN Feature maps work strongly enough to prevent this, even in SMMs?\n* At the end of page one you write that \"We believe... set prediction... requires distinguishing objects from each other\"- this reads as if you are saying that set prediction is a better measure of OCL competency than object discovery - I am not sure I follow this argument if so, as it seems that entanglement of representations should be less of a hindrance to object discovery than the production of e.g. object-wise masks?\n* At the start of section 3 you state that \"GRU [...]. takes current and previous slot representations as input and hidden states\" but in Algorithm 1 it seems that the GRU is fed only the means?\n* It seems as if figure 1 (left) might differ from Algorithm 1 in a few ways (at least, it was sufficiently unclear that perhaps labelling edges would be worthwhile). Most notably, the L2 difference for computing the covariance matrix is taken after the GRU-mean update in the algorithm, but before in the figure. Additionally, I don't think the MLP / LayerNorm are represented in the figure.\n\nNitpicks:\n* The penultimate sentence of the first paragraph in sec 2.1 is somewhat difficult to parse\n* The second paragraph in sec 2.1 - \"updating iteration as\" needs changing\n* In the first paragraph of sec 4.3 you describe the role of the 4th channel in the spatial broadcast decoder as being \"for the weights of the mixture component\" ; whilst this is correct in the context of the sentence, it is slightly confusing given that the slots are represented with mixture components $$\\pi$ within the SMM itself. It might be clearer to talk about the weights of _masks_.", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The paper is very well-written and easy to follow. The methodology is motivated and presented clearly, and the experiments are thorough.\n\nThe idea itself is a fairly simple and elegant extension of SA which has not been explored in the community before. As comparisons with SA are made using implicit gradients (which are commonly used for improved training stability of such models now) on various benchmarks, there are good reasons to expect that the proposed SMM model may replace SA in most use-cases (i.e. it improves upon the _practical_ state of the art in a simple way). Given the considerable popularity of SA and its derivatives, this means the paper should have considerable impact in the Object-Centric Learning community.\n\nThe ablations compare SA and SMM against their barebones counter-parts (which are closed to traditional k-Means and GMMs), bolstering the suggestion that the learning of the variance is truly beneficial for learning superior slot representations. In addition to comparing on reconstruction and property identification datasets, they show that the quality of edited images (formed through the manipulation of slots) is superior in the SMM model also, lending further credence to the former claim.", "weaknesses": "Given the emphasis on the method as an improvement over SA across numerous datasets and kinds of OCL tasks, there are no significant weaknesses in the methodology or choice of experiments. That being said, it would have been interesting to see an analysis of the properties of the learned representations (disentanglement, interpolation/extrapolation over generative factors, etc.), but the paper already contains sufficient information to demonstrate the value of SMMs - most notably the Concept Sampling experiments which show steerability at the level of distinct generative factors.\n\nOne minor note is that whilst the authors do take care to fairly compare SA against SMMs, it is still possible that SA with $2D$ slots would have been a \"fairer\" baseline (though, as they rightly point out, this would have been a model with more learnable parameters), in that the capacity in the slots may have been the limiting factor (unlikely); though I suspect the GMMs would still have performed more strongly (as $D$ is not so much the bottleneck, as the form of the distribution the models can capture). It would also be interesting to see how the behaviour of the two models varied as slots became very small, or sparse regularization was applied to the slot representations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "6LcrEGPrHP", "reviewer": "Reviewer_fszf", "summary": "The paper studies object-centric learning based on the existing slot attention method. Particularly, it incorporates Gaussian mixture model (GMM) rather than the soft k-means clustering used in the original slot attention to learn better slot representations, achieving promising results on three tasks: image reconstruction, set prediction and object discovery.", "questions": "See above.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1) The paper studies a very important problem of object-centric learning. \n\n2) The introduced GMM mechanism to learn distinguishable slot representations appears to be more effective than using a single mean vector to represent each slot.", "weaknesses": "1) The key contribution of this paper is the application of GMM mechanism for better slot representation learning, but the improvements are more about the performance in supervised tasks: image reconstruction and set prediction, instead of the more desirable unsupervised object segmentation. This makes the general contribution to be less appealing.\n\n2) In page 2, the claimed Contributions (2) and (3) are not very meaningful, or at least both can be combined, because both sentences describe the experimental results.\n\n3) In Equations (1)(2), it's suggested to use math symbols instead of English words. It’s also suggested to use bold symbols to represent vectors. \n\n4) In page 4, it is unclear how the function f_theta(x, u, diag) works. Obviously, the input has three elements instead of two as described in the text \"R^{NxD} x R^{KxD} -> R^{NxK}\". More details should be provided because this is the primary technique contribution of this paper.\n\n5) In Section 4.3, for the experiments of object discovery, the dataset(ClevrTex) is a bit too simple and the evaluation metric ARI is actually not suitable as well because the scores can easily achieve perfect numbers. It is advised to evaluate on more complex (real-world) datasets and use additional metrics such as AP scores, as also pointed out in the paper \"Promising or elusive? unsupervised object segmentation from real-world single images, NeurIPS 2022\". \n\n6) In Section 5 “Related Works”, in the field of object-centric learning, recently, there are a number of works using pretrained feature representations to discovery objects, such as DINOSAUR (Bridging the Gap to Real-World Object-Centric Learning, ICLR 2023), Odin (Object discovery and representation networks, ECCV 2022), and CutLER (Cut and Learn for Unsupervised Object Detection and Instance Segmentation, CVPR 2023). They should be discussed appropriately. \n\nMore other related works should be discussed as well, including (1) Invariant Slot Attention: Object Discovery with Slot-Centric Reference Frames, ICML 2023, (2) Spotlight Attention: Robust Object-Centric Learning With a Spatial Locality Prior, arxiv 2023.\n\nTo sum up, the paper can be further improved in the following aspects: 1) a better presentation about the technique parts, 2) more concrete experiments to demonstrate better performance in object discovery, 3) discussions about more related and recent works in the field of object-centric learning.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)