Original: O BJECT -C ENTRIC L EARNING WITH S LOT M IXTURE M ODELS
Revised: O BJECT -C ENTRIC L EARNING WITH S LOT M IXTURE M ODELS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Object-centric architectures usually apply some differentiable module on the whole feature map to decompose it into sets of entities representations called slots.
Revised: Object-centric architectures usually apply some differentiable module on the whole feature map to decompose it into sets of entity representations called slots.

Original: Some of these methods structurally resemble clustering algorithms, where the center of the cluster in latent space serves as  slot representation.
Revised: Some of these methods structurally resemble clustering algorithms, where the center of the cluster in latent space serves as a slot representation.

Original: Slot Attention is an example of such a method as a learnable analog of the soft k-Means algorithm.
Revised: Slot Attention is an example of such a method as a learnable analog of the soft K-Means algorithm.

Original: In our work, we use the learnable clustering method based on Gaussian Mixture Model, unlike other approaches we represent slots not only as centers of clusters but we also use information about the distance between clusters and assigned vectors, which leads to more expressive slots representations.
Revised: In our work, we use the learnable clustering method based on Gaussian Mixture Model, unlike other approaches we represent slots not only as centers of clusters but we also use information about the distance between clusters and assigned vectors, which leads to more expressive slot representations.

