[論文レビュー] BerDiff: Conditional Bernoulli Diffusion Model for Medical Image Segmentation

Tao Chen, Chenhui Wang+1 more|arXiv (Cornell University)|Apr 10, 2023

Radiomics and Machine Learning in Medical Imaging参考文献 26被引用数 4

ひとこと要約

本稿では、医用画像分類のための条件付きベルヌーイ拡散モデルBerDiffを提案する。ガウスノイズの代わりにベルヌーイノイズを拡散カーネルとして用いることで、より正確で多様な二値分類マスクの生成が可能になる。確率的逆過程を活用することで、顕著な領域を強調する複数の妥当なマスクを生成し、LIDC-IDRIおよびBRATS 2021データセットにおいてSOTA手法を上回る89.7%のDiceスコアを達成した。

ABSTRACT

Medical image segmentation is a challenging task with inherent ambiguity and high uncertainty, attributed to factors such as unclear tumor boundaries and multiple plausible annotations. The accuracy and diversity of segmentation masks are both crucial for providing valuable references to radiologists in clinical practice. While existing diffusion models have shown strong capacities in various visual generation tasks, it is still challenging to deal with discrete masks in segmentation. To achieve accurate and diverse medical image segmentation masks, we propose a novel conditional Bernoulli Diffusion model for medical image segmentation (BerDiff). Instead of using the Gaussian noise, we first propose to use the Bernoulli noise as the diffusion kernel to enhance the capacity of the diffusion model for binary segmentation tasks, resulting in more accurate segmentation masks. Second, by leveraging the stochastic nature of the diffusion model, our BerDiff randomly samples the initial Bernoulli noise and intermediate latent variables multiple times to produce a range of diverse segmentation masks, which can highlight salient regions of interest that can serve as valuable references for radiologists. In addition, our BerDiff can efficiently sample sub-sequences from the overall trajectory of the reverse diffusion, thereby speeding up the segmentation process. Extensive experimental results on two medical image segmentation datasets with different modalities demonstrate that our BerDiff outperforms other recently published state-of-the-art methods. Our results suggest diffusion models could serve as a strong backbone for medical image segmentation.

研究の動機と目的

• 医用画像分類における固有の曖昧さと不確実性、特に境界が不明瞭で複数の妥当なアノテーションが存在する場合に対処すること。
• 臨床意思決定支援のため、拡散モデルの確率的性質を活用して分類の正確性と多様性を向上させること。
• ガウスノイズの仮定を越えて、離散的二値分類タスクに特化した拡散モデルを構築すること。
• 逆拡散軌道からの部分シーケンスサンプリングにより、効率的な推論を可能とすること。
• ガウスベースの拡散モデルおよび既存のSOTA手法と比較して、医用画像分類において優れた性能を示すことを実証すること。

提案手法

• ガウスノイズに代えてベルヌーイノイズを拡散カーネルとして用いる、新しい条件付き拡散モデルを提案し、離散的分類マスクをよりよくモデル化する。
• 各ステップでスケジュールβ₁,…,βₜに従いベルヌーイノイズを追加するマルコフ連鎖の前向きプロセスを定義し、真値マスクを段階的に純粋なノイズに変換する。
• 入力画像を条件として、ノイズ除去U-Netアーキテクチャを用いて逆過程におけるノイズ推定を実行する。
• 初期ベルヌーイノイズおよび中間潜在変数を異なるランダムシードで再初期化することで、多様なサンプリングを可能にする。
• DDIMの原則に基づき、完全な逆拡散軌道から部分シーケンスをサンプリングする手法を導入し、品質を損なわずに推論を高速化する。
• 変分下界の目的関数を用いて、エンドツーエンドでモデルを訓練し、ノイズ除去スコア関数を最適化する。

実験結果

リサーチクエスチョン

• RQ1 ベルヌーイノイズを拡散カーネルとして用いることで、離散的二値医用画像分類タスクにおける拡散モデルの性能が向上するか？
• RQ2 ベルヌーイ拡散モデルの確率的逆過程は、臨床的根拠として妥当な多様な分類マスクをどのように生成するか？
• RQ3 逆軌道からの部分シーケンスサンプリングは、推論時間を顕著に短縮しつつ、分類品質を維持できるか？
• RQ4 実医用画像データセットにおいて、BerDiffはガウスベースの拡散モデルおよび他のSOTA手法と比較して、正確性と多様性の面で優れているか？
• RQ5 ベルヌーイノイズの使用は、特に訓練ステップ数が限られた状況下でも、収束性と性能を向上させるか？

主な発見

• BerDiffはLIDC-IDRIデータセットで89.7%のDiceスコアを達成し、SegDiff（89.3%）およびU-Net（89.2%）を上回る。
• BRATS 2021データセットでは、従来のU-NetおよびSwin UNETRのようなトランスフォーマーベースのモデルを上回り、複雑な腫瘍分類において優れた性能を示した。
• 顕著な領域を強調する多様な分類マスクを生成でき、特に小さな肺結節のような曖昧なケースでも、真値構造を効果的に捉えた顕著性マスクを生成した。
• アブレーションスタディの結果、十分な訓練ステップ数（例：86,500ステップ）がある場合、ベルヌーイベースの拡散がガウスベースの拡散を上回ることが示された。
• 部分シーケンスサンプリングにより、中間ステップをスキップすることで、推論時間を短縮し、高品質な出力を維持した。
• 可視化比較では、MRIスキャンにおける微細な腫瘍境界など、検出が難しい領域において、他の拡散モデルおよび非拡散モデルと比較してBerDiffがより正確なセグメンテーションを生成した。