GenMask: Adattamento di DiT per la Segmentazione tramite Maschera Diretta
GenMask: Adapting DiT for Segmentation via Direct Mask
March 25, 2026
Autori: Yuhuan Yang, Xianwei Zhuang, Yuxuan Cai, Chaofan Ma, Shuai Bai, Jiangchao Yao, Ya Zhang, Junyang Lin, Yanfeng Wang
cs.AI
Abstract
Gli approcci recenti per la segmentazione hanno sfruttato modelli generativi preaddestrati come estrattori di feature, trattando la segmentazione come un task di adattamento downstream tramite recupero indiretto delle caratteristiche. Questo utilizzo implicito soffre di un disallineamento fondamentale nella rappresentazione. Dipende inoltre fortemente da pipeline di estrazione indiretta delle feature, che complicano il flusso di lavoro e limitano l'adattamento. In questo articolo, sosteniamo che invece di un adattamento indiretto, i task di segmentazione dovrebbero essere addestrati direttamente in maniera generativa. Identifichiamo un ostacolo chiave per questa formulazione unificata: i latenti VAE delle maschere binarie hanno distribuzione netta, sono robusti al rumore e linearmente separabili, distinti dai latenti delle immagini naturali. Per colmare questa lacuna, introduciamo una strategia di campionamento dei timestep per le maschere binarie che enfatizza livelli di rumore estremi per la segmentazione e moderati per la generazione di immagini, consentendo un addestramento congiunto armonioso. Presentiamo GenMask, un DiT addestrato per generare maschere di segmentazione in bianco e nero così come immagini colorate nello spazio RGB sotto l'obiettivo generativo originale. GenMask preserva l'architettura originale del DiT eliminando al contempo la necessità di pipeline di estrazione di feature specifiche per i task di segmentazione. Empiricamente, GenMask raggiunge prestazioni state-of-the-art su benchmark di segmentazione referenziale e ragionata, e le ablation study quantificano il contributo di ciascun componente.
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.