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 entity 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 a 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 slot 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 objectcentric scenarios, achieving the 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. (2019a);
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: 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: 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-centric 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 slot 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)),
Revised: (2019b)),

Original: and also improves the generalization ability of the model for image reconstruction task.
Revised: and also improves the generalization ability of the model for image reconstruction task.

Original: The paper is structured as follows.
Revised: The paper is structured as follows.

Original: In Section 2, we provide background information about Slot
Revised: In Section 2, we provide background information about Slot

Original: Attention module and Mixture Models.
Revised: Attention module and Mixture Models.

Original: We also describe the process of their training.
Revised: We also describe the process of their training.

Original: In Section 3, we introduce a Slot Mixture Module — the modification of a Slot Attention Module which provides more expressive slot 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: and outperforms even highly specialized models
Revised: and outperforms even highly specialized models

Original: Zhang et al. (2019b).
Revised: 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
Revised: , ShapeStacks

Original: Groth et al.
Revised: 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
Revised: (2021)) and one with real life images

Original: COCO- (Lin et al.
Revised: COCO- (Lin et al.

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

Original: In
Revised: In

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

Original: Object Discovery task on the CLEVR10 dataset.
Revised: Object Discovery task on the CLEVR10 dataset.

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

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

Original: Mixture Model clustering is a better choice for object-centric learning.
Revised: Mixture Model clustering is a better choice for object-centric learning.

Original: We give a short overview of related works in Section 5.
Revised: We give a short overview of related works in Section 5.

Original: In Section 6, we discuss the obtained results, advantages, and limitations of our work.
Revised: 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).
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 slot representations as a Gaussian Mixture improve 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).

Original: The Slot Mixture Module shows a much faster convergence on the Object Discovery task compare to the Original Slot Attention (Section 4.3).
Revised: The Slot Mixture Module shows a much faster convergence on the Object Discovery task compare to the Original Slot Attention (Section 4.3).

Original:  B ACKGROUND 2. S LOT A TTENTION
Revised:  B ACKGROUND 2. S LOT A TTENTION

