SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziyi Wu $^{1,2}$  Jingyu Hu $^{1,*}$  Wuyue Lu $^{1,*}$  Igor Gilitschenski $^{1,2}$  Animesh Garg $^{1,2}$ $^{1}$ University of Toronto  $^{2}$ Vector Institute
Object-centric learning aims to represent visual data with a set of object entities (a.k.a. slots), providing structured representations that enable systematic generalization. Leveraging advanced architectures like Transformers, recent approaches have made significant progress in unsupervised object discovery. In addition, slot-based representations hold great potential for generative modeling, such as controllable image generation and object manipulation in image editing. However, current slot-based methods often produce blurry images and distorted objects, exhibiting poor generative modeling capabilities. In this paper, we focus on improving slot-to-image decoding, a crucial aspect for high-quality visual generation. We introduce SlotDiffusion – an object-centric Latent Diffusion Model (LDM) designed for both image and video data. Thanks to the powerful modeling capacity of LDMs, SlotDiffusion surpasses previous slot models in unsupervised object segmentation and visual generation across six datasets. Furthermore, our learned object features can be utilized by existing object-centric dynamics models, improving video prediction quality and downstream temporal reasoning tasks. Finally, we demonstrate the scalability of SlotDiffusion to unconstrained real-world datasets such as PASCAL VOC and COCO, when integrated with self-supervised pre-trained image encoders. Additional results and details are available at our website.
Humans perceive the world by identifying discrete concepts such as objects and events [88], which serve as intermediate representations to support high-level reasoning and systematic generalization of intelligence [24]. In contrast, modern deep learning models typically represent visual data with patch-based features [14, 28], disregarding the compositional structure of scenes. Inspired by the human perception system, object-centric learning aims to discover the modular and causal structure of visual inputs. This approach thus holds the potential to improve the generalizability and interpretability of AI algorithms [50, 80]. For example, explicitly decomposing scenes into conceptual entities facilitates visual reasoning [8, 12, 13] and causal inference [80, 109]. Also, capturing the compositional structure of the world proves beneficial for both image generation [83, 90] and video prediction [100, 103].
However, unsupervised object-centric learning from raw perceptual input is challenging, as it requires capturing the appearance, position, and motion of each object. Earlier attempts at this problem usually bake in strong scene or object priors in their frameworks [19, 40, 55, 56], limiting their applicability to synthetic datasets [5, 42]. Later works extend the Scaled Dot-Product Attention mechanism [94] and introduce a general Slot Attention operation [60], which eliminates domain-specific priors. With a simple input reconstruction objective, these models learn to segment objects from raw images [60] and videos [46]. To date, state-of-the-art unsupervised slot models [83, 84] have shown promising scene decomposition results on naturalistic data under specialized scenarios such as traffic videos.
Despite tremendous progress in object segmentation, the generative capability of slot-based methods has been underexplored. For example, while the state-of-the-art object-centric model STEVE [84]
is able to decompose complex data into slots, it reconstructs videos with severe object distortion and temporal inconsistency from slots [100]. Moreover, our own experiments reveal that all popular slot models [46, 60, 83] suffer from low slot-to-image decoding quality, preventing their adoption in visual generation such as image editing. Therefore, the goal of this work is to improve the generative capacity of object-centric models on complex data while preserving their segmentation performance.
We propose SlotDiffusion, an unsupervised object-centric model with a Latent Diffusion Model (LDM) [73] based slot decoder. We consider slots extracted from images as basic visual concepts (i.e., objects with different attributes), akin to text embeddings of phrases from sentences. The LDM decoder learns to denoise image feature maps guided by object slots, providing training signals for object discovery. Thanks to the expressive power of diffusion models, SlotDiffusion achieves a better trade-off between scene decomposition and visual generation. We evaluate our method on six synthetic and three real-world datasets consisting of complex images and videos. SlotDiffusion improves object segmentation results and achieves significant advancements in compositional generation. Crucially, we demonstrate that SlotDiffusion seamlessly integrates with recent development in object-centric learning, such as dynamics modeling [100] and the integration of pre-trained image encoders [81].
Our main contributions are as follows: (i) SlotDiffusion: a diffusion model based object-centric learning method. (ii) We apply SlotDiffusion to unsupervised object discovery and visual generation, where it achieves state-of-the-art results across both image and video datasets. (iii) We demonstrate that our learned slots can be directly utilized by the state-of-the-art object-centric dynamics model, improving future prediction and temporal reasoning performance. (iv) We further scale up SlotDiffusion to real-world datasets by integrating it with a self-supervised pre-trained image encoder.
In this section, we provide a brief overview of related works on unsupervised object-centric learning and diffusion models, which is further expanded in Appendix B.
Unsupervised object-centric learning aims to discover the underlying structure of visual data and represent it with a set of feature vectors (a.k.a. slots). To achieve unsupervised scene decomposition, most methods use input reconstruction as the training signal. Earlier works often perform iterative inference to extract object slots, followed by a CNN-based decoder for reconstruction [6, 17, 23, 40, 55, 56]. AIR [19] and SQAIR [48] use a patch-based decoder to decode each object in a canonical pose, and then transform them back to the original position. Slot Attention [60] and SAVi [46] adopt the spatial broadcast decoder [97] to predict the RGB images and segmentation masks from each slot, which are combined via alpha masking. However, the per-slot CNN decoder limits the modeling capacity of these methods to synthetic datasets [5, 42]. Recently, SLATE [83] and STEVE [84] proposed a Transformer-based slot decoder. They first pre-train a dVAE to tokenize the input, and then train the slot-conditioned Transformer decoder to reconstruct patch tokens in an autoregressive manner. The powerful Transformer architecture and the feature-level reconstruction target enable them to segment naturalistic images and videos [79, 81, 85]. However, as observed in [83, 100], the slot-to-image reconstruction quality of Transformer decoders is low on complex data, sometimes even underperform CNN decoders. In this work, we leverage diffusion models [34, 73] as the slot decoder, which runs iterative denoising conditioned on slots. Thanks to their expressive power, our method not only improves scene decomposition, but also generates results with higher visual quality.
Diffusion models (DMs) [34, 87] have recently achieved tremendous progress in image generation [11, 64, 71, 77], showing their great capability in sample quality and input conditioning. The generative process of DMs is formulated as an iterative denoising procedure from Gaussian noise to clean data, where the denoiser is typically implemented as a U-Net [74]. However, the memory and computation requirements of DMs scale quadratically with the input resolution due to the self-attention layers in the U-Net. To reduce the training cost, LDM [73] proposes to first down-sample images to feature maps with a pre-trained VAE encoder, and then run the diffusion process in this low-resolution latent space. LDM also introduces cross-attention as a flexible mechanism for conditional generation. For example, text-guided LDMs [73] perform cross-attention between text embeddings and U-Net's feature maps at multiple resolutions to guide the denoising process. In this work, we adopt LDM as the slot decoder due to its strong generation capacity, where the conditioning is achieved by cross-attention between the denoising feature maps and the object slots.
While conventional DMs excel in controllable generation such as image editing [31, 44] and compositional generation [57, 75], they often require specific supervision such as text to guide the generation
Figure 1: SlotDiffusion architecture overview. Given an initial frame of a video  $x_0$ , we initialize slots from a set of learnable vectors  $\tilde{S}_0$ , and perform Slot Attention with image features to update object slots  $S_0$ . During training, a U-Net denoiser predicts the noise  $\epsilon_t$  added to the image tokens  $z_0$  conditioned on slots via cross-attention. The entire model is applied recurrently to all video frames with a Predictor initializing future slots from current step.
process. In contrast, SlotDiffusion discovers meaningful visual concepts from data in an unsupervised manner, which can be used for controllable generation as will be shown in the experiments.
Concurrent work. Recently, Jiang et al. [41] also employ LDM in object-centric learning. However, they utilize an image tokenizer pre-trained on external datasets, while all components of our model are trained purely on target datasets. Moreover, their work only evaluates on images, while we test SlotDiffusion on both images and videos. In addition, we combine our model with recent development in object-centric learning, showing our potential in dynamics modeling and handling real-world data.
3 SlotDiffusion: Object-Oriented Diffusion Model
In this section, we describe SlotDiffusion by first reviewing previous unsupervised object-centric learning methods (Section 3.1). Then, we introduce our slot-conditioned diffusion model decoder (Section 3.2). Finally, we show how to perform compositional generation with the visual concepts learned by SlotDiffusion (Section 3.3). The overall model architecture is illustrated in Figure 1.
3.1 Background: Slot-based Object-Centric Learning
Unsupervised object-centric learning aims to represent a scene with a set of object slots without object-level supervision. Here, we review the SAVi family [46, 84] as they achieve state-of-the-art object discovery on videos. SAVi builds upon Slot Attention [60] and runs on videos in a recurrent encoder-decoder manner. Given  $T$  input frames  $\{\pmb{x}_t\}_{t=1}^T$ , SAVi first leverages an encoder to extract per-frame features, adds positional encodings, and flattens them into a set of vectors  $\pmb{h}_t = f_{\mathrm{enc}}(\pmb{x}_t) \in \mathbb{R}^{M \times D_{\mathrm{enc}}}$ . Then, the model initializes  $N$  slots  $\tilde{S}_t \in \mathbb{R}^{N \times D_{\mathrm{slot}}}$  from a set of learnable vectors  $(t = 1)$ , and updates them with Slot Attention as  $S_t = f_{\mathrm{SA}}(\tilde{S}_t, \pmb{h}_t)$ .  $f_{\mathrm{SA}}$  performs soft feature clustering, where slots compete with each other to capture certain areas of the input via iterative attention [94]. To achieve temporally aligned slots, SAVi leverages a Transformer-based predictor to initialize  $\tilde{S}_t$ $(t \geq 2)$  as  $\tilde{S}_t = f_{\mathrm{pred}}(S_{t-1})$ . Finally, the model utilizes a decoder  $f_{\mathrm{dec}}$  to reconstruct input  $\pmb{x}_t$  from slots  $S_t$  and uses the reconstruction loss as training signal. Below we detail the slot decoder  $f_{\mathrm{dec}}$  studied in previous works as we focus on improving it in this paper, and refer the readers to [46] for more information about other model components. See Appendix A for figures of different slot decoders.
Mixture-based decoder. Vanilla SAVi [46] adopts the Spatial Broadcast decoder [97] consisting of a stack of up-sampling deconvolution layers. It first broadcasts each slot  $S_{t}^{i}$  to a 2D grid to form a feature map. Each feature map is then decoded to an RGB image  $\pmb{y}_{t}^{i}$  and an alpha mask  $\pmb{m}_{t}^{i}$ , which are combined into the final reconstructed image  $\hat{\pmb{x}}_t$ . SAVi is trained with an MSE reconstruction loss:
$$
\left(\boldsymbol {y} _ {t} ^ {i}, \boldsymbol {m} _ {t} ^ {i}\right) = f _ {\operatorname {d e c}} ^ {\operatorname {m i x}} \left(\mathcal {S} _ {t} ^ {i}\right), \hat {\boldsymbol {x}} _ {t} = \sum_ {i = 1} ^ {N} \boldsymbol {m} _ {t} ^ {i} \odot \boldsymbol {y} _ {t} ^ {i}, \mathcal {L} _ {\text {i m a g e}} = \sum_ {t = 1} ^ {T} \left\| \boldsymbol {x} _ {t} - \hat {\boldsymbol {x}} _ {t} \right\| ^ {2}. \tag {1}
$$
Transformer-based decoder. The above mixture-based decoder has limited modeling capacity as it decodes each slot separately without interactions. Also, pixel-level reconstruction biases the model to low-level color statistics, which only proves effective on objects with uniform colors, and cannot
scale to complex data with textured objects. Current state-of-the-art model STEVE [84] thus proposes to reconstruct intermediate features produced by another network [83]. Given frame  $\boldsymbol{x}_t$ , STEVE leverages a pre-trained dVAE encoder to convert it into a sequence of patch tokens  $\boldsymbol{o}_t = \{\boldsymbol{o}_t^i\}_{i=1}^L$  which serve as the reconstruction targets for the autoregressive Transformer decoder:
$$
\boldsymbol {o} _ {t} = f _ {\text {e n c}} ^ {\mathrm {d V A E}} \left(\boldsymbol {x} _ {t}\right), \hat {\boldsymbol {o}} _ {t} ^ {l} = f _ {\text {d e c}} ^ {\text {t r a n s}} \left(\mathcal {S} _ {t}; \boldsymbol {o} _ {t} ^ {1}, \dots , \boldsymbol {o} _ {t} ^ {l - 1}\right), \mathcal {L} _ {\text {t o k e n}} = \sum_ {t = 1} ^ {T} \sum_ {l = 1} ^ {L} \operatorname {C r o s s E n t r o p y} \left(\boldsymbol {o} _ {t} ^ {l}, \hat {\boldsymbol {o}} _ {t} ^ {l}\right). \tag {2}
$$
Thanks to the cross-attention mechanism in the Transformer decoder and the feature-level reconstruction objective, STEVE works on naturalistic videos with textured objects and backgrounds. However, the generation quality of the Transformer slot decoder is still low on complex data [83, 100].
3.2 Slot-conditioned Diffusion Model
Object-centric generative models [90, 110] often decompose the generation process to first predicting object slots in the latent space, and then decoding slots back to the pixel space. Here, the image-to-slot abstraction and the slot-to-image decoding are often performed by a pre-trained object-centric model. Therefore, the generation quality is greatly affected by the scene decomposition result and the slot decoder capacity. As observed in [83, 100], methods with mixture-based decoders usually generate images with blurry objects, as they fail to capture objects from complex data. While models with Transformer-based decoders are better at scene decomposition, they still produce results of low visual quality. We attribute this to two issues: (i) treating images as sequences of tokens ignores their inherent spatial structure; (ii) autoregressive token prediction causes severe error accumulation. In this work, we overcome these drawbacks by introducing diffusion models as the slot decoder, which preserves the spatial dimension of images, and iteratively refines the generation results. From now on, we omit  $t$  as the video timestamp, and will only use  $t$  to denote the timestep in the diffusion process.
Diffusion model. Diffusion models (DMs) [34, 87] are probabilistic models that learn a data distribution  $p_{\theta}(\mathbf{X}_0)$  by gradually denoising a standard Gaussian distribution, in the form of  $p_{\theta}(\mathbf{X}_0) = \int p_{\theta}(\mathbf{X}_{0:T}) d\mathbf{X}_{1:T}$ , where  $\mathbf{X}_{1:T}$  are intermediate denoising results. The forward process of DMs is a Markov Chain that iteratively adds Gaussian noise to the clean data  $\mathbf{X}_0$ , which is controlled by a pre-defined variance schedule  $\{\beta_t\}_{t=1}^T$ . Let  $\alpha_t = 1 - \beta_t$  and  $\bar{\alpha}_t = \prod_{s=1}^t \alpha_s$ , we have:
$$
q \left(\boldsymbol {X} _ {t} \mid \boldsymbol {X} _ {t - 1}\right) = \mathcal {N} \left(\boldsymbol {X} _ {t} \mid \sqrt {1 - \beta_ {t}} \boldsymbol {X} _ {t - 1}, \beta_ {t} \boldsymbol {I}\right) \Rightarrow q \left(\boldsymbol {X} _ {t} \mid \boldsymbol {X} _ {0}\right) = \mathcal {N} \left(\boldsymbol {X} _ {t} \mid \sqrt {\bar {\alpha} _ {t}} \boldsymbol {X} _ {0}, (1 - \bar {\alpha} _ {t}) \boldsymbol {I}\right). \tag {3}
$$
During training, a denoising model  $\epsilon_{\theta}(\mathbf{X}_t,t)$  is trained to predict the noise applied to a noisy sample:
$$
\boldsymbol {X} _ {t} = \sqrt {\bar {\alpha} _ {t}} \boldsymbol {X} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon_ {t}, \quad \mathcal {L} _ {\mathrm {D M}} = | | \epsilon_ {t} - \epsilon_ {\theta} (\boldsymbol {X} _ {t}, t) | | ^ {2}, \text {w h e r e} \epsilon_ {t} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}). \tag {4}
$$
At inference time, we can start from a random Gaussian noise, and apply the trained denoiser to iteratively refine the sample. See Appendix C for a detailed formulation of diffusion models.
SlotDiffusion. Our model consists of the same model components as SAVi and STEVE, while only replacing the slot decoder with the DM-based one. Inspired by the success of feature-level reconstruction in STEVE, we adopt the Latent Diffusion Model (LDM) [73] as the decoder, which denoises features in the latent space. This improves the segmentation results with a higher-level reconstruction target, and greatly reduces the training cost as it runs at a lower resolution. Specifically, we pre-train a VQ-VAE [72] to extract feature maps  $\mathbf{z} \in \mathbb{R}^{h \times w \times D_{\mathrm{vq}}}$  from  $\mathbf{x}$  before training SlotDiffusion:
$$
\boldsymbol {z} = f _ {\mathrm {e n c}} ^ {\mathrm {V Q}} (\boldsymbol {x}), \hat {\boldsymbol {x}} = f _ {\mathrm {d e c}} ^ {\mathrm {V Q}} (\boldsymbol {z}), \mathcal {L} _ {\mathrm {V Q}} = \left\| \boldsymbol {x} - \hat {\boldsymbol {x}} \right\| ^ {2} + \operatorname {L P I P S} (\boldsymbol {x}, \hat {\boldsymbol {x}}), \tag {5}
$$
where  $\mathrm{LPIPS}(\cdot ,\cdot)$  is the VGG perceptual loss [108]. In preliminary experiments, we also tried KL-VAE [73] and VQ-GAN [20] as the image tokenizer, but did not observe clear improvements.
To condition the LDM decoder on slots  $S$ , we notice that slots are a set of  $N$  feature vectors, which are similar to text embeddings output by language models [69, 70]. Therefore, we follow recent text-guided LDMs [73] to guide the denoising process via cross-attention:
$$
\boldsymbol {c} = \operatorname {C r o s s A t t e n t i o n} (Q (\tilde {\boldsymbol {c}}), K (\mathcal {S}), V (\mathcal {S})). \tag {6}
$$
Here,  $Q$ ,  $K$ ,  $V$  are learnable linear projections,  $\tilde{c}$  is an intermediate feature map from the denoising U-Net  $\epsilon_{\theta}$ , and  $c$  is the feature map updated with slot information. An important property of crossattention is that the result in Eq. (6) is order-invariant to the conditional input  $S$ , thus preserving the permutation-equivariance of slots, which is a key property of object-centric models. In practice, we perform conditioning after several layers in the U-Net at multiple resolutions. Overall, our model is trained with a slot-conditioned denoising loss over VQ-VAE feature maps:
$$
\boldsymbol {z} _ {t} = \sqrt {\bar {\alpha} _ {t}} \boldsymbol {z} + (1 - \bar {\alpha} _ {t}) \epsilon_ {t}, \quad \mathcal {L} _ {\text {s l o t}} = \left\| \epsilon_ {t} - \epsilon_ {\theta} \left(\boldsymbol {z} _ {t}, t, \boldsymbol {S}\right) \right\| ^ {2}, \text {w h e r e} \epsilon_ {t} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}). \tag {7}
$$
3.3 Compositional Generation with Visual Concepts
Previous conditional DMs [71, 73, 77] rely on labels such as text to control the generation process. In contrast, SlotDiffusion is conditioned on slots output by Slot Attention, which is trained jointly with the DM. We consider each object slot extracted from an image as an instance of a basic visual concept (e.g., red cube). Therefore, we can discover visual concepts from unlabeled data, and build a library of slots for each of them. Then, we can compose concepts via slots to generate novel samples.
Visual concept library. We follow SLATE [83] to build visual concept libraries by clustering slots. Take image datasets for example, we first collect slots extracted from all training images, and then apply K-Means clustering to discover  $K$  clusters, each serves as a concept. All slots assigned to a cluster form a library for that concept. On video datasets, we discover concepts by clustering slots extracted from the first frame. As shown in our experiments, SlotDiffusion discovers semantically meaningful concepts, such as objects with different shapes, and different components of human faces.
To generate new samples, we first select  $N$  concepts from  $K$  candidates, pick one slot from each concept's library, and then decode them with our LDM-based slot decoder. To avoid severe occlusions, we reject slots with object segmentation masks overlapping greater than an mIoU threshold.
SlotDiffusion is a generic unsupervised object-centric learning framework that is applicable to both image and video data. We conduct extensive experiments to answer the following questions: (i) Can we learn object-oriented scene decomposition supervised by the denoising loss of DMs? (Section 4.2) (ii) Will our LDM-based decoder improve the visual generation quality of slot models? (Section 4.3) (iii) Is the object-centric representation learned by SlotDiffusion useful for downstream dynamics modeling tasks? (Section 4.4) (iv) Can we extend our method to handle real-world data? (Section 4.5) (v) Can SlotDiffusion benefit from other recent improvements in object-centric learning? (Section 4.6) (vi) What is the impact of each design choice on SlotDiffusion? (Section 4.7)
Datasets. We evaluate our method in unsupervised object discovery and slot-based visual generation on six datasets, namely, the two most complex image datasets CLEVRTex [43] and CelebA [59] from SLATE [83], and four video datasets MOVi-D/E/Solid/Tex [25] from STEVE [84]. Then, we show SlotDiffusion's capability for downstream video prediction and reasoning tasks on Physion [4]. Finally, we scale our method to unconstrained real-world images on PASCAL VOC 2012 [21] and MS COCO 2017 [54]. We briefly introduce each dataset below and provide more details in Appendix D.1.
CLEVRTex augments the CLEVR [42] dataset with more diverse object shapes, materials, and textures. The backgrounds also present complex textures compared to the plain gray one in CLEVR. We train our model on the training set with 40k images, and test on the test set with 5k samples.
CelebA contains over 200k real-world celebrity images. All images are mainly occupied by human faces, with roughly front-view head poses. This dataset is more challenging as real-world images typically have background clutter and complicated lighting conditions. We train our model on the training set with around 160k images, and test on the test split with around 20k images.
MOVi-D/E/Solid/Tex consist of videos generated using the Kubric simulator [25]. Their videos feature photo-realistic backgrounds and diverse real-world objects, where one or several objects are thrown to the ground to collide with other objects. We follow the official train-test split to evaluate our model.
Physion is a VQA dataset containing realistic simulations of eight physical scenarios, such as dropping and soft-body deformation. The goal of Physion is to predict whether two target objects will contact as the scene evolves. Following the official protocol, we first train models to learn scene dynamics and evaluate the video prediction results, and then perform VQA on the model's future rollout.