GenMask: Adaptando DiT para Segmentação via Máscara Direta

Resumo

Abordagens recentes para segmentação têm aproveitado modelos generativos pré-treinados como extratores de características, tratando a segmentação como uma tarefa de adaptação downstream via recuperação indireta de features. Este uso implícito sofre de um desalinhamento fundamental na representação. Também depende fortemente de pipelines indiretos de extração de características, que complicam o fluxo de trabalho e limitam a adaptação. Neste artigo, argumentamos que, em vez de adaptação indireta, as tarefas de segmentação devem ser treinadas diretamente de maneira generativa. Identificamos um obstáculo fundamental para esta formulação unificada: os latentes de VAE para máscaras binárias são distribuídos de forma abrupta, robustos ao ruído e linearmente separáveis, distintos dos latentes de imagens naturais. Para superar esta lacuna, introduzimos uma estratégia de amostragem de timesteps para máscaras binárias que enfatiza níveis extremos de ruído para segmentação e ruído moderado para geração de imagens, permitindo um treinamento conjunto harmonioso. Apresentamos o GenMask, um DiT treinado para gerar máscaras de segmentação em preto e branco, bem como imagens coloridas no espaço RGB, sob o objetivo generativo original. O GenMask preserva a arquitetura DiT original enquanto elimina a necessidade de pipelines de extração de características específicas para tarefas de segmentação. Empiricamente, o GenMask atinge desempenho state-of-the-art em benchmarks de segmentação por referência e por raciocínio, e ablations quantificam a contribuição 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.