SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Diffusion model for generation. Recently, diffusion models [87] (DMs) have achieved tremendous progress in generation tasks, including 2D images [11, 34, 64, 65, 71, 77], 2D videos [27, 35, 36, 82], 3D shapes [53, 58, 66], and scene editing [31, 44, 63, 75], showing their great capacity in generation quality and input conditioning. The generative process of DMs is formulated as an iterative denoising procedure with a denoising network, usually implemented as a U-Net [74]. However, the memory consumption of DMs scales quadratically with the input resolution due to the use of self-attention layers in the U-Net. To reduce the training cost, LDM [73] proposes to run the diffusion process in the latent space of a pre-trained image auto-encoder, e.g., VQ-VAE [72]. LDM also introduces a flexible conditioning mechanism via cross-attention between the U-Net feature maps and the conditional inputs. For example, recent text-guided LDMs [73] first extract text embeddings using pre-trained language models [69, 70], and then perform cross-attention between the language features and the U-Net denoising feature maps at multiple resolutions. In this work, we adopt the LDM as our slot decoder due to its strong generation capacity and low computation cost. To condition the decoder on 1D object slots, we draw an analogy with 1D text embeddings, and also apply cross-attention between slots and feature maps in the U-Net denoiser.
Diffusion model for representation learning. There have been a few works studying the representation learned by DMs. Baranchuk et al. [3] and Xu et al. [101] utilize the denoising feature maps of DMs for data-efficient segmentation, as the feature maps already capture object boundary information. Preechakul et al. [68] jointly trains a CNN encoder and a DM to perform auto-encoding for representation learning. DiffMAE [98] combines MAE [29] and DMs to learn better inpainting. Li et al. [52] leverages pre-trained text-guided DMs for zero-shot image classification. In contrast, SlotDiffusion does not rely on any supervision or pre-trained DMs. Instead, we incorporate it with the unsupervised Slot Attention framework to learn structured object-centric representations.
Conditional diffusion model without labels. Several works have explored learning conditional DMs in an unsupervised manner. Hu et al. [37] and Bao et al. [2] utilize self-supervised feature extractors to obtain image features, and then perform clustering to create pseudo labels such as cluster ids for training a conditional DM. Harvey and Wood [26] show that conditional DMs usually outperform unconditional DMs in terms of FID. Inspired by this observation, they leverage CLIP embeddings [69] of unlabeled images to train conditional DMs without manual annotations. Similarly, SlotDiffusion also learns a conditional DM without labels. Instead of relying on pre-trained feature extractors, we build upon the unsupervised Slot Attention framework, and can learn both object-centric representations and conditional DMs from scratch.
C Details on Diffusion Models
Diffusion models 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}$ . Here,  $\mathbf{X}_{1:T}$  are intermediate denoising results with the same shape as the clean data  $\mathbf{X}_0 \sim q(\mathbf{X})$ , and  $\theta$  are learnable parameters of the denoising U-Net model.
The joint distribution  $q(\mathbf{X}_{1:T}|\mathbf{X}_0)$  is called the forward process or diffusion process, which is a fixed Markov Chain that gradually adds Gaussian noise to  $\mathbf{X}_0$ . The noise is controlled by a pre-defined variance schedule  $\{\beta_t\}_{t=1}^T$ :
$$
q \left(\boldsymbol {X} _ {1: T} \mid \boldsymbol {X} _ {0}\right) = \prod_ {t = 1} ^ {T} q \left(\boldsymbol {X} _ {t} \mid \boldsymbol {X} _ {t - 1}\right), \quad 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). \tag {8}
$$
Thanks to the nice property of Gaussian distributions,  $\mathbf{X}_t$  can be sampled directly from  $\mathbf{X}_0$  in closed form without adding the noise  $t$  times. 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} _ {0}\right) = \mathcal {N} \left(\boldsymbol {X} _ {t} \mid \sqrt {\bar {\alpha} _ {t}} \boldsymbol {X} _ {0}, (1 - \bar {\alpha} _ {t}) \boldsymbol {I}\right) \Rightarrow \boldsymbol {X} _ {t} = \sqrt {\bar {\alpha} _ {t}} \boldsymbol {X} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \epsilon_ {t}, \text {w h e r e} \epsilon_ {t} \sim \mathcal {N} (\boldsymbol {0}, \boldsymbol {I}). \tag {9}
$$
We can now train a model to reverse this process and thus generate target data from random noise  $\mathbf{X}_T \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$ . The reverse process  $p_{\theta}(\mathbf{X}_{0:T})$  is also defined as a Markov Chain with a learned Gaussian transition:
$$
p _ {\theta} \left(\boldsymbol {X} _ {0: T}\right) = p \left(\boldsymbol {X} _ {T}\right) \prod_ {t = 1} ^ {T} p _ {\theta} \left(\boldsymbol {X} _ {t - 1} \mid \boldsymbol {X} _ {t}\right), \quad p _ {\theta} \left(\boldsymbol {X} _ {t - 1} \mid \boldsymbol {X} _ {t}\right) = \mathcal {N} \left(\boldsymbol {X} _ {t - 1} \mid \mu_ {\theta} \left(\boldsymbol {X} _ {t}, t\right), \Sigma_ {\theta} \left(\boldsymbol {X} _ {t}, t\right)\right) \tag {10}
$$
In practice, we do not learn the variance and usually set it to  $\Sigma_{t} = \beta_{t}\pmb{I}$  or  $\frac{1 - \bar{\alpha}_{t - 1}}{1 - \bar{\alpha}_t}\beta_t\pmb{I}$  since it leads to unstable training [34]. Also, instead of learning the mean  $\mu_{\theta}$  directly, we learn to predict the noise  $\epsilon_{t}$  in Equation (8). See Ho et al. [34] for how we can sample  $X_{t - 1}$  given  $X_{t}$  and the predicted  $\epsilon_{t}$  at the inference stage.
The training process of diffusion models is thus straightforward given Equation (9). At each training step, we sample a batch of clean data  $\mathbf{X}_0$  from the training set, timesteps  $t$  uniformly from  $\{1,\dots,T\}$ , and random Gaussian noise  $\epsilon_t \sim \mathcal{N}(\mathbf{0},\mathbf{I})$ . We then create the noisy version of data  $\mathbf{X}_t$  by applying Equation (9). A denoising model  $\epsilon_{\theta}$  is trained to predict the noise with an MSE loss:
$$
\mathcal {L} _ {D M} = \mathbb {E} _ {\boldsymbol {X}, t, \epsilon_ {t}} [ \| \epsilon_ {t} - \epsilon_ {\theta} (\boldsymbol {X} _ {t}, t) \| ^ {2} ] \tag {11}
$$
D Detailed Experimental Setup
In this section, we detail the datasets, evaluation metrics, baseline methods used in the experiments, and the implementation details of our model.
CLEVRtex [43]. This dataset augments the CLEVR [42] dataset with more diverse object shapes, materials, and textures. The backgrounds in CLEVRtex also present complex textures compared to the plain gray ones in CLEVR. Therefore, this dataset is visually much more complex than CLEVR. As shown in their paper, only 3 out of 10 benchmarked unsupervised object-centric models can achieve an mIoU higher than  $30\%$ . We train our model on the training set consisting of 40k images, and test on the test set with 5k samples. We use the same data pre-processing steps, i.e. first center-crop to  $192 \times 192$ , and then resize to  $128 \times 128$ .
CelebA [59]. This dataset contains over 200k real-world celebrity images. All images are mainly occupied by human faces, covering a large variation of poses and backgrounds. This dataset is more challenging compared to previous synthetic datasets 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. For data pre-processing, we simply resize all images to  $128 \times 128$ .
MOVi-D/E [25]. MOVi-D and MOVi-E are the 2 most challenging versions from the MOVi benchmark generated using the Kubric simulator. Their videos feature photo-realistic backgrounds and real-world objects from the Google Scanned Objects (GSO) repository [15], where one or several objects are thrown to the ground to collide with other objects. Compared to MOVi-D, MOVi-E applies linear camera motion. We follow the official train-test split to evaluate our model. For data pre-processing, we simply resize all frames to  $128 \times 128$ .
MOVi-Solid/Tex [84]. MOVi-Solid and MOVi-Tex share similar backgrounds and object motions with MOVi-D and MOVi-E. Besides, MOVi-Solid introduces more complex shapes, and MOVi-Tex maps richer textures to the object surfaces, making them harder to be distinguished from the background. We follow the train-test split in [84], and resize all frames to  $128 \times 128$ .
Physion [4]. This dataset contains eight categories of videos, each showing a common physical phenomenon, such as rigid- and soft-body collisions, falling, rolling, and sliding motions. The foreground objects used in the simulation vary in categories, textures, colors, and sizes. It also uses diverse background as the scene environment, and randomize the camera pose in rendering the videos. Physion splits the videos into three sets, namely, Training, Readout Fitting, and Testing. We truncate all videos by 150 frames as most of the interactions end before that, and sub-sample the videos by a factor of 3 for training the dynamics model. Following the official evaluation protocol, the dynamics models are first trained on videos from the Training set under future prediction loss. Then, they observe the first 45 frames of videos in the Readout Fitting and Testing set, and perform rollout to generate future scene representations (e.g., image feature maps or object slots). A linear readout model is trained on observed and rollout scene representations from the Readout Fitting set to classify whether the two cued objects (one in red and one in yellow) contact. Finally, the classification accuracy of the trained readout model on the Testing set scene representations is reported. Please refer to their paper [4] for detailed descriptions of the VQA evaluation. For video prediction, we evaluate on the test set videos. All video frames are resized to  $128 \times 128$  as pre-processing.
PASCAL VOC 2012 [21]. PASCAL VOC is a real-world image dataset containing natural scenes with diverse objects. We follow Seitzer et al. [81] to train on the "trainaug" variant with 10,582 images, and test on the official validation set with 1,449 images. For data pre-processing, we resize the smaller dimension of the image to 224, and perform random cropping to  $224 \times 224$  followed by random horizontal flip during training. For testing, we evaluate the  $224 \times 224$  center crop of images, and ignore pixels that are unlabeled.
Microsoft COCO 2017 [54]. Similar to VOC, MS COCO is a real-world benchmark commonly used in object detection and segmentation. We follow Seitzer et al. [81] to train on the training set with 118,287 images, and test on the validation set with 5,000 images. For data pre-processing, we resize the smaller dimension of the image to 224, and perform center crop to  $224 \times 224$  for both training and testing. We apply random horizontal flip during training. During evaluation, we filter out crowd instance annotations, and ignore pixels where objects overlap.
To evaluate the generation quality of SlotDiffusion, we compute the mean squared error (MSE) of the images reconstructed from the object slots. Following prior works, we scale the images to [0, 1], and sum the errors over channel and spatial dimensions. As pointed out by Zhang et al. [108], MSE does not align well with human perception of visual quality as it favors over-smooth results. Therefore, we additionally compute the VGG perceptual distance (LPIPS) [108]. For video generation, we compute the per-frame MSE and LPIPS, and then average them over the temporal dimension. Finally, to evaluate the compositional generation results, we adopt the Fréchet Inception Distance  $(FID)^{1}$  [32] and the Fréchet Video Distance  $(FVD)^{2}$  [92]. We compute these metrics between 5k generated images or 500 generated videos and the entire training set of that dataset.
<sup>1</sup>Implementation from https://github.com/mseitzer/pytorch-fid.
Implementation from https://github.com/universome/stylegan-v.
Table 5: Variations in model architectures and training settings on different datasets.
Dataset	CLEVRTex	CelebA	MOVi-D/E	MOVi-Solid/Tex	Physion	VOC / COCO
Image Encoder fenc	ResNet18	ResNet18	ResNet18	Plain CNN	ResNet18	ViT-S/8
Number of Slots N	11	4	15	12	8	6/7
Slot Size Dslot	192	192	192	192	192	192/256
Slot Attention Iteration	3	3	2	2	2	3
Max Learning Rate	2e-4	2e-4	1e-4	1e-4	1e-4	1e-4
Gradient Clipping	1.0	1.0	0.05	0.05	0.05	0.05
Batch Size	64	64	32	32	48	64
Training Epochs	100	50	30	30	10	500/100
To evaluate the scene decomposition results, we compute the FG-ARI, mIoU, and mBO of the object masks. FG-ARI is widely used in previous object-centric learning papers, which only consider foreground objects. As suggested in [17, 43], we should also compute the mIoU which evaluates background segmentation. For mBO, it first assigns each ground-truth object mask with the max-overlapping predicted mask, and then averages the IoU of all pairs of masks. It is worth noting that on video datasets, we flatten the temporal dimension into spatial dimensions when computing the metrics, thus taking the temporal consistency of object masks into consideration, i.e., slot ID swapping of the same object will degrade the performance. Being able to consistently track all the objects is an important property of video slot models.
Object-centric models. We adopt Slot Attention [60] and SAVi [46] as representative models which use a mixture-based CNN decoder, and SLATE [83] and STEVE [84] for models with a Transformer-based decoder. We use the unconditional version of these models, i.e., they are not trained with the initial frame supervision proposed in [46]. Since we augment SlotDiffusion with a stronger ResNet18 encoder compared to the previous stacked CNN encoder, we also re-train baselines with the same ResNet18 encoder, which achieves better performance than reported in their papers. In addition, we tried scaling up the slot decoder of baselines by adding more layers and channels, but observed severe training instability beyond a certain point. Instead, SlotDiffusion shows great scalability with model sizes. Its performance grows consistently with a larger LDM U-Net. On real-world datasets PASCAL VOC and COCO, we directly copy the numbers of DINOSAUR [81] from its paper, as they did not release their code until 15 days before the paper submission deadline. For Slot Attention and SLATE here, we replace the CNN encoder with the same DINO [7] pre-trained ViT encoder. We trained longer enough and achieved better results than reported in the DINOSAUR paper.
Video prediction and VQA. We adopt the same baselines as SlotFormer [100]. Please refer to their papers about the implementation details of these future prediction models. We train baselines on Physion with the default hyper-parameters. For VQA, we directly copy the numbers from SlotFormer.
D.4 Our Implementation Details
VQ-VAE [72]. We use the same architecture for all datasets, which is adopted from LDM [73]. We use 3 encoder and decoder blocks, resulting in  $4\mathrm{x}$  down-sampling of the feature maps  $\mathcal{Z}$  compared to input images  $\pmb{x}$ . We pre-train the VQ-VAE for 100 epochs on each dataset with a cosine learning rate schedule decaying from  $1\times 10^{-3}$ , and fix it during the object-centric model training.
SlotDiffusion. For the object-centric model, we only replace the decoder with LDM [73] compared to Slot Attention on images and SAVi on videos. We use a modified ResNet18 encoder [46] to extract image features<sup>3</sup>, except for real-world datasets PASCAL VOC and COCO, we use the DINO self-supervised pre-trained [7] encoder, where the ViT-S/8 variant is chosen following DINOSAUR<sup>4</sup> [81]. On video datasets, we train on small clips of 3 frames. For the LDM-based slot decoder, the training target is to predict the noise  $\epsilon$  added to the features  $z$  produced by a pre-trained VQ-VAE. Following prior works, the denoising network  $\epsilon_{\theta}(z_t,t,S)$  is implemented as a U-Net [74] with global self-attention and cross-attention layers in each block. We use the same noise schedule  $\{\beta_t\}_{t=1}^T$  and U-Net hyper-parameters as LDM [73]. The attention map from the last Slot Attention iteration is used as the predicted segmentation. See Table 5 for detailed slot configurations and training settings.
<sup>3</sup>Interestingly, we found that ResNet18 encoder leads to training divergence on MOVi-Solid and MOVi-Tex datasets for both SlotDiffusion and baselines. Therefore, we adopt the plain CNN encoder on these two datasets.
4Pre-trained weight obtained from the HuggingFace transformers library.
Table 6: Comparison of model complexity on the MOVi-D/E video datasets. We measure the training memory consumption, time per training step, and generation time of 100 videos at test stage. For training, we report the default settings (batch size 32 of length-3 video clips, frame resolution  $128 \times 128$ ) on NVIDIA A40 GPUs. For testing, we report the inference time on NVIDIA T4 GPUs (each video contains 24 frames).
SAVi	STEVE	SlotDiffusion
Train	Memory (GB)	32	51	24
Time (s)	0.57	0.87	0.77
Test	Time (min)	0.7	226	7
Figure 13: Visualization of image segmentation results on CLEVRTex (left) and CelebA (right).
Slot-to-image decoding. Diffusion models are notoriously slow in doing generation due to the iterative denoising process, where the diffusion step  $T$  is often 1,000. Fortunately, researchers have developed several methods to accelerate this procedure. We employ the DPM-Solver [62] which reduces the sampling steps to 20. Therefore, our reconstruction speed is even faster than models with Transformer-based decoders, as they need to forward their model number of tokens (typically 1,024) times for autoregressive token prediction, while we only need 20 times.
SlotFormer [100] on Physion. We use its online official implementation<sup>5</sup>. We first pre-train a SlotDiffusion model on Physion, and use it to extract object slots from all videos and save them offline. Then, we train SlotFormer to predict slots at future timesteps conditioned on slots from input timesteps. We follow the same model architecture and training settings as [100]. For video prediction, we use the pre-trained LDM-based slot decoder to decode the predicted slots back to frames, so that we can evaluate the generation quality. For VQA, we train additional readout models over the predicted slots to answer questions. Please refer to [100] for more information.
E.1 Computational Requirements
We empirically show the speed and GPU memory requirement at training time, as well as the time required to reconstruct 100 videos at test time of SAVi, STEVE, and SlotDiffusion in Table 6. Interestingly, our model requires the least GPU memory during training despite achieving the best performance. This is because we run the diffusion process at latent space, whose spatial dimension is only  $1/4$  of the input resolution. In contrast, SAVi applies CNN to directly reconstruct images at the original resolution, and STEVE uses an autoregressive Transformer to predict a long sequence of 1,024 tokens. Both of these designs consume large GPU memory. In terms of training and generation speed, SlotDiffusion ranks second. SAVi runs extremely fast since it decodes images in one-shot, while STEVE and our model both need to do iterative sampling. Thanks to the efficient DPM-Solver sampler, we only require 20 times forward pass, while STEVE requires 1,024 times.
Total training time. SAVi is the fastest to train which takes 30 hours on a single NVIDIA A40 GPU. STEVE and SlotDiffusion both need 20 hours to train the VAE image tokenizer. STEVE's second stage training requires 50 hours on 2 GPUs, while SlotDiffusion requires 40 hours on 1 GPU.
E.2 Object Segmentation
Image. We visualize our image segmentation results in Figure 13. SlotDiffusion is able to discover objects or semantically meaningful and consistent regions from images in an unsupervised way. For
5 https://github.com/pairlab/SlotFormer.
Figure 14: Visualization of video segmentation results on MOVi-D (left) and MOVi-E (right).
example, on CelebA, we decompose images into four parts: clothes, hairs, faces, and backgrounds. As shown in the section of compositional generation, we can generate novel samples by randomly composing these basic visual concepts. Notably, the object masks produced on CLEVRTex usually do not match the exact boundary of objects. This is because the segmentation masks are predicted in a low resolution (4x down-sampling in our case) and up-sampled to the original resolution. How to refine the object masks from a coarse initialization is an interesting future direction.
Video. We show more video segmentation results in Figure 14. SAVi typically generates blob-shape masks that do not follow the exact boundary of objects due to the low modeling capacity of the mixture-based decoder. Compared to STEVE, SlotDiffusion usually has more consistent tracking results and fewer object-sharing issues, especially on large objects. We also note that the shown examples of SlotDiffusion on MOVi-D have FG-ARI scores higher than  $80\%$ , although their visual quality is far from perfect. This indicates that the FG-ARI metric has been saturated, and we should focus more on better metrics such as mIoU [43].
E.3 Compositional Generation
Reconstruction. As in the main paper, we first show more reconstruction results in Figure 15. Compared to baselines that produce blurry frames, SlotDiffusion is able to preserve the local appearances such as textures on the objects and backgrounds. However, there is still a large room for improvement in finer details. For example, in the top example, we cannot reconstruct the texts on the object surfaces. Also, in the bottom example, we fail to retain the smooth surface of the brown shoes.
Generation. In Figure 16, we compare original dataset samples to novel results produced by baselines and SlotDiffusion. Due to the limited capacity, models with the mixture-based decoder generate blurry images. Although the Transformer-based decoder helps SLATE and STEVE decompose complex scenes, they typically produce objects with distorted attributes, such as deformed object shapes and unnatural facial features. Thanks to the iterative refinement process of DMs, SlotDiffusion is able to generate high-fidelity images and videos with fine details, coherent with the real samples.
Remark: the generation diversity of SlotDiffusion's LDM decoder. One good property of DMs is that they can generate diverse results from a single conditional input [11, 66, 73]. This is because the conditional inputs do not fully describe the output's content. For example, in text-guided generation, the input texts usually only describe the global style and layout of the images, while LDMs have the freedom to generate diverse local details. In contrast, SlotDiffusion conditions the LDM on slots, which are designed to capture all information (position, shape, texture, etc.) about objects. Therefore, we expect our LDM to faithfully reconstruct the original objects, instead of generating diverse results, which is also a desired property in our downstream applications. For example, when performing face replacement (Figure 8 (b) in the main paper), we want to maintain the facial attributes of the source face. Nevertheless, we note that SlotDiffusion is still able to generate novel content when necessary. For example, when removing an object from the scene in image editing (Figure 8 (a) in the main paper), it inpaints the background with plausible textures.