Source-linked AI summary
Diffuse, Attend, and Segment: Unsupervised Zero-Shot Segmentation using Stable Diffusion
Junjiao Tian, Lavisha Aggarwal, Andrea Colaco, Zsolt Kira, Mar Gonzalez-Franco
TL;DR
The paper targets segmentation that is both unsupervised and zero-shot, without annotations or prior knowledge of the target. It uses self-attention from a pre-trained stable diffusion model and iteratively merges attention maps with KL divergence. DiffSeg reports state-of-the-art performance, including absolute gains of 26% in pixel accuracy and 17% in mean IoU on COCO-Stuff-27, while requiring no language dependency or auxiliary images.
Problem
Unsupervised zero-shot segmentation without annotations or prior knowledge remains challenging because existing approaches often require target adaptation, extensive annotations, or external information.
Method
DiffSeg uses self-attention from a pre-trained stable diffusion model, aggregating and iteratively merging attention maps with KL divergence to form segmentation masks.
Results
26% in pixel accuracy and 17% in mean IoU are the absolute gains over a prior unsupervised zero-shot method on COCO-Stuff-27.
Takeaways & Limitations
DiffSeg segments images in a pure zero-shot manner without language dependency, auxiliary images, prior knowledge, or external resources, and generalizes to diverse image styles.
Takeaways & Limitations
DiffSeg generates high-quality masks but does not label each mask without an additional extension.
Abstract
from arXiv · showhide
Producing quality segmentation masks for images is a fundamental problem in computer vision. Recent research has explored large-scale supervised training to enable zero-shot segmentation on virtually any image style and unsupervised training to enable segmentation without dense annotations. However, constructing a model capable of segmenting anything in a zero-shot manner without any annotations is still challenging. In this paper, we propose to utilize the self-attention layers in stable diffusion models to achieve this goal because the pre-trained stable diffusion model has learned inherent concepts of objects within its attention layers. Specifically, we introduce a simple yet effective iterative merging process based on measuring KL divergence among attention maps to merge them into valid segmentation masks. The proposed method does not require any training or language dependency to extract quality segmentation for any images. On COCO-Stuff-27, our method surpasses the prior unsupervised zero-shot SOTA method by an absolute 26% in pixel accuracy and 17% in mean IoU. The project page is at \url{https://sites.google.com/view/diffseg/home}.
1. Introduction
DiffSeg addresses the challenge of unsupervised, zero-shot segmentation without annotations, prior knowledge, language, or external resources. It uses self-attention from a pre-trained stable diffusion model and iterative merging to produce masks, achieving large gains on COCO-Stuff-27.
- Semantic segmentation partitions images into semantically consistent regions and supports downstream applications including image editing, medical imaging, and autonomous driving.
- Zero-shot transfer segmentation for images with unknown categories is challenging because existing strong approaches often depend on extensive annotations or target-data adaptation.SAM uses 1.1B segmentation annotations, while many unsupervised methods adapt on target data.
- The target setting removes annotations and prior knowledge of the target image.
- DiffSeg uses self-attention maps from a pre-trained stable diffusion model to produce segmentation masks without training, language dependency, or external information.Its pipeline combines attention aggregation, iterative attention merging, and non-maximum suppression.
- 26% in pixel accuracy and 17% in mean IoU are the absolute gains DiffSeg reports over a prior unsupervised zero-shot method on COCO-Stuff-27.
2. Related Works
Related work spans diffusion-based visual representations, unsupervised segmentation, and zero-shot transfer segmentation. DiffSeg differs by using stable-diffusion self-attention directly, without target training, text, auxiliary images, or image synthesis.
- Diffusion Models: Stable diffusion features have supported zero-shot classification, supervised and label-efficient segmentation, semantic correspondence, and open-vocabulary segmentation, typically with additional training.
- Unsupervised Segmentation: Unsupervised segmentation methods generate dense masks without annotations but generally require unsupervised training on the target dataset.
- Zero-Shot Transfer Segmentation: Zero-shot segmentation methods can segment without training, but some use auxiliary images or text inputs to facilitate transfer.
- Zero-Shot Transfer Segmentation: DiffSeg uses a diffusion model to generate segmentation without synthesizing or querying multiple images and without knowing object categories.
3. Method
DiffSeg extracts object-grouping information from stable-diffusion self-attention, aligns maps across resolutions, and iteratively merges similar maps into segmentation proposals. The method uses multi-scale receptive fields and KL divergence, then normalizes and suppresses overlapping proposals.
- Stable Diffusion Model Review: Stable diffusion compresses images into latent space and processes them with an encoder-decoder U-Net containing self-attention and cross-attention mechanisms.The paper investigates self-attention because it learns global attention across the image.
- Stable Diffusion Model Review: Each self-attention tensor encodes correlations between a spatial location and all locations, while unconditional self-attention is hypothesized to contain object-grouping information without text.
- Stable Diffusion Model Review: Intra-attention similarity reinforces locations within an object group, while inter-attention similarity links maps whose reference locations belong to the same group.
- Stable Diffusion Model Review: Lower-resolution attention maps provide larger receptive fields for grouping whole objects, whereas higher-resolution maps give finer-grained grouping of components and small objects.
- Non-Maximum Suppression: Non-maximum suppression converts the resulting probability-map proposals into a segmentation mask by assigning each pixel to the proposal with the largest probability.
- Iterative Attention Merging: The method samples M × M anchor maps and iteratively merges similar maps using KL divergence between valid probability distributions.Intra-attention and inter-attention similarities guide reinforcement and expansion of object activations during merging.
4. Experiments
DiffSeg is evaluated on COCO-Stuff-27 and Cityscapes using pre-trained stable diffusion features, with comparisons spanning clustering baselines, input resolutions, and auxiliary-resource requirements. Its proportional aggregation balances segmentation detail and consistency, while performance remains sensitive to resolution and selected hyper-parameters.
- Experimental setup: COCO-Stuff-27 evaluation uses pixel accuracy and mean IoU, with Hungarian matching assigning predicted masks to ground-truth masks.The experiments use the diffusion model's native 512 × 512 resolution and a lower 320 × 320 resolution following prior work.
- Implementation and sensitivity: The experiments use pre-trained Stable Diffusion V1.4 and report results with the hyper-parameters listed in Table 3.The sensitivity study examines aggregation weights and other hyper-parameters used in the main evaluations.
- COCO-Stuff-27: DiffSeg significantly outperforms K-Means-C, K-Means-S, and DBSCAN on COCO-Stuff-27 using the same self-attention tensors.K-Means-S uses image-specific cluster counts, whereas DiffSeg avoids requiring the number of clusters beforehand.
- Cityscapes: On Cityscapes, DiffSeg is on par with prior works at 320-resolution and outperforms them in accuracy and mIoU at 512-resolution.Input resolution affects Cityscapes more severely than COCO because Cityscapes contains smaller classes such as light poles and traffic signs.
- Resource requirements: DiffSeg achieves its performance in a pure zero-shot manner without language dependency or auxiliary images.The method is benchmarked without requiring additional images during inference.
- Aggregation weights: Proportional aggregation assigns higher importance to higher-resolution maps, balancing detailed but fractured outputs against coherent but coarse segmentation.Very low resolutions can merge the entire image into one object under the tested settings.
5. Adding Semantics
DiffSeg produces high-quality segmentation masks but does not label each mask. The paper describes a simple extension for producing labeled segmentation masks.
- Adding Semantics: DiffSeg generates high-quality segmentation masks but does not label each mask.A simple extension for labeled segmentation masks is proposed in Appendix 9.8.
6. Visualization
DiffSeg demonstrates generalization across images with diverse visual styles, including sketches, paintings, smartphone photographs, satellite images, CT scans, and synthetic diffusion-generated images.
- DiffSeg segments images spanning sketches, paintings, smartphone photographs, satellite imagery, CT scans, and synthetic diffusion-generated content.Examples also include Cityscapes and SUN-RGBD images.
7. Conclusion
The paper frames unsupervised zero-shot segmentation as a difficult setting because existing approaches commonly require adaptation or external resources. DiffSeg addresses this with a training-free, resource-free approach based on a pretrained stable diffusion model and reports strong benchmark performance and broad generalization.
- Unsupervised and zero-shot segmentation is challenging, and existing methods commonly require unsupervised adaptation or external resources.
- DiffSeg segments images without prior knowledge, external resources, or additional training by using a pretrained stable diffusion model.Its attention maps are iteratively merged using Intra-Attention Similarity and Inter-Attention Similarity.
- DiffSeg achieves state-of-the-art performance on popular benchmarks and generalizes to images with diverse styles.
9. Appendix
The appendix documents Stable Diffusion’s multi-resolution self-attention structure and explains how DiffSeg averages and iteratively merges attention maps using KL-divergence thresholds.
- Stable Diffusion contains 16 Transformer blocks, each producing a 4D self-attention tensor at a different resolution.
- Averaging attention across multi-head channels is supported because channel differences are much smaller than the paper’s KL merge threshold.The resulting attention maps are therefore very similar across channels.
- DiffSeg’s iterative attention-merging algorithm initializes proposals, identifies attention maps within a KL-divergence threshold, and repeatedly merges them.The process runs for N − 1 iterations after initial proposal construction.
9.4. Comparisons to K-Means and DBSCAN Baselines
The appendix compares DiffSeg with clustering baselines and examines its transfer to DINO features and sensitivity to inference hyper-parameters. DiffSeg performs better than DinoSeg, while most settings are broadly stable except the KL threshold.
- Comparisons to K-Means and DBSCAN Baselines: DiffSeg’s iterative merging avoids k-means’s need to specify the cluster count and its stochastic dependence on initialization.
- Comparisons to K-Means and DBSCAN Baselines: DBSCAN removes the need to specify cluster count but requires tuning eps, which controls neighborhood membership.
- DinoSeg: DiffSeg is theoretically applicable to transfer features from any backbone and is evaluated with a DINO-Base backbone.
- DinoSeg: DinoSeg performs worse than DiffSeg, potentially because DINO uses smaller pretraining data, lower input resolution, and 28 × 28 attention outputs.
- Additional Ablation study: The KL threshold τ is the most sensitive hyper-parameter and is tuned on a small 180-image training subset.
- Additional Ablation study: Using three merging iterations preserves similar accuracy afterward while improving the trade-off between system latency and performance.
- Additional Ablation study: A reasonable τ range is 0.9–1.1, with τ = 1.0 suggested for wild-image segmentation and dataset-specific tuning preferred for benchmark results.
9.7. Limitations
DiffSeg generalizes broadly but remains constrained on specialized self-driving imagery and is not yet real-time. Its main limitations arise from attention-map resolution, training-domain exposure, and iterative computation.
- DiffSeg performs poorly on specialized datasets such as Cityscapes despite its broad zero-shot segmentation capability.
- The largest attention map is 64×64, which may be too coarse for segmenting small objects in self-driving scenes.
- Limited stable-diffusion exposure to vehicle-centric scenes may negatively affect downstream performance on self-driving imagery.
- DiffSeg is not real-time because stable diffusion is large and attention aggregation and merging are iterative.
9.8. Adding Semantics
Semantic DiffSeg extends DiffSeg with labels by generating an image vocabulary, extracting noun-grounded cross-attention, and assigning labels to generated masks. It then merges masks sharing the same predicted label.
- Overview: Semantic DiffSeg adds vocabulary generation, cross-attention extraction, and aggregation and prediction to produce labeled segmentation masks.
- Vocabulary Generation: An off-the-shelf captioning model generates a caption, whose nouns define the output label space while retaining their caption positions.
- Cross-Attention Extraction: The method uses the generated caption as input and extracts cross-attention maps corresponding to noun tokens for concept grounding.
- Aggregation and Prediction: Cross-attention maps are spatially upsampled and aggregated across resolutions into a normalized tensor.
- Aggregation and Prediction: Each DiffSeg mask receives the label with the maximum elementwise overlap against aggregated cross-attention maps, then same-label masks are merged.
- Examples: Examples show DiffSeg outputs alongside semantic segmentation with mask merging for captions describing racing and parked cars.
9.9. Additional Visualization
The paper visualizes segmentation outputs across indoor, natural, synthetic, and autonomous-driving imagery. Comparisons use overlays and inputs, with ground truth shown for COCO-Stuff-27 and Cityscapes.
- SUN-RGBD: SUN-RGBD visualizations arrange overlay, input, and segmentation images from left to right.
- COCO-Stuff-27: COCO-Stuff-27 visualizations arrange overlay, ground truth, and segmentation images from left to right.
- Synthetic Images: Synthetic-image visualizations arrange overlay, input, and segmentation images from left to right.
- Cityscapes: Cityscapes visualizations arrange overlay, ground truth, and segmentation images from left to right.