GenMask: Adaptación de DiT para Segmentación mediante Máscara Directa
GenMask: Adapting DiT for Segmentation via Direct Mask
March 25, 2026
Autores: Yuhuan Yang, Xianwei Zhuang, Yuxuan Cai, Chaofan Ma, Shuai Bai, Jiangchao Yao, Ya Zhang, Junyang Lin, Yanfeng Wang
cs.AI
Resumen
Los enfoques recientes para la segmentación han aprovechado modelos generativos preentrenados como extractores de características, tratando la segmentación como una tarea de adaptación posterior mediante la recuperación indirecta de características. Este uso implícito adolece de una desalineación fundamental en la representación. También depende en gran medida de pipelines de extracción de características indirectos, lo que complica el flujo de trabajo y limita la adaptación. En este artículo, sostenemos que, en lugar de una adaptación indirecta, las tareas de segmentación deberían entrenarse directamente de manera generativa. Identificamos un obstáculo clave para esta formulación unificada: los espacios latentes de las máscaras binarias en un VAE presentan una distribución abrupta, son robustos al ruido y linealmente separables, a diferencia de los espacios latentes de imágenes naturales. Para salvar esta brecha, introducimos una estrategia de muestreo de pasos de tiempo para máscaras binarias que enfatiza niveles de ruido extremos para la segmentación y ruido moderado para la generación de imágenes, permitiendo un entrenamiento conjunto armonioso. Presentamos GenMask, un DiT que se entrena para generar máscaras de segmentación en blanco y negro, así como imágenes coloridas en el espacio RGB, bajo el objetivo generativo original. GenMask preserva la arquitectura DiT original mientras elimina la necesidad de pipelines de extracción de características específicos para tareas de segmentación. Empíricamente, GenMask alcanza un rendimiento de vanguardia en benchmarks de segmentación referencial y de razonamiento, y los estudios de ablación cuantifican la contribución de cada 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.