I MPROVING O BJECT - CENTRIC L EARNING WITH Q UERY O PTIMIZATION
A BSTRACT
The ability to decompose complex natural scenes into meaningful object-centric abstractions lies at the core of human perception and reasoning.
In the recent culmination of unsupervised object-centric learning, the Slot-Attention module has played an important role with its simple yet effective design and fostered many powerful variants.
These methods, however, have been exceedingly difficult to train without supervision and are ambiguous in the notion of object, especially for complex natural scenes.
In this paper, we propose to address these issues by investigating the potential of learnable queries as initializations for Slot-Attention learning, uniting it with efforts from existing attempts on improving Slot-Attention learning with bi-level optimization.
With simple code adjustments on Slot-Attention, our model, Bi-level Optimized Query Slot Attention, achieves state-of-the-art results on 3 challenging synthetic and 7 complex real-world datasets in unsupervised image segmentation and reconstruction, outperforming previous baselines by a large margin.
We provide thorough ablative studies to validate the necessity and effectiveness of our design.
Additionally, our model exhibits great potential for concept binding and zero-shot learning.
Our work is made publicly available at https://bo-qsa.github.io .
I NTRODUCTION
Objects, and their interactions, are the foundations of human cognition (Spelke & Kinzler, 2007).
The endowment on making abstractions from perception and organizing them systematically empowers humans the ability to accomplish and generalize across a broad range of tasks, such as scene modeling (Bear et al., 2020), visual reasoning (Yi et al., 2020) and simulating interactions (Bear et al., 2020).
The key to such success lies in the emergence of symbol-like mental representations of object concepts (Whitehead, 1928).
However, important as it is, disentangling object-centric concepts from visual stimuli is an exceedingly difficult task to accomplish with limited supervision (Greff et al., 2020) and requires proper inductive biases (Schölkopf et al., 2021).
Motivated by the development of symbolic thought in human cognition, slot-based representations, instance (Greff et al., 2017; 2019; Locatello et al., 2020), sequential (Gregor et al., 2015; Burgess et al., 2019; Engelcke et al., 2021; Goyal et al., 2021), or spatial (Crawford & Pineau, 2019; Lin et al., 2020; Jiang et al., 2019), have been the key inductive bias to recent advances in unsupervised object-centric learning.
Among them, the Slot-Attention module has received tremendous focus given its simple yet effective design (Locatello et al., 2020).
By leveraging the iterative attention mechanism,
Slot-Attention learns to compete between slots for explaining parts of the input, exhibiting a softclustering effect on visual signals.
It is later proven to be more memory and training efficient as a plug-and-play module for unsupervised object-centric learning (Locatello et al., 2020) and fostered powerful variants in understanding images (Singh et al., 2021; Xu et al., 2022), 3D scenes (Yu et al., 2022; Sajjadi et al., 2022a) and videos (Kipf et al., 2022; Elsayed et al., 2022; Singh et al., 2022).
However, as revealed by recent studies, the Slot-Attention module comes with innate discrepancies for object-centric representation learning.
First, with slots randomly initialized each time, the objectcentric representations obtained by these models do not necessarily bind to object concepts (Kipf et al., 2022).
Intuitively, such randomness leads to undesired scenarios where slots with similar initializations compete for objects on different images.
Such randomness challenges the iterative refinement procedure as it now needs to project sets of potentially similar representations to independent constituents of the input.
As discovered by Chang et al.
(2022), differentiating through such recurrences contributes to various training instabilities with growing spectral norm of Slot-Attention weights.
This leads to the second and perhaps least desired property of Slot-Attention; it relies heavily on hyper-parameter tuning, including gradient clipping, learning rate warm-up, etc ., and further hurts the flexibility of Slot-Attention in adapting to broader applications with more complex signals.
To this end, we propose an extension of the Slot-Attention module, Bi-level Optimized Query Slot
Attention (BO-QSA), to tackle the aforementioned problems.
First, we follow the bi-level optimization framework proposed by Chang et al.
for easing the training difficulty in Slot-Attention.
More importantly, instead of sampling from a learnable Gaussian distribution, we propose to directly learn the slot initializations as queries.
With these learnable representations, we eliminate the ambiguous competitions between slots and provide a better chance for them to bind to specific object concepts.
We improve the training of query-initialized Slot-Attention with a straight-through gradient estimator (STE) by connecting our method with first-order approaches (Finn et al., 2017; Nichol &
Schulman, 2018; Geng et al., 2021) in solving bi-level optimization problems.
The experimental results show that the proposed BO-QSA can achieve state-of-the-art results on both synthetic and real-world image datasets with simple code adjustments to the original Slot-Attention module.
With our model significantly outperforming previous methods in both synthetic and real domains, we provide thorough ablative studies demonstrating the effectiveness of our model design.
We later show that our BO-QSA possesses the potential of binding object concepts to slots.
To validate this potential, we design zero-shot transfer learning experiments to show the generalization power of our model on unsupervised object-centric learning.
As the experiments suggest (see Sec. 5), our model could potentially be a principle approach for unsupervised object-centric learning and serve as a general plug-and-play module for a broader range of modalities where variants of Slot-Attention prosper.
We hope these efforts can help foster new insights in the field of object-centric learning.
Contributions In summary, our main contributions are three-fold:
• We propose BO-QSA, a query-initialized Slot-Attention model that unites straight-through gradient updates to learnable queries with methods on improving Slot-Attention with bi-level optimization.
• We show that, with simple code adjustments on Slot-Attention, the proposed BO-QSA achieves state-of-the-art results on several challenging synthetic and real-world image benchmarks, outperforming previous methods by a large margin.
• We show the potential of our BO-QSA being a better approach to concept binding and learning generalizable representations with qualitative results and zero-shot transfer learning experiments.
P RELIMINARIES 2. O BJECT -C ENTRIC R EPRESENTATION L EARNING WITH S LOT -A TTENTION
Slot-Attention (Locatello et al., 2020) takes a set of N input feature vectors x P R N ˆ D input and maps them to a set of K output vectors ( i.e ., slots)
s P R K ˆ D slots .
It leverages an iterative attention mechanism to first map inputs and slots to the same dimension D with linear transformations k p¨q , q p¨q and v p¨q parameterized by ϕ attn .
At each iteration, the slots compete to explain part of the visual input by computing the attention matrix A with softmax function over slots and updating slots with the weighted average of visual values:
[Equation]
The slots are initialized from a learnable Gaussian distribution with mean µ and variance σ .
They are refined iteratively within the Slot-Attention module by passing the updates into a Gated Recurrent
Unit (GRU) (Cho et al., 2014) and MLP parameterized by ϕ update for T iterations: [Equation]
[Equation]
The final prediction ˆ s can be treated as the learned object-centric representation w.r.t.
to input features x .
In the image domain, we take as input a set of images I and encode them with f ϕ enc to obtain features x P R HW ˆ D input .
After obtaining ˆ s through the iterative refinement procedure with h ϕ update , images could be decoded from these object-centric representations with a mixture-based decoder or autoregressive transformer-based decoder.
We refer the readers to Appendix A.1 for details on different decoder designs and their ways of visualizing learned object concepts.
2. I MPROVING S LOT -A TTENTION WITH B I - LEVEL O PTIMIZATION
The problem of bi-level optimization embeds the optimization of an inner objective within the outer objective.
Normally, a bi-level optimization problem can be formulated as:
[Equation]
where we call f p θ, ϕ q the outer objective function and g p θ, ϕ q the inner objective function.
To jointly optimize both objectives w.r.t.
parameters θ and ϕ , a straightforward approach to solving Eq.
(2) is to represent the inner solution of θ as a function of ϕ , i.e ., θ ˚ p ϕ q “ arg min θ 1 g p θ 1 , ϕ q .
Then we can optimize the outer objective with gradient descent by approximating ∇ ϕ f p θ ˚ p ϕ q , ϕ q as a function of ϕ .
When the inner optimization objective could be solved by a fixed point iteration θ “ F ϕ p θ q (Amos & Kolter, 2017; Bai et al., 2019), the bi-level optimization problem could be solved by
[Equation]
For efficiency concerns, recent methods often use the first-order approximation of the infinite
Neumann’s series (Shaban et al., 2019; Geng et al., 2021) for updating ϕ .
Given that Slot-Attention is, in essence, an iterative refinement method that falls into the same framework, Chang et al.
(2022) adapted this technique to improve Slot-Attention training and obtained significant improvement both in model performance and training stability.
We provide more discussions on this in Sec. 3.2 and also other bi-level optimization methods for approximating ∇ ϕ f p θ ˚ p ϕ q , ϕ q in Appendix A.2.
M ETHOD 3. Q UERY S LOT A TTENTION
As mentioned in Sec. 1, the Slot-Attention module adopts a random initialization of slots and conducts iterative refinement to obtain object-centric representations
ˆ s as in Eq.
However, as argued by Kipf et al.
(2022), such random initializations provide no hint on the notion of object and no means for controllably probing concepts from the model.
As shown by Chang et al.
, this random initialization plays a minimal role and could be detached from training.
This indicates that the estimation of ˆ s relies heavily on the task-specific iterative refining of slots over data, leaving a limited possibility for slots to bind to specific concepts and be leveraged as generalizable representations.
To address this issue, we focus on the Query Slot Attention (QSA), which initializes the slots in the
Slot-Attention module with learnable queries s 0 “ ϕ init .
Such a design is motivated by the success of recent query-based networks (Van Den Oord et al., 2017; Jaegle et al., 2021b).
It facilitates an objectcentric model to learn general symbolic-like representations that could be quickly adapted by refining over task-specific requirements, as discussed in Sec.
and Kipf et al.
Meanwhile, in contrast to the use of learnable queries in other encoder-decoder structures ( e.g . discrete VAE (dVAE)), the slot initializations s 0 are not necessarily required to encode image features since they were designed for separating them.
This resembles recent discoveries in query networks (Carion et al., 2020; Yang et al., 2021) where queries could be generalizable probes for input properties.
Despite the good properties and potentials QSA presents, it is shown detrimental to initialize slots independently in
Slot-Attention under unsupervised settings (Locatello et al., 2020).
3. R ETHINKING B I - LEVEL O PTIMIZATION M ETHODS FOR Q UERY S LOT A TTENTION
To improve the learning of QSA, we rewind to the idea of improving the learning of the vanilla Slot-Attention module with bi-level optimization (Chang et al., 2022).
Under this formulation,
Slot-Attention could be treated as solving the following objectives:
[Equation]
where x i and s i denote the input feature from the i -th image and its corresponding slots, and Φ “ t ϕ init , ϕ attn , ϕ update u denotes parameters for assigning input features x to different slots.
Under this setting, the outer objective L is usually a reconstruction objective and the inner objective could be viewed as a soft-clustering objective (Locatello et al., 2020).
Next, the inner objective is solved by iterative refinement, which could be formulated as solving for fixed-points (Chang et al., 2022) of
[Equation]
where F Φ p¨ , ¨q is an fixed-point operation.
As introduced by Chang et al. (2022) in Implicit SlotAttention (I-SA), with Eq.
(3), the instabilities through the iterative updates could be avoided by detaching gradients, treating slots in the final iteration as an approximation of s ˚ i , and computing first-order gradient approximations for updating Φ with s ˚ i .
However, we demonstrate in Tab.that this design is only beneficial for randomly initialized slots and detrimental for query-initialized Slot-Attention architectures since it relies heavily on the good approximation of the solution to the inner objective.
With no randomness in slot initializations or gradient during training, starting from a fixed set of initialization points puts challenges on the learning of Slot-Attention update F Φ as it will be difficult to provide a good approximation of s ˚ i with only a fixed number of iterations (see in Appendix B.2).
This urges the need for information flow to the slot initialization queries.
3. B I - LEVEL O PTIMIZED Q UERY S LOT A TTENTION
[Table]
return
[Table]
We propose BO-QSA to address the learning problem of QSA.
As shown in Algorithm 1, we initialize slots with learnable queries in BO-QSA and perform T steps of
Slot-Attention update to obtain an approximation of s ˚ i .
These near-optimal solutions of the inner objective are passed into one additional Slot-Attention step where gradients to all previous iterations are detached.
In contrary to I-SA, we use a STE (Bengio et al., 2013; Van
Den Oord et al., 2017) to backpropagate gradients and also to slot initialization queries.
Such designs help find good starting points for the inner optimization problem on clustering, alleviating the problem of bi-level optimization with QSA mentioned in Sec. 3.2.
Similar to dVAE, the STE adds bias to the gradient of the initialization queries.
However, since these learnable queries are meant for disentangling image features, they do not have to maintain information about the approximated s ˚ .
Such bias could lead to learned queries which are better pivots for separating different image features, similar to anchors, or filter queries learned for different tasks (Carion et al., 2020; Zhang et al., 2021).
Note that we do not add constraints on the consistency between s 0 and ˆ s ( e.g . ||
sg p ˆ s q ´ s
|| 2 ) as done in dVAE since we find such constraints lead to a mean-representation of datasets that forbids better concept binding (see in Appendix B.3).
As shown in Tab. 7 and Fig.
3, our learned slot initialization queries do fulfill this goal by providing a more separable initialization space and can significantly facilitate model learning.
R ELATED W ORK
Unsupervised Object-Centric Learning Our work falls into the recent line of research on unsupervised object-centric learning on images (Greff et al., 2016;
Eslami et al., 2016; Greff et al., 2017; 2019; Burgess et al., 2019; Crawford & Pineau, 2019; Engelcke et al., 2020; Lin et al., 2020; Bear et al., 2020;
Locatello et al., 2020; Zoran et al., 2021).
A thorough review and discussion on this type of method can be found in Greff et al.
One critical issue of these methods is on handling complex natural scenes.
Singh et al.
; Lamb et al.
(2021) leverages a transformer-based decoder with Slot-Attention for addressing this problem.
Similar attempts have also been made by exploiting self-supervised contrastive learning (Choudhury et al., 2021; Caron et al., 2021; Wang et al., 2022;
Hénaff et al., 2022) and energy-based models (Du et al., 2021; Yu et al., 2022).
Our work builds upon
Slot-Attention by extending it with learnable queries and a novel optimization method for learning.
Our compelling experimental suggests our model could potentially serve as a general plug-and-play module for a wider range of modalities where variants of Slot-Attention prosper (Kipf et al., 2022;
Elsayed et al., 2022;
Singh et al., 2022; Yu et al., 2022; Sajjadi et al., 2022a;b).
Query Networks Sets of latent queries are commonly used in neural networks.
These methods leverage permutation equivariant network modules ( e.g . GNNs (Scarselli et al., 2008) and attention modules (Vaswani et al., 2017)) in model design for solving set-related tasks such as clustering (Lee et al., 2019), outlier detection (Zaheer et al., 2017; Zhang et al., 2019), etc .
These learned latent queries have been shown to have good potential as features for tasks like contrastive learning (Caron et al., 2020), object detection (Carion et al., 2020), and data compression (Jaegle et al., 2021a;b).
In contrast to the recent success of query networks in supervised or weakly-supervised learning (Carion et al., 2020; Zhang et al., 2021; Kipf et al., 2022; Elsayed et al., 2022; Xu et al., 2022), Locatello et al.
(2020) demonstrates the detrimental effect of using independently initialized slots in Slot-Attention learning.
However, we show that our BO-QSA method successfully overcomes this issue and generalizes the success of query networks to the domain of unsupervised object-centric learning.
Bi-level Optimization Our work is closely related to bi-level optimization methods with iterative fixed update rules for solving the inner objective.
Specifically, methods are designed with implicit differentiation (Amos & Kolter, 2017; Bai et al., 2019) to stabilize the iterative update procedure.
Similar formulations are also found when combined with meta-learning where Madan et al.
(2021) train queries through recurrence in a meta-learning fashion and Rajeswaran et al.
(2019) provides a unified view of the optimization problem with implicit gradients.
Concurrent work from Chang et al.
(2022) formulate the Slot-Attention learning from an implicit gradient perspective with gradient stopping derived from first-order hyper-gradient methods (Geng et al., 2021).
However, they ignore the important role of slot initializations in generalization and concept binding.
As our experiments suggest, such gradient-stopping methods do not guarantee superior performance compared to the original Slot-Attention.
We leave the details to Sec.
5.3 for an in-depth discussion.
E XPERIMENTS
In this section, we aim to address the following questions with our experimental results:
• How good is our proposed BO-QSA on both synthetic and complex natural scenes?
• How important is the query and the optimization method in BO-QSA?
• Does BO-QSA possess the potential for concept binding and zero-shot transfer?
We provide details in the following sections with thorough comparative and ablative experiments and leave the details on model implementation and hyperparameter selection to Appendix A.3.
Here we clarify the datasets and metrics selected for evaluating our model on each domain:
Synthetic Domain For the synthetic domain, we select three well-established challenging multiobject datasets Shapestacks (Groth et al., 2018),
ObjectsRoom (Kabra et al., 2019), and CLEVRTEX for evaluating our BO-QSA model.
Specifically, we consider three metrics to evaluate the quality of object segmentation and reconstruction.
Adjusted Rand Index (ARI) (Hubert & Arabie, 1985) and Mean Segmentation Covering (MSC)
(Engelcke et al., 2020) for segmentation and Mean Squared
Error (MSE) for reconstruction.
Following the evaluation setting of recent works, we report the first two segmentation metrics over foreground objects (ARI-FG and MSC-FG).
Additionally, we conduct extra experiments on more datasets and leave the discussion to Appendix B.1.
Real-world Images For the real image domain
, we use two tasks (1) unsupervised foreground extraction and (2) unsupervised multi-object segmentation for evaluating our method.
Specifically, we select Stanford Dogs (Khosla et al., 2011), Stanford Cars (Krause et al., 2013), CUB200 Birds (Welinder et al., 2010), and Flowers (Nilsback & Zisserman, 2010) as our benchmarking datasets for foreground extraction and YCB (Calli et al., 2017), ScanNet (Dai et al., 2017), COCO (Lin et al., 2014) proposed by Yang & Yang (2022) for multi-object segmentation.
We use mean Intersection over Union (mIoU) and Dice as metrics for evaluating the quality of foreground extraction and use the evaluation metrics adopted by Yang & Yang (2022) for multi-object segmentation.
5. O BJECT D ISCOVERY ON S YNTHETIC D ATASETS
Experimental Setup We explore our proposed BO-QSA with two types of decoder designs, mixture-based and transformer-based, as discussed in Sec.
2.1 and Appendix A.1.
We follow the decoder architecture in Slot-Attention (Locatello et al., 2020) for mixture-based decoders and
[Table] [Table]
SLATE (Singh et al., 2021) for transformer-based decoders.
For both types of models, we use the
Slot-Attention module with a CNN image encoder and initialize slots with learnable embeddings.
Results We report multi-object segmentation results on synthetic datasets in Tab. 1 and visualize qualitative results in Fig.
As shown in Tab. 1, our BO-QSA achieves the state-of-the-art results with large improvements over previous object-centric learning methods on all metrics in ShapeStacks and ObjectsRoom.
We also observe more stable model performance, i.e . smaller variances in results, across different trials of experiments.
Our model with mixture-based decoders obtains the best overall performance on all datasets.
More specifically, our mixture-based BO-QSA significantly outperforms the vanilla Slot-Attention model ( „ 15%) with minimal architectural differences.
This validates the importance of the learnable queries and our optimization method.
We will continue this discussion in Sec.
As shown in Tab. 2, our model also achieves state-of-the-art results on the unsupervised object segmentation task in CLEVRTEX with consistent improvement over Slot-Attention on the
CAMO and OOD generalization split.
Interestingly, our model (1) shows larger reconstruction errors, (2) generalizes well in out-of-distribution scenarios, and (3) shows marginal improvement in camouflaged images.
We attribute (1) and (3) to the simple architecture of encoders/decoders currently adopted and provide insights on (2) in Sec. 5.4.
[Table]
Mixture-based vs. Transformer-based Decoder We observe inferior segmentation but superior reconstruction performance of transformer-based variants of Slot-Attention on synthetic datasets.
Specifically, we compare the MSE of models on ShapeStacks and ObjectsRoom.
As shown in Tab. 3, transformer-based methods provide better reconstruction results.
We attribute the low segmentation performance to mask prediction in these methods, which relies on the attention matrix computed over input features.
This leads to coarse object masks as a result of image tokenization.
Nonetheless, we observe consistent improvement by applying our slot encoder to both mixture and transformer decoders.
5. O BJECT D ISCOVERY ON R EAL D ATASETS
Experimental Setup For real-world experiments, we use the same slot encoder design used in Sec.
5.1 with a 4-layer CNN image encoder and initialize slots with learnable queries.
For
[Table] [Table]
- - Slot-Attention (Locatello et al., 2020) 35.
45.9 Voynov et al.
- DRC (Yu et al., 2021) 56.
- - Melas-Kyriazi et al.
- SLATE (Singh et al., 2021)
79.1 I-SA (Chang et al., 2022) 63.
Ours (mixture) 25.
51.6 Ours (transformer) 71.
[Figure] Dogs
[Figure] Cars
CUB200 Birds
[Figure]
unsupervised foreground extraction, we follow Yu et al.
and report the best model performance on all datasets.
During the evaluation, we select the slot’s mask prediction that has a maximum intersection with the ground-truth foreground mask as our predicted foreground.
For unsupervised multi-object segmentation, we follow Yang & Yang (2022) and report the models’ performance on all datasets across trials with different random seeds.
[Table]
Results We show quantitative experimental results in
Tab. 5 and Tab. 4.
We also visualize qualitative results in Fig.
For multi-object segmentation, as shown in Tab. 4, our model outperforms existing object-centric learning baselines by a large margin, especially on the
YCB dataset where the segmented objects have clear semantic meanings.
For foreground extraction, as shown in Tab. 5, our method significantly outperforms all existing baselines on the task of foreground extraction, achieving new state-of-the-art on all datasets.
We recognize the dis-
[Table]
[Figure]
crepancy of mixture-based decoders in both Slot-Attention and our mixture-based design in modeling real-world images, reflecting similar discoveries from recent works (Singh et al., 2021) that mixturebased decoder struggles in modeling real-world images.
On the other hand, our transformer-based model shows significant improvements over the vanilla version.
Notably, our method outperforms a broad range of models, including GAN-based generative models ( i.e . OneGAN, Voynov et al.
(2020)), and large-scale pre-trained contrastive methods ( i.e . MoCo-v2, BYOL, R2O).
As shown in Tab. 6, our method achieves comparable results with state-of-the-art self-supervised contrastive learning methods without large-scale pre-training and data augmentation.
This result sheds light on the potential of object-centric learning as a pre-training task for learning general visual representations.
5. A BLATIVE S TUDIES
Experimental Setup We perform ablative studies over our designs by comparing them with different design variants on ShapeStacks and Stanford Dogs.
For slot initialization, we consider (1) the original Slot-Attention module’s sampling initialization (SA), and (2) initializing with learnable queries (QSA).
For optimization, we consider (1) the original optimization in Slot-Attention ( i.e .
w/o detach or STE), (2) the I-SA optimization where gradients to slots in iterative updates are detached ( i.e . w/ detach only), and (3) our optimization where we both detach the gradients into iterative refinement, and pass gradient to the initialization queries with STE ( i.e . w/ detach and STE).
For simplicity, we term these variants with prefixes (I-) for I-SA and (BO-) for our full method.
We run all ablations on each dataset with the same encoder-decoder architecture.
Results We show experimental results in Tab. 7 and Fig.
First, from Tab. 7, we observe that
BO-QSA significantly outperforms other variants.
For sample-based slot initializations, our method shows a similar effect compared with I-SA on improving Slot-Attention learning.
For query-based slot initializations, we validate the difficulty in training query-based Slot-Attention with its inferior performance.
We further show the ineffectiveness of I-SA for query-based Slot-Attention.
The experiments on query-based Slot-Attention prove that both of our design choices are necessary and effective for superior performance.
To study the effect of learned queries, we visualize in Fig.where we set different numbers of iterative updates of Slot-Attention during inference on the Stanford
S h a p e S t ac k s
C LE V R
TE X
Transfer
B i r d s
D og s
C a r s
Transfer
Y CB
S ca n N e t
C O C O
Transfer
Dogs dataset.
We can see that our BO-QSA significantly outperforms other variants with only one iteration.
This indicates that our query-based design can help ease training difficulties.
In Fig.
3, we further visualize the learned initializations and post-iteration slots in the same feature space using t-SNE (Van der Maaten & Hinton, 2008).
Our initializers provide a more separable space when differentiating image features, which validates the desired model behaviors mentioned in Sec. 3.3.
5. A DDITIONAL A NALYSES
In this section, we provide additional analyses on the potential of our BO-QSA as a concept binder for generalizing to new examples.
First, we qualitatively visualize our learned content for each slot (without additional clustering) in ShapeStacks, Birds, and YCB in Fig.
We observe high similarity within the learned content of each slot, indicating similar concepts learned by specific slots.
This shows the potential of the slots in our BO-QSA for binding specific concepts on object properties ( e.g . colors, contours, and spatial positions).
Although we can not control which concepts to learn, these results are important indicators that our learned initialization queries could potentially be generalizable concept probes.
We further
[Table]
provide quantitative evaluations where we use models trained on dataset X for zero-shot inference on dataset Y. We term this transfer as (X Ñ Y).
As shown in Tab. 8, when adapting models trained on YCB to zero-shot inference on ScanNet and COCO, our method outperform I-SA and also the majority of fine-tuned methods shown in Tab. 4.
Due to the page limit, we show in Appendix B.1 that this superior transfer capability is general across datasets when compared to Slot-Attention variants.
C ONCLUSIONS
We introduce BO-QSA for unsupervised object-centric representation learning.
We initialize
Slot-Attention with learnable queries, and combine bi-level optimization and straight-through gradient estimators to ease the difficulty in query-based Slot-Attention learning.
With simple code adjustments on Slot-Attention, we obtain state-of-the-art model for unsupervised object segmentation in both synthetic and natural image domains, outperforming previous baselines by a large margin.
More importantly, our learned model exhibits concept-binding effects where visual concepts are attached to specific slot queries.
With a fixed number of initialized slots, our model is limited to handling a fixed maximum number of objects in the inputs.
However, our queries could be learned to bind object attributes, which leads to meaningful segmentation of images by grouping similar properties ( e.g . color, position, etc .).
As a future direction, this connects our method with weakly-supervised contrastive learning methods that learn grounded visual representations with language.
A CKNOWLEDGEMENT
We gratefully thank all colleagues from BIGAI for fruitful discussions.
We would also like to thank the anonymous reviewers for their constructive feedback.
This work reported herein was supported by National Key R&D Program of China (2021ZD0150200)
A M ODEL A RCHITECTURE AND D ESIGN A. D ESIGN OF D ECODERS
In this section, we follow the notations used in Sec.
and describe two common approaches, mixture-based and transformer-based, for decoding images from the learned slot representations.
Mixture-based Decoder
The mixture-based decoder (Watters et al., 2019) decodes each slot
ˆ s i into an object image x
i and mask m i with decoding functions g img ϕ dec and g mask ϕ dec , which are implemented using CNNs.
The decoded images and masks are calculated by:
[Equation]
During training, a reconstruction objective is employed for supervising model learning.
Despite its wide usage, mixture-based decoders showed limited capability at handling natural scenes with high visual complexity (Singh et al., 2021).
Autoregressive Transformer Decoder Recently, Singh et al.
(2021; 2022) reveal the limitations of mixture decoder and leverage transformers and dVAEs (Van Den Oord et al., 2017; Ramesh et al., 2021) for decoding slot-based object-centric representations.
To obtain decoded images ˆ I , they learn a separate dVAE for first encoding I into a sequence of L tokens z “ t z 1 , ¨ ¨ ¨ , z L u with dVAE encoder f dVAE ϕ enc .
Next, they use a transformer decoder g transformer ϕ
dec to auto-regressively predict image tokens with learned slot representation ˆ s :
[Equation]
To train the entire model, we have the reconstruction objective supervising the learning of z with dVAE decoder g dVAE ϕ dec .
Next, the objective for object-centric learning relies on the correct prediction from the auto-regressive transformer for predicting correct tokens:
[Equation]
where L dVAE “
[Equation]
g
[Equation]
dec p [Equation]
I [Equation] 22 , L CE “
L ÿ l “
CrossEntropy p
z l , o l q
Under this setting, the model does not predict additional masks and relies on the attention A within the Slot-Attention module for obtaining slot-specific object masks.
Although such models can achieve competitive results on real-world synthetic datasets, as our experiments suggest, they can be inferior to mixture-based decoders on segmentation in synthetic datasets.
We suspect that this originates from the low resolution when discretizing images into tokens.
A. B I - LEVEL O PTIMIZATION AND M ETA -L EARNING
Recall the bi-level optimization problem we introduced in Sec. 2.2.
[Equation]
where we call f p θ, ϕ q the outer objective function and g p θ, ϕ q the inner objective function.
To jointly optimize both objectives w.r.t.
parameters θ and ϕ , a straightforward approach to solving Eq.
(6) is to represent the inner solution of θ as a function of ϕ , i.e ., θ ˚ p ϕ q “ arg min θ 1 g p θ 1 , ϕ q .
Then we can optimize the outer objective with gradient descent:
[Equation]
However, the difficulty of this method lies in the calculation of ∇ ϕ θ ˚ p ϕ q where we need to solve linear equation from implicit gradient theorem:
[Equation]
If ∇ 2 , 2 g p θ ˚ , ϕ q is invertible, we can solve for ∇ ϕ θ ˚ p ϕ q and obtain the gradient update on ϕ :
[Equation]
f k ´ p ∇ 1 , 2 g k q J p ∇ 2 , 2 g k q ´ 1 ∇ 1 f k ˘
[Figure]
where ∇ 1 f k “ ∇ 2 f p θ ˚ p ϕ k q , ϕ k q and ∇ 1 f k “ ∇ 1 f p θ ˚ p ϕ k q , ϕ k q .
Various methods have been proposed to approximate the solution (Pedregosa, 2016; Lorraine et al., 2020), and we refer the authors to Ye et al.
for a thorough review of related methods.
Bi-level optimization is closely related to meta-learning.
In meta-learning, we have meta-training tasks which comes in as N different collections of datasets D “ t D i “ D tr i Y D val i u Ni “ 1 .
The inner and outer objectives in Eq.
(6) are substituted by averaging training and validation errors over multiple tasks (Franceschi et al., 2018):
[Equation]
where L i represents task-dependent error on D i .
The final goal of meta-learning aims at seeking the meta-parameter ϕ that is shared between tasks which later enables few-shot learning and fast adaptation.
With its connections with bi-level optimization, the previously mentioned optimization methods are broadly adapted for solving meta-learning problems (Finn et al., 2017; Nichol &
Schulman, 2018; Rajeswaran et al., 2019).
From the meta-learning perspective, our attempt shares similar insights with first-order meta-learning methods (Finn et al., 2017; Nichol & Schulman, 2018), where we use the gradient at some task-specific optimal solution s ˚ i of the inner optimization for optimizing slot initialization queries which are shared across datasets on the outer objective.
This meta-learning perspective also indicates the potentials of our BO-QSA for fast adaptation and generalization.
A. I MPLEMENTATION D ETAILS
We provide a visualization of our designed slot-encoder in Fig.
5 and discuss the implementation details for different experimental settings in the following sections.
A.3. S LOT I NITIALIZATION
We initialize all models with the number of slots shown in Tab. 13.
During training, we add a small perturbation to the queries by sampling from a zero-mean distribution with variance σ as we found it empirically helpful for better performance.
We perform annealing over σ to gradually eliminate the effect of this random perturbation during training.
We adopt the cosine annealing strategy such that σ starts from 1 and gradually anneals to 0 after N σ training steps, where N σ is a hyperparameter that controls the annealing rate of σ .
In our experiments, we use N σ “ 0 on Cars and Flowers and
N σ “ 30000 on the rest of the datasets.
[Table] [Table]
A.3. BO-QSA WITH M IXTURE -B ASED D ECODERS
For mixture-based decoders
[Table]
we use the same Slot-Attention architecture as in Locatello et al.
(2020) with slots initialized by learnable queries.
Given an input image, Slot-Attention uses a CNN encoder to extract image features.
After adding positional embedding
[Table]
these features are input into the
Slot-Attention module slot updates.
Finally, these slots are decoded by the mixture decoder to reconstruct the input image.
We provide the details of our image encoder in Tab. 9.
For the mixturebased decoder, we use six transposed convolutional layers with ReLU activations following Locatello et al.
We visualize the details of our mixture-based decoder design in Tab. 10.
We train our model for 250k steps with a batch size of 128 and describe all training configurations and hyperparameter selection Tab. 11. [Table]
[Table]
A.3. BO-QSA WITH T RANSFORMER -B ASED D ECODER
For transformer-based decoders, we adopt the transformer architecture proposed by SLATE (Singh et al., 2021).
For the transformer-based BO-QSA, unlike SLATE, we use the same CNN as in mixture-based BO-QSA (instead of the dVAE encoder) to extract features from the image as input to the Slot-Attention module as we find such changes help solve the problem on coarse object boundary prediction mentioned in Sec. 5.1.
Next, we use the same overall architecture of dVAE as mentioned in SLATE Singh et al.
However, we change the kernel size of the dVAE encoder from 1 tosince we find that such changes can help increase model performance when decomposing scenes.
We train our model for 250k steps with a batch size of 128, and all the training configuration in our experiments is described in Tab.
A.3. B ASELINES
The reproduction of Slot-Attention and SLATE follows the architecture and hyperparameter selection mentioned in their paper.
Similar to our models, we train all baseline models with 250K steps on all datasets.
For SLATE, we use the input image size of 96 on the ShapeStacks dataset as we find that the image size of 128 will cause all objects to be divided into the same slot, resulting in low
[Table]
ARI and MSC.
For a fair comparison with numbers reported in SLATE’s paper, we report the MSE of models by first computing per-pixel errors and then multiplying it by the total number of pixels. For CLEVRTEX
[Table]
we follow the same experimental setting of (BO-QSA+mixture) for ShapeStacks and set the number of slots to 11.
For YCB, ScanNet
[Table]
and COCO, we follow the same experimental setting of (BO-QSA+transformer) for birds and set the number of slots to 6.
[Table]
B A DDITIONAL E XPERIMENTS B. Z ERO - SHOT T RANSFER
In this section, we continue the discussion in Sec.
5.4 and provide additional zero-shot transfer results.
Similarly, we use the notation ( X Ñ Y ) to denote the zero-shot adaptation of models trained unsupervisedly on dataset X to new datasets Y .
For unsupervised multi-object segmentation, we report transfer results from ScanNet and COCO to all other real-image multi-object segmentation datasets in addition to the results on YCB (mentioned in Sec. 5.4).
As shown in Tab. 14, our model shows consistent improvement over Slot-Attention and
I-SA during zero-shot transfer.
[Table]
For unsupervised foreground extraction, we report transfer results from Stanford Dogs and CUB
Birds to all other real-image foreground extraction datasets.
As we can see from Tab. 15, our model achieves the overall best results compared with other powerful Slot-Attention variants (models that achieve best or second-best results in our ablation studies as in Tab. 7) except for (Birds Ñ Cars).
However, our optimization method still helps improve zero-shot transfer for randomly initialized
Slot-Attention.
[Table]
B. A NALYSIS N UMBER OF S LOT -A TTENTION I TERATIONS
As described in Sec. 3.
[Table]
we study whether a fixed point s ˚ could be reached by a fixed number of iterations during training.
Since we hypothesized that the low performance of I-QSA in Sec. 5.
originated from the insufficient number of starting points for fixed-point approximation
[Table]
we conduct experiments on increasing the number of Slot-Attention iterations during training for I-QSA on the
Dog dataset.
As shown in Tab.
[Table]
increasing the number of Slot-Attention iterations during training for I-QSA significantly improves its performance.
However
[Table]
we found that adding more iterations after a threshold ( i
[Table]
e . 7 in this case) does not further improve the overall performance.
This verifies the need for learning slot initialization vectors for better approximating the fixed point solution of the inner soft-clustering objective in Slot-Attention
[Table]
[Table]
B. D ESIGN C HOICES ON S LOT I NITIALIZATION
As described in Sec. 3.3, our method is connected with recent works on dVAE.
However, we do not require the initialization queries to maintain information about the post-iteration slots ˆ s as we found such constraints lead to the learning of the mean representation of datasets which forbids disentanglement and concept binding.
In this section, we provide experimental results to verify this argument.
Specifically, we consider three different ways to update slot initialization queries in addition to our proposed method: 1) using the running mean of the post-iteration slots as initialization queries (RunningMean), 2) running K-Means clustering on post-iteration slots and updating the initialization queries using re-clustered centers by Hungarian matching (KMeans), 3) adding consistency loss between initialization queries and post-iteration slots as done in VQ-VAE (VQ-constraint).
For (1) and (2)
[Table]
we empirically found such designs to be suffering from frequent updates and therefore use momentum updates to stabilize their training.
We term these variants with the suffix (-M).
[Table]
As shown in Tab.
[Table]
our model achieves the best overall performance compared to other initialization methods.
Specifically, we found that using the running mean of post-iteration slots or K-Means cluster centers re-clustered from post-iteration slots to be harmful to model performance.
We attribute this [Figure] effect to the learning of the mean-representation of datasets.
This is further proved in experiments with VQ-VAE loss on consistency between slot initializations and post-iteration slots ( i.e . ||
sg p ˆ s q ´ s 0 || 2 ), where the VQ-constraint variant showed inferior performance.
We also found that the weight of this additional loss needs to be carefully tuned for the model to decompose objects.
Empirically, most configurations of this hyperparameter will lead to bad reconstructions except for certain small weights ( e.g . 0.01 reported here).
Above all, we believe these experimental results verify the effectiveness of our design choices on initialization query learning.
We provide additional visualizations on the learned contents of slots for each update method in Fig. 6.
B. E XPERIMENTS ON A DDITIONAL D ATASETS
In addition to datasets considered in Sec. 5, we conduct experiments on other synthetic datasets and visualize qualitative results.
More specifically, we test our model on PTR (Hong et al., 2021).
PTR is a synthetic dataset of 3D objects from PartNet with rendering variations.
We run our BO-QSA with the same configuration mentioned in Appendix A.3 previously.
We compare our method with the vanilla Slot-Attention module on multi-object segmentation.
We report ARI-FG and MSC-FG scores of our model compared with the vanilla Slot-Attention on the PTR validation set.
[Table]
As we can see from Tab. 18, our model achieves similar performance compared with Slot-Attention on ARI-FG and significantly outperforms it on MSC-FG.
We attribute this result to the capability of precisely segmenting objects.
As ARI-FG applies masks to each slot prediction for calculating results, it does not require models to precisely segment the object from the background.
However, MSC-FG uses a mIoU-like measure that requires the model to precisely predict the object boundaries.
This indicates that our model is better at precisely segmenting objects without noise.
Similarly, we observe the binding of certain slots to scene backgrounds, but with more complex concepts, the binding of slots to concepts is not as straightforward as in ShapeStacks and CUB200 Birds.
To further investigate the effectiveness and generality of our method, we adapt BO-QSA to the recent 3D object-centric learning model, uORF (Yu et al., 2022), and test it on 3D datasets including CLEVR567, Room-Chair, and Room-Diverse.
uORF can decompose complex 3D scenes from a single image by combining NeRF (Mildenhall et al., 2021) with Slot-Attention.
We only modify the initialization and optimization method of the Slot-Attention module in uORF, leaving all other hyperparameters unchanged.
As we can see from Tab. 19, with our method, the uORF model that trained withepochs can achieve a similar or even superior result compared to the original model trained withepochs.
Additionally, when the dataset complexity increases ( e.g ., in Room-Diverse), our method demonstrates significant improvement.
Please refer to uORF (Yu et al., 2022) for more details about the model, datasets, and evaluation metrics [Table] [Table] C L IMITATIONS AND F UTURE W ORK
We discuss all limitations of our work found in the experiments.
First, we observed a strong correlation between the powerfulness of encoder-decoder architectures and model performance.
However, in contrast to supervised learning, more powerful encoders/decoders do not guarantee superior performance.
Gaining insights from how contrastive learning methods have shown the effect of concept emergence with large-scale pretraining, we can also incorporate such representations learned by self-supervised learning into object-centric learning to unite the best of both worlds.
Second, our work is primarily limited by the fixed number of slot initialization vectors.
In contrast to the vanilla Slot-Attention that could generalize to a new number of objects, our model can not easily generalize to scenarios with new concepts since our model learns a fixed set of separating spaces that best disentangle different parts of the image.
This problem is also frequently met in semantic segmentation and object classification, where we can only use existing concepts to interpret novel objects/semantic entities.
Although solutions to this close-vocabulary problem have been proposed in supervised classification and segmentation, we leave the exploration of this problem in object-centric learning to future work.
Finally, the current learned slot initialization vectors do not explicitly bind towards concepts and need to be mined by humans.
We believe this is an important next step in our current work to combine unsupervised object-centric learning with semantic alignments from language for concept grounding.
This opens future research directions on learning finer-level organization of object concepts under more complex scenarios ( e.g . hierarchical grouping) with weak supervision of correspondence.
D A DDITIONAL V ISUALIZATIONS
We provide more qualitative results of our model on different datasets in the following pages.