PaperlayerView exact revision on arXiv

Source-linked AI summary

Diffusion Model as a Generalist Segmentation Learner

Haoxiao Wang, Antao Xiang, Haiyang Sun, Peilin Sun, Changhao Pan, Yifu Chen, Minjie Hong, Weijie Wang, Shuang Chen, Yue Chen, Zhou Zhao

arXiv:2604.24575v1cs.CV

TL;DR

Segmentation systems remain fragmented across tasks, vocabularies, and domains, motivating a unified model. DiGSeg repurposes a pretrained diffusion backbone with latent conditioning and text guidance, achieving state-of-the-art segmentation, open-vocabulary generalization, and cross-domain transfer.

  • Problem

    Segmentation systems remain fragmented across tasks and domains, lacking a unified conditioned interface that generalizes across vocabularies and visual settings.

  • Method

    DiGSeg fine-tunes a pretrained diffusion denoising U-Net to generate segmentation-consistent latents, using visual conditioning and multi-scale CLIP-aligned language features.

  • Results

    DiGSeg achieves state-of-the-art performance with strong open-vocabulary generalization and cross-domain transfer without task-specific architectures.

  • Takeaways & Limitations

    Diffusion backbones can serve as generalist segmentation learners through a unified interface connecting generative modeling with dense visual understanding.

  • Takeaways & Limitations

    DiGSeg is slower at inference than feed-forward segmenters, reflecting a quality-efficiency trade-off.

Abstract

from arXiv · show

Diffusion models are primarily trained for image synthesis, yet their denoising trajectories encode rich, spatially aligned visual priors. In this paper, we demonstrate that these priors can be utilized for text-conditioned semantic and open-vocabulary segmentation, and this approach can be generalized to various downstream tasks to make a general-purpose diffusion segmentation framework. Concretely, we introduce DiGSeg (Diffusion Models as a Generalist Segmentation Learner), which repurposes a pretrained diffusion model into a unified segmentation framework. Our approach encodes the input image and ground-truth mask into the latent space and concatenates them as conditioning signals for the diffusion U-Net. A parallel CLIP-aligned text pathway injects language features across multiple scales, enabling the model to align textual queries with evolving visual representations. This design transforms an off-the-shelf diffusion backbone into a universal interface that produces structured segmentation masks conditioned on both appearance and arbitrary text prompts. Extensive experiments demonstrate state-of-the-art performance on standard semantic segmentation benchmarks, as well as strong open-vocabulary generalization and cross-domain transfer to medical, remote sensing, and agricultural scenarios-without domain-specific architectural customization. These results indicate that modern diffusion backbones can serve as generalist segmentation learners rather than pure generators, narrowing the gap between visual generation and visual understanding.

1 Introduction

The introduction identifies fragmentation and unreliable attention-map repurposing as barriers to unified segmentation, then presents DiGSeg as a fine-tuned diffusion framework with visual and text conditioning for generalist segmentation.

  • Motivation: Segmentation systems remain fragmented across tasks and domains because they use different architectures, label spaces, and training pipelines.This fragmentation spans natural-scene semantic and instance segmentation as well as medical imaging, remote sensing, and agriculture.
  • Limitations of prior work: Diffusion repurposing methods often produce noisy, low-resolution, layer-inconsistent attention maps that yield fragmented masks requiring substantial post-processing.High-resolution maps may lack coherence, whereas low-resolution maps preserve semantic consistency but lose small-object and boundary information.
  • Method: DiGSeg explicitly converts diffusion representations into segmentation by encoding images and ground-truth maps into latent space while fine-tuning only the denoising U-Net.Almost all other pretrained diffusion components remain frozen, and the U-Net is trained to produce outputs aligned with segmentation-consistent latents.
  • Contributions: A visual latent pathway and CLIP-aligned text conditioner inject language across multiple denoising scales, binding textual queries to evolving mask latents.The framework is intended to support standard segmentation, zero-shot and open-vocabulary capabilities, and cross-domain transfer without task-specific architectures.

2 Related Work

Prior segmentation work progressed from CNNs to unified transformer architectures, while open-vocabulary methods aligned visual features with language supervision. Recent studies also repurposed diffusion models’ spatial and semantic priors for segmentation and other dense prediction tasks.

  • Semantic segmentation: Semantic segmentation assigns pixel-level segments and class labels, with CNNs such as DeepLabv3+ capturing multi-scale context.Transformer frameworks later unified panoptic, instance, and semantic segmentation, including Mask2Former and OneFormer.
  • Open-vocabulary segmentation: Open-vocabulary segmentation recognizes and segments arbitrary categories beyond the training set by leveraging language supervision.Early approaches aligned visual features with pretrained text embeddings, while CLIP helped establish vision–language methods in this direction.
  • Diffusion-based segmentation: Recent diffusion-based methods repurposed pretrained text-to-image models by extracting internal attention maps or latent features for segmentation.These studies showed that diffusion backbones encode strong spatial and semantic organization.
  • Diffusion-based segmentation: Diffusion models are being developed from generative priors into generalized segmentation learners while preserving visual knowledge and enabling explicit semantic control.This direction motivates using diffusion backbones beyond pure generation for structured prediction.
  • Dense prediction: Diffusion models have also shown promise for dense prediction, particularly monocular depth estimation through iterative denoising and fine-tuned diffusion U-Nets.DiffusionDepth and Marigold exemplify image-conditioned and pretrained-image-diffusion approaches to depth prediction.

