直接生成分割掩码,让DiT模型同时学会画图和分图。
GenMask: Adapting DiT for Segmentation via Direct Mask Generation
- 用特殊采样策略让DiT直接生成黑白分割图,不依赖特征提取
- 在多个分割任务上达到顶尖水平,比现有方法更准确
- 适合想统一图像生成与分割的开发者或研究者
当前分割方法多依赖预训练生成模型提取特征,通过间接方式完成任务,存在表征错位问题,且流程复杂。本文主张应直接以生成方式训练分割任务。我们发现二值分割掩码的VAE隐变量分布集中、抗噪性强、线性可分,与自然图像显著不同。为此,提出针对掩码的步数采样策略,强调高噪声用于分割,适度噪声用于图像生成,实现协同训练。提出GenMask:基于DiT架构,在原始生成目标下同时生成黑白分割掩码与彩色图像,无需为分割设计额外特征提取流程。实验表明,GenMask在参考与推理分割基准上达到领先性能,消融实验证明各组件的有效性。
原文摘要 · Abstract (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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。