Source-linked AI summary
GenMask: Adapting DiT for Segmentation via Direct Mask Generation
Yuhuan Yang, Xianwei Zhuang, Yuxuan Cai, Chaofan Ma, Shuai Bai, Jiangchao Yao, Ya Zhang, Junyang Lin, Yanfeng Wang
TL;DR
Existing diffusion-based segmentation methods rely on indirect feature extraction and face a mismatch between natural-image and binary-mask representations. GenMask trains a DiT to generate masks and images jointly, using distinct timestep sampling strategies to bridge this gap while preserving the original architecture. It achieves state-of-the-art performance on referring and reasoning segmentation benchmarks and supports joint training with generated data.
Problem
Diffusion-based segmentation relies on indirect feature extraction and mismatched representations because natural-image VAE features differ from compact binary-mask labels.
Method
GenMask trains a DiT under a unified conditional generative objective, using extreme-noise sampling for masks, moderate-noise sampling for images, and VAE image shortcuts for segmentation.
Results
GenMask achieves state-of-the-art performance on referring and reasoning segmentation benchmarks while ablations identify timestep sampling, generative-data mixing, and VAE shortcuts as important components.
Takeaways & Limitations
Direct mask generation removes bespoke feature-extraction pipelines and allows generated image data to be incorporated into segmentation training.
Abstract
from arXiv · showhide
Recent approaches for segmentation have leveraged pretrained generative models as feature extractors, treating segmentation as a downstream adaptation task via indirect feature retrieval. This implicit use suffers from a fundamental misalignment in representation. It also depends heavily on indirect feature extraction pipelines, which complicate the workflow and limit adaptation. In this paper, we argue that instead of indirect adaptation, segmentation tasks should be trained directly in a generative manner. We identify a key obstacle to this unified formulation: VAE latents of binary masks are sharply distributed, noise robust, and linearly separable, distinct from natural image latents. To bridge this gap, we introduce timesteps sampling strategy for binary masks that emphasizes extreme noise levels for segmentation and moderate noise for image generation, enabling harmonious joint training. We present GenMask, a DiT trains to generate black-and-white segmentation masks as well as colorful images in RGB space under the original generative objective. GenMask preserves the original DiT architecture while removing the need of feature extraction pipelines tailored for segmentation tasks. Empirically, GenMask attains state-of-the-art performance on referring and reasoning segmentation benchmarks and ablations quantify the contribution of each component.
1. Introduction
Existing diffusion-based segmentation methods use pretrained generative models indirectly, creating representation mismatch and cumbersome feature-extraction pipelines. GenMask instead trains a DiT directly to generate segmentation masks within a unified generative framework.
- Text-based segmentation predicts binary masks from natural-language descriptions, increasingly using pretrained visual representations for downstream adaptation.
- Existing diffusion-based methods extract hidden features through denoising or inversion before producing segmentation masks.
- Diffusion-model segmentation faces a mismatch between low-level VAE features and compact semantic labels, while indirect extraction pipelines complicate adaptation.
- GenMask directly generates black-and-white segmentation masks in RGB space using the original DiT architecture and generative objective.
- Generation and segmentation share one conditional generation framework, with separate timestep strategies for intermediate image denoising and extreme-noise segmentation training.
- The model combines a pretrained DiT with a vision-language encoder and injects the input image’s VAE latent to provide low-level cues for pixel-level segmentation.
2. Method
GenMask unifies text-to-image generation and language-guided segmentation in one DiT framework, adapting the timestep schedule to the distinct latent behavior of binary masks. The method emphasizes high-noise training for segmentation and supports one-step mask decoding without architectural changes.
- Architecture Overview: GenMask integrates text-to-image generation and language-guided segmentation through the same diffusion training process, differing only in timestep sampling.Segmentation uses an aggressively long-tailed schedule focused on high-noise regions.
- Latent Distribution for Binary Masks: Binary-mask latents are noise-robust and occupy a narrow, biased region unlike natural-image latents, creating a distributional mismatch for generative modeling.At high noise, mask shape and boundaries remain recognizable while natural-image content is largely destroyed.
- Latent Distribution for Binary Masks: VAE representations of binary masks are effectively linearly separable, with a single PCA component producing labels highly similar to the input masks.The toy analysis uses N = 100 masks and VAE latent dimension d = 16.
- Latent Distribution for Binary Masks: Low-noise perturbations preserve mask-embedding separability, whereas only high-intensity noise substantially degrades it.This motivates concentrating segmentation training on high-noise samples rather than using the image-generation schedule.
- Time Shift for Segmentation: The segmentation timestep distribution concentrates 90% of samples in the high-noise region, with a 13% peak that is over 8× higher than generation’s peak.Generation emphasizes intermediate noise levels with a comparatively mild, near-uniform schedule.
- One-step Inference for Segmentation: One-step mask decoding preserves the original diffusion architecture and avoids additional training parameters, while segmentation inference remains a single forward pass.The generated latent is decoded into the final mask after fixing t = 1.
3. Experiments
GenMask is evaluated on referring and reasoning segmentation benchmarks, with ablations examining sampling, supervision format, mixed training, and low-level VAE inputs. The model also generates colorful images and binary masks.
- Benchmark evaluation: GenMask is compared on RefCOCO, RefCOCO+, and Ref-COCOg using mIoU and oIoU.
- Benchmark evaluation: GenMask is evaluated on ReasonSeg, including a finetuned variant marked with an asterisk.
- Ablation studies: The most extreme long-tailed sampling strategy, a = 0.05, achieves the best segmentation performance, while smoother distributions degrade results.
- Ablation studies: MSE loss in VAE space performs best because it most closely matches the original DiT training objective; raw BCE is harder to optimize through the VAE.
- Ablation studies: Adding text-to-image data at a 1:1 ratio improves segmentation performance in joint training.
- Ablation studies: Removing the VAE input causes a significant decline in segmentation performance, underscoring the importance of low-level information for pixel-level prediction.
- Visualization: GenMask simultaneously generates colorful images and binary masks, with masks overlaid on original images for visualization.
4. Related Work
Related work has used latent diffusion models both as generative systems and as sources of transferable representations. Segmentation methods commonly extract diffusion features indirectly and decode them with external components.
- Latent diffusion models perform diffusion in compressed latent spaces, while transformer denoisers support scalable text-to-image, text-to-video, and unified generation-editing systems.
- Diffusion-model representations have been used for discrimination, hybrid generation-recognition architectures, and distillation into downstream tasks.
- Diffusion-based segmentation methods use frozen feature decoding, inversion, activation aggregation, or feature distillation with external decoders.
5. Conclusion
GenMask treats segmentation as conditional generation and bridges mask-image latent differences with separate timestep sampling for joint training. Experiments report state-of-the-art performance on referring and reasoning segmentation benchmarks.
- GenMask directly treats mask production as a conditional generation problem rather than using indirect feature extraction.
- A separate timestep sampling strategy bridges binary-mask and natural-image VAE latents for joint segmentation and generation training.
- GenMask removes bespoke feature-extraction pipelines, closes the optimization gap with pretraining, and incorporates generated data into segmentation training.
- Experiments on referring and reasoning segmentation benchmarks demonstrate GenMask’s state-of-the-art performance.
- Future work includes scaling to larger DiT backbones and extending the paradigm to medical and audio-visual segmentation.