3 Method

DiGSeg repurposes a pretrained diffusion model as a unified segmentation learner by modeling segmentation maps through conditional latent denoising. Visual latents, multi-scale CLIP text conditioning, and diffusion inference support semantic, open-vocabulary, and cross-domain segmentation.

  • Framework Overview: DiGSeg fine-tunes a pretrained diffusion model into a unified learner for semantic and open-vocabulary segmentation across diverse domains.The framework uses diffusion fine-tuning rather than treating the model solely as a feature extractor.
  • Conditional Diffusion Formulation: Segmentation is formulated as conditional latent diffusion, modeling pθ(y|x) and reversing the usual image-generation conditioning to generate segmentation-consistent latents from images.At inference, denoising reconstructs the segmentation latent from Gaussian noise conditioned on the input image.
  • Visual Latent Pathway: The Visual Latent Pathway encodes the RGB image and segmentation map into compact latents, replicating the single-channel mask across three channels for compatibility with the pretrained VAE.The VAE decoder reconstructs image and segmentation representations back to pixel space.
  • CLIP-Aligned Text Conditioner: A frozen CLIP text encoder injects class-name or natural-language embeddings into multiple U-Net scales through cross-attention, jointly refining global semantics and local spatial cues.The denoiser is conditioned on both the image latent and text embedding, enabling text-controllable open-vocabulary segmentation.
  • Inference and Thresholding: A fixed threshold τ = 0.7 consistently degrades performance across categories without post-processing, while threshold preferences vary with object size, texture, and spatial sparsity.Small or fine-grained objects typically favor lower thresholds, whereas larger homogeneous objects prefer higher thresholds.

4 Experiments

Experiments evaluate DiGSeg on semantic, open-vocabulary, and task-specific segmentation benchmarks using mIoU, with Dice additionally reported for medical segmentation. Results show strong generalization without architectural changes, data efficiency under limited supervision, and improved diffusion inference speed through DDIM-trailing.

  • Semantic and open-vocabulary segmentation: DiGSeg achieves state-of-the-art results on almost all evaluated semantic and open-vocabulary segmentation datasets against thirteen open-vocabulary baselines.Comparisons use mIoU and include MAFTP w/ MaskAdapter, SemLA, and OPMap-related baselines.
  • Task-specific segmentation: DiGSeg maintains significant performance across general, agricultural, and remote-sensing domains without changing its architecture.Task-specific evaluations include DeepGlobe, while Monuseg results are provided in the supplementary materials.
  • Data efficiency under limited supervision: Using half of ADE20K training data yields nearly identical performance to the full dataset, while quarter-data training remains impressively strong.The results indicate effective use of pretrained visual priors and generalizable representations in low-data scenarios.
  • Test-time ensembling: Ten test-time ensemble predictions improve ADE20K by 1.2% and COCO by 1.9% over a single prediction.Performance improves consistently with ensemble size, but the gain begins tapering after 10 predictions per sample.
  • Inference Speed: DiGSeg is slower than feed-forward segmenters, but DDIM-trailing makes it noticeably faster than previous diffusion-based segmentation models.This reflects the quality-efficiency trade-off typical of diffusion-based methods, while leaving room for further acceleration techniques.

5 Conclusion

DiGSeg repurposes pretrained diffusion models into generalized segmentation learners through latent-conditioned fine-tuning and a multi-scale CLIP-aligned text pathway. Evaluations show state-of-the-art performance, exceptional open-vocabulary generalization, and a unified bridge between generative modeling and dense visual understanding.

  • Ablations: Increasing the E-Step range improves segmentation accuracy but reduces inference speed.The ablation identifies an accuracy–speed trade-off in E-Step scheduling.
  • Contributions: DiGSeg repurposes pretrained diffusion models from pure generators into generalized segmentation learners using latent-conditioned fine-tuning and a multi-scale CLIP-aligned text pathway.The framework channels diffusion backbones’ visual priors through a unified segmentation interface.
  • Results: DiGSeg achieves state-of-the-art performance on standard benchmarks and exceptional open-vocabulary generalization across specialized domains.The evaluations span standard benchmarks and specialized domains.
  • Appendix: Additional implementation details, extended ablations, and further quantitative and qualitative results are provided in Sections A–D.Implementation details appear in Section A, extended ablations in Section B, and quantitative and qualitative results in Sections C and D, respectively.

