GenMask: DiT aanpassen voor segmentatie via directe masker
GenMask: Adapting DiT for Segmentation via Direct Mask
March 25, 2026
Auteurs: Yuhuan Yang, Xianwei Zhuang, Yuxuan Cai, Chaofan Ma, Shuai Bai, Jiangchao Yao, Ya Zhang, Junyang Lin, Yanfeng Wang
cs.AI
Samenvatting
Recente benaderingen voor segmentatie maken gebruik van voorgetrainde generatieve modellen als feature-extractors, waarbij segmentatie wordt behandeld als een downstream-aanpassingstaak via indirecte feature-retrieval. Dit impliciete gebruik lijdt onder een fundamentele misalignering in de representatie. Het is ook sterk afhankelijk van indirecte feature-extractiepijplijnen, wat de workflow compliceert en de aanpassing beperkt. In dit artikel beargumenteren wij dat segmentatietaken, in plaats van indirecte aanpassing, direct op een generatieve manier getraind moeten worden. Wij identificeren een belangrijk obstakel voor deze geünificeerde formulering: VAE-latents van binaire maskers zijn scherp verdeeld, ruisrobuust en lineair scheidbaar, onderscheiden van latents van natuurlijke afbeeldingen. Om deze kloof te overbruggen, introduceren wij een timestep-samplingstrategie voor binaire maskers die extreme ruisniveaus benadrukt voor segmentatie en gematigde ruis voor beeldgeneratie, waardoor harmonieuze gezamenlijke training mogelijk wordt. Wij presenteren GenMask, een DiT die getraind wordt om zwart-wit segmentatiemaskers evenals kleurrijke afbeeldingen in RGB-ruimte te genereren onder het oorspronkelijke generatieve doel. GenMask behoudt de oorspronkelijke DiT-architectuur terwijl de noodzaak van feature-extractiepijplijnen specifiek voor segmentatietaken wordt geëlimineerd. Empirisch behaalt GenMask state-of-the-art prestaties op refererende en redenerende segmentatiebenchmarks, en ablatiestudies kwantificeren de bijdrage van elke component.
English
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.