[論文レビュー] SegDiff: Image Segmentation with Diffusion Probabilistic Models

Tomer Amit, Tal Shaharbany+2 more|arXiv (Cornell University)|Dec 1, 2021

Advanced Neural Network Applications参考文献 47被引用数 13

ひとこと要約

この論文は、事前学習されたバックボーンを必要とせず、エンドツーエンドで画像セグメンテーションを実行する新しい拡散確率的モデル、SegDiffを紹介する。この手法は、入力画像からのエンコード特徴量と現在のセグメンテーション予測からの特徴量を合算してノイズ除去プロセスを条件づけ、段階的に出力を改善する。複数の確率的生成とモデル平均化を活用することで、Cityscapes、Vaihingen、MoNuSegベンチマークで最先端の性能を達成した。

ABSTRACT

Diffusion Probabilistic Methods are employed for state-of-the-art image generation. In this work, we present a method for extending such models for performing image segmentation. The method learns end-to-end, without relying on a pre-trained backbone. The information in the input image and in the current estimation of the segmentation map is merged by summing the output of two encoders. Additional encoding layers and a decoder are then used to iteratively refine the segmentation map, using a diffusion model. Since the diffusion model is probabilistic, it is applied multiple times, and the results are merged into a final segmentation map. The new method produces state-of-the-art results on the Cityscapes validation set, the Vaihingen building segmentation benchmark, and the MoNuSeg dataset.

研究の動機と目的

• 拡散確率的モデルを、主に画像生成に使われてきたものから、真のラベルが一意で決定的である画像セグメンテーションタスクへと拡張すること。
• 事前学習されたバックボーンに依存しないように、モデル全体をスクラッチからエンドツーエンドで学習すること。
• 入力画像と現在のセグメンテーション推定値を特徴量の加算によって統合する、新しい条件づけ機構を考案すること。
• 複数の確率的出力を生成して平均化することで、モデルのキャリブレーションと精度を向上させること。
• 特に小規模データセットにおいても、多様なセグメンテーションベンチマークで最先端の性能を示すこと。

提案手法

• モデルは、ノイズ除去ネットワークとしてU-Netアーキテクチャを用い、入力画像 I 用と現在のセグメンテーション推定値 xₜ 用の2つの別々のエンコーダーを持つ。
• 2つのエンコーダーからの特徴量が、U-Netのエンコーダー経路の初期段階で合算され、画像とセグメンテーション情報の共同モデリングが可能になる。
• 合算された特徴量を用いて、ノイズ除去ネットワークは次の状態 xₜ₋₁ を予測し、T 個の拡散ステップにわたり段階的にセグメンテーションマップを改善する。
• 変分推論フレームワークを用いて学習し、逆過程の損失を最小化することで、ノイズ除去の目的関数を最適化する。
• 1つの入力に対して複数の独立した生成を実行し、最終的なセグメンテーションは結果を平均化することで得られ、ロバストネスと精度が向上する。
• リーマン型の条件づけ機構が採用されており、入力画像をエンコードして現在の予測表現に加算する形で、結合ではなく加算による統合が行われる。

実験結果

リサーチクエスチョン

• RQ1 真のラベルが一意で決定的である画像セグメンテーションタスクに、拡散確率的モデルを効果的に適用できるか。これは、複数の妥当な出力が得られる生成タスクとは対照的である。
• RQ2 事前学習されたバックボーンを避ける場合、入力画像を拡散モデルのセグメンテーションタスクに効果的に条件づける方法は何か。
• RQ3 複数の確率的サンプルを生成して平均化することで、単一の決定的予測と比較して、セグメンテーションの精度とモデルキャリブレーションが向上するか。
• RQ4 特徴量統合機構の選択（例：加算対比連結）が、セグメンテーション拡散設定における性能に与える影響は何か。
• RQ5 この新しいフレームワークにおいて、推論速度（拡散ステップ数）とセグメンテーション精度のトレードオフは何か。

主な発見

• SegDiffはCityscapesの検証セットで最先端の性能を達成し、特に「バス」のような困難なカテゴリにおいて、先行手法よりもmIoUが向上した。
• Vaihingenの建物セグメンテーションベンチマークでは、特徴量の加算を連結よりも用いることで、既存のアプローチを大きく上回った。
• MoNuSegデータセットでは、核のセグメンテーションにおいて強力な結果を達成し、多様な医療画像分野への汎用性を示した。
• 拡散ステップ数を100から25に減らしても、性能は1〜2 mIoUの低下にとどまり、最大4倍の高速化が可能で、精度の損失は最小限だった。
• RRDBブロックの数は性能にほとんど影響せず、最適な設定と最悪の設定との間でmIoU差は1未満にとどまった。
• 9つの生成サンプルを平均化することで、精度とモデルキャリブレーションが顕著に向上し、決定的タスクにおける確率的サンプリングの利点が裏付けられた。