[論文レビュー] Masked Diffusion as Self-supervised Representation Learner

Zixuan Pan, Jianxu Chen+1 more|arXiv (Cornell University)|Aug 10, 2023

Generative Adversarial Networks and Image Synthesis被引用数 6

ひとこと要約

本論文は、拡散モデルにおけるガウスノイズの代わりにマスキングを導入することで、セマンティックセグメンテーションのための自己教師付き表現学習モデル「マスキング拡散モデル（MDM）」を提案する。構造的類似度（SSIM）損失を用いて事前学習と下流の密予測タスクの整合性を高め、特に少サンプル設定において最先端の性能を達成した。生成能力が強力な表現学習に不可欠であるという仮定に反し、生成能力がなくても強力な表現学習が可能であることを示した。

ABSTRACT

Denoising diffusion probabilistic models have recently demonstrated state-of-the-art generative performance and have been used as strong pixel-level representation learners. This paper decomposes the interrelation between the generative capability and representation learning ability inherent in diffusion models. We present the masked diffusion model (MDM), a scalable self-supervised representation learner for semantic segmentation, substituting the conventional additive Gaussian noise of traditional diffusion with a masking mechanism. Our proposed approach convincingly surpasses prior benchmarks, demonstrating remarkable advancements in both medical and natural image semantic segmentation tasks, particularly in few-shot scenarios.

研究の動機と目的

• 拡散モデルの生成能力と表現学習の潜在的性能を分離し、ノイズ除去が強力な特徴を生成するために不可欠であるという仮定に挑戦する。
• 特に軽量モデルを用いた少サンプル設定において、再構成事前学習と下流の密予測タスクとの間の不一致を是正する。
• ノイズの代わりにマスキングを用いることで、スケーラブルな自己教師付き学習を可能にする、画期的な事前学習パラダイム「マスキング拡散モデル（MDM）」を提案する。
• MSE損失の代わりにSSIM損失を用いることで、事前学習の目的とセグメンテーション性能との整合性を高める。
• MDMがDDPMおよびMAEを含む、医療画像および自然画像セグメンテーションベンチマークにおいて、実証的に優れた性能を示すことを検証する。

提案手法

• 時間ステップに応じて確率的に画像パッチをマスキングする動的マスキング機構を、従来の加法的ガウスノイズの代わりに導入する。
• Tステップの拡散プロセスを用い、マスキングされた入力を元に元の画像を再構成する時間に依存するU-Netを学習する。
• 事前学習と下流のセグメンテーションタスクの分布的乖離を最小化するため、平均二乗誤差（MSE）の代わりに構造的類似度（SSIM）損失を用いる。
• 各ステップで元の画像（またはマスキングされた領域）を予測するようにモデルを学習し、意味的整合性を最適化する再構成ターゲットを採用する。
• DDPMの変種ではノイズ予測戦略を採用し、腐食プロセス中に追加されたノイズを予測することで、表現品質を向上させる。
• 主にmIoUを指標として、下流のセマンティックセグメンテーションタスクにおける微調整による学習済み表現の評価を実施する。

実験結果

リサーチクエスチョン

• RQ1 拡散モデルの生成能力は、その表現学習性能と分離可能か？
• RQ2 拡散モデルにおけるガウスノイズの代わりにマスキングを用いることで、セグメンテーションに適した強力な意味的表現が得られるか？
• RQ3 事前学習の目的と下流の密予測タスクとの整合性を高めるために、SSIM損失はMSE損失よりも有効か？
• RQ4 MDMは、DDPMおよびMAEと比較して、少サンプルセグメンテーション設定でどのように性能を発揮するか？
• RQ5 再構成ターゲットとしてノイズ（εt）と元の画像（x0）のどちらを選ぶかが、下流のセグメンテーション精度に顕著な影響を与えるか？

主な発見

• MDMは、10%のラベル付きデータでGlaS（84.51 ± 1.15）およびFFHQ-34（59.18 ± 0.11）で、少サンプル設定において最先端のmIoUを達成し、DDPMおよびMAEを上回った。
• SSIM損失を用いることで、MDMはGlaSで10%ラベルデータの状況で84.51 ± 1.15のmIoUを達成し、DDPM（82.32 ± 0.77）およびMAE（79.47 ± 1.08）を顕著に上回った。
• SSIM損失の使用はMDMおよびDDPMの性能を向上させるが、MAEでは向上しなかったことから、その有効性は腐食メカニズムに依存するコンテキスト依存的であることが示唆された。
• MSE損失を用いたMDM（82.70 ± 0.79）は、既にMAE（79.47 ± 1.08）を上回っており、さらにSSIM損失を用いることで性能が向上（84.51 ± 1.15）した。
• 単一ステップの形でも、MDMおよびDDPMはMAEを上回ったことから、自己符号化の単純な拡張を超えて、拡散プロセスが表現学習を強化していることが示唆された。
• ノイズ（εt）を予測するように学習したモデルは、元の画像（x0）を予測するモデルよりも高い下流の精度を達成しており、従来の拡散モデルの文献と整合的であった。