SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Jindong Jiang
jindong.jiang@rutgers.edu
Gautam Singh
singh.gautam@rutgers.edu
The recent success of transformer-based image generative models in object-centric learning highlights the importance of powerful image generators for handling complex scenes. However, despite the high expressiveness of diffusion models in image generation, their integration into object-centric learning remains largely unexplored in this domain. In this paper, we explore the feasibility and potential of integrating diffusion models into object-centric learning and investigate the pros and cons of this approach. We introduce Latent Slot Diffusion (LSD), a novel model that serves dual purposes: it is the first object-centric learning model to replace conventional slot decoders with a latent diffusion model conditioned on object slots, and it is also the first unsupervised compositional conditional diffusion model that operates without the need for supervised annotations like text. Through experiments on various object-centric tasks, including the first application of the FFHQ dataset in this field, we demonstrate that LSD significantly outperforms state-of-the-art transformer-based decoders, particularly in more complex scenes, and exhibits superior unsupervised compositional generation quality. In addition, we conduct a preliminary investigation into the integration of pre-trained diffusion models in LSD and demonstrate its effectiveness in real-world image segmentation and generation. Project page is available at https://latentslotdiffusion.github.io
The fundamental structure of the physical world is compositional and modular. While this structure is naturally revealed in some data modalities like language in the form of tokens or words, in other modalities such as images, it is elusive how one may discover this structure. However, this representational compositionality and modularity is essential for various applications that require the systematic manipulation of knowledge pieces to achieve high-level cognitive abilities. This includes reasoning [46, 5], causal inference [68], and out-of-distribution generalization [2, 22].
Object-centric learning [27] aims to discover the latent compositional structure from unstructured observation by learning to bind relevant features, thereby forming useful tokens in an unsupervised way. For images, one of the most popular approaches is to auto-encode the image via the Slot Attention [53] encoder. Slot Attention applies competitive spatial attention to partition the image into separate local areas and then obtain a representation, called a slot, from each area. Then, a decoder generates the image from the slots with the aim of minimizing the reconstruction error. Due to the limited capacity and competition between slots, each slot is encouraged to capture a reusable and compositional entity, such as an object.
The primary challenge that remains in the framework of unsupervised object-centric learning is making it work for complex naturalistic scene images. Until recently, most object-centric models have adopted a special type of decoder known as the mixture decoder. While a weak slot-wise decoder [81] is predominantly employed in this mixture decoder due to its efficacy in promoting the emergence of object-centric representation in relatively simple scene images, further studies [70, 73] have shown that this strong prior can make handling complex naturalistic scene images more challenging. Contrary to conventional belief, Singh et al. [70] recently proposed departing from this low-capacity mixture-decoder approach and suggested to use an expressive transformer-based autoregressive image generator in object-centric learning [70, 73]. It was shown that increasing the decoder capacity is crucial for handling complex and naturalistic scenes in this framework [8, 7, 69, 84, 71].
The success of transformer-based image generative modeling in object-centric learning naturally leads to the question: can diffusion models, another pillar of modern deep generative modeling known for their highly expressive generation capabilities, also benefit object-centric learning? Diffusion models [75, 31] are based on a stochastic denoising process and have demonstrated impressive performance in a variety of image generation tasks [61, 58, 65, 14, 63, 66, 56], sometimes surpassing transformer-based autoregressive models. Moreover, diffusion models possess unique modeling capabilities that transformer-based autoregressive models cannot provide [75]. However, despite their potential, the applicability of diffusion models to unsupervised object-centric learning remains largely unexplored. Consequently, it is crucial to examine the feasibility of this approach and to identify the associated benefits and limitations.
In this paper, we address this question by introducing a novel model called Latent Slot Diffusion (LSD). The LSD model can be interpreted from two perspectives. From the perspective of object-centric learning, LSD can be viewed as the first model substituting conventional slot decoders with a conditional latent diffusion model, in which the conditioning is object-centric slots provided by Slot Attention. From the diffusion model perspective, our approach is the first unsupervised compositional conditional diffusion model. While traditional conditional diffusion models [61, 63, 65, 52] require supervised annotations, such as a text description of an image for compositional generation, LSD is a diffusion model that enables the construction of such compositional descriptions in terms of visual concepts extracted from images through unsupervised object-centric learning.
In our experiments, we evaluate the proposed model across various object-centric tasks, including unsupervised object segmentation, downstream property prediction, compositional generation, and image editing. We show that the LSD model delivers significantly superior performance compared to the state-of-the-art model, namely, the transformer-based autoregressive generative model. A notable attribute of the proposed model is that LSD's performance advantage over autoregressive transformers increases as the scene complexity increases. In particular, LSD enables exploring, for the first time, the applicability of object-centric model to the FFHQ dataset [41], a collection of high-resolution and high-quality face images that surpasses the generative capabilities of existing object-centric models. Additionally, we discuss the overfitting issue faced by LSD on very simple scene images, such as those in the CLEVR dataset [37], and offer suggestions for addressing this problem.
2 Latent Slot Diffusion
In this section, we outline the auto-encoding framework of our proposed model, Latent Slot Diffusion or LSD $^2$ , beginning with our object-centric encoder and subsequently describing our proposed decoder.
2.1 Object-Centric Encoder
Given an input image  $\mathbf{x} \in \mathbb{R}^{H \times W \times C}$ , our object-centric encoder seeks to decompose and represent it as a collection of  $N$  vectors or slots  $\mathbf{S} \in \mathbb{R}^{N \times D}$ , where each slot (denoted as  $\mathbf{s}_n \in \mathbb{R}^D$ ) is encouraged to represent a compositional entity in the image. For this, we adopt Slot Attention, an architecture that is also used in the current state-of-the-art object-centric learning approaches [53, 70, 43]. We now describe how Slot Attention works in our model.
In Slot Attention, we first encode the input image  $\mathbf{x}$  as a set of  $M$  input features  $\mathbf{E} \in \mathbb{R}^{M \times D_{\mathrm{input}}}$  via a backbone network  $f_{\phi}^{\mathrm{backbone}}$ , i.e.,  $\mathbf{E} = f_{\phi}^{\mathrm{backbone}}(\mathbf{x})$ . The network  $f_{\phi}^{\mathrm{backbone}}$  is implemented as a CNN
$^{2}$ Code will be made available at https://github.com/JindongJiang/latent-slot-diffusion
Figure 1: Method. Left: In training, we encode the given image as image latent and as slots. We then add noise to the image latent and we train a denoising network to predict the noise given the noisy latent and the slots. Right: Given the trained model, we can generate novel images by composing a slot-based concept prompt and decoding it using the trained latent slot diffusion decoder.
(detailed in Appendix D.2) whose final output feature map is flattened to form a set. Next, the features in  $\mathbf{E}$  are grouped into  $N$  spatial groupings and the information in each grouping is aggregated to produce a slot. The grouping is achieved via an iterative slot refinement procedure. At the start of the refinement procedure, the slots  $\mathbf{S}$  are filled with random Gaussian noise. Then, they are refined via competitive attention over the input features, where the  $N$  slots act as the queries and the  $M$  input features act as the keys and values. The queries and keys undergo a dot-product to produce  $N\times M$  attention proportions. Next, on these attention proportions, the softmax function is applied along the axis  $N$  to produce attention weights  $\mathbf{A}$  that capture the soft assignment of each input feature to a slot. Then, for each  $n$ , all input features are sum-pooled weighted by their attention weights  $\mathbf{A}_{n,1},\ldots ,\mathbf{A}_{n,M}$  to produce an attention readout  $\mathbf{u}_n\in \mathbb{R}^D$ . These steps can be formally described as follows:
$$
\mathbf {A} = \underset {N} {\mathrm {s o f t m a x}} \left(\frac {q (\mathbf {S}) \cdot k (\mathbf {E}) ^ {T}}{\sqrt {D}}\right) \Rightarrow \mathbf {A} _ {n, m} = \frac {\mathbf {A} _ {n , m}}{\sum_ {m = 1} ^ {M} \mathbf {A} _ {n , m}} \Rightarrow \mathbf {u} _ {n} = \sum_ {m = 1} ^ {M} v (\mathbf {E} _ {m}) \mathbf {A} _ {n, m},
$$
where,  $q,k,v$  are linear projections that map the slots and input features to a common dimension  $D$ . Using the bottom-up information captured by the readout  $\mathbf{u}_n$ , the slots are updated by an RNN as  $\mathbf{s}_n = f_{\phi}^{\mathrm{RNN}}(\mathbf{s}_n,\mathbf{u}_n)$ . In practice, competitive attention and RNN update are performed iteratively several times and slots from the last iteration are considered the final slot representation  $\mathbf{S}$ .
2.2 Latent Slot Diffusion Decoder
In this section, we describe our proposed decoding approach called Latent Slot Diffusion Decoder or LSD decoder for reconstructing the image given the slot representation S. The design of the LSD decoder takes advantage of the recent advances in generative modeling based on diffusion [63, 31]. An overview of our decoding approach is provided in Figure 1.
2.2.1 Pre-Trained Image Auto-Encoder
One of the key design components of the LSD decoder is a pre-trained auto-encoder (AE) [21, 63], which provides a way to map an image  $\mathbf{x}$  to a lower-dimensional latent representation  $\mathbf{z}_0$  via an encoder  $f_{\phi}^{\mathrm{AE}}$ . This enables LSD to reduce the computational burden of reconstructing high-resolution images by using the lower-dimensional latent  $\mathbf{z}_0$  as an intermediate reconstruction target. It also allows us to later obtain the original resolution image without compromising image content and fidelity by decoding the latent  $\mathbf{z}_0$  using the AE decoder  $g_{\theta}^{\mathrm{AE}}$ . These can be summarized as:
$$
\mathbf {z} _ {0} = f _ {\phi} ^ {\mathrm {A E}} (\mathbf {x}), \qquad \hat {\mathbf {x}} = g _ {\theta} ^ {\mathrm {A E}} (\mathbf {z} _ {0}), \qquad \text {w h e r e} \mathbf {z} _ {0} \in \mathbb {R} ^ {H _ {\mathrm {A E}} \times W _ {\mathrm {A E}} \times D _ {\mathrm {A E}}}  , \text {a n d} \hat {\mathbf {x}} \in \mathbb {R} ^ {H \times W \times C}.
$$
In our model, we employ an auto-encoder pre-trained on OpenImages<sup>3</sup>.
<sup>3</sup>We use the pre-trained weights from https://ommer-lab.com/files/latent-diffusion/k1-f8.zip
2.2.2 Slot-Conditioned Diffusion
In LSD, we leverage diffusion modeling to reconstruct the image latent  $\mathbf{z}_0$  conditioned on the slots S. This modeling approach has been primarily explored in supervised contexts, e.g., for text-to-image generation in Latent Diffusion Models (LDM) [63]. However, unlike LDM, in this work, instead of conditioning the decoder on embeddings of supervised labels, we condition it on slots where the process of obtaining the slots themselves, i.e., Slot Attention, is jointly trained with the decoder without supervision. Following LDM, our decoder works by training a decoding distribution  $p_{\theta}(\mathbf{z}_0|\mathbf{S})$  to maximize the log-likelihood  $\log p_{\theta}(\mathbf{z}_0|\mathbf{S})$  of the image latent  $\mathbf{z}_0$  given the slots S. This decoding distribution  $p_{\theta}(\mathbf{z}_0|\mathbf{S})$  is modeled as a  $T$ -step denoising process:
$$
p _ {\theta} (\mathbf {z} _ {0} | \mathbf {S}) = \int p (\mathbf {z} _ {T}) \prod_ {t = T, \dots , 1} p _ {\theta} (\mathbf {z} _ {t - 1} | \mathbf {z} _ {t}, t, \mathbf {S}) \mathrm {d} \mathbf {z} _ {1: T},
$$
where  $p(\mathbf{z}_T) = \mathcal{N}(\mathbf{0},\mathbf{I}), p_\theta (\mathbf{z}_{t - 1}|\mathbf{z}_t,t,\mathbf{S})$  is a one-step denoising distribution, and  $\mathbf{z}_T,\mathbf{z}_{T - 1},\ldots ,\mathbf{z}_0$  is a sequence of progressively denoised latents. The one-step denoising distribution  $p_{\theta}(\mathbf{z}_{t - 1}|\mathbf{z}_t,t,\mathbf{S})$  is parametrized via a neural network  $g_{\theta}^{\mathrm{LSD}}$  in the following manner:
$$
p _ {\theta} (\mathbf {z} _ {t - 1} | \mathbf {z} _ {t}, t, \mathbf {S}) = \mathcal {N} \left(\frac {1}{\sqrt {\alpha_ {t}}} \left(\mathbf {z} _ {t} - \frac {\beta_ {t}}{\sqrt {1 - \bar {\alpha} _ {t}}} \hat {\epsilon} _ {t}\right), \beta_ {t} \mathbf {I}\right), \quad \mathrm {w h e r e} \hat {\epsilon} _ {t} = g _ {\theta} ^ {\mathrm {L S D}} (\mathbf {z} _ {t}, t, \mathbf {S}),
$$
$\beta_{1}, \ldots, \beta_{T}$  is a linearly increasing variance schedule,  $\alpha_{t} = 1 - \beta_{t}$ , and  $\bar{\alpha}_{t} = \prod_{i=1}^{t} (1 - \beta_{i})$ .
Sampling Procedure. To sample a  $\mathbf{z}_0\sim p_\theta (\mathbf{z}_0|\mathbf{S})$ , we adopt an iterative denoising procedure as in [63, 31]. The sampling process starts with a latent representation  $\mathbf{z}_T\sim \mathcal{N}(\mathbf{0},\mathbf{I})$  filled with random Gaussian noise. Next, conditioned on the slots, we denoise it  $T$  times by sampling sequentially from the one-step denoising distribution  $\mathbf{z}_{t - 1}\sim p_{\theta}(\mathbf{z}_{t - 1}|\mathbf{z}_t,t,\mathbf{S})$  for  $t = T,\dots ,1$ . This produces a sequence of latents  $\mathbf{z}_T,\mathbf{z}_{T - 1},\ldots ,\mathbf{z}_0$  that become progressively cleaner. Finally,  $\mathbf{z}_0$  can be considered as the reconstructed latent representation.
Training Procedure. Following LDM [63], the training of  $p_{\theta}(\mathbf{z}_0|\mathbf{S})$  can be cast to a simple procedure for training  $g_{\theta}^{\mathrm{LSD}}$  as follows. Given an image  $\mathbf{x}$ , its slot representation  $\mathbf{S}$ , and its image latent  $\mathbf{z}_0$ , we first randomly choose a noise level  $t \in \{1, \dots, T\}$  from a uniform distribution. Given the  $t$ , we corrupt the clean latent  $\mathbf{z}_0$  and obtain a noised latent  $\mathbf{z}_t$  as follows:
$$
\mathbf {z} _ {t} = \sqrt {\bar {\alpha} _ {t}} \mathbf {z} _ {0} + \sqrt {1 - \bar {\alpha} _ {t}} \boldsymbol {\epsilon} _ {t}, \quad \text {w h e r e} \boldsymbol {\epsilon} _ {t} \sim \mathcal {N} (\mathbf {0}, \mathbf {I}), \quad \bar {\alpha} _ {t} = \prod_ {i = 1} ^ {t} (1 - \beta_ {i}).
$$
The noised latent  $\mathbf{z}_t$  is then given as input to  $g_{\theta}^{\mathrm{LSD}}$  along with the slots  $\mathbf{S}$  and denoising time-step  $t$  to predict the noise  $\pmb{\epsilon}_{t}$ . The network  $g_{\theta}^{\mathrm{LSD}}$  is trained by minimizing the mean squared error between the predicted noise  $\hat{\pmb{\epsilon}}_t$  and the true noise  $\pmb{\epsilon}_{t}$ :
$$
\hat {\boldsymbol {\epsilon}} _ {t} = g _ {\theta} ^ {\mathrm {L S D}} (\mathbf {z} _ {t}, t, \mathbf {S}) \quad \Longrightarrow \quad \mathcal {L} (\phi , \theta) = | | \hat {\boldsymbol {\epsilon}} _ {t} - \boldsymbol {\epsilon} _ {t} | | ^ {2}.
$$
2.2.3 Denoising Network
We implement the denoising network  $g_{\theta}^{\mathrm{LSD}}$  as a variant of the conventional UNet architecture adapted to incorporate slot-conditioning. Our denoising network consists of a stack of  $L$  layers where each layer  $l$  is a UNet-style CNN layer followed by a slot-conditioned transformer:
$$
\tilde {\mathbf {h}} _ {l} = \operatorname {C N N} _ {\theta} ^ {l} \left(\left[ \mathbf {h} _ {l - 1}, \mathbf {h} _ {\operatorname {s k i p} (l)} \right], t\right), \quad \Longrightarrow \quad \mathbf {h} _ {l} = \operatorname {T r a n s f o r m e r} _ {\theta} ^ {l} \left(\tilde {\mathbf {h}} _ {l} + \mathbf {p} _ {l}, \operatorname {c o n d} = \mathbf {S}\right),
$$
where  $\mathbf{h}_0 = \mathbf{z}_t$  is the input,  $\mathbf{h}_1,\dots ,\mathbf{h}_{L - 1}$  are the hidden states and  $\hat{\pmb{\epsilon}}_t = \mathbf{h}_L$  is the output.
CNN Layers. Following UNet [64], the convolutional layers  $\mathbb{C}\mathbb{N}\mathbb{N}_{\theta}^{1},\ldots ,\mathbb{C}\mathbb{N}_{\theta}^{L}$  downsample the feature map via the first  $\frac{L}{2}$  layers and then upsample it back to the original resolution via the remaining layers. Following standard UNet, these CNN layers also receive inputs via skip connection from an earlier layer denoted by  $\mathrm{skip}(l)$ . This network design has also been explored in LDM [63].
Slot-Conditioned Transformer. The role of the transformer is to incorporate the information from the slots into the UNet-based denoising process. For this, in each layer  $l$ , the intermediate feature map  $\tilde{\mathbf{h}}_l$  produced by the CNN layer is flattened into a set of features. To this, positional embeddings  $\mathbf{p}_l$  are added and the resulting features are provided as input to the transformer. Within the transformer, these features interact with each other and with the slots  $\mathbf{S}$ , thus incorporating the information from the slots into the denoising process. The transformer output is then reshaped back to a feature map  $\mathbf{h}_l$ .
3 Compositional Image Synthesis
In this section, we describe how a trained LSD model can be used to compose and synthesize novel images. The conventional approach to composing novel images is commonly supervised and rely on text prompts. This involves using words from a vocabulary to create a sentence, which is then given to a text-to-image model to synthesize the desired image [62, 61, 65, 63, 58]. However, in a fully unsupervised setting like ours, we first need to build a library of visual concepts by simply observing a large set of unlabelled images. Then, similarly to composing a sentence prompt using words, we compose a concept prompt by selecting concepts from the visual concept library. By providing this concept prompt to the LSD decoder, we can then synthesize a desired novel image. This approach of unsupervised compositional image synthesis was also explored in [70].
Unsupervised Visual Concept Library. To build a library of visual concepts from unlabelled images, we first take a large batch of  $B$  images  $\mathbf{x}_1, \ldots, \mathbf{x}_B$ . We then apply slot attention to obtain slots for these images  $\mathbf{S}_1, \ldots, \mathbf{S}_B$ . Next, we gather all these slots into a single set  $S$  and perform  $K$ -means on it. The  $K$ -means procedure assigns each slot in  $S$  to one of the  $K$  clusters. We consider the set of slots assigned to a  $k$ -th cluster as a visual concept library  $\mathcal{V}_k$ . With  $K$  as the number of clusters, this method results in  $K$  visual concept libraries  $\mathcal{V}_1, \ldots, \mathcal{V}_K$ . Our experiments will demonstrate that this basic  $K$ -means technique can generate semantically meaningful concept libraries. For example, on a dataset of human face images like FFHQ [41], the  $K$  libraries correspond to practical concept categories such as hair style, face, clothing, and background.
Novel Image Synthesis. Given libraries  $\nu_{1},\dots,\nu_{K}$ , we can compose a concept prompt  $\mathbf{S}_{\mathrm{compose}}$  by picking  $K$  slots, each from the corresponding  $k$ -th library, and stacking them together:
$$
\mathbf {S} _ {\text {c o m p o s e}} = \left(\mathbf {s} _ {1}, \dots , \mathbf {s} _ {K}\right), \quad \text {w h e r e} \mathbf {s} _ {k} \sim \operatorname {U n i f o r m} \left(\mathcal {V} _ {k}\right)
$$
We then give the composed prompt  $\mathbf{S}_{\mathrm{compose}}$  to the LSD decoder to generate the image latent:  $\mathbf{z}_{\mathrm{compose}} \sim p_{\theta}(\mathbf{z}_0|\mathbf{S}_{\mathrm{compose}})$ . Applying the image decoder on  $\mathbf{z}_{\mathrm{compose}}$  generates the desired novel scene image  $\mathbf{x}_{\mathrm{compose}} = g_{\theta}^{\mathrm{AE}}(\mathbf{z}_{\mathrm{compose}})$ . For instance, in the FFHQ dataset, the concept prompt can be a collection of chosen hair style, face, clothing, and background. Decoding this prompt would generate a face image that conforms to this prompt.
Unsupervised Object-Centric Learning. Object-centric representation learning aims to decompose multi-object scenes into meaningful object entities. A common approach to this is by auto-encoding. In this line, the focus has been to design an appropriate decoder that supports good decomposition. The most widely used decoders include the mixture-decoder [6, 25, 53, 26, 19, 78, 1, 17, 18, 39, 89, 80], spatial transformer decoder [20, 12, 50, 35, 13, 9, 36, 49, 72, 86], Neural Radiance Fields (NeRF) [77, 83, 74], transformer decoder [70, 73, 71, 84, 7, 67, 23], energy-based models [16] and complex-valued functions [55]. Among these, the mixture decoder has been shown to struggle on complex scene images [70, 73] while the spatial transformer decoder requires careful tuning of hyperparameters which limits their applicability in realistic scenes. Neural Radiance Fields require camera poses, and the learning setting involves 3D scenes unlike ours. Another line of work seeks object-centric scene decomposition without reconstruction. This includes works such as [69, 82, 28, 79, 3, 54, 87]. However, unlike ours, these approaches lack the ability to generate images. Preceding object-centric learning, several works pursued disentanglement within a single-vector representation of the scene and use it for compositional image generation [11, 30, 42, 45, 10]. However, lacking spatial binding mechanisms, these methods struggle in multi-object scenes [25] unlike ours.
Diffusion Models. Diffusion models (DMs) are a recent class of generative models that can produce high-quality images by reversing a stochastic process that gradually adds noise to an image [31, 76]. DMs have been applied to various tasks in computer vision, such as class-conditioned generation, text-to-image generation, image editing, super-resolution, and inpainting [14, 33, 61, 58, 65, 56, 32, 66]. Recently, [63] proposed Latent Diffusion Models (LDM). By virtue of operating on a low-dimensional latent space, LDM reduces the computational demands of DMs significantly. [51] introduced a method for multi-object scene generation by combining signals from multiple text-conditioned denoising networks. However, to achieve controllable generation, many of these existing DMs require additional labels such as text descriptions to train and control the generation process. In contrast, our model can generate images compositionally using object concepts directly extracted from images. Moreover,
Figure 2: Visualization of Unsupervised Object Segmentation. We show visualizations of predicted segments on CLEVR, CLEVRTex, MOVi-E, and FFHQ datasets.
DMs have also been used for representation learning. [4] demonstrated that the intermediate features of a learned DM are useful representations for label-efficient semantic segmentation. [60] introduced an image encoder that compresses an image into a latent representation, jointly trained with the denoising objective. Nevertheless, these representations are unstructured and not modular like ours.