GenMask: Anpassung von DiT für Segmentierung durch direkte Maskierung
GenMask: Adapting DiT for Segmentation via Direct Mask
March 25, 2026
Autoren: Yuhuan Yang, Xianwei Zhuang, Yuxuan Cai, Chaofan Ma, Shuai Bai, Jiangchao Yao, Ya Zhang, Junyang Lin, Yanfeng Wang
cs.AI
Zusammenfassung
Aktuelle Ansätze für Segmentierung nutzen vortrainierte generative Modelle als Merkmalsextraktoren und behandeln Segmentierung als Downstream-Anpassungsaufgabe durch indirekte Merkmalsgewinnung. Diese implizite Verwendung leidet unter einer grundlegenden Fehlausrichtung der Repräsentation. Sie ist zudem stark von indirekten Merkmalsextraktions-Pipelines abhängig, was den Arbeitsablauf verkompliziert und die Anpassung einschränkt. In diesem Beitrag argumentieren wir, dass Segmentierungsaufgaben direkt auf generative Weise trainiert werden sollten, anstatt eine indirekte Anpassung vorzunehmen. Wir identifizieren ein Haupthindernis für diese vereinheitlichte Formulierung: VAE-Latenzvariablen von binären Masken sind scharf verteilt, rauschrobust und linear separierbar, was sie von Latenzvariablen natürlicher Bilder unterscheidet. Um diese Lücke zu schließen, führen wir eine Zeitschritt-Abtaststrategie für binäre Masken ein, die extreme Rauschpegel für die Segmentierung und moderate Rauschpegel für die Bildgenerierung betont, um harmonisches gemeinsames Training zu ermöglichen. Wir stellen GenMask vor, ein DiT, das darauf trainiert wird, schwarz-weiße Segmentierungsmasken sowie farbige Bilder im RGB-Raum unter dem ursprünglichen generativen Ziel zu erzeugen. GenMask bewahrt die ursprüngliche DiT-Architektur, während es den Bedarf an speziell für Segmentierungsaufgaben entwickelten Merkmalsextraktions-Pipelines beseitigt. Empirisch erzielt GenMask state-of-the-art Leistung auf Referenz- und Reasoning-Segmentierungs-Benchmarks, und Ablationstudien quantifizieren den Beitrag jeder Komponente.
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.