Improved Object-Centric Diffusion Learning with Registers and Contrastive Alignment

AI summary

Overview

Research area: Computer Vision — unsupervised object-centric learning (OCL), slot-based scene decomposition, and latent diffusion models.

Technical level: Advanced. The paper assumes familiarity with Slot Attention, latent diffusion models, cross-attention conditioning, contrastive objectives, and mutual-information arguments.

Scope: The paper introduces CODA (Contrastive Object-centric Diffusion Alignment), a training framework that combines input-independent "register slots" and a contrastive alignment loss to improve object discovery, property prediction, and compositional image generation in diffusion-based object-centric learning.

What This Paper Is About

Object-centric learning tries to decompose an image into a small set of slot vectors, each representing one object or concept, without any human labels. Recent methods use a pretrained diffusion model as the slot decoder, which generates good images but suffers from two failures: slots get tangled together (one slot encodes parts of several objects), and slots do not reliably correspond to distinct image regions. CODA attacks both problems by adding register slots that soak up residual attention and a contrastive loss that explicitly pushes each image and its slots to correspond.

Key Contributions

1. Register-augmented slot diffusion. Input-independent register slots are added alongside semantic slots. They carry no semantic content but act as attention sinks, absorbing residual attention mass so that object slots stay focused on meaningful object-concept associations and entanglement is reduced.

2. Mitigating text-conditioning bias. Rather than adding new adapter layers, CODA finetunes only the key, value, and output projections inside the diffusion model's cross-attention layers, so the pretrained Stable Diffusion backbone aligns better with slot conditioning instead of expecting text embeddings.

3. Contrastive alignment objective. A contrastive loss is defined as the negative of the denoising loss evaluated on mismatched ("negative") slots. Combined with the denoising loss, the total objective is argued to be a tractable surrogate for maximizing the mutual information between slots and inputs, formalized in Theorem 1 and Corollary 1.

4. Comprehensive evaluation. CODA is benchmarked against unsupervised diffusion-based OCL baselines on synthetic (MOVi-C, MOVi-E) and real-world (VOC, COCO) datasets for object discovery, property prediction, and compositional generation. Code and pretrained models are released at https://github.com/sony/coda.

Main Findings

• Synthetic object discovery: On MOVi-C, CODA improves FG-ARI by +7.15% and mIoU by +7.75% over the strongest baseline. On MOVi-E, which has more objects per scene, it improves FG-ARI by +2.59% and mIoU by +3.36%. SLATE and LSD struggle to produce accurate object segmentations on these benchmarks.

• Real-world object discovery: On COCO, CODA reaches FG-ARI of 47.54 against 41.4 for SlotAdapt, a +6.14% improvement. On VOC it reaches FG-ARI 32.23, mBO_i 55.38, mBO_c 61.32, mIoU_i 50.77, and mIoU_c 56.30, improving instance-level discovery by +3.88% mBO_i and +3.97% mIoU_i, and semantic-level discovery by +5.72% mBO_c and +7.00% mIoU_c over SlotAdapt.

• Property prediction: With frozen slot representations and a 2-layer MLP predictor (hidden dimension 786), CODA improves all metrics except 3D bounding box. On MOVi-C it reaches Position MSE 0.01, 3D B-Box MSE 2.11, and Category accuracy 74.12. On MOVi-E it reaches Position MSE 0.01, 3D B-Box MSE 4.22, and Category accuracy 78.06. The authors attribute weaker 3D bounding box results to DINOv2 features lacking fine-grained geometric detail.

• Image generation: On COCO reconstruction, CODA achieves KID 0.35 and FID 10.65, versus 0.39 / 10.86 for SlotAdapt, 5.85 / 19.45 for SlotDiffusion, and 19.09 / 35.54 for LSD. In compositional generation it reaches KID 30.44 and FID 31.03, against 34.38 / 40.57 for SlotAdapt, 57.31 / 64.21 for SlotDiffusion, and 103.48 / 167.23 for LSD.

• Ablations on VOC: The frozen-Stable-Diffusion baseline is improved moderately by finetuning cross-attention key, value, and output projections, substantially by adding registers (especially in mBO), and further in mIoU by the contrastive loss. Applying the contrastive loss without stopping gradients in the diffusion model degrades performance. All components together give the best results.

• Negative slot design: Hard negatives are built by replacing half of the slots from one image with slots from another image, with slot initialization shared between the two images to keep the mixed combinations semantically plausible. The paper notes naive mixing can produce invalid combinations, such as dropping background slots.

