Spatial Symmetry in Slot Attention
Editors: Sophia Sanborn, Christian Shewmake, Simone Azeglio, Arianna Di Bernardo, Nina Miolane
Abstract
Automatically discovering composable abstractions from raw perceptual data is a longstanding challenge in machine learning.
Slot-based neural networks have recently shown promise at discovering and representing objects in visual scenes in a self-supervised fashion.
While they make use of permutation symmetry of objects to drive learning of abstractions, they largely ignore other spatial symmetries present in the visual world.
In this work, we introduce a simple, yet eﬀective, method for incorporating spatial symmetries in attentional slot-based methods.
We incorporate equivariance to translation and scale into the attention and generation mechanism of Slot Attention solely via translating and scaling positional encodings.
Both changes result in little computational overhead, are easy to implement, and can result in large gains in data eﬃciency and scene decomposition performance.
1. Introduction
Slot-based neural networks learn to represent inputs using a discrete number of latent vectors, often referred to as “slots”.
These are a promising class of architectures for learning object representations (Greﬀ et al., 2020).
In Slot Attention (Locatello et al., 2020), slots learn to describe the individual objects in an image through an iterative clustering procedure that leverages the permutation equivariance of objects.
However, other inductive biases, such as equivariance to the location and scale of objects is absent, and thus must be learned in a potentially sample- and parameter-ineﬃcient manner from input data alone.
This is diﬀerent from humans, who are believed to attach reference frames to objects to facilitate translation symmetric reasoning about objects and their parts (Hinton, 1981; Hawkins et al., 2019).
Spatial symmetries were successfully incorporated as an inductive bias to improve sample eﬃciency, generalization and consistency of predictions of neural networks (Thomas et al., 2018;
Wang et al., 2020; Han et al., 2022).
These advances have, however, had only a limited impact in object discovery.
Prior object discovery methods often use monolithic encoders (Eslami et al., 2016; Kosiorek et al., 2018) to process images and populate latent slots.
Limited equivariance to translation is present in the case of convolutional encoders with output anchors (Crawford and Pineau, 2019; Lin et al., 2020).
Some works employ spatial symmetries in the decoder (Eslami et al., 2016; Lin et al., 2020) using the Spatial
Transformer (Jaderberg et al., 2015), which is equivariant to 2 D − aﬃne transformations.
Yet another popular choice, the Spatial Broadcast Decoder (Watters et al., 2019), breaks symmetry by appending absolute positions to pixels.
In this work, we explore the symmetry of object translation and object scale in Slot
Attention (Locatello et al., 2020) applied to object discovery.
We equip each slot with an explicit representation of position and a scale, and ensure that the same model weights can be used to detect and reconstruct objects at diﬀerent positions and scales.
Equivariance is achieved by encoding pixel positions relative to each slot both in Slot Attention and in the Spatial Broadcast Decoder.
Although our model is not end-to-end equivariant (Cohen and Welling, 2016, 2017), as we use a standard convolutional encoder to allow for some ﬂexibility in encoding absolute positions and scales of objects (Park et al., 2022), we ﬁnd that it has better sample eﬃciency and generalization properties.
Additionally, we ﬁnd that the equivariant model is more likely to converge to favorable solutions, instead of collapsing to failure modes, such as always predicting Voronoi tessellated segmentation masks.
2. Equivariant Slot Attention
The key observation is that many slot-based models (Locatello et al., 2020; Singh et al., 2021; Sajjadi et al., 2022) and other scene representation approaches (Mildenhall et al., 2020; Sajjadi et al., 2021) append absolute (2 D or 3 D ) positions to latent representations in order to encode and reconstruct images.
These models are sensitive to positions and have to re-learn spatial symmetries from data.
By giving slots explicit positions and scales, we can make position encodings relative to slots, making the model symmetric.
Speciﬁcally, we propose translation and scale equivariant Slot Attention and Spatial Broadcast Decoder, but the same technique could be used with other models and symmetries.
Slot Attention (SA) (Locatello et al., 2020) computes cross attention between input tokens ( inputs ∈ R N × D inputs ) and latent slots ( slots ∈ R K × D slots ).
The input tokens have an absolute coordinate grid, abs grid ∈ R N × 2 , attached to them.
This makes the cross attention sensitive to positions.
Keys and values for this are computed as follows using learned linear projections / MLPs f, k, g, v 1 :
[Equation]
In contrast, in equivariant Slot Attention (Algorithm 1, see Appendix), we equip the K slots with randomly sampled positions, S p ∈ R K × 2 , and scales, S s ∈ R K × 2 and use these to create a separate relative coordinate grid for each slot:
[Equation]
[Figure]
( y , sy )
[Figure]
[Figure]
By attaching this relative grid to the input tokens we eﬀectively center and scale the input tokens into each slot’s own coordinate frame.
This achieves the desired spatial symmetry.
In more detail, we replace (1) with the following:
[Equation]
After the cross attention step, slots are updated as in standard SA.
S p and S s are replaced by the center of mass and spread of the attention mask ( attn ∈ R N × K ), respectively.
[Equation]
Similarly, we make the Spatial Broadcast Decoder (Watters et al., 2019) translation and scale equivariant: we compute the ﬁnal slot positions and scales using the attention map of the last iteration of SA, create relative coordinate grids as in (2), and then add them to the broadcasted slots after applying a learned linear transformation.
This ensures that an object can be decoded using the same weights at arbitrary sizes and scales (Figure 1).
3. Experiments
We evaluate equivariant Slot Attention across four synthetic datasets: Tetrominoes (Greﬀ et al., 2019), CLEVRTex (Karazija et al., 2021), ObjectsRoom (Eslami et al., 2018)
(in the
Appendix), and CLEVR (Johnson et al., 2017) (in the Appendix).
These datasets cover simple backgrounds with simple objects (Tetrominoes, CLEVR, ObjectsRoom) as well as fully-textured backgrounds/objects (CLEVRTex).
We test (1) whether equivariant Slot
Attention generalizes out of distribution if the data is fully symmetric to translations, and (2) whether incorporating spatial symmetries leads to overall better scene decomposition on standard multi-object benchmark tasks.
Generalization and sample eﬃciency in Tetrominoes: A proof of concept The
Tetris-like objects in the Tetrominoes dataset have the same appearance regardless of their position (no occlusion, lighting or perspective changes); hence, Slot Attention should beneﬁt from the inductive bias of translation equivariance. It achieves above 90% FG-ARI using Method CLEVRTex
CLEVRTex CAMO
CLEVRTex OOD [Table] [Table] only one eighth of the dataset size required by the baseline SA model (256 vs. 4096).
In
Figure 2 (left), we perform a generalization experiment wherein we ﬁlter the training set for images with objects only appearing on the left side.
The validation set is unchanged.
We ﬁnd that non-equivariant Slot Attention is less likely to detect objects on the right side of the image (FG-ARI 80 .
6 ± 6 . 8% compared to 94 .
8 ± 1 .
5% for T-SA), likely because it does not have any in-built inductive biases to promote generalization to unseen spatial conﬁgurations in the input.
Comparison to state of the art on CLEVRTex
CLEVRTex is a challenging dataset with textured foreground objects and backgrounds.
Previously, it was understood that
Slot Attention cannot handle textures, as the FG-ARI score of 62.4% for the original Slot
Attention (Table 1) is close to a naive Voronoi tessellation baseline (around 52% FG-ARI).
Our main ﬁnding is that the results of Slot Attention (SA) can be signiﬁcantly improved by adding translation equivariance (T-SA).
A further beneﬁt can be observed by adding both translation and scale equivariance (TS-SA).
This version of Slot Attention is trained using a simple 4-layer CNN backbone as in Locatello et al.
We further ﬁnd that replacing the simple CNN encoder of Slot Attention with a ResNet-34 (He et al., 2016) backbone signiﬁcantly improves scene decomposition performance (Table 1, SA (ours)).
Adding translation equivariance further improves Slot Attention’s ability to segment objects.
Here, adding scale equivariance does not lead to signiﬁcant further improvement, it does however enable explicit control of slot scales when decoding the learned slot representations.
ResNet T-SA and TS-SA outperforms all baselines reported in Karazija et al.
(2021) without pre-training.
Sauvalle and de La Fortelle (2022) reported around 95% FG-ARI with a pretrained SegFormer backbone (Xie et al., 2021) and a background model, which could be further combined with our approach.
4. Conclusion
We have introduced translation- and scale-equivariant Slot Attention.
Our method enables incorporation of spatial symmetries with little computational overhead via simple changes to the positional encoding used both in the attention mechanism and the decoder of Slot Attention.
We are excited about the potential of incorporating additional symmetries through similar mechanisms to a broader class of slot-based neural architectures
Appendix A. Limitations
• Our method struggles with small spatial input grids.
We guess that this is because a sparsely sampled position grid results in a poor learning signal and unstable gradients.
• Computing keys and values per slot scales linearly with the number of slots.
For large scale applications with 100s of slots this could cause memory issues.
• The real world is 3D and we only model 2D translation and scale symmetries in this work.
Thus the input domain does not respect spatial symmetries 100%.
We intentionally allow for global positions to leak through the CNN encoder so that the model may leverage these when necessary.
We interpret this as a feature not a limitation of our model.
A better solution might be to explicitly model 3D symmetries but that is beyond the scope of 2D slot attention.
• Our experiments are entirely on synthetic data.
This is a limitation of this paper and not of the model itself and an exciting direction for future work.
Appendix B. Related Work
Prior object discovery approaches use an encoder-decoder framework with few exceptions (Greﬀ et al., 2019;
Kipf et al., 2020; Huang et al., 2020).
The encoder processes images and populates latent slots, and the decoder uses the information about object present in latent slots to reconstruct the input.
Some works employ spatial symmetries in the decoder (Eslami et al., 2016; Kosiorek et al., 2018; Crawford and Pineau, 2019; Lin et al., 2020; Jiang and Ahn, 2020; Monnier et al., 2021; Smirnov et al., 2021; Sauvalle and de La Fortelle, 2022) using the Spatial Transformer (Jaderberg et al., 2015), which is equivariant to 2 D − aﬃne transformations.
In contrast, most prior works use monolithic encoders (Greﬀ et al., 2016;
Eslami et al., 2016; Greﬀ et al., 2017; Kosiorek et al., 2018; Burgess et al., 2019;
Engelcke et al., 2020; Jiang and Ahn, 2020; Engelcke et al., 2021; Monnier et al., 2021; Smirnov et al., 2021; Emami et al., 2021) with only limited equivariance to translation in the case of convolutional encoders with output anchors (Crawford and Pineau, 2019; Lin et al., 2020).
Alternatively, models iteratively process the encoded inputs to reﬁne object detections (Locatello et al., 2020; Huang et al., 2020).
Slot Attention (Locatello et al., 2020) appends absolute coordinates to feature maps in the encoder, thus making object detections sensitive to locations.
Appendix C. Theoretical analysis of translation equivariance
We show that Slot Attention is equivariant to joint translation of the input features and the initial slot positions.
We do not show equivariance to translations of individual objects due to occlusions, but a similar line of reasoning would otherwise work.
We formalize Slot Attention as a function with four inputs and two outputs:
[Equation]
Here, inputs : Z 2 → R D inputs map feature coordinates to input vectors, abs grid : Z 2 →
R 2 is a linear function that maps feature coordinates to real-valued coordinate encodings,
S ∈ R K × D slots are the initial latent slots and S p ∈ R K × 2 are the initial slot positions.
Correspondingly, S (cid:48) and S (cid:48) p are the ﬁnal latent slots and slot positions after T rounds of slot attention.
Next, we show that the ﬁnal slot positions are equivariant to a joint translation of the input and initial slot positions, and that the ﬁnal latent slots are invariant to said transformation:
[Equation]
Here, t belongs to the group of translations over Z 2 , L t ( x ) = x + t , x ∈ Z 2 , is a group action that translates an integer-valued coordinate and R t ( y ) =
y + abs grid( t ), y ∈ R 2 , is a translation in the space of real-valued coordinates.
We use the inverse of R t because slot positions are used to re-center feature coordinate grids by inverting the translations applied to the input features.
We formalize the keys and values for the k th slot as mapping an integer-valued feature coordinate x to a vector in R D .
[Equation]
Next, we show that a joint translation of the inputs and the slot positions is equivalent to the translation of the key coordinates.
The same can be shown for values.
[Equation]
We use the assumption that abs grid is linear.
Functions f, g, k, v are applied per-position and we do not require linearity.
The cross attention between keys and slots (Algorithm 1, line 12) is computed for each pixel coordinate separately.
Hence, it is trivially translation equivariant given the result for keys obtained above.
The slot updates (line 13) are an attention-weighted sum over values.
The sum is invariant to joint translations of both the attention mask and the values.
The renormalization of the attention mask for each slot on line 14 is also invariant to translations of the attention mask.
We assume the sum to be ﬁnite.
Formally, we have the following properties hold for the attention mask computed on lines 12 and 14, and the updates computed on line 13:
[Equation]
Next, we compute the updated slot positions based on the attention mask (line 17):
[Equation] [Equation]
We show that translation equivariance holds for the updated positions.
[Equation]
Finally, since updates are invariant to the input transformation, lines 21 and 22 are also invariant.
Hence, the updated latent slots S (cid:48) are invariant to the input transformation.
The equivariant of S kp and invariance of S (cid:48) holds over multiple iterations of Algorithm 1.
Appendix D. Pseudo-Code
Algorithm 1 gives the self-explanatory pseudo-code of our method.
In line 7 we scale the relative grid using δ after adjusting it using S s as otherwise for small objects, rel grid will have numerically large values which we found made model training diﬃcult.
Appendix E. Additional results E.1. Out-of-distribution generalization on ObjectsRoom
Slot Attention uses the same mechanism to segment the foreground and the background.
We test the interaction between translation and scale equivariance and multi-segment backgrounds in ObjectsRoom (Eslami et al., 2018), Figure 3.
We ﬁnd that both T-SA and
TS-SA are more likely to learn the correct segments of the background (two walls, ground and ceiling), leading to between 10 and 15% absolute improvement in ARI.
We ﬁnd that equivariant Slot Attention is robust to all out-of-distribution test sets, whereas the baseline deteriorates in the Empty Room OOD set due to increased oversegmentation of the background, see Table 2.
E.2. Robustness to data augmentation on CLEVR
Given enough parameters and a training dataset that covers all spatial conﬁgurations a powerful deep learning model could potentially learn to be equivariant to spatial transformations at the object level.
Data augmentation is typically used to augment the training set do
slots = LayerNorm(slots)  # Computes relative grids per slot, and associated key, value embeddings.
rel grid =
(abs grid − S p ) / S s × δ keys =
f ( k (inputs) + g (rel grid)) values =
f ( v (inputs) + g (rel grid))  # Inverted dot production attention.
attn = softmax( 1 √ K keys ∗ q (slots) T , axis = “slots”) updates = WeightedMean(weights = attn , values = values)  # Updates S p , S s and slots.
S p =
WeightedMean(weights = attn , values = abs grid)
S s = (cid:112) WeightedMean(weights
= attn + (cid:15), values =
(abs grid − S p ) 2 ) if t < T + 1 then slots =
GRU(state = slots prev , inputs = updates) slots += MLP(LayerNorm(slots)) end 22 end
Method
ARI (11 slots)
Validation
Six Objects
Empty Room
Identical Colors
SA 71.2 ± 16.
71.3 ± 16.
65.9 ± 14.
69.6 ± 16.
T-SA
[Table]
[Figure]
[Table]
with all possible spatial variations at the image level hoping to achieve some of this eﬀect.
We analyze, on the CLEVR-10 dataset, the eﬀect of data augmentation during training and evaluation.
Baseline SA achieves close to 99% FG-ARI on this dataset.
However, as seen in Table 3, under the column “CLEVR Augm.
Eval.”
, this model is not robust to perturbations in the translation and scale of input images at test time.
We hypothesize that SA tends to overﬁt to the objects appearing in the central portion of the images as well as to the constrant background.
In column 4, “CLEVR Augm.
Train. & Eval.”
, we ﬁnd that model performance is not restored simply by using data augmentation during training.
Details of the augmentation used are discussed in Appendix F.
On the other hand, our TS-SA model is relatively robust to test time augmentation and additionally beneﬁts from global image level data augmentation during training suggesting that these two changes are somewhat orthogonal.
We conclude that the inductive biases facilitated by equivariance cannot be supplanted by data augmentation alone.
Appendix F. Datasets and data preprocessing
We use the standard pre-processing pipeline in CLEVR (e.g. Locatello et al.
In the data augmentation experiment, we sample random square crops that cover at least 25% of the original unprocessed image; these crops are then resized to 64 × 64, as in the original data processing.
For CLEVRTex, we use the data processing from Karazija et al.
We do not perform any data preprocessing for Tetrominoes and objects room.
RGB values in all datasets are scaled to [0 , 1].
In the Tetrominoes dataset, in Figure 2 (right), we sample random square crops with a minimum area coverage tuned to be optimal for the baseline.
The same setting is them used for our method.
Appendix G. Model architectures and hyper-parameters
We use the same encoder (Table 6) and decoder (Table 7) on objects room, CLEVR and
CLEVRTex.
In CLEVRTex, we also use a ResNet-34 encoder.
The ResNet is not pre-trained and the downsampling before its ﬁrst stage is removed (the ﬁrst convolutional layer’s stride is set to 1 and the max-pooling layer is removed).
In Tetrominoes, we do not downsample in the encoder (Table 5) and we use a per-pixel decoder (Table 8), similarly to Kabra et al.
Diﬀerent from the original Slot Attention, we use learnable initial latent slots (Kipf et al., 2022; Elsayed et al., 2022), which usually lead to better results.
Initial slot positions are randomly sampled from U ( − 1 , 1) in both the x and y axes, and initial slot scales are sample from N (0 . 1 , 0 . 1) and clipped between 0.01 and 5.
We did not perform hyper-parameters search for the initialization.
Hyper-parameters are reported in Table 4.
All Slot Attention models are trained for 500k steps without early stopping or model selection.
In the Tetrominoes experiment with limited dataset sizes, we found it suﬃcient to train for only 20k steps.
[Table]
[Table] [Table]
[Table] [Table]