A More Implementation Details · A.1 Multi-Resolution Noise Construction

The implementation constructs a multi-resolution noise pyramid for diffusion, using progressively lower-resolution noise components combined with exponentially decaying weights. This design controls the global influence of low-frequency noise across diffusion steps.

  • A.1 Multi-Resolution Noise Construction: DiGSeg constructs a noise pyramid with levels corresponding to progressively lower spatial resolutions.The pyramid is denoted {ϵ(i)}L, as stated in the passage.
  • A.1 Multi-Resolution Noise Construction: The construction follows a multi-scale noising strategy commonly used in latent diffusion models.The passage presents the noise pyramid as an adaptation of this established diffusion-noising pattern.
  • A.1 Multi-Resolution Noise Construction: Each noise level is generated by downsampling a base Gaussian tensor, sampling independent low-resolution Gaussian noise, and upsampling it to latent resolution.This procedure preserves a common latent resolution while introducing scale-specific noise.
  • A.1 Multi-Resolution Noise Construction: The multi-resolution formulation combines noise signals with exponentially decaying weights s^i.The weighting assigns progressively smaller contributions to higher-index levels.
  • A.1 Multi-Resolution Noise Construction: The parameter s satisfies 0 < s < 1 and controls the global influence of low-frequency noise components.Lower-frequency components are therefore governed by the selected decay factor.
  • A.1 Multi-Resolution Noise Construction: The noise formulation is defined over diffusion steps, with T denoting the total number of diffusion steps.The supplied passage identifies T as the total diffusion-step count.

A.2 Training Architecture Details · A.3 Inference Details

DiGSeg fine-tunes selected attention and projection components of a frozen SDv2 latent-diffusion backbone, then performs text-conditioned candidate selection, class-specific mask generation, and dense aggregation. Its inference protocol combines CLIP vocabulary filtering with diffusion denoising and pixel-wise background-aware labeling.

  • A.2 Training Architecture Details: DiGSeg uses the SDv2 U-Net while keeping the VAE encoder and decoder frozen for latent-space consistency.The pretrained multi-scale convolutional and attention structures are retained.
  • A.2 Training Architecture Details: Training fine-tunes cross-attention blocks at 1/4 and 1/8 resolutions, mid-block attention, and an MLP projecting CLIP ViT-L/14 text embeddings.Other convolutional layers remain frozen under the original SD configuration.
  • A.3 Inference Details: Inference constructs candidate class lists, injects text embeddings, and aggregates per-class predictions into a final dense segmentation map.These are the three stated components of the full inference protocol.
  • A.3 Inference Details: A CLIP open-vocabulary classifier proposes likely classes from a fixed benchmark vocabulary using the prompt "A photo of a class".Class text embeddings are normalized, precomputed, and cached.
  • A.3 Inference Details: Classes with image-text cosine similarity above the default threshold τ = 0.25 become initial candidates, with patch-level refinement recovering localized objects.Image embeddings are compared against every cached class embedding.
  • A.3 Inference Details: For each candidate class, DiGSeg injects its text embedding into diffusion cross-attention, runs DDIM-trailing denoising, and decodes a continuous class-specific mask logit map.The predicted mask is formally represented as Mĉ = fθ(x, tc).
  • A.3 Inference Details: The final segmentation assigns each pixel the candidate class with the highest predicted mask logit.All candidate mask logits are collected before pixel-wise aggregation.
  • A.3 Inference Details: Pixels whose maximum logit falls below τbg are assigned the background label instead of a candidate class.The passage states that this rule works well across all benchmarks.

B More Ablation Experiments

The study evaluates whether the Stable Diffusion VAE can faithfully encode segmentation masks in latent space. Results indicate that the VAE preserves boundaries and label semantics without degradation.

  • SD VAE with Segmentation Masks: The Stable Diffusion VAE is evaluated for encoding segmentation masks on ADE20K and COCO-Stuff using 1,000 randomly sampled ground-truth masks per dataset.Single-channel label maps are expanded to three channels, normalized to the VAE input range, decoded to RGB, and converted back to discrete labels through nearest-neighbor palette matching.
  • SD VAE with Segmentation Masks: The VAE is expressive enough to encode segmentation masks without degrading boundaries or label semantics, enabling effective learning in the latent space.Channel-wise consistency after VAE decoding remains extremely small across the evaluated datasets.

C More Quantitative results · C.1 Results on Pheno-Bench Dataset

