[Paper Review] Improving Object-centric Learning with Query Optimization

Baoxiong Jia, Yu Liu+1 more|arXiv (Cornell University)|Oct 17, 2022

Advanced Image and Video Retrieval Techniques4 citations

TL;DR

This paper proposes Bi-level Optimized Query Slot Attention (BO-QSA), a novel method that improves unsupervised object-centric learning by learning slot initializations as queries instead of random initialization, enabling more stable and conceptually meaningful object representations. The approach achieves state-of-the-art performance on 10 diverse datasets—synthetic and real-world—by integrating straight-through gradient estimation with bi-level optimization, significantly outperforming prior Slot-Attention variants in segmentation and reconstruction tasks.

ABSTRACT

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

Motivation & Objective

• Address the instability and ambiguity in object-centric representation learning caused by random slot initialization in Slot-Attention.
• Reduce reliance on hyperparameter tuning and training heuristics (e.g., gradient clipping, learning rate warm-up) in unsupervised object-centric models.
• Improve the binding of learned slots to actual object concepts, enabling better generalization and zero-shot transfer potential.
• Enhance training efficiency and convergence by integrating learnable queries with bi-level optimization and straight-through gradient estimation.
• Enable plug-and-play deployment of improved object-centric learning across diverse vision tasks, including 3D scene reconstruction and image segmentation.

Proposed method

• Replace random slot initialization in Slot-Attention with learnable query vectors that are optimized via bi-level optimization.
• Introduce a straight-through gradient estimator (STE) to enable end-to-end backpropagation through the discrete slot assignment process.
• Formulate the optimization as a bi-level problem: inner loop optimizes slot features via attention, outer loop updates the learnable query initializations.
• Maintain the original Slot-Attention architecture but modify only the initialization and optimization procedure, ensuring backward compatibility.
• Use first-order approximation to reduce computational cost in the outer-loop optimization, enabling efficient training.
• Apply the method as a plug-in to existing models like uORF, improving performance without architectural changes.

Experimental results

Research questions

• RQ1 Can learnable query initialization stabilize training and improve performance in unsupervised object-centric learning?
• RQ2 Does replacing random initialization with optimized queries lead to better object concept binding in complex scenes?
• RQ3 To what extent does the proposed method reduce dependency on hyperparameter tuning in Slot-Attention?
• RQ4 Can the model generalize to zero-shot scenarios and transfer to new datasets without fine-tuning?
• RQ5 How does the method perform across diverse benchmarks, including synthetic and real-world image datasets?

Key findings

• BO-QSA achieves state-of-the-art performance on 3 synthetic and 7 real-world datasets, outperforming previous SOTA methods in both reconstruction and segmentation.
• On the CLEVR-567 dataset, uORF+BO-QSA with 600 training epochs achieves 74.5 NV-ARI, surpassing the original uORF model trained for 1200 epochs (84.4 NV-ARI).
• In the complex Room-Diverse benchmark, uORF+BO-QSA with 240 epochs reaches 63.0 NV-ARI, compared to 56.6 for the original uORF at 240 epochs.
• The model demonstrates improved generalization in zero-shot transfer learning, indicating strong potential for concept binding.
• Ablation studies confirm that learnable queries are essential for performance gains, with random initialization leading to inferior results.
• The method reduces training instability and dependency on hyperparameter tuning, as evidenced by consistent convergence across diverse datasets.