[論文レビュー] Diffusion Models for Implicit Image Segmentation Ensembles

Julia Wolleb, Robin Sandkühler+3 more|arXiv (Cornell University)|Dec 6, 2021

Medical Image Segmentation Techniques被引用数 71

ひとこと要約

本論文は拡散確率モデルを用いて脳腫瘍のセグメンテーションマスクを生成し、確率的サンプリングを活用して暗黙のアンサンブルを形成し、不確実性マップを生成し、BRATS2020でのセグメンテーション精度を向上させる。

ABSTRACT

Diffusion models have shown impressive performance for generative modelling of images. In this paper, we present a novel semantic segmentation method based on diffusion models. By modifying the training and sampling scheme, we show that diffusion models can perform lesion segmentation of medical images. To generate an image specific segmentation, we train the model on the ground truth segmentation, and use the image as a prior during training and in every step during the sampling process. With the given stochastic sampling process, we can generate a distribution of segmentation masks. This property allows us to compute pixel-wise uncertainty maps of the segmentation, and allows an implicit ensemble of segmentations that increases the segmentation performance. We evaluate our method on the BRATS2020 dataset for brain tumor segmentation. Compared to state-of-the-art segmentation models, our approach yields good segmentation results and, additionally, detailed uncertainty maps.

研究の動機と目的

• 医療画像セグメンテーションにおける不確実性の定量化を動機づけ、病変セグメンテーション性能を向上させる。
• 訓練時およびサンプリング時に入力画像をアトラス事前として利用する拡散モデルベースの手法を提案する。
• 確率的拡散サンプリングを通じてセグメンテーションマスクの暗黙的なアンサンル生成を可能にする。
• 複数のセグメンテーションサンプルのピクセル単位分散を通じて不確実性マップを提供する。

提案手法

• 入力MR画像を前提として結合を取り入れつつ、グラウンドトゥルースマスクにノイズを加えることでセグメンテーションマスクを生成するDDPMを訓練する。
• DDPMのサンプリング方程式を画像事前 X_t = b 直前(x_b,t) に条件付けするよう変更し、平均/分散のセグメンテーションアンサンブルを生成する。
• 線形ノイズスケジュールを用いT=1000ステップ、デノイジングステップ epsilon_theta を予測するU-Netアーキテクチャを使用する。
• 複数のセグメンテーションマスクをサンプリングし、それらの分散を計算してピクセル単位の不確実性マップを求める。
• 複数のサンプリングマスクを平均化してアンサンルを作成し、最終的なセグメンテーションを得る。
• BRATS2020でnnU-NetとSegNetと比較し、Dice、Jaccard、HD95 の指標を報告する。

実験結果

リサーチクエスチョン

• RQ1 拡散モデルを対応する画像を事前として条件付けることで医療画像の意味セグメンテーションに適用できるか。
• RQ2 確率的拡散サンプリングによる複数セグメンテーションサンプルの生成は、単一サンプル予測と比べてセグメンテーション指標を改善するか。
• RQ3 BRATS2020における拡散ベースのセグメンテーションは、精度と不確実性視覚化の観点で最先端手法（nnU-Net, SegNet）とどう比較されるか。

主な発見

Dice係数	HD95	Jaccard係数	空欄
我々（1回のサンプリング実行）	0.866 [0.892]	6.052	0.795 [0.819]	31
我々（5回のサンプリングによるアンサンブル）	0.881 [0.909]	5.178	0.819 [0.845]	34
nnU-Net（5-fold交差検証のアンサンブル）	0.891 [0.905]	5.004	0.831 [0.845]	17
SegNet（1回実行）	0.839 [0.867]	7.190	0.761 [0.786]	34
ベイズ SegNet（5回の実行によるアンサンブル）	0.838 [0.841]	13.707	0.747 [0.749]	3

• 単一サンプルDDPMセグメンテーションはDice 0.866、Jaccard 0.795、HD95 6.052を達成。
• 5つの拡散サンプルをアンサンブルするとDiceが0.881、Jaccardが0.819、HD95が5.178へ改善。
• nnU-NetはDice 0.891、Jaccard 0.831、HD95 5.004を達成。
• SegNetはDice 0.839、Jaccard 0.761、HD95 7.190。
• ベイズ SegNet（MCドロップアウト、5サンプル）ではDice 0.838、Jaccard 0.747、HD95 13.707。
• アンサンブル平均マップはセグメンテーション性能を改善し、詳細な不確実性可視化を可能にする。