Diffusion Model as a Generalist Segmentation Learner

論文「Diffusion Model as a Generalist Segmentation Learner (DiGSeg)」の詳細な技術的サマリーを以下に示す。

1. 問題定義

現在のセグメンテーションシステムは高度に断片化されている。モデルは通常、特定のタスク（セマンティック、インスタンス、パンオプティック）、特定の語彙（クローズドセット対オープンボキャブラリー）、または特定のドメイン（自然画像、医療、リモートセンシング）に合わせて設計されている。これらの設定間を移行するには、アーキテクチャの再設計やゼロからの再学習が必要となることが多い。

最近の研究では、アテンションマップを抽出することで事前学習済みのテキストから画像への拡散モデルをセグメンテーションに転用しようとする試みがあるが、これらのアプローチには重大な限界がある：

• ノイズと不整合： 生のアテンションマップは解像度が低く、ノイズが多く、層間で不整合であることが多い。
• ポストプロセッシングへの依存： 使用可能なマスクを生成するために、重たいヒューリスティックなポストプロセッシングを頻繁に必要とする。
• 汎化性の欠如： 既存の手法は単一のタスクに焦点を当てることが多く、アーキテクチャのカスタマイズなしに多様なドメインやオープンボキャブラリークエリにわたる堅牢な汎化を実証できていない。

本論文が扱う核心的な問いは以下の通りである：単一の統合された拡散ベースモデルが、ドメイン固有のアーキテクチャ変更なしに、セマンティック、オープンボキャブラリー、およびクロスドメインタスクを処理できる汎用セグメンテーション学習器として機能し得るか？

2. 手法：DiGSeg フレームワーク

著者らは、事前学習済みの拡散モデル（具体的には Stable Diffusion v2）を、生成事前知識から、統合された微調整プロトコルを通じて明示的なセグメンテーション学習器へと変換するDiGSegを提案する。

コアアーキテクチャ

このフレームワークは主に 3 つのコンポーネントで構成される：

1. 視覚潜在経路：

   • 画像を生ピクセルとして扱うのではなく、入力 RGB 画像 ($x$) とグランドトゥルースのセグメンテーションマスク ($y$) を、Stable Diffusion から取得した凍結された VAE エンコーダを用いてコンパクトな潜在空間に符号化する。
   • 3 チャンネル VAE との互換性を維持するため、単一チャンネルのマスクは 3 チャンネルに複製される。
   • これにより、空間構造を保持しつつ、潜在空間における微細な対応関係の効率的な学習が可能になる。

2. CLIP 整合テキスト条件付け：

   • オープンボキャブラリー機能を実現するため、凍結された CLIP テキストエンコーダがクラス名や自然言語記述の埋め込み ($t_{clip}$) を生成する。
   • これらの埋め込みは、複数のノイズ除去スケールにおいてクロスアテンションを介して拡散 U-Net に注入される。これにより、モデルはテキストクエリを、局所的かつ全球的に発展する視覚表現と整合させることができる。

3. セグメンテーション整合ノイズ除去 U-Net：

   • 学習目的： モデルは、画像潜在 ($z_x$) とテキスト埋め込み ($t_{clip}$) に条件付けられたセグメンテーション潜在 ($z_y$) に加えられたノイズを予測するように訓練される。
   • 生成定式化： 拡散を特徴抽出器として使用する以前の研究とは異なり、DiGSeg はセグメンテーションを条件付きノイズ除去生成タスク：$p_\theta(y|x)$ として扱う。U-Net は、ランダムノイズからセグメンテーション整合的な潜在表現へとノイズを除去することを学習する。
   • ノイズ戦略： 新しいアニーリング型マルチスケールノイズスケジューリングが採用される。高周波の摂動は初期段階で支配的となり微細な構造を学習し、低周波成分は後段で影響力を持ち、意味的整合性を確保する。

推論プロセス

