GenMask: Адаптация DiT для сегментации с помощью прямого маскирования
GenMask: Adapting DiT for Segmentation via Direct Mask
March 25, 2026
Авторы: Yuhuan Yang, Xianwei Zhuang, Yuxuan Cai, Chaofan Ma, Shuai Bai, Jiangchao Yao, Ya Zhang, Junyang Lin, Yanfeng Wang
cs.AI
Аннотация
В последних подходах к сегментации используются предобученные генеративные модели в качестве экстракторов признаков, рассматривая сегментацию как задачу адаптации через косвенное извлечение признаков. Такое неявное использование страдает от фундаментального несоответствия репрезентаций. Оно также сильно зависит от косвенных конвейеров извлечения признаков, что усложняет рабочий процесс и ограничивает адаптацию. В данной статье мы утверждаем, что вместо косвенной адаптации задачи сегментации следует обучать напрямую генеративным способом. Мы выявляем ключевое препятствие для такой унифицированной формулировки: латентные переменные VAE для бинарных масок имеют резкое распределение, устойчивы к шуму и линейно разделимы, что отличает их от латентных переменных натуральных изображений. Чтобы преодолеть этот разрыв, мы представляем стратегию сэмплирования по временным шагам для бинарных масок, которая подчеркивает экстремальные уровни шума для сегментации и умеренный шум для генерации изображений, обеспечивая гармоничное совместное обучение. Мы представляем GenMask — модель DiT, которая обучается генерировать черно-белые маски сегментации, а также цветные изображения в пространстве RGB в рамках исходной генеративной цели. GenMask сохраняет исходную архитектуру DiT, устраняя необходимость в конвейерах извлечения признаков, специфичных для задач сегментации. Экспериментально GenMask достигает передовых результатов на бенчмарках референциальной и логической сегментации, а ablation-исследования количественно оценивают вклад каждого компонента.
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.