SlowGuess/NeurIPS_2023_OCR huggingface dataset:
PASCAL VOC 2012/MS COCO 2017 are real-world datasets commonly used in object detection and segmentation. They are more challenging than CelebA as the images capture unconstrained natural scenes with multiple objects of varying sizes. We will denote them as VOC and COCO for short.
Baselines. We compare SlotDiffusion with state-of-the-art fully unsupervised object-centric models. On image datasets, we adopt Slot Attention (SA for short) [60] and SLATE [83]. On video datasets, we adopt SAVi [46] and STEVE [84]. They are representative models which use the mixture-based
Figure 2: Unsupervised object segmentation measured by FG-ARI (left), mIoU (middle), and mBO (right). Mixture, Transformer, and Diffusion stand for SA/SAVi, SLATE/STEVE, and SlotDiffusion, respectively.
Figure 3: Visualization of video segmentation results on MOVi-D (left) and MOVi-E (right).
decoder and the Transformer-based decoder. We will introduce other baselines in each task below. See Appendix D.3 for implementation details of all baselines adopted in this paper.
Our Implementation Details. We use the same image encoder, Slot Attention module, and transition predictor as baselines, while only replacing the slot decoder with the conditional LDM. We first pre-train VQ-VAE on each dataset, and then freeze it and train the object-centric model with Eq. (7). Please refer to Appendix D.4 for complete training details, and Appendix E.1 for a comparison of computation requirements of SlotDiffusion and baseline methods.
4.2 Evaluation on Object Segmentation
Evaluation Metrics. Following previous works, we use foreground Adjusted Rand Index (FG-ARI) to measure how well the predicted object masks match the ground-truth segmentation. As pointed out by [17, 43], FG-ARI is sub-optimal since it only considers foreground pixels and ignores false segmentation of backgrounds. Therefore, we also adopt the mean Intersection over Union ( $mIoU$ ) and mean Best Overlap ( $mBO$ ) computed over all pixels. See Appendix D.2 for metric implementations.
Results. Figure 2 presents the results on unsupervised object segmentation. We do not evaluate on CelebA as it does not provide object mask annotations. SlotDiffusion outperforms baselines in three metrics across all datasets, proving that the denoising loss in LDMs is a good training signal for scene decomposition. Compared to CLEVRTex, our method achieves larger improvements on video datasets since SlotDiffusion also improves the tracking consistency of objects. Figure 3 shows a qualitative result of video segmentation. As observed in previous work [84], MOVi-E's moving camera provides motion cues compared to static cameras in MOVi-D, making object discovery easier. Indeed, SAVi degenerates to stripe patterns on MOVi-D, but is able to produce meaningful masks on MOVi-E. Compared to STEVE, our method has fewer object-sharing issues (one object being captured by multiple slots), especially on large objects. See Appendix E.2 for more qualitative results.
4.3 Evaluation on Generation Capacity
Evaluation Metrics. To generate high-quality results, object-centric generative models need to first represent objects with slots faithfully, and then compose them in novel ways. Hence, we first evaluate slot-to-image reconstruction, which inspects the expressive power of object slots. We adopt MSE and LPIPS [108] following previous works. As discussed in [100, 108], LPIPS better aligns with human perception, while MSE favors blurry images. We thus focus on comparing LPIPS, and report MSE
Figure 4: Slot-based image and video reconstruction measured by MSE (left) and LPIPS (right).
Figure 5: Reconstruction results on CLEVRTex (left) and CelebA (right). Our method better preserves the textures on objects and backgrounds in CLEVRTex, and the human hairs and facial features in CelebA.
(a) Visual concepts discovery and composition
(b) Quantitative results on compositional generation
Figure 6: The mechanism of slot-based compositional generation (a) and the quantitative results (b). Our method first discovers visual concepts from images in the form of slots, such as objects and background maps in CLEVRTex. Then, we can synthesize novel samples by composing concept slots from different images.
Figure 7: Compositional generation samples on CLEVRtex (left), CelebA (middle), and MOVi-E (right). We randomly select slots from the visual concept libraries built from training data, and decode them into images or videos. SlotDiffusion synthesizes higher-fidelity results with more details compared to baselines.
for completeness. In addition to reconstruction, we conduct compositional generation with the visual concepts learned in Section 3.3. FID [32] and FVD [92] are used to measure the generation quality.
Results on reconstruction. Figure 4 presents the reconstruction results. SlotDiffusion outperforms both baselines with a sizeable margin in LPIPS on all datasets, and achieves the second best MSE. As shown in the qualitative results in Figure 5, SA with a mixture-based decoder produces blurry images despite a lower MSE. Also, the Transformer-based decoder reconstructs distorted objects, and loses all the textures on object surfaces and backgrounds. In contrast, our LDM-based decoder is able to iteratively refine the results, leading to accurate textures and details such as human hairs.
Results on compositional generation. As shown in Figure 6 (a), we discover visual concepts by clustering slots extracted from images, and randomly composing slots from each concept's library to generate novel samples. Figure 6 (b) summarizes the quantitative results. SlotDiffusion significantly improves the FID and FVD scores, since it synthesizes images and videos with coherent contents and much better details compared to baselines. This is verified by the qualitative results in Figure 7. See Appendix E.3 for more visualizations of compositional generation results.
(a) Scene editing on CLEVRTex
(b) Face replacement on CelebA
Figure 8: Slot-based image editing. Left: we can manipulate objects, retrieve and swap background maps in a scene. Right: we can replace the human faces in the first column with the faces from the first row.
Table 1: Video prediction results on Physion.
Method	MSE	LPIPS	FVD
PredRNN	286.3	0.38	756.7
VQFormer	619.0	0.33	694.5
STEVE + SlotFormer	832.0	0.43	930.6
Ours + SlotFormer	477.7	0.26	582.2
Table 2: Physion VQA accuracy. We report observation (Obs.), observation plus rollout (Dyn.), and the improvement  $(\uparrow)$ .
Method	Obs. (%)	Dyn. (%)	↑(%)
RPIN	62.8	63.8	+1.0
pDEIT-lstm	59.2	60.0	+0.8
STEVE + SlotFormer	65.2	67.1	+1.9
Ours + SlotFormer	67.5	69.7	+2.2
Figure 9: Qualitative results of video prediction on Physion. The first row shows ground-truth frames, while the other rows show the model rollout results. See Appendix E.4 for visualizations of all baselines.
Slot-based image editing. In Figure 8, we apply SlotDiffusion to image editing by manipulating slots. On CLEVRTex, we can remove an object by segmenting it out, and insert it to another scene. We can also extract background maps from an image and swap it with another image. On CelebA, SlotDiffusion decomposes images to faces, hairs, clothes, and backgrounds. By swapping the face slot, our method generates realistic images with new faces, while keeping other components unchanged.
4.4 Downstream Tasks: Video Prediction and VQA
In this subsection, we evaluate the object slots learned by SlotDiffusion in downstream tasks. We select Physion dataset as it requires capturing the motion and physical properties of objects for accurate dynamics simulation. We train the state-of-the-art object-centric dynamics model SlotFormer [100] over slots extracted by our model, and unroll it to perform video prediction and question answering.
Video prediction baselines. We adopt three baselines that build upon different image feature representations: PredRNN [96] that uses global image feature maps, VQFormer which utilizes image patch tokens encoded by VQ-VAE [72], and STEVE + SlotFormer that runs the same SlotFormer dynamics model over object-centric slot representations extracted by STEVE.
Results on video prediction. Table 1 presents the results on the visual quality of the generated videos. SlotDiffusion with SlotFormer outperforms baselines clearly in LPIPS and FVD, and achieves a competitive result on MSE. As discussed in Section 4.3 and [100], MSE is a poor metric in generation. Figure 9 visualizes the model rollouts. Compared to STEVE with SlotFormer, our method generates videos with higher fidelity such as the consistent boxes and the red object.
Table 3: Unsupervised object segmentation results on real-world datasets: VOC (left) and COCO (right).
PASCAL VOC	FG-ARI	mBOi	mBOc	MS COCO	FG-ARI	mBOi	mBOc
SA + DINO ViT	12.3	24.6	24.9	SA + DINO ViT	21.4	17.2	19.2
SLATE + DINO ViT	15.6	35.9	41.5	SLATE + DINO ViT	32.5	29.1	33.6
DINOSAUR	23.2	43.6	50.8	DINOSAUR	34.3	32.3	38.8
Ours + DINO ViT	17.8	50.4	55.3	Ours + DINO ViT	37.2	31.0	35.0
Figure 10: Segmentation results on real-world images: COCO (left) and VOC (right). SlotDiffusion is able to discover unannotated but meaningful regions highlighted by the red arrows, such as trees and keyboards.
Method	FG-ARI	mIoU	mBO
SlotDiffusion	68.39	57.56	61.94
SlotDiffusion + Slot-TTA	69.66	58.67	62.82
Table 4: Unsupervised object segmentation results of SlotDiffusion with other improvements on CLEVRTex.
Method	FG-ARI	mIoU	mBO
BO-QSA	80.47	64.26	66.27
SlotDiffusion	68.39	57.56	61.94
BO-SlotDiffusion	78.50	64.35	68.68
VQA baselines. We select two baselines following [100].  $RPIN$  is an object-centric dynamics model with ground-truth object bounding box information.  $pDEIT$ -lstm builds LSTM over frame features extracted by an ImageNet [10] pre-trained DeiT [91]. We also compare with  $STEVE + SlotFormer$ . Results on VQA. Table 2 summarizes the VQA accuracy on observation only (Obs.) and observation plus rollout frames (Dyn.). SlotDiffusion with SlotFormer achieves the best result in both cases, with the highest rollout gain. This proves that our method learns informative slots which capture correct features for learning object dynamics, such as the motion of objects falling down in Figure 9.
4.5 Scaling to Real-World Data
To show the scalability of our method to real-world data, we follow [81] to leverage pre-trained models. We replace the image encoder in SlotDiffusion with a self-supervised DINO [7] pre-trained ViT encoder. The model is still trained with the LDM denoising loss over VQ-VAE image tokens.
Baselines. To directly compare the effect of the slot decoders, we adopt SA and SLATE with the same DINO pre-trained ViT image encoder. We also compare with DINOSAUR [81] with the same image encoder, which reconstructs the ViT feature maps instead of input images as training signal.
Results. Table 3 presents the segmentation results on VOC and COCO. We follow [81] to report  $\mathrm{mBO}^i$  and  $\mathrm{mBO}^c$ , which are mBOs between predicted masks and ground-truth instance and semantic segmentation masks. SlotDiffusion consistently outperforms SA and SLATE over all metrics, proving the effectiveness of our LDM-based slot decoder. Our model is also competitive with the state-of-the-art method DINOSAUR. The performance variations may be because COCO has more instances per image than VOC. As shown in the qualitative results in Figure 10, SlotDiffusion is able to segment the major objects in an image, and discover unannotated but semantically meaningful regions.
4.6 Combining SlotDiffusion with Recent Improvements in Slot-based Models
As an active research topic, there have been several recent improvements on slot-based models [39, 67, 85]. Since these ideas are orthogonal to our modifications to the slot decoder, we can incorporate them into our framework to further boost the performance.
Slot-TTA [67] proposes a test-time adaptation method by optimizing the reconstruction loss of slot models on testing images. Similarly, we optimize the denoising loss of SlotDiffusion at test time. As shown in Table 4 left, Slot-TTA consistently improves the segmentation results of our method.
BO-QSA [39] applies bi-level optimization in the iterative slot update process. We directly combine it with SlotDiffusion (dubbed BO-SlotDiffusion). As shown in Table 4 right, BO-SlotDiffusion greatly
Figure 11: Ablation study on the number of diffusion steps  $T$  of SlotDiffusion. We show the reconstruction performance (MSE, LPIPS) and segmentation results (FG-ARI, mIoU) on MOVi-D and MOVi-E. We also plot the best baseline results on each dataset as dashed lines. We do not plot MSE as it favors blurry generations.
improves the segmentation results compared to vanilla SlotDiffusion, and achieves comparable performance with BO-QSA. Note that the generation quality of BO-QSA is close to vanilla Slot Attention, and thus both our methods are significantly better than it.
We study the effect of the number of diffusion process steps  $T$  in our LDM-based decoder. We plot both the reconstruction and segmentation results on MOVi-D and MOVi-E in Figure 11. As expected, using more denoising steps leads to better generation quality. Interestingly, smaller  $T$  results in better segmentation performance. This indicates that there is a trade-off between the reconstruction and segmentation quality of our model, and we select  $T = 1000$  to strike a good balance. Literature on self-supervised representation learning [29, 38] suggests that a (reasonably) more difficult pretext task usually leads to better representations. Analogously, we hypothesize that a smaller  $T$  makes the denoising pretext task harder, resulting in better object-centric representations.
Overall, SlotDiffusion outperforms baselines consistently across different  $T$  in LPIPS, mIoU, and FG-ARI, proving that our method is robust against this hyper-parameter.
In this paper, we propose SlotDiffusion by introducing diffusion models to object-centric learning. Conditioned on object slots, our LDM-based decoder performs iterative denoising over latent image features, providing strong learning signals for unsupervised scene decomposition. Experimental results on both image and video datasets demonstrate our advanced segmentation and visual generation performance. Moreover, SlotDiffusion can be combined with recent progress in object-centric models to achieve state-of-the-art video prediction and VQA results, and scale up to real-world data. We briefly discuss the limitations and future directions here, which are further extended in Appendix F.
Limitations. With the help of pre-trained encoders such as DINO ViT [7], SlotDiffusion is able to segment objects from naturalistic images [21, 54]. However, we are still unable to decode natural images faithfully from slots. In addition, objects in real-world data are not well-defined. For example, the part-whole hierarchy causes ambiguity in the segmentation of articulated objects, and the best way of decomposition depends on the downstream task. Finally, our learned slot representations are still entangled, while many tasks require disentanglement between object positions and semantics.
Future Works. From the view of generative modeling, we can explore better DM designs to improve the modeling capacity of local details. To generate high-quality natural images, one solution is to incorporate pre-trained generative models such as the text-guided Stable Diffusion [73]. From the view of object-centric learning, we want to apply SlotDiffusion to more downstream tasks such as reinforcement learning [95, 105] and 3D modeling [86, 104]. It is interesting to see if task-specific supervision signals (e.g., rewards in RL) can lead to better scene decomposition.
Unsupervised object-centric learning is an important task in machine learning and computer vision. Equipping machines with a structured latent space enables more interpretable and robust AI algorithms, and can benefit downstream tasks such as scene understanding and robotics. We believe this work benefits both the research community and the society.
Potential negative societal impacts. We do not see significant risks of human rights violations or security threats in our work. However, since our method can be applied to generation tasks, it might generate unpleasant content. Future research should also avoid malicious uses such as social surveillance, and be careful about the training cost of large diffusion models.
We would like to thank Gautam Singh, Maximilian Seitzer for the help with data preparation and discussions about technical details of baselines, Xuanchi Ren, Tianyu Hua, Wei Yu, all colleagues from TISL for fruitful discussions about diffusion models, and Xuanchi Ren, Liquan Wang, Quincy Yu, Calvin Yu, Ritviks Singh for support in computing resources. We would also like to thank the anonymous reviewers for their constructive feedback. Animesh Garg is supported by CIFAR AI chair, NSERC Discovery Award, University of Toronto XSeed Grant, and NSERC Exploration grant. We would like to acknowledge Vector Institute for computation support.
[1] Federico Baldassarre and Hossein Azizpour. Towards self-supervised learning of global and object-centric representations. In ICLR2022 Workshop on the Elements of Reasoning: Objects, Structure and Causality, 2022.
[2] Fan Bao, Chongxuan Li, Jiacheng Sun, and Jun Zhu. Why are conditional generative models better than unconditional ones? In NeurIPS 2022 Workshop on Score-Based Methods, 2022.
[3] Dmitry Baranchuk, Andrey Voynov, Ivan Rubachev, Valentin Khrulkov, and Artem Babenko. Label-efficient semantic segmentation with diffusion models. In ICLR, 2022.
[4] Daniel Bear, Elias Wang, Damian Mrowca, Felix Jedidja Binder, Hsiao-Yu Tung, RT Pramod, Cameron Holdaway, Sirui Tao, Kevin A Smith, Fan-Yun Sun, et al. Physion: Evaluating physical prediction from vision in humans and machines. In NeurIPS Datasets and Benchmarks Track, 2021.
[5] Chris Burgess and Hyunjik Kim. 3d shapes dataset. https://github.com/deepmind/3dshapes-dataset/, 2018.
[6] Christopher P Burgess, Loic Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matt Botvinick, and Alexander Lerchner. Monet: Unsupervised scene decomposition and representation. arXiv preprint arXiv:1901.11390, 2019.
[7] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In ICCV, 2021.
[8] Zhenfang Chen, Jiayuan Mao, Jiajun Wu, Kwan-Yee Kenneth Wong, Joshua B Tenenbaum, and Chuang Gan. Grounding physical concepts of objects and events through dynamic visual reasoning. In ICLR, 2020.
[9] Antonia Creswell, Rishabh Kabra, Chris Burgess, and Murray Shanahan. Unsupervised object-based transition models for 3d partially observable environments. NeurIPS, 2021.
[10] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
[11] Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. NeurIPS, 2021.
[12] David Ding, Felix Hill, Adam Santoro, Malcolm Reynolds, and Matt Botvinick. Attention over learned object embeddings enables complex visual reasoning. NeurIPS, 2021.
[13] Mingyu Ding, Zhenfang Chen, Tao Du, Ping Luo, Josh Tenenbaum, and Chuang Gan. Dynamic visual reasoning by learning differentiable physics models from video and language. NeurIPS, 2021.
[14] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2020.
Figure 12 shows the general pipeline of video Slot Attention models, and compare two existing slot decoders with our proposed Latent Diffusion Model (LDM) [73] based decoder.
(a) General pipeline of video-based Slot Attention models.
(b) Mixture-based CNN decoder.
(c) Transformer-based autoregressive decoder.
(d) Latent diffusion model based decoder.
Figure 12: Illustration of (a) the training pipeline of video Slot Attention models, (b) SAVi's mixture-based decoder [46], (c) STEVE's Transformer-based decoder [84], and (d) our LDM-based decoder.
Unsupervised object-centric learning from images. Our work is directly related to research that aims at learning to represent visual data with a set of feature vectors without explicit supervision. Earlier attempts start from synthetic datasets with naive shape primitives, uniformly colored objects, and simple backgrounds [6, 17, 18, 19, 22, 23, 60]. They typically perform iterative inference to extract object-centric features from images, followed by a mixture-based CNN decoder applied to each slot separately for reconstruction. For example, MONet [6] and Slot Attention [60] both adopt the spatial broadcast decoder [97] to predict an RGB image and an objectness mask from each slot, and combine them via alpha masking. Recently, SLATE [83] challenges the traditional design with a Transformer-based decoder, which helps the model scale to more complex data such as textured objects and backgrounds. Such decoder has also been employed in later works [79, 85], proving its effectiveness on naturalistic inputs. Another line of works [1, 30, 81, 99] focuses on feature-level reconstruction or contrastive learning for object discovery without pixel-level decoding. However, they are not applicable for generation tasks as they cannot reconstruct images from slots. With the recent progress in 3D reconstruction, there are also methods leveraging Neural Radiance Fields (NeRF) [89, 104] and light fields [79, 86] as the slot decoder. However, they usually require multiview images and camera poses as training data, which is incompatible with our setting. Compared to previous works, SlotDiffusion is the first attempt introducing diffusion models to unsupervised object-centric learning, achieving state-of-the-art results in both object discovery and generation.
Unsupervised object-centric learning from videos. Compared to images, videos provide additional information such as motion cues. Our work also builds upon recent efforts in decomposing raw videos into temporally aligned object slots [40, 46, 48, 56, 93, 106]. These works usually integrate a dynamics module to model the interactions between objects compared to their image-based counterparts. For example, STOVE [49] adopts a Graph Neural Network (GNN), while OAT [9], SAVi [46], and PARTS [110] leverage the powerful Transformer architecture. However, the limited modeling capacity of their mixture-based decoders still prevents them from scaling to more realistic videos. Some works [16, 102] thus introduce additional supervisions such as optical flow, depth maps, and initial object positions. Recently, STEVE [84] pushes the limit of fully unsupervised object-centric learning on videos with the Transformer slot decoder from SLATE and the Transformer dynamics module from SAVi. Nevertheless, the autoregressive generation mechanism of Transformer degrades its generation capacity, which we aim to address in this paper. There are also methods learning scene decomposition from videos by contrasting feature [45, 61], which cannot generate videos. Similar to SAVi and STEVE, SlotDiffusion can also be applied to video data with the Transformer-based dynamics module. In addition, we generate videos with much better temporal consistency.