• Efficiency: Register slots add negligible overhead. For Stable Diffusion v1.5, 77 padding tokens passed through the frozen text encoder yield 77 fixed register slots. Learnable registers were also explored, and fixed registers performed best. The overall gain figures for MOVi improvements come from comparisons the paper states in text; ablation values beyond the VOC table are described as being in appendices that are not included in the truncated content.

Methodology in Plain English

CODA keeps the standard pipeline: a DINOv2 encoder extracts image features, Slot Attention turns those features into N slot vectors, and a pretrained Stable Diffusion v1.5 U-Net decodes the image by denoising latents conditioned on those slots through cross-attention.

Three changes are layered onto this pipeline. First, register slots — embeddings obtained by encoding an all-padding token sequence with the frozen CLIP ViT-L/14 text encoder — are concatenated with the semantic slots and fed to cross-attention. Since softmax forces cross-attention weights to sum to one, queries that match no semantic slot cause attention to smear across slots; registers give that excess attention somewhere harmless to go, similar to attention sinks in language models.

Second, instead of adding new adapter modules or retraining the diffusion model from scratch, only the key, value, and output projections in cross-attention are finetuned, which removes the model's expectation of text conditioning while preserving generative quality.

Third, a contrastive term is added to training. The denoising loss is minimized for the correct slots, while its negative is minimized with respect to mismatched slots — meaning the model is pushed to denoise well under the true slots and poorly under wrong ones. During this contrastive term the diffusion decoder is frozen and only Slot Attention is updated, so the model cannot cheat by changing how it generates. The Theory section shows that the denoising gap between matched and mismatched slots is a practical stand-in (labeled Δ) for the mutual information between images and slots, which direct optimization cannot compute.

Why This Matters

Impact on research. CODA shows that the two dominant failure modes of diffusion-based OCL — slot entanglement and weak slot-image alignment — can be addressed with a small, architecturally simple set of changes rather than a new model or extra supervision. It also supplies a mutual-information justification for a contrastive denoising objective, connecting practical OCL training to an information-theoretic target. Because it avoids generated captions or other pseudo-labels used by some prior methods, it stays fully unsupervised.

Real-world applications:

• Image and video editing: CODA's slot manipulation supports removing objects by discarding their slots and swapping objects across scenes, enabling fine-grained, controllable edits.
• Controllable content creation: Faithful single-slot generation lets creators compose novel scenes from recombined object concepts rather than regenerating whole images.
• Robotics and embodied agents: Object-centric representations with accurate boundaries and disentangled attributes support manipulation and planning over discrete entities.
• Visual reasoning and causal inference: Cleaner slot-level attributes (position, category) provide structured inputs for downstream reasoning and world-modeling tasks.

Industry relevance. The work comes from Sony AI with Stanford University and Sony Group Corporation, and the released code and pretrained models point to practical deployment. Register slots add negligible overhead and the finetuning is restricted to a few cross-attention projections, so the approach stays compatible with off-the-shelf diffusion models — an attractive property for teams that cannot afford retraining from scratch or large new adapter architectures.

Future Directions

• Remaining weakness in geometric attributes. The 3D bounding box results are the one place CODA underperforms its baselines on MOVi-C and MOVi-E, which the authors tie to the lack of fine-grained geometry in DINOv2 features. Swapping in features with stronger spatial detail is an open question.
• Better negative construction. Negative slots are currently built by replacing half of the slots with those from another image under shared initialization, and the paper acknowledges this can produce invalid combinations such as missing background slots. More principled ways of generating hard but valid negatives could improve the contrastive signal.
• Scaling the sensitivity analyses. The trade-off parameter λ_cl, the replacement ratio for negative slots, and register design (fixed versus learnable) were studied in appendices that the truncated content does not report, leaving room to characterize how far these choices generalize to harder data.
• Beyond static images. The introduction motivates OCL for video editing, and the related-work discussion notes prior contrastive work on temporal consistency for video. Extending CODA's register and alignment mechanisms to video and to larger, more cluttered real-world scenes is a natural next step.

Target Audience

Researchers and graduate students working on object-centric learning, unsupervised segmentation, and diffusion-based generative models. It is also relevant to engineers building compositional image-editing or controllable-generation systems who want a lightweight way to condition a pretrained diffusion model on structured object representations, and to readers interested in the connection between contrastive denoising objectives and mutual information.

Authors’ 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. Code and pretrained models are available at https://github.com/sony/coda.