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 잠재 공간이 자연 이미지 잠재 공간과 달리 급격한 분포, 노이즈 강건성, 선형 분리 가능성을 보인다는 점입니다. 이러한 차이를 해소하기 위해 우리는 분할에는 극단적 노이즈 수준을, 이미지 생성에는 중간 수준의 노이즈를 강조하는 이진 마스크용 타임스텝 샘플링 전략을 도입하여 조화로운 공동 훈련을 가능하게 합니다. 우리는 원래의 생성 목표 하에서 흑백 분할 마스크와 RGB 공간의 컬러 이미지를 모두 생성하도록 훈련된 DiT 기반 GenMask를 제시합니다. 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.