# GenMask: Adapting DiT for Segmentation via Direct Mask Generation

arXiv: https://arxiv.org/abs/2603.23906 · HTML: https://mlsanity.com/paper/2603.23906/

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 78 of its 89 references that are in the map.
- First posted on arXiv 25 Mar 2026.

## Builds on (most cited first, 30 of 78)

- [Learning Transferable Visual Models From Natural Language Supervision](https://mlsanity.com/paper/2103.00020/) (2021)
- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [Scalable Diffusion Models with Transformers](https://mlsanity.com/paper/2212.09748/) (2022)
- [Classifier-Free Diffusion Guidance](https://mlsanity.com/paper/2207.12598/) (2022)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [Qwen2.5-VL Technical Report](https://mlsanity.com/paper/2502.13923/) (2025)
- [SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis](https://mlsanity.com/paper/2307.01952/) (2023)
- [Scaling Rectified Flow Transformers for High-Resolution Image Synthesis](https://mlsanity.com/paper/2403.03206/) (2024)
- [Flow Straight and Fast: Learning to Generate and Transfer Data with Rectified Flow](https://mlsanity.com/paper/2209.03003/) (2022)
- [Wan: Open and Advanced Large-Scale Video Generative Models](https://mlsanity.com/paper/2503.20314/) (2025)
- [Semantic Understanding of Scenes through the ADE20K Dataset](https://mlsanity.com/paper/1608.05442/) (2016)
- [FLUX.1 Kontext: Flow Matching for In-Context Image Generation and Editing in Latent Space](https://mlsanity.com/paper/2506.15742/) (2025)
- [Grounded SAM: Assembling Open-World Models for Diverse Visual Tasks](https://mlsanity.com/paper/2401.14159/) (2024)
- [COCO-Stuff: Thing and Stuff Classes in Context](https://mlsanity.com/paper/1612.03716/) (2016)
- [PixArt-α: Fast Training of Diffusion Transformer for Photorealistic Text-to-Image Synthesis](https://mlsanity.com/paper/2310.00426/) (2023)
- [LISA: Reasoning Segmentation via Large Language Model](https://mlsanity.com/paper/2308.00692/) (2023)
- [Open-Vocabulary Semantic Segmentation with Mask-adapted CLIP](https://mlsanity.com/paper/2210.04150/) (2022)
- [Segment Everything Everywhere All at Once](https://mlsanity.com/paper/2304.06718/) (2023)
- [GroupViT: Semantic Segmentation Emerges from Text Supervision](https://mlsanity.com/paper/2202.11094/) (2022)
- [Scaling Open-Vocabulary Image Segmentation with Image-Level Labels](https://mlsanity.com/paper/2112.12143/) (2021)
- [GLaMM: Pixel Grounding Large Multimodal Model](https://mlsanity.com/paper/2311.03356/) (2023)
- [Emergent Correspondence from Image Diffusion](https://mlsanity.com/paper/2306.03881/) (2023)
- [Step1X-Edit: A Practical Framework for General Image Editing](https://mlsanity.com/paper/2504.17761/) (2025)
- [DiffusionDB: A Large-scale Prompt Gallery Dataset for Text-to-Image Generative Models](https://mlsanity.com/paper/2210.14896/) (2022)
- [BLIP3-o: A Family of Fully Open Unified Multimodal Models-Architecture, Training and Dataset](https://mlsanity.com/paper/2505.09568/) (2025)
- [CRIS: CLIP-Driven Referring Image Segmentation](https://mlsanity.com/paper/2111.15174/) (2021)
- [LAVT: Language-Aware Vision Transformer for Referring Image Segmentation](https://mlsanity.com/paper/2112.02244/) (2021)
- [PixArt-Sigma: Weak-to-Strong Training of Diffusion Transformer for 4K Text-to-Image Generation](https://mlsanity.com/paper/2403.04692/) (2024)
- [Your Diffusion Model is Secretly a Zero-Shot Classifier](https://mlsanity.com/paper/2303.16203/) (2023)

## Built on by (most cited first, 3 of 3)

- [ChangeFlow -- Latent Rectified Flow for Change Detection in Remote Sensing](https://mlsanity.com/paper/2605.15375/) (2026)
- [Generative Semantic Segmentation via an Observable Semantic-Image Interface and Hierarchical Generator Evidence Alignment](https://mlsanity.com/paper/2608.11537/) (2026)
- [TRACE-Bench: Decomposing and Diagnosing Multi-Reference Image Generation](https://mlsanity.com/paper/2608.16765/) (2026)

## Abstract (arXiv)

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.