GenMask : Adaptation de DiT pour la segmentation via un masque direct
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
Résumé
Les approches récentes en segmentation ont exploité des modèles génératifs pré-entraînés comme extracteurs de caractéristiques, traitant la segmentation comme une tâche d'adaptation en aval via une extraction indirecte de caractéristiques. Cette utilisation implicite souffre d'un désalignement fondamental dans la représentation. Elle dépend également fortement de pipelines d'extraction de caractéristiques indirects, qui complexifient le flux de travail et limitent l'adaptation. Dans cet article, nous soutenons qu'au lieu d'une adaptation indirecte, les tâches de segmentation devraient être entraînées directement de manière générative. Nous identifions un obstacle clé à cette formulation unifiée : les latents des masques binaires dans les VAE ont une distribution abrupte, sont robustes au bruit et linéairement séparables, distincts des latents des images naturelles. Pour combler cet écart, nous introduisons une stratégie d'échantillonnage des pas de temps pour les masques binaires qui privilégie les niveaux de bruit extrêmes pour la segmentation et un bruit modéré pour la génération d'images, permettant un entraînement conjoint harmonieux. Nous présentons GenMask, un DiT entraîné pour générer des masques de segmentation en noir et blanc ainsi que des images colorées dans l'espace RGB selon l'objectif génératif original. GenMask préserve l'architecture DiT originale tout en supprimant le besoin de pipelines d'extraction de caractéristiques spécifiques aux tâches de segmentation. Empiriquement, GenMask atteint des performances state-of-the-art sur des benchmarks de segmentation par référence et raisonnement, et des ablations quantifient la contribution de chaque composant.
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.