Learning to Compose: Improving Object Centric Learning by Injecting Compositionality

In this approach, we randomly sample $N$ slots among $2N$ slots i.e., $\mathbf{S}^{c}\overset{N}{\sim}(\mathbf{S}^{1}\cup\mathbf{S}^{2})$. As it explores over all of the possible combinations, this composition function encourages the slot representation itself to be highly composable to generate valid images for any combinations. On the other hand, it may produce invalid combinations of slots on rare occasions, e.g., omitting the background slots or sampling two objects placed in the same location.

One way to mitigate such suspicious compositions is to constrain $\mathbf{S}^{c}$ to be valid composition of the scene. However, strictly ensuring this constraint is non-trivial due to the stochastic nature of slot attention i.e., each slot is sampled stochastically from its underlying distribution and the association between the slots and scenes varies depending on the initialization. Instead, we adopt a rather simple approach that employs the identical slot initialization $\mathbf{S}^{(0)}$ in Equation 3 for two images, and sample the exclusive set of slots. Formally, let $I_{1}$ and $I_{2}$ be a random partition of slot indices i.e., $I_{1}\cup I_{2}=\{1,...,N\},I_{1}\cap I_{2}=\emptyset$. Then we construct the composite slot by $\mathbf{S}^{c}=\mathbf{S}^{1}_{I_{1}}\cup\mathbf{S}^{2}_{I_{2}}$, where $\mathbf{S}^{1}$ and $\mathbf{S}^{2}$ are slots extracted by Equation 3 from $\mathbf{x}^{1}$ and $\mathbf{x}^{2}$, respectively, which are initialized with the same $\mathbf{S}^{(0)}$.The underlying intuition is that the slot initialization is reasonably correlated with the objects it captures (Figure 7), hence sampling from exclusive slots is likely to be valid scenes than the random sampling.

To model the generative prior $p(\mathbf{x}^{c})$, we opt for a diffusion model (Ho et al., 2020) due to its excellence in generation quality and mode coverage (Xiao et al., 2022). The latter is especially important in our framework since the model evaluates the prior over potentially out-of-distribution samples generated by the composition (Section 3.1). Instead of introducing an additional pre-trained diffusion model, we employ the diffusion-based decoder in the auto-encoding path (Section 2), and reuse it as a generative prior. This way, our decoder $D_{\phi}$ is trained by minimizing the reconstruction loss by denoising objective in Equation 4, while serving as a generative prior in the composition path. It greatly improves the parameter-efficiency and memory, and the need for pre-trained generative prior per dataset.

Given the generative prior, we maximize the likelihood $p(\mathbf{x}^{c})$ with respect to $\mathbf{x}^{c}$ in the composition path. Since $\mathcal{L}_{\text{Diff}}$ in Equation 4 is minimizing the upper bound of negative log likelihood of $x^{c}$ (Ho et al., 2020), minimizing $\mathcal{L}_{\text{Diff}}$ with respect to $\mathbf{x}^{c}$ leads to the maximization of the likelihood $p(\mathbf{x}^{c})$. However, computing the gradient of $\mathcal{L}_{\text{Diff}}$ requires expensive computation of Jacobian maxtrix of the decoder and it often degrades the overall training stability. Following (Poole et al., 2022), the gradient of $\mathcal{L}_{\text{Diff}}$ with respect to $\theta$ can be approximated as:

where $\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})$ is a noise, $t\sim\mathcal{U}(t_{\text{min}},t_{\text{max}})$ is a timestep, respectively, $w(t)$ is a weighting function dependent to $t$, and $\mathbf{x}_{t}^{c}=\sqrt{\bar{\alpha}_{t}}\mathbf{x}^{c}+\sigma_{t}\epsilon$ is a corrupted image of $\mathbf{x}^{c}$ from forward diffusion process. By updating the encoder parameters $\theta$ with $\nabla_{\theta}\mathcal{L}_{\text{Prior}}$, $\mathbf{x}^{c}$ is guided toward high probability density region following the diffusion prior. Note that optimization of the Equation 6 is with only respect to the encoder parameter while fixing the decoder. It prevents suspicious collaboration between the encoder and decoders in generating composite images from suboptimal slots.

As discussed earlier, our framework exploits the diffusion model $D_{\phi}$ as a decoder and generative prior in the auto-encoding and composition paths, respectively. One drawback is that the diffusion decoder requires an iterative denoising process to generate the composite image $\mathbf{x}^{c}$, which takes significant time and makes the backpropagation through the decoder non-trivial. To address this problem, we employ a one-shot decoder $D_{\psi}$ as a surrogate for $D_{\phi}$ to support fast and differentiable decoding of $\mathbf{x}^{c}$. ¹¹1We also consider one-step denoising result of the diffusion decoder using Tweedie’s formula (Stein, 1981; Robbins, 1992) but observe severe degradation in performance due to its inferior quality.

