[論文レビュー] GenMask: Adapting DiT for Segmentation via Direct Mask Generation

Yuhuan Yang, Xianwei Zhuang+7 more|arXiv (Cornell University)|Mar 25, 2026

Generative Adversarial Networks and Image Synthesis被引用数 0

ひとこと要約

GenMask は拡散トランスフォーマー（DiT）を用いてRGB空間で直接バイナリセグメンテーションマスクを生成する生成的な統一目標の下で学習し、間接的な特徴抽出パイプラインに依存せず、 referring および reasoning セグメンテーションのベンチマークで最先端の結果を達成する。

ABSTRACT

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.

研究の動機と目的

• 生成モデルを用いた下流適応タスクとしてのセグメンテーションを動機付け、既存の間接的な特徴ベースアプローチのミスアラインメントを強調する。
• DiT バックボーンの下で直接バイナリマスクと RGB 画像を生成する統一的生成トレーニングフレームワーク（GenMask）を提案する。
• Binary マスクの VAE 潜在変数と自然画像とのギャップを、特定の timestep サンプリング戦略によって埋める。
• 生成とセグメンテーションの両方の条件付けを可能にするため、低レベルのテクスチャ手掛かりを VAE 潜在とビジョン-言語モデルを通じて組み込む。
• セグメンテーションと生成データの共同トレーニングが referring および reasoning セグメンテーションタスクで最先端の結果を導くことを経験的に示す。

提案手法

• 元の DiT アーキテクチャを保持し、同じ生成的目標の下で RGB 空間でバイナリマスクを生成するセグメンテーション目的を追加する。
• セグメンテーション用に高ノイズタイムステップを強調し、画像生成には中程度/ノイズ少なめのタイムステップを用いる専門的なタイムステップサンプリング戦略を導入する。
• 入力画像の VAE 潜在表現を低レベルのショートカットとして組み込み、セグメンテーションのためのテクスチャ/カラーの手掛かりを提供する。
• 指示をエンコードし、生成とセグメンテーションタスクの条件付けを統一するためにビジョン-言語モデルを使用する。
• 三つのセグメンテーション監視形式（VAE 空間での MSE、VAE をデコップルした RGB 空間での BCE、学習済み線形射影後の BCE）を試し、DiT 目的と最も整合するのは VAE 空間での MSE であることを確認する。
• セグメンテーションを t=1 で実行し VAE で復号してバイナリマスクを得ることで、1ステップ推論を実現する。

実験結果

リサーチクエスチョン

• RQ1 生成特徴抽出パイプラインに依存せず、生成フレームワーク内で直接セグメンテーションを効果的に訓練できるか。
• RQ2 バイナリマスク表現と自然画像表現のギャップを埋めて共同生成とセグメンテーションを可能にするにはどうするか。
• RQ3 ディフュージョンベースのモデルを訓練する際、タスク特有のタイムステップサンプリング戦略はセグメンテーション性能を改善するか。
• RQ4 低レベルの VAE 潜在情報とビジョン-言語 conditioning モデルの組み込みがセグメンテーション品質に与える影響は何か。
• RQ5 生成データをセグメンテーション訓練に補助として用いることで referring および reasoning セグメンテーションベンチマークの性能は向上するか。

主な発見

• GenMask は referring セグメンテーションベンチマーク（例：RefCOCO 系列および ReasonSeg）で最先端の性能を達成する。
• セグメンテーション専用のタイムステップサンプリング戦略が高ノイズ領域を強調することでセグメンテーション性能を大幅に改善する。
• VAE 潜在空間での MSE 損失をセグメンテーション監視に用いる形式が、試した形式の中で最良の結果を生む。
• セグメンテーションのための低レベル入力として VAE 潜在を組み込むとピクセルレベルの正確さが大幅に改善される。
• 生成データと jointly training（1:1 ミックス）を行うことでセグメンテーション性能が向上する。
• セグメンテーションのための one-step 推論（t=1）は、統一的な目標の下で正確なマスク生成に十分である。