SlowGuess/NeurIPS_2023_OCR huggingface dataset:
As illustrated by the LSD samples in Figure 6, certain parts of the background, such as the upper regions of the images, are split and assigned to multiple object slots. This behavior leads to information mixing between objects and the background, which substantially reduces the performance of slot segmentation and representation. We conjecture that this issue may stem from overfitting in the model. Due to the simplicity of the texture and the absence of objects in these areas, the powerful decoder of LSD can easily memorize the rendering process of those regions independently of the slot-conditioning. As a result, it does not provide a sufficient training signal for the encoder to group those background parts into a meaningful slot.
In light of these findings, we tested a straightforward approach to address this problem. We propose introducing additional complexity to the background by incorporating data samples from the CLEVRTex dataset into the training process. By integrating the CLEVRTex samples, we introduce 60 different background textures, which makes implicit memorization difficult. This effectively provides the learning signal to capture the background component into the slot-conditioning. The results of this approach, labeled as LSD(Mix), can be seen in Figure 6. As clearly illustrated, including CLEVRTex data samples in the training results in cleaner background segmentation and alleviates the information mixing issue. As a result, LSD(Mix) achieves approximately  $22\%$  gains of improvement in both mIoU and mBO compared to the original model trained solely on CLEVR samples, as shown in Table 3. However, we also see that the FID score indicated a decrease in the generation quality of
Figure 6: Visualization of Unsupervised Object Segmentation on CLEVR dataset. The LSD(Mix) variant corresponds to the model trained on both CLEVR and CLEVRTex training images.
LSD(Mix) compared to the original LSD, likely due to the increased demand of model capacity to model both datasets. Overall, this experiment confirms that introducing complexity into the training dataset can mitigate the overfitting problem and allow LSD to effectively perform representation learning on simple datasets.
C LSD with Pre-Trained Diffusion Models
One of the main focuses of recent research in diffusion models have been on adapting pre-trained models for additional control signals [47, 57, 88]. Specifically, these studies propose the integration of a learnable adapter network into large-scale pre-trained diffusion models, such as Stable Diffusion, to enable controllable generation based on various types of conditioning such as sketches, semantic layouts, and depth maps. An inherent advantage of these approaches is their ability to bypass the resource-intensive processes of training the entire generative model, while still enabling the production of high-quality controllable images through minimal training of the adapter network.
We note that such a model design can also be utilized for slot-based conditioning. However, unlike existing works that primarily focus on learning the adapter to integrate the provided conditioning into the pre-trained diffusion model, our focus lies in training an adapter to achieve two objectives simultaneously: (1) we aim to learn how to extract object representations (slots) from input images in an unsupervised manner, and (2) we aim to learn how to incorporate these slot representations into the pre-trained diffusion model to facilitate slot-based conditional generation. This leads to an intriguing question: can learning through such a model, which has not been specifically trained on multi-object
Table 3: Quantitative Evaluation on CLEVR dataset. We evaluate the models by reporting mBO, mIoU and FG-ARI scores for segmentation quality, property prediction scores for representation quality, and FID scores for compositional generation quality.
SLATE	SLATE+	LSD	LSD (Mix)
mBO (↑)	64.86	67.42	38.49	61.09
mIoU (↑)	63.96	66.62	37.49	59.76
FG-ARI (↑)	69.20	88.56	76.32	76.30
Shape (↑)	95.66	95.70	75.16	80.48
Material (↑)	97.62	97.96	94.21	96.24
Position (↓)	0.59	0.51	0.86	0.99
FID (↓)	52.96	20.93	16.22	19.45
Figure 7: Qualitative Results of Control-LSD on CLEVR, CLEVRTex, MOVi-E, and COCO. The CLEVR(Mix) corresponds to the model trained on both CLEVR and CLEVRTex samples.
scenarios, provide the necessary learning signal to acquire object-centric representations? In this section, we provide a preliminary exploration on this question by introducing two LSD variants: Control-LSD and Latent Slot Stable Diffusion (Stable-LSD).
C.1 Control-LSD
Model Design Control-LSD utilizes a similar design of the adapter structure proposed in ControlNet [88] and consists of three main components: (1) a pre-trained Stable Diffusion model that is frozen throughout the training, (2) a "control block" that is a trainable copy of the unet encoder from Stable Diffusion, and (3) a object-centric encoder from the LSD model. Similar to LSD, the object-centric encoder takes an image as input and generates a set of slots, which are then conditioned on using a Slot-Conditioned Transformer. However, instead of providing the slots to the diffusion decoder, Control-LSD feeds the slots into the trainable control block while solely employing null text conditioning in the Stable Diffusion. To integrate the slot control signal, the outputs of each intermediate control block are element-wise added to the intermediate blocks and skip-connections within the frozen Stable Diffusion model. Following LSD, we use a input resolution of  $256 \times 256$  and attention map resolution of  $64 \times 64$  for all datasets. We also use the same number of slots as in LSD for multi-object dataset, and 10 slots for the COCO dataset.
Results We evaluate Control-LSD on three multi-object datasets: CLEVR, CLEVRTex, and MOVi-E, as well as one real-world dataset, COCO [48]. Following the proposed LSD, we also include a CLEVR(Mix) variant for the CLEVR dataset which was trained on both CLEVR and CLEVRTex samples. The segmentation results are presented in Figure 7. We can see that Control-LSD struggles to achieve satisfactory object segmentation on the CLEVR dataset. The model has difficulties in distinguishing between the background and the objects. We also note that the model demonstrate
more reasonable results on CLVERTex and introducing samples from CLEVRTex to the training process of CLEVR (i.e., CLEVR(Mix)) also led to an improvement on CLEVR, but the results are still inferior to those achieved by LSD. Interestingly, we also observe that Control-LSD achieve better segmentation quality when applied to more realistic scenes, such as those in the MOVi-E dataset. We conjecture that this is attributed to the fact that the complexity of MOVi-E aligns more closely with real-world images used for pre-training the Stable Diffusion model, as compared to the samples from CLEVR and CLEVRTex datasets. This indicates that the distribution gap between the data used for pre-training Stable Diffusion and the multi-object datasets used for learning object-centric representation might be a potential contributing factor to the lower performance of Control-LSD on the multi-object datasets.
Finally, we also evaluate Control-LSD on a real-world multi-object dataset COCO [48]. Notably, Control-LSD showcases its ability to capture semantically meaningful entities in the slots, as illustrated in the segmentation results in Figure 7. However, we find it challenging to obtain accurate object segmentations aligned with COCO annotations. In this regard, we would like to share two key observations when applying the model on real-world scenes: (1) Without human annotations, unsupervised models have difficulty distinguishing between discrete parts of an object and the entire object. This is known as the part-whole ambiguity. For instance, while COCO annotates an entire human as one segment, Control-LSD might segment it into two distinct parts: the upper body and the legs. It is important to emphasize that neither of the two segmentation results are incorrect, as each could be beneficial for specific downstream tasks. (2) The segmentation granularity is sensitive to the number of slots. This sensitivity can be attributed to the model's tendency to always utilize the entire set of slots to segment the image. Consequently, if the slot count is less than the number of entities in an image, the model tends to produce semantic masks rather than object-level segments. Conversely, if the slots outnumber the entities, the model tends to over-segment the image, i.e., splitting some of the entity into multiple slots. This behavior limits the applicability of Control-LSD in real-world data, where the true number of objects is not known to the model. Therefore, we believe further investigation is needed for applying the model to real-world applications.
C.2 Stable-LSD
While Control-LSD bypasses the need to train a Diffusion decoder, the computational cost of training the adapter network—which includes both the control block and the object encoder—is still substantial. As an illustration, the memory demands are so significant that even with four 48GB GPUs, our experiments could only handle a training batch size of 32, which also leads to significant training latency. To address this problem, we introduce the second LSD variant which further reduce trainable modules, we name it Latent Slot Stable Diffusion, or Stable-LSD for short. We show that despite its simple architecture, Stable-LSD demonstrate strong performance on real-world scenarios. We will start with describing the model design.
Model Design Firstly, following the DINOSAUR model [69], we employ a pre-trained image encoder, DINOv2 [59], as the CNN backbone for the object encoder. This encoder remains frozen during training. It is important to emphasize that unlike DINOSAUR, Stable-LSD is a generative model capable of generating images. Secondly, the learned slots are provided directly to the cross attention layers of pre-trained DMs as special text tokens. This means that the slot attention module acts dual purposes — as an object encoder extracting object-centric representations and as a feature adapter converting input tokens from the image to the text domain. Notably, within the Stable-LSD framework, only the slot attention module requires training. The resulting design significantly reduces the computational demands observed with Control-LSD. In our experiment, we were able to train the model with a batch size of 128 using two 48GB GPUs.
Real-World Segmentation Results Despite its simplicity, Stable-LSD demonstrate its effectiveness in real-world object-centric learning. To evaluate its object-centric learning ability, we compare Stable-LSD with DINOSAUR on instance-level FG-ARI and mBO. The result can be found in Table 4. The results show that Stable-LSD achieves comparative performance to DINOSAUR with relatively superior FG-ARI value and a diminished mBO value. We show the mask visualization on Figure 8.
4The Stable-LSD implementation will be available in the official release: https://github.com/JindongJiang/latent-slot-diffusion
Table 4: Quantitative Evaluation of Stable-LSD on COCO dataset. We compare Stable-LSD with DINOSAUR [69] on instance-level FG-ARI and mBO. The results of DINOSAUR is copied from the original paper.
Stable-LSD (Ours)	DINOSAUR
mBO (↑)	30.38	31.60
FG-ARI (↑)	35.02	34.10
Figure 8: Unsupervised Image Segmentation and Generation Results of Stable-LSD on COCO dataset. For image generation, we apply different CFG values on the learned slots with  $\text{cfg} = 1.0$  equivalent to image reconstruction.
Real-World Image Generation Results We further evaluate the image generation quality of Stable-LSD. In Figure 8, we show the image reconstruction results (denoted by  $\mathrm{cfg} = 1.0$ ). The results demonstrate that the model output suffers from visual artifacts which constrained both the image reconstruction and generation capabilities. To further improve the generation results, we explore the technique of classifier-free guidance which is widely employed in text-to-image diffusion models.
Classifier-free guidance [34] uses a mixing weight (abbreviated as  $cfg$ ) to control the weighted sum of noise predictions between a conditional and an unconditional diffusion model. With  $cfg = 1$ , it operates similarly to standard conditional generation. When  $cfg > 1$ , the model generates images that are more align to the text input, but comes with a cost of reduced image diversity. Practically speaking, an optimal  $cfg > 1$  can yield images that not only better match the conditions but also exhibit enhanced quality. Since Stable-LSD provides object slots as text tokens to the pre-trained DMs, classifier-free guidance can be applied directly to the input slot, akin to how it's applied to input text in standard DMs, without any model modifications. And our experiment shows that a  $cfg > 1$  for slots can also significantly improve the image generation quality.
In Figure 8, we tested various CFG values with Stable-LSD to optimize image generation quality. As we can see in the figure, from  $\mathrm{cfg} = 1.0$  to 1.3, the images appear cleaner with less visual artifacts,
Figure 9: Real-World Image Generations of Stable-LSD. Each image is conditioned on the same set of learned slots and with a different initial noise map in the denoising process. We use  $\text{cfg} = 1.3$  for all samples.
delivering conceptually more consistent images with the input slots. Nevertheless, we also observe that higher CFG values strip away intricate visual details, leading to overly simplified images. We settled on a CFG value of 1.3 for following experiments.
In Figure 9, we show image samples generated based on the same set of learned slots. The variance in these samples comes from differing the initial noise maps in the denoising process. These results demonstrate that with a proper CFG value, Stable-LSD achieves slot-based conditional image generation with unconstrained real-world objects, which is for the first time in object-centric generative models. In addition, beyond showing the diversity of the generated images, these results also underscore a crucial point from the representation learning prospective: the learned slots contain rich semantic information rather than merely performing texture-based clustering. An evidence to this is the model's ability to maintain overall semantic consistency while exhibiting diverse presentations of the same semantic content across samples. For instance, even if the background of the stop sign images shows different tree shapes, they all clearly have trees in the background. Achieving this level of consistency would be challenging if the slots lacked an abstract semantic understanding of the objects in the scene. We show additional real-world segmentation and generation results in Figures 18 - 23.
We believe the preliminary attempts detailed in this section to be valuable explorations toward large-scale object-centric learning, and we hope our insights pave the way for further research in this area.
Table 5: Full Table of Segmentation Performance and Representation Quality. Extended results from the main text, now including standard deviation values.
(a) CLEVRTex
Segmentation	SLATE	SLATE+	LSD (Ours)	Representation	SLATE	SLATE+	LSD (Ours)
mBO (↑)	50.88 ± 0.70	54.90 ± 1.15	63.93 ± 1.93	Shape (↑)	74.24 ± 2.41	71.63 ± 4.94	80.23 ± 1.59
mIoU (↑)	49.54 ± 0.70	52.96 ± 1.40	62.52 ± 1.88	Material (↑)	69.73 ± 2.04	63.61 ± 2.40	75.56 ± 1.00
FG-ARI (↑)	44.19 ± 1.80	70.71 ± 3.35	64.41 ± 8.53	Position (↓)	1.28 ± 0.10	1.26 ± 0.15	1.13 ± 0.04
(b) MOVi-C
Segmentation	SLATE	\( SLATE^+ \)	LSD (Ours)	Representation	SLATE	\( SLATE^+ \)	LSD (Ours)
mBO (↑)	39.37 ± 0.80	38.17 ± 1.29	45.57 ± 0.80	Position (↓)	1.37 ± 0.06	1.28 ± 0.07	1.14 ± 0.03
mIoU (↑)	37.75 ± 0.71	36.44 ± 1.35	44.19 ± 0.91	3D B-Box (↓)	1.48 ± 0.03	1.44 ± 0.11	1.44 ± 0.02
FG-ARI (↑)	49.54 ± 1.44	52.04 ± 0.45	51.98 ± 3.53	Category (↑)	42.45 ± 0.13	45.32 ± 1.13	46.11 ± 0.59
(c) MOVi-E
Segmentation	SLATE	\( SLATE^+ \)	LSD (Ours)	Representation	SLATE	\( SLATE^+ \)	LSD (Ours)
mBO (↑)	30.17 ± 1.70	22.17 ± 0.38	38.96 ± 0.47	Position (↓)	2.09 ± 0.13	2.15 ± 0.07	1.85 ± 0.05
mIoU (↑)	28.59 ± 1.66	20.63 ± 0.35	37.64 ± 0.45	3D B-Box (↓)	3.36 ± 0.12	3.37 ± 0.27	2.94 ± 0.003
FG-ARI (↑)	46.06 ± 3.32	45.25 ± 0.77	52.17 ± 0.89	Category (↑)	38.93 ± 0.16	38.00 ± 0.36	42.96 ± 0.21
To ensure reproducibility, we are sharing additional implementation details of LSD in this section. Our implementation will be made available at https://github.com/JindongJiang/ latent-slot-diffusion
D.1 Experiment Setup
Datasets. In our experiments, we create three data subsets for each dataset: training, validation, and testing. For CLEVR, we utilize the official split for training and validation sets. However, for evaluation, we do not use the official test set as CLEVR does not provide ground-truth object segmentation masks. Instead, we use the "clevr_with_mask" subset from the Multi-Object Datasets [38] for evaluation. This subset is designed to have the same object distribution as CLEVR while providing ground-truth masks and attributes annotations for each object in the image, which allows us to evaluate the segmentation and representation quality of each model. We use the full subset as the test set for evaluation. For CLEVRTex, there is no official split provided, so we allocate  $80\%$  of the data for training,  $10\%$  for validation, and  $10\%$  for testing. Regarding MOVi-C and MOVi-E, we use  $90\%$  of the training set data for training and reserve  $10\%$  for validation. We use the official validation set for testing. It is important to note that the test sets in MOVi-C and MOVi-E are designed for out-of-distribution (OOD) evaluation, so we do not use them for test time evaluation. For the FFHQ dataset, we use  $86\%$  of the dataset ( $\sim 60\mathrm{K}$  images) for training and  $7\%$  ( $\sim 5\mathrm{K}$  images) for validation. To ensure consistency across all models, we shuffled the dataset and utilized a fixed random seed when splitting the data.
Segmentation Metrics. To evaluate the unsupervised object segmentation, we use three metrics: the foreground adjusted rand index (FG-ARI), the mean intersection over union (mIoU), and the mean best overlap (mBO). We predict the object mask by using the argmax along the slot dimension in the attention map. In computing the mIoU, we use Hungarian matching to obtain the ground-truth and slots assignment. Following [15], we use the negative cosine similarity as the matching loss. In mBO, we assign the ground-truth segment to the slot with the largest overlap. Note that the mBO metric is generally simpler than mIoU since it does not require a strict one-to-one mapping between the ground-truth and predicted masks.
We will provide an overview of the implementation details in this section. The hyperparameters used in our approach are listed in Table 6.
Auto-Encoding CNN backbone. The object-centric encoder in our model includes a CNN backbone, which transforms the raw image input into feature vectors that can be processed by the slot attention
Module	Hyperparameter	Dataset
CLEVR	CLEVRTex	MOVi-C/E	FFHQ
General	Batch Size	64	64	128	128
Training Steps	200K	200K	200K	200K
# K-means Clusters	5	5	18	18
CNN Backbone	Input Resolution	256
Output Resolution	64
Self Attention	Middle Layer
Base Channels	128
Channel Multipliers	[1,1,2,4]
# Heads	8
# Res Blocks / Layers	2
Output Channels	128	128	192	192
Learning Rate	0.0001	0.00003	0.0001	0.0001
Slot Attention	Input Resolution	64
# Iterations	3
Slot Size	128	128	192	192
# Slots	11	11	24	4
Learning Rate	0.0001	0.00003	0.0001	0.0001
Auto-Encoder	Model	KL-8
Input Resolution	256
Output Resolution	32
Output Channels	4
LSD Decoder	Input Resolution	32
Input Channels	4
β scheduler	Linear
Mid Layer Attention	Yes
# Res Blocks / Layers	2
Image Latent Scaling	0.18215
Learning Rate	0.0001
# Heads	4	4	8	8
Base Channels	128	128	192	192
Attention Resolution	[1,2,4]	[1,2,4]	[1,2,4,4]	[1,2,4,8]
Channel Multipliers	[1,2,4]	[1,2,4]	[1,2,4,4]	[1,2,4,4]
Table 6: Hyperparameters of our model used in our experiments.
module. Previous studies have used multi-layer CNN networks as image encoders, as demonstrated by [53, 70, 73]. However, in our experiments with Latent Diffusion Models, we have found that such encoders tend to produce low-level features that mainly represent pixel position or color information. This leads to the slot attention module grouping pixels based on these low-level features, rather than on high-level object information. We hypothesize that this problem may arise from the early stages of training, where the model lacks a semantic understanding of the input image. This results in the LSD Decoder relying only on low-level and local information from slots to denoise the noisy input, without encouraging the object-centric encoder to learn semantic-level pixel grouping. As a result, without additional guidance, the model may become trapped in a suboptimal solution. We have observed this issue only when a Diffusion Decoder is used.
To overcome this issue, we have incorporated a UNet architecture [64] as the image encoder in the object-centric encoder. The main advantage of this approach is that the auto-encoding structure of UNet allows high-level global context information to be blended into the output features and slots, facilitating the learning process and helping the model escape from suboptimal pixel groupings. Additionally, the skip connections in UNet allow the output features to retain fine-grained low-level details from early CNN layers, ensuring that the resulting representations are rich in both high-level object information and low-level texture information. These improvements to the object-centric
encoder result in more accurate object-centric representations and enable better image generation results. We provide detailed design information in Table 6. Prior to applying the UNet structure, we first employ a single CNN layer with a kernel size of 4 and stride of 4 to downsample the image from  $256 \times 256 \times 3$  to  $64 \times 64 \times S$ , where  $S$  is the base channel size in UNet. In Appendix E.1, we provide an additional ablation study for the UNet encoder.
Slot-Conditioned Cross Attention. We implement the conditioning mechanism between the object slots and decoder features through a transformer network as in LDM [63] and SLATE [70]. This conditioning mechanism consists of two transformer layers, including a self-attention layer that computes self-attention within the features generated by the UNet in the LSD Decoder, followed by a cross-attention layer that computes cross-attention between the object slots and UNet features.
In the cross-attention layer, we use the object slots  $\mathbf{S}$  as key-value pairs, and the queries are intermediate layers of the diffusion network augmented with positional embedding  $(\tilde{\mathbf{h}}_l + \mathbf{p}_l)$ . We then apply multi-head cross attention for each head as Attention  $(\mathbf{Q},\mathbf{K},\mathbf{V}) = \text{softmax}\left(\frac{\mathbf{Q}\mathbf{K}^T}{\sqrt{d}}\right)\cdot V$ , where
$$
\mathbf {Q} = \mathbf {W} _ {Q} ^ {(i)} \cdot (\tilde {\mathbf {h}} _ {l} + \mathbf {p} _ {l}), \mathbf {K} = \mathbf {W} _ {K} ^ {(i)} \cdot \mathbf {S}, \mathbf {V} = \mathbf {W} _ {V} ^ {(i)} \cdot \mathbf {S}.
$$
Here,  $i$  is the head index,  $\tilde{\mathbf{h}}_l$  represents the intermediate layers of the diffusion network,  $\mathbf{p}_l$  is the positional embedding,  $\mathbf{S}$  is the set of object slots computed from the object-centric encoder, and  $\mathbf{W}_V^{(i)},\mathbf{W}_Q^{(i)},\mathbf{W}_K^{(i)}$  are learnable projection matrices.