1. 入力画像が潜在表現に符号化される。
2. セグメンテーション潜在がガウスノイズとして初期化される。
3. テキスト条件付けの下で、U-Net が潜在を反復的にノイズ除去する（効率化のため DDIM サンプリングに「trailing」タイムステップ戦略を使用）。
4. 最終的なクリーンな潜在は、凍結された VAE デコーダを用いてピクセル空間に復号され、セグメンテーションマスクを生成する。
5. テスト時アンサンブル： 異なるノイズシードによる複数の推論パスを潜在空間で平均化し、空間的整合性を向上させ、アーティファクトを低減する。

3. 主要な貢献

• 統合フレームワーク： DiGSeg は、事前学習済みの拡散バックボーンを、タスク固有のアーキテクチャ変更なしにクローズドボキャブラリー、オープンボキャブラリー、およびクロスドメインタスクを処理する汎用セグメンテーション学習器として成功裏に転用した最初のフレームワークである。
• 直接潜在学習： ノイズのあるアテンションマップに依存するのではなく、ノイズ除去 U-Net を微調整してセグメンテーション整合的な潜在を直接生成させ、重たいポストプロセッシングなしに高品質なマスクを生成する。
• マルチスケールテキスト条件付け： 複数のノイズ除去スケールにわたって言語特徴を注入するメカニズムを導入し、任意のテキストプロンプトと視覚特徴間の堅牢な整合を可能にする。
• 新規ノイズ戦略： 微細な境界学習とグローバルな意味的整合性のバランスを取るアニーリング型マルチスケールノイズスケジューリングを提案する。

4. 実験結果

著者らは DiGSeg を標準ベンチマークおよび専門ドメインで評価し、最先端（SOTA）のパフォーマンスを実証した。

• オープンボキャブラリーセグメンテーション：
  • A-847, PC-459, A-150, PC-59, および Cityscapesにおいて、13 のベースライン（ODISE, OVSeg, CAT-Seg などを含む）を上回った。
  • 2 番目に良い手法に対して大幅な改善（例：PC-59 で +2.8 mIoU、PC-459 で +2.2）を達成した。
• 標準セマンティックセグメンテーション：
  • COCO (50.8 mIoU) およびADE20K (58.6 mIoU) において SOTA を達成し、Mask2Former や OneFormer などの強力なトランスフォーマーベースモデルを凌駕した。
• クロスドメイン転移：
  • アーキテクチャ変更なしに、医療 (REFUGE-2)、農業 (PhenoBench)、およびリモートセンシング (DeepGlobe) データセットへの強力なゼロショット転移を実証した。
  • 注記： 医療ドメインでの性能は自然画像に比べて低かったが、これは CLIP の事前学習データにおけるドメインギャップに起因するものであり、それでも競争力のある結果であった。
• 効率性とデータ利用：
  • モデルは非常にデータ効率的であり、トレーニングデータの**50%**のみでほぼ同等の性能を達成した。
  • 推論速度は DDIM-trailing により向上し（1 ステップで飽和に達する）、拡散の性質上フィードフォワードモデルより遅いものの改善された。

5. 意義

この研究は、視覚生成と視覚理解の間のギャップを根本的に埋めるものである。

• パラダイムシフト： 拡散モデルは合成のみに特化したものであるという概念に挑戦し、適切に条件付けられれば強力な汎用識別学習器として機能し得ることを証明した。
• 汎化性： 大規模生成モデルに符号化された豊かな空間的および意味的事前知識が、多様なドメイン（自然、医療、農業）にわたる密な予測タスクへ転移可能であることを示した。
• 単純性： 複雑な多段階パイプラインや専用ヘッドを回避することで、DiGSeg はセグメンテーションのためのストリームライン化された統合インターフェースを提供し、単一の生成バックボーンがコンピュータビジョンの広範なタスクを処理できる未来を示唆している。

結論として、DiGSeg は、潜在条件付きプロトコルとテキスト整合性を用いて微調整された現代の拡散バックボーンが、オープンボキャブラリーおよびクロスドメインシナリオにおいて優れた柔軟性を提供しつつ、専門的なセグメンテーションモデルを上回る性能を発揮し得ることを確立している。