GenMask: Adapting DiT for Segmentation via Direct Mask

2026-03-24

AI Summary

This work addresses the limitations of existing approaches that repurpose pretrained generative models as feature extractors for segmentation, which suffer from representation misalignment and reliance on complex, indirect pipelines. Instead, the authors propose directly training a segmentation model in a generative manner by extending the DiT architecture to jointly generate color images and binary segmentation masks under the original generative objective, eliminating the need for customized feature extraction modules. The key innovation lies in a mask-aware differential noise-step sampling strategy, combined with coordinated scheduling of VAE latent space and diffusion timesteps, enabling unified generation of both images and masks in RGB space. The method achieves state-of-the-art performance on referring expression and reasoning-based segmentation benchmarks, with ablation studies confirming the effectiveness of each proposed component.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Machine Learning: Deep Generative Models & Autoencoders
• Natural Language Processing: Generation

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Retrieval-Augmented Generation (RAG) and multi-modal RAG
• Social Networks and Social Media: Generative AI / large language models and their impact on social systems

Abstract

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.

Problem

• segmentation
• generative models
• feature extraction
• representation alignment
• downstream adaptation

Innovation

• GenMask
• DiT
• direct mask generation
• timesteps sampling strategy
• generative segmentation