Improved Object-Centric Diffusion Learning with Registers and Contrastive Alignment

TL;DR

CODA addresses unsupervised object-centric learning with diffusion decoders by tackling slot entanglement and weak slot–image alignment. It introduces input-independent register slots, finetuning of cross-attention keys/queries, and a contrastive alignment objective that serves as a tractable surrogate for maximizing mutual information between inputs and slots. Across MOVi-C/E, VOC, and COCO, CODA yields consistent improvements in object discovery, property prediction, and compositional image generation, including notable gains on COCO FG-ARI. The method remains efficient and drop-in compatible with off-the-shelf diffusion backbones, offering a practical pathway toward robust OCL in real-world scenes.

Abstract

Slot Attention (SA) with pretrained diffusion models has recently shown promise for object-centric learning (OCL), but suffers from slot entanglement and weak alignment between object slots and image content. We propose Contrastive Object-centric Diffusion Alignment (CODA), a simple extension that (i) employs register slots to absorb residual attention and reduce interference between object slots, and (ii) applies a contrastive alignment loss to explicitly encourage slot-image correspondence. The resulting training objective serves as a tractable surrogate for maximizing mutual information (MI) between slots and inputs, strengthening slot representation quality. On both synthetic (MOVi-C/E) and real-world datasets (VOC, COCO), CODA improves object discovery (e.g., +6.1% FG-ARI on COCO), property prediction, and compositional image generation over strong baselines. Register slots add negligible overhead, keeping CODA efficient and scalable. These results indicate potential applications of CODA as an effective framework for robust OCL in complex, real-world scenes.

Figures (12)

• Figure 1: Image generation from individual slots. Top: slot masks. Bottom: generated images. Both methods can reconstruct the full scene when conditioned on all slots (last column). However, Stable-LSD (without register slots) fails to generate images from individual slots. Our method yields faithful single-concept generations, demonstrating disentangled and well-aligned slots.
• Figure 2: Overview of CODA. The input image is encoded with DINOv2 and processed by Slot Attention (SA) to produce slot representations. The semantic slots ${\mathbf{s}}$, together with register slots $\overline{{\mathbf{r}}}$, serve as conditioning inputs for the cross-attention layers of a pretrained diffusion model. SA is trained jointly with the key, value, and output projections of the cross-attention layers using a denoising objective that minimizes the mean squared error between the true and predicted noise. In addition, a contrastive loss is applied to align each image with its corresponding slot representations.
• Figure 3: Illustration of compositional editing. CODA can compose novel scenes from real-world images by removing (left) or swapping (right) the slots, shown as masked regions in the images.
• Figure 4: Illustration of the ablation study on VOC. We start from the pretrained diffusion model as a slot decoder (Baseline), adding register slots (Reg), finetuning the key, value, and output projections in the cross-attention layers (CA), adding contrastive alignment (CO).
• Figure 5: Segmentation masks learned by CODA on COCO

Theorems & Definitions (3)

• Theorem 1
• Corollary 1
• Remark 1