Original: Slot Attention (SA) module (Locatello et al.
Revised: Slot Attention (SA) module (Locatello et al.

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

Original: Randomly initialized slots from a
Revised: Randomly initialized slots from a

Original: Gaussian distribution with trainable parameters are used to get q projections of slots.
Revised: Gaussian distribution with trainable parameters are used to get q projections of slots.

Original: Feature map vectors with corresponding projections serve as k and v vectors.
Revised: Feature map vectors with corresponding projections serve as k and v vectors.

Original: Dot-product attention between q and k vectors with the softmax across q dimension implies competition between slots for explaining parts of the input.
Revised: Dot-product attention between q and k vectors with the softmax across q dimension implies competition between slots for explaining parts of the input.

Original: Attention coefficients are used to assign v vectors to slots via a weighted mean.
Revised: Attention coefficients are used to assign v vectors to slots via a weighted mean.

Original: [Equation]
Revised: [Equation]

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

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

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

Original: 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: 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.

Original: These properties make the Slot
Revised: These properties make the Slot

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

Original: 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: 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.

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

Original: 2. M IXTURE M ODELS
Revised: 2. M IXTURE M ODELS

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 a 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 = 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 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.
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.

Original:  S LOT M IXTURE M ODULE
Revised:  S LOT M IXTURE M ODULE

Original: For the purposes of the object-centric learning we propose a modified Gaussian 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: 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.
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: This set of slots is further used in the downstream task.
Revised: This set of slots is further used in the downstream task.

Original: We also use the same additional neural network update step for the mean values before updating the covariance 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 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 (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.
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 the Algorithm 1.
Revised: The full pseudocode is presented in the Algorithm 1.

Original: A function f θ ( x , µ , Σ diag ) :
Revised: A function f θ ( x , µ , Σ diag ) :

Original: R N × D × R K × D
Revised: R N × D × R K × D

Original: → R N × K stands for a log-Gaussian density function with additional k and q projections for input and slot vectors.
Revised: → R N × K stands for a log-Gaussian density function with additional k and q projections for input and slot vectors.

Original: 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.
Revised: 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 are initialized from Gaussian distributions with trainable parameters.
Revised: π is initialized as uniform categorical distribution, µ and Σ diag are initialized from Gaussian distributions with trainable parameters.

Original: Input feature added positional embeddings; µ , Σ diag ∈ R K × D , π ∈ R K — SMM initialization parameters.
Revised: Input feature added positional embeddings; µ , Σ diag ∈ R K × D , π ∈ R K — SMM initialization parameters.

Original: Output slots
Revised: Output slots

Original: R
Revised: R

Original: 1 x = MLP(LayerNorm( x )) 2 for t = 0 ...
Revised: 1 x = MLP(LayerNorm( x )) 2 for t = 0 ...

Original: T do logits =
Revised: T do logits =

Original: f θ ( x , µ , Σ diag ) gammas =
Revised: f θ ( x , µ , Σ diag ) gammas =

Original: SoftMax(logits + log π , dim= 1 ) π
Revised: SoftMax(logits + log π , dim= 1 ) π

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

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

Original: Σ diag =
Revised: Σ diag =

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

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

Original: E XPERIMENTS
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: 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 ).
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 a Hungarian algorithm (Kuhn (1955)) and the learning signal is provided by Huber
Revised: Sets of predicted and target vectors are matched using a Hungarian algorithm (Kuhn (1955)) and the learning signal is provided by Huber

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

Original: Slot Mixture Module is suitable for operating with sets as it preserves permutation equivariance regarding 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.

Original: The scheme of the model is shown in the Fig.
Revised: The scheme of the model is shown in the Fig.

Original: GRU
Revised: GRU

Original: N times
Revised: N times

Original: M a t M u l f ( x , μ , σ ) σ π
Revised: M a t M u l f ( x , μ , σ ) σ π

Original: S o f t M a x x
Revised: S o f t M a x x

Original: M e a n ( d i m = x )
Revised: M e a n ( d i m = x )

Original: P a i r w
Revised: P a i r w

Original: i s e L 2 d
Revised: i s e L 2 d

Original: i s t μ π μ
Revised: i s t μ π μ

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

Original: Empty Slots
Revised: Empty Slots

Original: MLP
Revised: MLP

Original: S 1 S
Revised: S 1 S

Original: S 3 S
Revised: S 3 S

Original: Encoder
Revised: Encoder

Original: [Figure]
Revised: [Figure]

Original: M L P
Revised: M L P

Original: HungarianMatching
Revised: HungarianMatching

Original: Obj Obj
Revised: Obj Obj

Original: Obj ...
Revised: Obj ...

Original: N
Revised: N

Original: Concat
Revised: Concat

Original: [Figure]
Revised: [Figure]

Original: Setup We use the CLEVR dataset with rescaled images to a resolution of 128 × 128 as the data source.
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 (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 .
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).
Revised: Precision (AP) with a certain distance threshold ( ∞ means we do not use distance threshold).

Original: A predicted vector of properties and coordinates is correct if there is an object with the same properties within the threshold distance.
Revised: A predicted vector of properties and coordinates is correct if there is an object with the same properties within the threshold distance.

Original: The lower the distance threshold, the more difficult the task.
Revised: The lower the distance threshold, the more difficult the task.

Original: 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.
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.

Original: 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.
Revised: 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: [Table]
Revised: [Table]

Original: For comparison in an unsupervised image-to-image task
Revised: For comparison in an unsupervised image-to-image task

Original: [Table]
Revised: [Table]

Original: we use the SLATE model (Singh et al.
Revised: 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 (Chen et al.
Revised: SLATE uses an Image GPT (Chen et al.

Original: decoder conditioned on slot representations to reconstruct in the autoregressive manner the discrete visual tokens from a discrete VAE (dVAE)
Revised: decoder conditioned on slot representations to reconstruct in the autoregressive manner the discrete visual tokens from a discrete VAE (dVAE)

Original: (Im et al.
Revised: (Im et al.

Original: treating pre-computed slot representations from a dVAE output as query vectors and latent code-vectors of the image as key/value vectors.
Revised: 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.
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 (Lin et al.
Revised: (2018b)), and COCO-2017 (Lin et al.

Original: 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.
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
Revised: For ShapeStacks, ClevrTex and

Original: COCO we used images rescaled to the resolution of 96 × 96 , CLEVR-Mirror images are rescaled to 64 × 64 .
Revised: 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)),
Revised: (2021a)),

Original: except that we use a batch size equal to 64 and 2 .
Revised: 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 shows the 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)
Revised: We use mean squared error (MSE) and Frechet Inception Distance (FID)

Original: (Heusel et al.
Revised: (Heusel et al.

Original: (2017)) computed with the PyTorch-Ignite library (Fomin et al.
Revised: (2017)) computed with the PyTorch-Ignite library (Fomin et al.

Original: as measure of quality of the generated images.
Revised: as measure of quality of the generated images.

Original: 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.
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.

Original: [Table]
Revised: [Table]

Original: Fig.
Revised: Fig.

Original: 2 indicates the advantage of SMM in terms of cross-entropy during training across all the datasets we use.
Revised: 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.
Revised: The largest increase in quality was obtained for the ShapeStacks dataset.

Original: Examples of the images from ShapeStacks, ClevrTex and
Revised: Examples of the images from ShapeStacks, ClevrTex and

Original: CLEVR-Mirror datasets and their reconstructions are presented in the Fig.
Revised: CLEVR-Mirror datasets and their reconstructions are presented in the Fig.

Original: The model trained with
Revised: The model trained with

Original: 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).
Revised: 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. 5).
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.

Original: 4. O BJECT D ISCOVERY
Revised: 4. O BJECT D ISCOVERY

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

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

Original: (2019b)).
Revised: (2019b)).

Original: The resulting reconstruction is estimated as a mixture of decoded slots, where the
Revised: The resulting reconstruction is estimated as a mixture of decoded slots, where the

Original: [Figure]
Revised: [Figure]

Original: Setup To compare with the Slot Attention module, we consider the same training setup from the original work (Locatello et al.
Revised: Setup To compare with the Slot Attention module, we consider the same training setup from the original work (Locatello et al.

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

Original: 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: 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.

Original: Fig.
Revised: Fig.

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

Original: 4. C OMPARING VANILLA CLUSTERING
Revised: 4. C OMPARING VANILLA CLUSTERING

Original: 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.
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.

Original: Table 4 shows results of training set prediction model for CLEVR dataset using these vanilla clustering methods.
Revised: Table 4 shows results of training set prediction model for CLEVR dataset using these vanilla clustering methods.

Original: Our experiments demonstrate that GMM clustering is a better choice for object-centric learning, even without trainable layers.
Revised: Our experiments demonstrate that GMM clustering is a better choice for object-centric learning, even without trainable layers.

Original: [Figure]
Revised: [Figure]

Original: [Table]
Revised: [Table]

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

Original: [Figure]
Revised: [Figure]

Original: Set prediction Neural network models for sets are applied to different machine learning tasks, such as point cloud generation (Achlioptas et al.
Revised: 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
Revised: There are some approaches that are built to reflect the unordered nature of

Original: [Table]
Revised: [Table]

Original: sets, e.g., the Deep Set Prediction Network (DSPN) (Zhang et al.
Revised: 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.
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 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: MONET 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 models an image as a spatial Gaussian mixture model to jointly infer the object representation and segmentation.
Revised: IODINE models an image as a spatial Gaussian mixture model to jointly infer the object representation and segmentation.

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

Original: 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.
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 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 the
Revised: (2021)), a development of the

Original: GENESIS model, 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 slot 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: 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))
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: , we achieved improved reconstruction performance.
Revised: , we achieved improved reconstruction performance.

Original: 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.
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.

Original: We also show that Slot Mixture Module outperforms original Slot Attention on the Object Discovery task on the CLEVR10 dataset.
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.

Revised: A A RCHITECTURE D ETAILS

Revised: Tables 5, 6, 7 describe hyperparameters for our experiments.

Revised: In the case of using SMM instead of SA we use an additional dimensionality reduction for slots via trainable matrix multiplication.

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: Transformer

Revised: Hidden Dim.

Revised: SA/SMM

Revised: Num.

Revised: slots

Revised: SA/SMM

Revised: Iterations

Revised: SA/SMM

Revised: Slot dim.