We employ a bidirectional Transformer (Devlin et al., 2019) that takes the composite slot $\mathbf{S}^{c}$ and the learnable mask tokens $\mathbf{m}\in\mathbb{R}^{HW\times C}$ as input, and produces the composite image by a single forward process by $\mathbf{x}^{c}=D_{\psi}(\mathbf{m},\mathbf{S}^{c})$. The decoder is trained along with the auto-encoding path by:

Note that the generation quality of the one-shot decoder $D_{\psi}$ is behind the powerful diffusion decoder $D_{\phi}$, and serves only to compute the $\mathbf{x}^{c}$ in Equation 6. We observe that such weak decoder is sufficient to compute the meaningful gradient through the Equation 6, presumably because the gradients are accumulated over various noise levels $t$.

Since object-centric learning assumes no prior knowledge on data, the mismatch between the number of objects and slots is inevitable in practice. To evaluate such robustness, we vary the number of slots from 11 to 17. Figure 4(a) presents the result. It shows that the performance of the baselines is highly sensitive to the number of slots. Specifically, SLATE tends to deteriorate more as the number of slots increases. Compared to the baseline, our method achieves more robust performance by encoding an object into a slot while leaving excess slots empty.

To identify the effect of slot encoder, we consider two popular architectures in the literature; a multi-layer CNN encoder (Singh et al., 2022b) and UNet-based encoder (Ronneberger et al., 2015). Figure 4(b) summarizes the result. It shows that employing the weaker encoder generally deteriorates the performance of the baselines significantly, indicating that architectural bias in the encoder is critical in the auto-encoding objective. Interestingly, the performance of our method is hardly affected by such drastic modifications, showing great robustness.

It is widely observed that the choice of decoder is also crucial in object-centric learning, since the highly expressive decoder can often bypass the object representation to minimize the reconstruction loss (Singh et al., 2022a). To examine such effect, we gradually increase the feature dimensions of the decoder to 133$\%$, 166$\%$, and 200$\%$. Figure 4(c) summarizes the result. It shows that increasing the decoder capacity hampers the performance in SLATE. LSD exhibits the opposite trends showing a large improvement in FG-ARI, although its performance drops significantly in mIoU (Figure 6). Compared to the baselines, our method is much less sensitive to the decoder capacity, while the performance tends to improve slightly with increased capacity in all measures.

Overall, the results indicate that the quality of object-centric representation is significantly influenced by various factors in the auto-encoding-based methods. Conversely, our model consistently delivers outstanding performance across all configurations, even with major alterations to the encoder architecture. It demonstrates that our regularization through the composite path can directly encourage the model to learn compositional representation, greatly enhancing robustness to architectural biases.

To identify the contributions of each component in our framework, we conduct an ablation study and present the result in Table 1. The first row corresponds to our model with only the auto-encoding path, while the last row is the complete version of our model. Comparing the first row with the others shows that incorporating the composition path significantly improves overall quality. Adding $\mathcal{L}_{\text{Prior}}$, we observe a substantial improvement in all three metrics. Considering that FG-ARI measures the correct cluster membership of pixels within the objects, increased FG-ARI indicates that the generative prior encourages the encoder to capture more holistic object representations. This is because the generative prior penalizes the encoder for fragmenting the objects, thereby discouraging the generation of unrealistic partial objects in the composite image. Comparing the second and the third rows, we observe that sharing the slot initialization $\mathbf{S}^{(0)}$ slightly enhances mIoU and mBO scores. This improvement is likely attributed to the increased training stability by avoiding invalid slot combinations as shown in Figure 7. Incorporating regularization $\mathcal{L}_{\text{reg}}$ alone in the composition path does not improve the performance (fourth row), while combined with generative prior, it leads to significant improvement.

We present the compositional generation results to further investigate the impact of our composition path. Figure 5 presents the results. Given two images, we construct the composite representation by replacing one object slot from the first image (red arrow) to another from the second image (blue arrow), and producing the image by the decoder. Based on visualization of the learned slots, we observe that the baselines often fail to learn compositional slot representation, by separating objects into multiple slots or encoding background with an object. It leads to failures in object-level manipulation, such as retaining an object after the removal (LSD in MultiShapeNet and PTR), altering the content of the added object (SLATE in MultiShapeNet), or transforming background with the object (SLATE in PTR and LSD in Super-CLEVR). In contrast, our method produces both semantically meaningful and realistic images from composite slot representations, supporting our claim that we can regularize object-centric learning through the proposed compositional path.