The paper reports downstream results on Pheno-Bench and REFUGE-2, while noting unsatisfactory medical-domain performance potentially linked to CLIP’s limited medical understanding. On Pheno-Bench, DiGSeg achieves competitive baseline performance and transfers effectively to agricultural and biological imaging, especially for highly variable categories.

  • C More Quantitative results: The downstream evaluation includes Pheno-Bench and REFUGE-2 datasets.Pheno-Bench results appear in Tab. 8(a), while REFUGE-2 results appear in Tab. 8(b).
  • C More Quantitative results: Medical-domain performance is reported as unsatisfactory.The passage attributes this possibly to CLIP’s insufficient understanding of medical data.
  • C.1 Results on Pheno-Bench Dataset: Pheno-Bench evaluates fine-grained plant phenotyping.The task requires distinguishing subtle morphological variations among leaves, stems, and fruits.
  • C.1 Results on Pheno-Bench Dataset: DiGSeg achieves competitive performance against existing Pheno-Bench baselines.The results are presented in Tab. 8(a).
  • C.1 Results on Pheno-Bench Dataset: The learned representations transfer effectively to agricultural and biological imaging domains.This transferability is identified as a key implication of the Pheno-Bench results.
  • C.1 Results on Pheno-Bench Dataset: Improvement is particularly notable in categories with high intra-class variability.The passage connects this observation to the Pheno-Bench category-level results.

C.2 Results on REFUGE-2 Dataset · D More Qualitative results

On REFUGE-2, DiGSeg segments optic-disc and optic-cup regions and outperforms several strong baselines, but its absolute medical-domain performance remains limited. Additional qualitative results cover open-vocabulary and semantic segmentation tasks.

  • C.2 Results on REFUGE-2 Dataset: REFUGE-2 evaluates optic-disc and optic-cup segmentation in retinal fundus images for glaucoma diagnosis.The dataset targets two clinically relevant retinal structures.
  • C.2 Results on REFUGE-2 Dataset: DiGSeg outperforms several strong baselines on REFUGE-2 medical image segmentation.The comparison is reported in Table 8(b).
  • C.2 Results on REFUGE-2 Dataset: Despite this advantage, absolute performance on REFUGE-2 is comparatively limited relative to the method’s results in other domains.The paper characterizes the medical-domain results as limited despite baseline improvements.
  • C.2 Results on REFUGE-2 Dataset: The paper attributes the REFUGE-2 performance gap primarily to domain mismatch between medical imagery and CLIP’s natural-image training.CLIP was pretrained on web-scale natural image-text pairs.
  • C.2 Results on REFUGE-2 Dataset: Table 8(b) presents a quantitative comparison for REFUGE-2 medical image segmentation.The medical-scenario-specific module is visually distinguished in gray.
  • D More Qualitative results: The paper includes additional qualitative results for open-vocabulary segmentation tasks.These results are presented in the qualitative-results section.
  • D More Qualitative results: The additional qualitative results also cover semantic segmentation tasks.The section explicitly presents both open-vocabulary and semantic segmentation examples.

D.1 Open-Vocabulary Segmentation · D.2 Semantic Segmentation · E Discussion

DiGSeg delivers accurate, boundary-preserving open-vocabulary and semantic segmentation across diverse benchmarks, including fine-grained, cluttered, occluded, and scale-varying scenes. These results support the broader view that diffusion-model visual priors transfer from generation to general visual understanding.

  • D.1 Open-Vocabulary Segmentation: On A-847, DiGSeg segments rare and visually similar categories with sharp boundaries and fewer confusions, including for small or occluded objects.A-847 contains 847 categories and poses fine-grained recognition challenges.
  • D.1 Open-Vocabulary Segmentation: On PC-459, DiGSeg remains robust in cluttered scenes, segments complex shapes, and captures contextual relationships between foreground and background regions.PC-459 requires differentiation among 459 categories with diverse appearances.
  • D.1 Open-Vocabulary Segmentation: On A-150 and PC-59, DiGSeg produces clean masks, smooth boundaries, accurate region assignments, and predictions aligned with object contours.The A-150 and PC-59 results indicate effective use of semantic and spatial cues across commonly seen categories and diverse indoor and outdoor scenes.
  • D.2 Semantic Segmentation: On ADE20K, DiGSeg handles complex compositions and precisely segments both stuff and thing classes across object scales and lighting conditions.Examples include sky, road, wall, person, car, and furniture.
  • D.2 Semantic Segmentation: On COCO, DiGSeg produces sharp, accurately assigned segmentation maps despite heavy occlusion and significant scale variation, corroborating its quantitative findings and generalization capability.The semantic-segmentation evaluation covers two large-scale closed-set benchmarks: ADE20K and COCO.
  • E Discussion: DiGSeg aligns with related work showing that diffusion U-Net priors learned for visual generation can transfer to visual understanding, dissolving the generative–discriminative boundary.The discussion positions these approaches as part of a broader trajectory in which large-scale visual-generation priors serve as substrates for general visual understanding.
Loading 2604.24575v1…