Original: Our experiments demonstrate that using this approach instead of Slot Attention improves performance in different scenarios achieving  state-of-the-art performance in the set property prediction task.
Revised: Our experiments demonstrate that using this approach instead of Slot Attention improves performance in different objectcentric scenarios, achieving the state-of-the-art performance in the set property prediction task.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In recent years, interest in object-centric representations has greatly increased (Greff et al.
Revised: In recent years, interest in object-centric representations has greatly increased (Greff et al.

Original: Burgess et al.
Revised: Burgess et al.

Original: ; Li et al.
Revised: ; Li et al.

Original: Engelcke et al.
Revised: Engelcke et al.

Original: ; Locatello et al.
Revised: ; Locatello et al.

Original: Such representations have the potential to improve the generalization ability of machine learning methods in many domains, such as reinforcement learning Keramati et al.
Revised: Such representations have the potential to improve the generalization ability of machine learning methods in many domains, such as reinforcement learning (Keramati et al.

Original: ; Watters et al. 
Revised: ; Watters et al. (2019a);

Original: Kulkarni et al.
Revised: Kulkarni et al.

Original: Berner et al.
Revised: Berner et al.

Original: ; Sun et al.
Revised: ; Sun et al.

Original: (2019), scene representation and generation El-Nouby et al.
Revised: scene representation and generation (El-Nouby et al.

Original: ; Matsumori et al.
Revised: ; Matsumori et al.

Original: ; Kulkarni et al.
Revised: ; Kulkarni et al.

Original:  reasoning Yang et al.
Revised: , reasoning (Yang et al.

Original: (2020), object-centric visual tasks Groth et al.
Revised: (2020)), object-centric visual tasks (Groth et al.

Original: (2018a);
Revised: (2018a);

Original: Yi et al.
Revised: Yi et al.

Original: ; Singh et al.
Revised: ; Singh et al.

Original: (2021b), and planning (Migimatsu & Bohg (2020)).
Revised: (2021b)), and planning (Migimatsu & Bohg (2020)).

Original: Automatic segmentation of objects on the scene and the formation of a structured latent space can be carried out in various ways Greff et al.
Revised: Automatic segmentation of objects on the scene and the formation of a structured latent space can be carried out in various ways (Greff et al.

Original: (2020): augmentation of features with grouping information, the use of a tensor product representation, engaging ideas of attractor dynamics, etc.
Revised: augmentation of features with grouping information, the use of a tensor product representation, engaging ideas of attractor dynamics, etc.

Original: however, the most effective method for learning object-centric representations is Slot Attention Locatello et al.
Revised: However, the most effective method for learning object-centric representations is Slot Attention Locatello et al.

Original: Slot Attention maps the input feature vector received from the convolutional encoder to a fixed number of output feature vectors, which are called Slots.
Revised: Slot Attention maps the input feature vector received from the convolutional encoder to a fixed number of output feature vectors, which are called Slots.

Original: As a result of training, each object is assigned a corresponding slot.
Revised: As a result of training, each object is assigned a corresponding slot.

Original: If the number of slots is greater than the number of objects, then some of the slots remain empty (do not contain objects).
Revised: If the number of slots is greater than the number of objects, then some of the slots remain empty (do not contain objects).

Original: This approach showed significant results in such object-centir tasks as set property prediction and object discovery.
Revised: This approach showed significant results in such object-centric tasks as set property prediction and object discovery.

Original: The iterative slot-based approach can be considered as a variant of the soft k-Means algorithm
Revised: The iterative slot-based approach can be considered as a variant of the soft K-Means algorithm

Original: Bauckhage (2015b), where the key/value/query projections are replaced with the identity function and updating via the recurrent neural network is excluded Locatello et al.
Revised: (Bauckhage (2015b)), where the key/value/query projections are replaced with the identity function and updating via the recurrent neural network is excluded (Locatello et al.

Original: In our work, we propose another version of the generalization, when the k-Means algorithm is considered as a
Revised: In our work, we propose another version of the generalization, when the K-Means algorithm is considered as a

Original: Gaussian Mixture Model Bauckhage (2015a).
Revised: Gaussian Mixture Model (Bauckhage (2015a)).

Original: We represent slots not only as centers of clusters but we also use information about the distance between clusters and assigned vectors, which leads to more expressive slots representations.
Revised: We represent slots not only as centers of clusters, but we also use information about the distance between clusters and assigned vectors, which leads to more expressive slot representations.

Original: Representing slots in this way improves the quality of the model in object-centric problems, achieving the state-of-the-art results in the set property prediction task, even compared to highly specialized models Zhang et al.
Revised: Representing slots in this way improves the quality of the model in object-centric problems, achieving the state-of-the-art results in the set property prediction task, even compared to highly specialized models (Zhang et al.

Original: (2019b), and also improves the generalization ability of the model for image reconstruction task.
Revised: (2019b)), and also improves the generalization ability of the model for image reconstruction task.

Original: The paper is structured as follows: in Section 2 we provide background information about Mixture
Revised: The paper is structured as follows. In Section 2, we provide background information about Slot

Original: Models and their process of  training.
Revised: Attention module and Mixture Models. We also describe the process of their training.

Original: In Section 3.1 we introduce a Slot Mixture Module – the modification of a Slot Attention Module which provides more expressive slots representations.
Revised: In Section 3, we introduce a Slot Mixture Module — the modification of a Slot Attention Module which provides more expressive slot representations.

Original: In Section 4.1 by extensive experiments we show that the proposed Slot Mixture Module riches the state-of-the-art performance in the set property prediction task on the CLEVR dataset Johnson et al.
Revised: In Section 4.1, by extensive experiments we show that the proposed Slot Mixture Module riches the state-of-the-art performance in the set property prediction task on the CLEVR dataset Johnson et al.

Original: (2017) and outperforms even highly specialized models Zhang et al. (2019b).
Revised: and outperforms even highly specialized models Zhang et al. (2019b).

Original: In Section 4.2 we provide experimental results for the image reconstruction task on four datasets, three with synthetic images (CLEVR-Mirror Singh et al.
Revised: In Section 4.2, we provide experimental results for the image reconstruction task on four datasets: three with synthetic images (CLEVR-Mirror Singh et al.

Original: (2021a)
Revised: (2021a)

Original: , ShapeStacks Groth et al.
Revised: , ShapeStacks Groth et al.

Original: (2018b),
Revised: (2018b),

Original: ClevrTex Karazija et al.
Revised: ClevrTex Karazija et al.

Original: (2021)) and one with real life images COCO-2017 Lin et al.
Revised: (2021)) and one with real life images COCO- (Lin et al.

Original: The proposed Slot Mixture Module improves reconstruction performance. 
Revised: The proposed Slot Mixture Module improves reconstruction performance. In

Revised: Section 4.3, we demonstrate that Slot Mixture Module outperforms original Slot Attention on the

Revised: Object Discovery task on the CLEVR10 dataset.

Revised: In Section 4.4, we compare K-Means and Gaussian

Revised: Mixture Model clustering approaches on the Set Property Prediction task and show that Gaussian

Original:  In Section 6 we discuss the obtained results, advantages, and limitations of our work.
Revised: Mixture Model clustering is a better choice for object-centric learning. We give a short overview of related works in Section 5. In Section 6, we discuss the obtained results, advantages, and limitations of our work.

Original: The main contributions of our paper are as follows:
Revised: The main contributions of our paper are as follows:

Original: We proposed a generalization of slot-based approach for object-centric representations as a Gaussian Mixture Model (Section 3.1).
Revised: We proposed a generalization of slot-based approach for object-centric representations as a Gaussian Mixture Model (Section 3).

Original: Such a representation allows state-of-the-art performance in the set property prediction task, even in comparison with specialized models (Section 4.1), which are not aimed at building disentangled representations.
Revised: Such a representation allows state-of-the-art performance in the set property prediction task, even in comparison with specialized models (Section 4.1), which are not aimed at building disentangled representations.

Original: The representation of slots as a Gaussian Mixture improves the generalization ability of the model in other object-oriented tasks (Section 4.2). 
Revised: The slot representations as a Gaussian Mixture improve the generalization ability of the model in other object-oriented tasks (Section 4.2). The Slot Mixture Module shows a much faster convergence on the Object Discovery task compare to the Original Slot Attention (Section 4.3).  B ACKGROUND 2. S LOT A TTENTION Slot Attention (SA) module (Locatello et al.

Revised: (2020)) is an iterative attention mechanism that is designed to map a distributed feature map to a set of K slots.

Revised: Randomly initialized slots from a

Original: M IXTURE M ODELS
Revised: Gaussian distribution with trainable parameters are used to get q projections of slots. Feature map vectors with corresponding projections serve as k and v vectors. Dot-product attention between q and k vectors with the softmax across q dimension implies competition between slots for explaining parts of the input. Attention coefficients are used to assign v vectors to slots via a weighted mean.

Revised: [Equation]

Revised: Gated Recurrent Unit (GRU) (Cho et al.

Revised: (2014)) is used for addition slots update.

Revised: It takes slot representations before update iteration as a hidden state and updated slots as inputs.

Revised: The important property of Slot Attention is that it has permutation invariance with respect to input vectors of the feature map and permutation equivariance with respect to slots.

Revised: These properties make the Slot

Revised: Attention module suitable for operating with sets and object-centric representations.

Revised: Technically, Slot Attention is a learnable analogue of K-Means clustering algorithm with an additional trainable GRU update step and dot product (with trainable q , k , v ) projections instead of Euclidean distance as the measure of similarity between the input vectors and cluster centroids.

Revised: At the same time, K-Means clustering can be considered as a simplified version of the Gaussian Mixture Model.

Original: Mixture Models (MM) is a class of parametric probabilistic models, in which it is assumed that each x i from some observations
Revised: Mixture Models (MM) is a class of parametric probabilistic models, in which it is assumed that each x i from some observations

Original: X = { x 1 , ..., x N } ∈ R N × D is sampled from the mixture distribution with K mixture components and
Revised: X = { x 1 , ..., x N } ∈ R N × D is sampled from the mixture distribution with K mixture components and

Original: prior mixture weights π ∈ R K :
Revised: prior mixture weights π ∈ R K :

Original: [Equation]
Revised: [Equation]

Original: These models can be seen as the models with latent variables z
Revised: These models can be seen as the models with latent variables z

Original: i,k ∈ { z 1 , ..., z K } that indicate which component x i came from.
Revised: i,k ∈ { z 1 , ..., z K } that indicate which component x i came from.

Original: The problem is to find such K groups of component parameters θ k and component assignments of each sample x i that will maximize the likelihood of the model P ( X | θ ) .
Revised: The problem is to find such K groups of component parameters θ k and component assignments of each sample x i that will maximize the likelihood of the model P ( X | θ ) .

Original: The Expectation Maximization (EM) algorithm is an iterative algorithm that addresses this problem and includes two general steps.
Revised: The Expectation Maximization (EM) algorithm is an iterative algorithm that addresses this problem and includes two general steps.

Original: The Expectation (E) step : evaluate the expected value of the complete likelihood P ( X , Z | θ ∗ ) with respect to the conditional distribution of P ( Z | X , θ ) :
Revised: The Expectation (E) step : evaluate the expected value of the complete likelihood P ( X , Z | θ ∗ ) with respect to the conditional distribution of P ( Z | X , θ ) :

Original: [Equation]
Revised: [Equation]

Original: where I ( ∗ ) is an indicator function.
Revised: where I ( ∗ ) is an indicator function.

Original: The Maximization (M) step : find θ ∗ , π ∗ that maximize Q ( θ ∗ , π ∗ | θ , π ) :
Revised: The Maximization (M) step : find θ ∗ , π ∗ that maximize Q ( θ ∗ , π ∗ | θ , π ) :

Original: [Equation]
Revised: [Equation]

Original: 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 i | θ k ) =
Revised: One of the most widely used models of this kind is Gaussian Mixture Model (GMM), where each mixture component is modeled as a Gaussian distribution parameterized with its mean values and covariance matrix, which is diagonal in the simplest case: P ( x i | θ k ) =

Original: N ( x i | µ k , Σ k ) , Σ k [Equation] diag ( σ 2 k ) .
Revised: N ( x i | µ k , Σ k ) , Σ k = diag ( σ 2 k ) .

Original: In this case, EM algorithm is reduced to the following calculations.
Revised: In this case, EM algorithm is reduced to the following calculations.

Original: E step :
Revised: E step :

Original: [Equation]
Revised: [Equation]

Original: M step :
Revised: M step :

Original: [Equation]
Revised: [Equation]

Original: The goal of object-centric learning using slots is to distribute low-level features from the neural network to the slots in a trainable way so that the resulting slot representations correspond to higherlevel abstractions.
Revised: The key difference between the Gaussian Mixture Model and K-Means clustering is that GMM considers not only the centers of clusters, but also the distance between clusters and assigned vectors with the prior probabilities of each cluster.  S LOT M IXTURE M ODULE

Original: To address this problem in the case of visual tasks, we use a modified mixture model approach and call it the Slot Mixture Module (SMM).
Revised: For the purposes of the object-centric learning we propose a modified Gaussian Mixture Model approach and call it the Slot Mixture Module (SMM).

Original: The essence of the approach is as follows: a sequence of vectorsobtained from an image with some convolutional neural network is considered as observations, concatenated µ k , Σ k are considered as slots representations.
Revised: This module uses GMM E and M steps (Section 2.2) to map feature map vectors from the convolutional neural network (CNN) encoder to the set of slot representations, where slots are concatenation of mean values and diagonal of covariance matrix.

Original: The density function of each component is modeled using function f θ ( x , µ , Σ diag ) : R N × D × R K × D → R N × K with trainable parameters θ . This function involves some input pre-processing ( q, k, v projections in our case) and pairwise measure ofsimilarity between input vectors x and slots representations µ , Σ .
Revised: This set of slots is further used in the downstream task.

Original: The remaining calculations are the same as in the example with GMM, except for an extra step after updating the mean  values:
Revised: We also use the same additional neural network update step for the mean values before updating the covariance values:

Original: [Equation]
Revised: [Equation]

Original: These two steps, along with the function f θ , serve the needs of the downstream task by linking the external and internal models.
Revised: These two steps serve the needs of the downstream task by linking the external and internal models.

Original: The internal model (SMM) tries to update its parameters µ , Σ so that the input vectors x are assigned to slots with the maximum likelihood, while external model takes these parameters as input.
Revised: The internal model (E and M steps in SMM) tries to update its parameters µ , Σ so that the input vectors x are assigned to slots with the maximum likelihood, while external model takes these parameters as input.

Original: The full pseudocode is presented in algorithm 1.
Revised: The full pseudocode is presented in the Algorithm 1.

Original: The Slot Attention (SA) module Locatello et al. (2020) produces a set of task-dependent slot representations.
Revised: A function f θ ( x , µ , Σ diag ) : R N × D × R K × D

Original: From a technical point of view, Slot Attention is a trainable analogue ofk-Means clustering algorithm in thesame way that our SMM is a trainable analogue of Gaussian Mixture Model. The main difference between approaches is that Slot Attention does not explicitly use information about the distance of vectors to cluster centers when forming a slot representation, slots are  the cluster centers themselves.
Revised: → R N × K stands for a log-Gaussian density function with additional k and q projections for input and slot vectors. Slot Mixture Module can be seen as an extension of the Slot Attention Module 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.

Original: Algorithm 1: The Slot Mixture Module pseudocode.
Revised: Algorithm 1: The Slot Mixture Module pseudocode.

Original: π is initialized as uniform categorical distribution, µ and Σ diag 
Revised: π is initialized as uniform categorical distribution, µ and Σ diag are initialized from Gaussian distributions with trainable parameters. Input feature added positional embeddings; µ , Σ diag ∈ R K × D , π ∈ R K — SMM initialization parameters. Output slots R

Original: E XPERIMENTS
Revised: 1 x = MLP(LayerNorm( x )) 2 for t = 0 ... T do logits = f θ ( x , µ , Σ diag ) gammas = SoftMax(logits + log π , dim= 1 ) π

Revised: = gammas.mean(dim= 0 ) µ ∗ =

Revised: WeightedMean(weights=gammas, values= x ) µ = GRU(input= µ ∗ , hidden= µ ) µ += MLP(LayerNorm( µ ))

Revised: Σ diag =

Revised: WeightedMean(weights=gammas, values= ( x − µ ) 2 ))

Revised: 10 slots = concat([ µ , Σ diag ]) 11 return slots

Revised: E XPERIMENTS

Original: Since our model can be seen as an extension of the Slot Attention, we use it as our main competitor.
Revised: Since our model can be seen as an extension of the Slot Attention, we use it as our main competitor.

Original: In each experiment, we trained two versions of the same model: the one with the Slot Attention module and another one with our SMM, keeping the same training conditions.
Revised: In each experiment, we trained two versions of the same model: the one with the Slot Attention module and another one with our SMM, keeping the same training conditions.

Original: We also make use of the technique from Chang et al.
Revised: We also make use of the technique from (Chang et al.

Original: detaching slots from the gradient computational graph at the last iteration of the algorithm in every experiment.
Revised: detaching slots from the gradient computational graph at the last iteration of the algorithm in every experiment.

Original: [Table] With equal dimensionality of input vectors [Table] SMM implies twice the dimensionality of the slots in comparison with SA (SA represents slots with µ , while SMM uses concatenation of µ , Σ diag ).
Revised: With equal dimensionality of input vectors, SMM implies twice the dimensionality of the slots in comparison with SA (SA represents slots with µ , while SMM uses concatenation of µ , Σ diag ).

Original: To ensure a fair comparison of approaches, we added matrix multiplication after the SMM module, reducing the dimensionality by a factor of two.
Revised: To ensure a fair comparison of approaches, we added matrix multiplication after the SMM module, reducing the dimensionality by a factor of two.

Original: 4. S ET P ROPERTY P REDICTION
Revised: 4. S ET P ROPERTY P REDICTION

Original: Neural networks for sets are involved in various applications across many data modalities Carion et al.
Revised: Neural networks for sets are involved in various applications across many data modalities (Carion et al.

Original: Achlioptas et al.
Revised: Achlioptas et al.

Original: ; Simonovsky & Komodakis (2018); Fujita et al.
Revised: ; Simonovsky & Komodakis (2018); Fujita et al.

Original: Set
Revised: Set

Original: Property Prediction is a supervised task that requires the model to predict an unordered set of vectors representing the properties of objects from the input image.
Revised: Property Prediction is a supervised task that requires the model to predict an unordered set of vectors representing the properties of objects from the input image.

Original: Sets of predicted and target vectors are matched using  Hungarian algorithm.
Revised: Sets of predicted and target vectors are matched using a Hungarian algorithm (Kuhn (1955)) and the learning signal is provided by Huber

Revised: Loss (Zhang et al. (2019a)) between matched vectors.

Original: Slot Mixture Module is suitable for operating with sets as it preserves permutation equivariance with respect to mixture components (slots) and initializes them randomly. 
Revised: Slot Mixture Module is suitable for operating with sets as it preserves permutation equivariance regarding mixture components (slots) and initializes them randomly. The scheme of the model is shown in the Fig. GRU N times

Original: Setup We use the CLEVR dataset with rescaled images to a resolution of 128 × 128 as the data source.
Revised: M a t M u l f ( x , μ , σ ) σ π S o f t M a x x M e a n ( d i m = x ) P a i r w

Revised: i s e L 2 d

Revised: i s t μ π μ

Revised: M a t M u l σ Positional Embedding flatten μ σ sample

Revised: Empty Slots

Revised: MLP

Revised: S 1 S

Revised: S 3 S

Revised: Encoder

Revised: [Figure]

Revised: M L P

Revised: HungarianMatching

Revised: Obj

Revised: Obj

Revised: Obj ...

Revised: N

Revised: Concat

Revised: [Figure]

Revised: Setup We use the CLEVR dataset with rescaled images to a resolution of 128 × 128 as the data source.

Original: All modules except SA/SMM are the same as in the Locatello et al.
Revised: All modules except SA/SMM are the same as in the (Locatello et al.

Original: Each model is trained with Adam optimizer  for 1 . 5 × 10 5 iterations with OneCycleLR  learning rate scheduler at a maximum learning rate of 4 × 10 − 4 , we use a batch size of 512 .
Revised: Each model is trained with Adam optimizer (Kingma & Ba (2015)) for 1 . 5 × 10 5 iterations with OneCycleLR (Smith & Topin (2019)) learning rate scheduler at a maximum learning rate of 4 × 10 − 4 , we use a batch size of 512 .

Original: SA/SMM number of iterations is set to 5 during training and to 7 during evaluation.
Revised: SA/SMM number of iterations is set to 5 during training and to 7 during evaluation.

Original: The number of slots is equal to 10 since CLEVR images contain 10 or fewer objects.
Revised: The number of slots is equal to 10 since CLEVR images contain 10 or fewer objects.

Original: Results Quantitative results of the experiments are presented in Table 1.
Revised: Results Quantitative results of the experiments are presented in Table 1.

Original: We compute Average
Revised: We compute Average

Original: Precision (AP) with a certain distance threshold ( ∞ means we do not use distance threshold), a predicted vector of properties and coordinates isconsidered correct if there is an object with the same properties within the threshold distance, the lower the distance threshold the more difficult the task.
Revised: Precision (AP) with a certain distance threshold ( ∞ means we do not use distance threshold). A predicted vector of properties and coordinates is correct if there is an object with the same properties within the threshold distance. The lower the distance threshold, the more difficult the task.

Original: Our experiments demonstratethat it is possible tosignificantly improve  original Slot Attention performance via detaching slots before the last iteration of slots refinement and rescaling coordinates to a wider range of values, but it is still worse than the current state-of-the-art model iDSPN, while our SMM confidently outperforms iDSPN.
Revised: Our experiments demonstrate it is possible to improve significantly the original Slot Attention performance via detaching slots before the last iteration of slots refinement and rescaling coordinates to a wider range of values. With these modifications, the original Slot Attention performance is still worse than the current state-of-the-art model iDSPN, while our SMM confidently outperforms iDSPN.

Original: 4. I MAGE R ECONSTRUCTION U SING T RANSFORMER
Revised: [Table]

Original: For comparison in an unsupervised image-to-image task, we use the SLATE model (Singh et al.
Revised: For comparison in an unsupervised image-to-image task [Table] we use the SLATE model (Singh et al.

Original: (2021a))
Revised: (2021a))

Original: replacing the Slot Attention module with our SMM.
Revised: replacing the Slot Attention module with our SMM.

Original: Unlike the pixel-mixture decoders,
Revised: Unlike the pixel-mixture decoders,

Original: SLATE uses an Image GPT encoder conditioned on slots representations toautoregressively reconstruct  the discrete visual tokens from a discrete VAE (dVAE)  treating pre-computed slots representations from a discrete VAE output as query vectors and latent code-vectors of the image as key/value vectors.
Revised: SLATE uses an Image GPT (Chen et al. decoder conditioned on slot representations to reconstruct in the autoregressive manner the discrete visual tokens from a discrete VAE (dVAE) (Im et al. treating pre-computed slot representations from a dVAE output as query vectors and latent code-vectors of the image as key/value vectors.

Original: SLATE demonstrates an impressive ability to capture complex interactions between the slots and pixels of the synthetic images, but exhibits poor performance for real-world data.
Revised: SLATE demonstrates an impressive ability to capture complex interactions between the slots and pixels of the synthetic images, but exhibits poor performance for real-world data.

Original: dVAE encoder, decoder, and latent discrete tokens are receiving training signals from MSE between the input image and the reconstructed one with dVAE only.
Revised: dVAE encoder, decoder, and latent discrete tokens are receiving training signals from MSE between the input image and the reconstructed one with dVAE only.

Original: SA/SMM modules and Image GPT are trained with a cross-entropy using compressed image into dVAE tokens as a target distribution, these gradients are blocked from the rest of the model (i.e. dVAE), but both parts of the system are trained simultaneously.[Table]
Revised: SA/SMM modules and Image GPT are trained with a cross-entropy using compressed image into dVAE tokens as a target distribution, these gradients are blocked from the rest of the model (i.e. dVAE), but both parts of the system are trained simultaneously.

Original: Setup We consider the following datasets: CLEVR-Mirror Singh et al.
Revised: Setup We consider the following datasets: CLEVR-Mirror (Singh et al.

Original: (2021a),
Revised: (2021a))

Original:  ClevrTex Karazija et al.
Revised: , ClevrTex (Karazija et al.

Original: , ShapeStacks Groth et al.
Revised: , ShapeStacks (Groth et al.

Original: (2018b), and COCO-2017 ? .
Revised: (2018b)), and COCO-2017 (Lin et al.

Original: CLEVR-Mirror is an extension of the standard CLEVR dataset, which requires capturing global relations between local components due to the presence of a mirror, ShapeStacks tests the ability of the model to describe complex local interactions (multiple different objects stacked on each other), ClevrTex examines the capabilities of the model in the textural-rich scenes.
Revised: CLEVRMirror is an extension of the standard CLEVR dataset, which requires capturing global relations between local components due to the presence of a mirror, ShapeStacks tests the ability of the model to describe complex local interactions (multiple different objects stacked on each other), ClevrTex examines the capabilities of the model in the textural-rich scenes.

Original: For ShapeStacks, ClevrTex and COCO we used images rescaled to the resolution of 96 × 96 , CLEVR-Mirror images are rescaled to 64 × 64 .
Revised: For ShapeStacks, ClevrTex and COCO we used images rescaled to the resolution of 96 × 96 , CLEVR-Mirror images are rescaled to 64 × 64 .

Original: Training conditions with hyperparameters corresponding to a certain dataset are taken from Singh et al.
Revised: Training conditions with hyperparameters corresponding to a certain dataset are taken from (Singh et al.

Original: (2021a), except that we use a batch size equal to 64 and 2 .
Revised: (2021a)), except that we use a batch size equal to 64 and 2 .

Original: 5 × 10 5 training iterations for all experiments.
Revised: 5 × 10 5 training iterations for all experiments.

Original: Results Table 2 demonstrates metrics of the image reconstruction performance for the test-parts of different datasets.
Revised: Results Table 2 shows the metrics of the image reconstruction performance for the test-parts of different datasets.

Original: We use mean squared error (MSE) and Frechet Inception Distance (FID)  computed with the Pytorch-Ignite library Fomin et al.
Revised: We use mean squared error (MSE) and Frechet Inception Distance (FID) (Heusel et al. (2017)) computed with the PyTorch-Ignite library (Fomin et al.

Original: (2020) as measure of quality of the generated images.
Revised: as measure of quality of the generated images.

Original: We also evauate cross-entropy between tokens from discrete VAE (as one-hot distributions) and predicted distributions by image GPT that is conditioned by certain slots representations as MSE and FID can be limited by the discrete VAE. 
Revised: We also evaluate cross-entropy between tokens from discrete VAE (as one-hot distributions) and predicted distributions by Image GPT that are conditioned by certain slot representations, as MSE and FID can be limited by the discrete VAE. [Table]

Original: Figure 2 shows advantage of SMM in terms of cross-entropy during training across all the datasets we use.
Revised: Fig. 2 indicates the advantage of SMM in terms of cross-entropy during training across all the datasets we use.

Original: But as can be seen from the Table 2, this advantage does not always translate into much better image quality metrics of generated images due to dVAE limitations.
Revised: But as can be seen from the Table 2, this advantage does not always translate into much better image quality metrics of generated images due to dVAE limitations.

Original: The largest increase in quality was obtained for the ShapeStacks dataset, the model trained with
Revised: The largest increase in quality was obtained for the ShapeStacks dataset. Examples of the images from ShapeStacks, ClevrTex and

Original:  SMM is much more likely to reconstruct the correct order of objects from the original image, as can be seen inone of the examples in the Figure 1.
Revised: CLEVR-Mirror datasets and their reconstructions are presented in the Fig. The model trained with SMM is much more likely to reconstruct the correct order of objects from the original image, as can be seen in examples from the random batch of 64 images from ShapeStacks dataset (see Fig. 4).

Original: Even though we were able to improve the quality of real-world object-centric image generation, the final quality in general is quite poor and attention maps barely reflect human object-centered vision (see Fig. 3).
Revised: Even though we were able to improve the quality of real-world object-centric image generation, the final quality in general is quite poor and attention maps barely reflect human object-centered vision (see Fig. 5).

Original: Scaling and extending such visual models to handle a wide range of complex real-world images is an area of the future research. 
Revised: Scaling and extending such visual models to handle a wide range of complex real-world images is an area of the future research. 4. O BJECT D ISCOVERY

Revised: Another unsupervised object-centric image-to-image task is the Object Discovery.

Revised: In this task, each slot representation is decoded into the 4-channel image using a Spatial Broadcast decoder (Watters et al.

Revised: (2019b)).

Original: Images in the violet borders are from the model with the Slot Attention module, and images in orange borders are generated using slots from the Slot Mixture Module.
Revised: The resulting reconstruction is estimated as a mixture of decoded slots, where the [Figure] Setup To compare with the Slot Attention module, we consider the same training setup from the original work (Locatello et al.

Revised: for the CLEVR10 dataset but with the decreased number of training steps (300k instead of 500k).

Revised: Results Table 3 shows similarity between ground truth segmentation masks of objects (excluded background) and mixture coefficients estimated via the Adjusted Rand Index (ARI) score.

Revised: Fig.

Revised: demonstrates a much faster convergence of the model that uses SMM instead of the Slot Attention, which results in higher ARI score.

Original: Bordeaux for 
Revised: 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. Table 4 shows results of training set prediction model for CLEVR dataset using these vanilla clustering methods. Our experiments demonstrate that GMM clustering is a better choice for object-centric learning, even without trainable layers.

Original: [Figure]
Revised: [Figure]

Revised: [Table]

Original:  R ELATED W ORKS
Revised:  R ELATED W ORKS

Original:  Set prediction Neural network models for sets are applied to different machine learning tasks, such as point cloud generation (Achlioptas et al.
Revised: [Figure] Set prediction Neural network models for sets are applied to different machine learning tasks, such as point cloud generation (Achlioptas et al.

Original: (2017)), object detection (Carion et al.
Revised: (2017)), object detection (Carion et al.

Original: speaker diarization (Fujita et al.
Revised: speaker diarization (Fujita et al.

Original: and molecule generation (Simonovsky & Komodakis (2018)).
Revised: and molecule generation (Simonovsky & Komodakis (2018)).

Original: Albeit the set structure is suitable in many cases, traditional deep learning models are not inherently suitable for representing sets.
Revised: Albeit the set structure is suitable in many cases, traditional deep learning models are not inherently suitable for representing sets.

Original: There are some approaches that are built to reflect the unordered nature of  sets, e.g., the Deep Set Prediction Network (DSPN) (Zhang et al.
Revised: There are some approaches that are built to reflect the unordered nature of [Table] sets, e.g., the Deep Set Prediction Network (DSPN) (Zhang et al.

Original: (2019a)) reflects permutation symmetry by running an inner gradient descent loop that changes a set to encode more similarly to the input.
Revised: (2019a)) reflects permutation symmetry by running an inner gradient descent loop that changes a set to encode more similarly to the input.

Original: An improved version of DSPN — iDSPN Zhang et al. (2019b) with approximate implicit differentiation provides better optimizations with faster convergence and state-of-the-art performance on the CLEVR dataset Johnson et al.
Revised: An improved version of DSPN — iDSPN (Zhang et al. (2019b)) with approximate implicit differentiation provides better optimizations with faster convergence and state-of-the-art performance on the CLEVR dataset.

Original: Such models as Slot Attention(Locatello et al. and TSPN (Kosiorek et al.
Revised: Such models as Slot Attention and TSPN (Kosiorek et al.

Original: (2020)) use set-equivariant self-attention layers to represent the structure of sets.
Revised: (2020)) use set-equivariant self-attention layers to represent the structure of sets.

Original: Object-centric representation
Revised: Object-centric representation

Original: The discovery of objects in a scene in an unsupervised manner is a crucial aspect of representation learning especially for object-centric tasks.
Revised: The discovery of objects in a scene in an unsupervised manner is a crucial aspect of representation learning, especially for object-centric tasks.

Original: In recent years, many approaches have been presented to solve this problem (Greff et al.
Revised: In recent years, many approaches have been presented to solve this problem (Greff et al.

Original: ; Burgess et al.
Revised: ; Burgess et al.

Original: (2019); Engelcke et al.
Revised: (2019); Engelcke et al.

Original: (2020; 2021) and others).
Revised: (2020; 2021) and others).

Original: Such methods as IODINE (Greff et al.
Revised: Such methods as IODINE (Greff et al.

Original: MONET (Burgess et al.
Revised: MONET (Burgess et al.

Original: (2019)), and GENESIS (Engelcke et al.
Revised: (2019)), and GENESIS (Engelcke et al.

Original: (2020)) are built upon the Variational
Revised: (2020)) are built upon the Variational

Original: Autoencoder (VAE) framework (Kingma & Welling (2014);
Revised: Autoencoder (VAE) framework (Kingma & Welling (2014);

Original: Rezende et al.
Revised: Rezende et al.

Original: MONETBurgess et al. (2019)) uses the attention network that generates masks and conditions VAE on these masks.
Revised: MONET uses the attention network that generates masks and conditions VAE on these masks.

Original: IODINE(Greff et al. (2019)) models an image as a spatial Gaussian mixture model to jointly infer the object representation and segmentation.Compared to MONET Burgess et al. and IODINE (Greff et al.
Revised: IODINE models an image as a spatial Gaussian mixture model to jointly infer the object representation and segmentation.

Original: , GENESIS(Engelcke et al. explicitly models dependencies between scene components that allow the sampling of novel scenes.MONET Burgess et al.
Revised: Compared to MONET and IODINE, GENESIS explicitly models dependencies between scene components that allow the sampling of novel scenes.

Original: , IODINE (Greff et al. (2019)), and GENESIS(Engelcke et al. (2020)) use multiple steps to encode and decode an image, while Slot Attention(Locatello et al. (and its sequential extension for video Kipf et al.
Revised: MONET, IODINE, and GENESIS use multiple steps to encode and decode an image, while Slot Attention (and its sequential extension for video (Kipf et al.

Original: (2021)) uses one step but performs an iterative procedure inside this step.
Revised: (2021))) uses one step but performs an iterative procedure inside this step.

Original: The useful property of Slot Attention(Locatello et al. (2020)) is that it produces the set of output vectors (slots) with permutation symmetry.
Revised: The useful property of Slot Attention is that it produces the set of output vectors (slots) with permutation symmetry.

Original: Slots group input information and could be used in unsupervised tasks (object discovery) and supervised tasks (set prediction).
Revised: Slots group input information and could be used in unsupervised tasks (object discovery) and supervised tasks (set prediction).

Original: GENESIS-v2 (Engelcke et al.
Revised: GENESIS-v2 (Engelcke et al.

Original: (2021)), a development of theGENESIS model (Engelcke et al.
Revised: (2021)), a development of the

Original: (2020)), uses attention masks similarly to Locatello et al.
Revised: GENESIS model, uses attention masks similarly to (Locatello et al.

Original:  C ONCLUSION AND D ISCUSSION
Revised:  C ONCLUSION AND D ISCUSSION

Original: In this paper, we propose a new slot-based approach for object-centric representations – the Slot
Revised: In this paper, we propose a new slot-based approach for object-centric representations — the Slot

Original: Mixture Module.
Revised: Mixture Module.

Original: Our module is a generalization of the k-Means clustering algorithm to the Gaussian
Revised: Our module is a generalization of the K-Means clustering algorithm to the Gaussian

Original: Mixture Model Bauckhage (2015a).
Revised: Mixture Model (Bauckhage (2015a)).

Original: Unlike other approaches we represent slots not only as centers of clusters but we also use information about the distance between clusters and assigned vectors, which leads to more expressive slots representations.
Revised: Unlike other approaches, we represent slots not only as centers of clusters, but we also use information about the distance between clusters and assigned vectors, which leads to more expressive slot representations.

Original: We have demonstrated on the CLEVR dataset Johnson et al.
Revised: We have demonstrated on the CLEVR dataset (Johnson et al.

Original: (2017) that using this module achieves the best results in the set property prediction task, even compared to highly specialized models.
Revised: that using this module achieves the best results in the set property prediction task, even compared to highly specialized models.

Original: Also, the use of Slot Mixture Module shows considerable results in the image reconstruction task.
Revised: Also, the use of Slot Mixture Module shows considerable results in the image reconstruction task.

Original: On synthetic datasets, such as CLEVR-Mirror Singh et al.
Revised: On synthetic datasets, such as CLEVR-Mirror (Singh et al.

Original: (2021a), ShapeStacks Groth et al.
Revised: (2021a)) , ShapeStacks (Groth et al.

Original: (2018b),
Revised: (2018b))

Original:  ClevrTex Karazija et al.
Revised: , ClevrTex (Karazija et al.

Original: (2021), we achieved improved reconstruction performance.
Revised: , we achieved improved reconstruction performance.

Original: However, it is worth noting that modern object-centric models still do not perform well enough on real-life images, such as in COCO-17 ? .
Revised: However, it is worth noting that modern objectcentric models still do not perform well enough on real-life images, such as in COCO-17 (Lin et al.

Original: The generalization of such methods to real data is an important scientific problem.
Revised: The generalization of such methods to real data is an important scientific problem.

Revised: We also show that Slot Mixture Module outperforms original Slot Attention on the Object Discovery task on the CLEVR10 dataset.

Original: Also, in models using the slot representation, the number of slots is a hyperparameter, the same as the number of clusters for the k-Means clustering algorithm.
Revised: Also, in models using the slot representation, the number of slots is a hyperparameter, the same as the number of clusters for the K-Means clustering algorithm.

Original: In object-centric methods, this parameter is usually chosen equal to the maximum number of objects in the image plus one (assuming one slot is reserved for the background).
Revised: In object-centric methods, this parameter is usually chosen equal to the maximum number of objects in the image plus one (assuming one slot is reserved for the background).

Original: This value varies from image to image, making it difficult to transfer from one dataset to another.
Revised: This value varies from image to image, making it difficult to transfer from one dataset to another.

Original: It should be noted that, in general, the issue of transfer learning for slot-based models has not been sufficiently studied.
Revised: It should be noted that, in general, the issue of transfer learning for slot-based models has not been sufficiently studied.

Original: Another limitation of modern slot-based models is that they work with a small number of slots (about 10), which means a small number of objects in the image.
Revised: Another limitation of modern slot-based models is that they work with a small number of slots (about 10), which means a small number of objects in the image.

Original: Therefore, the issue of scalability of such models to a larger number of objects is considered as a direction for further research.
Revised: Therefore, the issue of scalability of such models to a larger number of objects is considered as a direction for further research.