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

cs.CVMar 25, 2026

GenMask: Adapting DiT for Segmentation via Direct Mask Generation

Authors: Yuhuan Yang, Xianwei Zhuang, Yuxuan Cai, Chaofan Ma, Shuai Bai, Jiangchao Yao, Ya Zhang, Junyang Lin,

Organizations: Cooperative Medianet Innovation Center, Shanghai Jiao Tong University · Alibaba Group · School of Artificial Intelligence, Shanghai Jiao Tong University · Institute of Artificial Intelligence for Medicine, Shanghai Jiao Tong University School of Medicine

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.

Figures & tables

Figure 1: Overall architecture of our model. Here brown and green denote the segmentation and generation data flow respectively. Generation task follows standard diffusion training process, and its timesteps sampling strategy is similar to Stable Diffusion 3 [ 16 ] , emphasizing on intermediate denoising steps. For segmentation task, we use an extreme long tailed sampling strategy. We also add VAE representation of the input image into DiT to supplement low-level information for segmentation.

Figure 2: The process of adding noise to natural image and binary mask. Compared with natural image, binary mask is much more robust to noise.

Panel of Figure 3 — Figure 3: VAE features for binary segmentation masks are linearly separable. (a) The input segmentation masks. (b) PCA label of its VAE representation. (c) The difference between the input masks and the PCA label after histogram normalization.

Figure 4: SVM validation accuracy on VAE embeddings of binary masks under different noise levels. The embeddings remain linearly separable under low noise, while only high-intensity perturbations substantially degrade separability.

Panel of Figure 5 — Figure 5: (Up) Importance resampling function for segmentation task with different hyperparameter $a$ . More extreme $a$ value means the distribution is more concentrated in high noise intensity regime. (Down) Different sampling strategies for segmentation and generation task separately. Generation task uses a relatively uniform sampling strategy, only emphasize the intermediate denoising steps. While segmentation task uses a extreme long tailed sampling strategy, with peak value 8× higher.

Figure 6: Three variants of segmentation supervision formats. (a) MSE loss in VAE space or BCE loss in RGB space with the help of VAE encoder-decoder. (b) By replacing the VAE decoder with a linear layer, we can apply BCE loss without backpropagation through frozen VAE.

Method	RefCOCO	RefCOCO+	RefCOCO-g
Method	test A / B	test A / B	val / test
Metric: oIoU
CRIS [ 71 ]	71.5 / 62.2	64.1 / 48.4	56.6 / 57.4
VLT [ 15 ]	76.0 / 69.6	68.4 / 56.9	63.5 / 66.2
LAVT [ 87 ]	75.8 / 68.8	68.4 / 55.1	61.2 / 62.1
BKINet [ 14 ]	76.4 / 69.4	69.9 / 53.4	64.2 / 63.8

Table 1: Performance comparison on RES benchmarks. The results are evaluated on RefCOCO, RefCOCO+ [ 26 ] , and RefCOCOg [ 51 ] (UMD partition) using mIoU and oIoU.

Model	Val Set		Test Set
Model	gIoU	cIoU	gIoU	cIoU
SEEM [ 95 ]	25.5	21.2	24.3	18.7
Grounded SAM [ 58 ]	26.0	14.5	21.3	16.4
OVSeg [ 33 ]	28.5	18.6	26.1	20.8
GLaMM [ 56 ]	47.4	47.2	–	–
SAM4MLLM [ 10 ]	46.7	48.1	–	–

Table 2: Performance comparison on ReasonSeg benchmarks. Here ∗ denotes the model finetuned on ReasonSeg training dataset.

Figure 7: Visualization results of our model. We are able to generate both colorful images (up) and binary masks (down) simultaneously. For segmentation, we apply the output binary mask on the original image for better visualization.

$a$ value	RefCOCO		RefCOCO+		RefCOCO-g
$a$ value	mIoU	oIoU	mIoU	oIoU	mIoU	oIoU
0.05	82.2	81.3	75.8	73.5	77.7	76.0
0.1	78.1	77.6	69.3	68.1	73.7	72.3
0.5	66.0	66.0	52.7	53.3	57.5	56.6

Table 3: Performance comparison of different sampling strategies for segmentation task. Adjusting segmentation sampling strategy to differ from generation is essential. A relatively smooth distribution (larger

a

) leads to degraded model performance, whereas the most extreme long-tailed distribution (

a=0.05

) achieves the best results.

Loss	RefCOCO		RefCOCO+		RefCOCO-g
Loss	mIoU	oIoU	mIoU	oIoU	mIoU	oIoU
MSE	82.2	81.3	75.8	73.5	77.7	76.0
BCE	78.1	77.6	73.1	70.3	75.0	72.2
BCE w/ linear	81.3	80.4	74.8	72.4	76.9	75.0

Table 4: Performance comparison of different segmentation supervision formats. MSE performs best since it’s the most closely aligned with the original DiT training objective. Raw BCE performs poorly because it requires backpropagation through VAE, making optimization more difficult. Replacing the VAE decoder with a linear layer alleviates the optimization difficulty.

Gen data	RefCOCO		RefCOCO+		RefCOCO-g
Gen data	mIoU	oIoU	mIoU	oIoU	mIoU	oIoU
$\checkmark$	82.2	81.3	75.8	73.5	77.7	76.0
$\times$	81.0	80.6	74.2	72.4	76.7	75.0

Table 5: Ablation study on mix training. Adding generation data in an 1:1 ratio is beneficial to segmentation performance.

VAE	RefCOCO		RefCOCO+		RefCOCO-g
VAE	mIoU	oIoU	mIoU	oIoU	mIoU	oIoU
$\checkmark$	82.2	81.3	75.8	73.5	77.7	76.0
$\times$	74.1	73.2	68.8	67.6	72.4	71.1

Table 6: Ablation study on VAE input for segmentation task. VAE input is crucial for segmentation task, it provides low-level information which is essential for accurate pixel-level prediction.

Image Segmentation

Image segmentation assigns pixels or voxels to objects and regions, supporting measurement, scene understanding, and decisions across image domains.

See all 399 papers

GET: Generative Embedding Translation for Medical Image Segmentation
Aug 23, 2026Md Maklachur Rahman, Md Hasan Al Banna, Saraf Anjum +2Image Segmentation Medical Image Segmentation
Texas A&M University, College Station, TX, USA · Independent Researcher · Alfa Laval
Diffusion Model as a Generalist Segmentation Learner
Apr 27, 2026Haoxiao Wang, Antao Xiang, Haiyang Sun +8Image Segmentation Text-Guided Image Segmentation
Zhejiang University · South China University of Technology · Nanjing University +1
FocusDiT: Masking Queries in Diffusion Transformers for Fine-grained Image Generation
Jun 1, 2026Xueji Fang, Liyuan Ma, Jianhao Zeng +3Diffusion Transformer Image Generation
Zhejiang University · Westlake University · Hangzhou, China

Aug 23, 2026J/K move · Enter open · S save

GET: Generative Embedding Translation for Medical Image Segmentation

Md Maklachur Rahman, Md Hasan Al Banna, Saraf Anjum, Mahmudul Hasan+1

Texas A&M University, College Station, TX, USA · Independent Researcher · Alfa Laval

Image Segmentation Medical Image Segmentation Diffusion Models for Medical Image Segmentation

Open paper AskPDF arXiv

Generative segmentation provides an alternative to direct pixel-wise prediction by operating on learned latent representations, but effective image-to-mask translation must preserve target structure while remaining computationally efficient. We propose Generative Embedding Translation (GET), a structured embedding-translation framework that progressively transforms image embeddings into mask embeddings within the frozen latent space of a Stable Diffusion VAE. GET uses a U-Net-style Embedding Translation Network with 1.07M trainable parameters, combining Mobile Bottleneck Convolutions, Subsampled Self-Attention, and Multi-scale Feature Enrichment for local modeling, global context, and multi-scale refinement. Across five medical segmentation datasets, GET outperforms generative, CNN, and Transformer baselines. Compared with the strongest generative baseline, GMS, GET improves average Dice and IoU by 0.93% and 1.26%, reduces HD95 by 0.81 pixels, and uses 31.41% fewer trainable parameters. Under bidirectional BUS-BUSI domain shift, GET further improves Dice and IoU by 3.51% and 3.39%, while reducing HD95 by 27.37 pixels. Our code is available at: https://github.com/maklachur/GET.

Abstract

Image Segmentation

Explore similar work

GET: Generative Embedding Translation for Medical Image Segmentation

Diffusion Model as a Generalist Segmentation Learner

FocusDiT: Masking Queries in Diffusion Transformers for Fine-grained Image Generation