[論文レビュー] Ambiguous Medical Image Segmentation using Diffusion Models

Aimon Rahman, Jeya Maria Jose Valanarasu+2 more|arXiv (Cornell University)|Apr 10, 2023

Radiomics and Machine Learning in Medical Imaging被引用数 7

ひとこと要約

本稿では、拡散モデルに基づくフレームワークであるCIMDを提案する。CIMDは、拡散サンプリングの固有の確率性を活用することで、1つの入力から複数の多様で現実的なセグメンテーションマスクを生成する。従来の方法が別個の事前ネットワークを必要とするのに対し、CIMDは複数の階層的レベルにわたり確率性を統合し、精度（Diceスコア：0.915）と多様性（CIスコア：0.759）の両面で最先端の性能を達成するとともに、臨床的に意味のあるばらつきを保持する。

ABSTRACT

Collective insights from a group of experts have always proven to outperform an individual's best diagnostic for clinical tasks. For the task of medical image segmentation, existing research on AI-based alternatives focuses more on developing models that can imitate the best individual rather than harnessing the power of expert groups. In this paper, we introduce a single diffusion model-based approach that produces multiple plausible outputs by learning a distribution over group insights. Our proposed model generates a distribution of segmentation masks by leveraging the inherent stochastic sampling process of diffusion using only minimal additional learning. We demonstrate on three different medical image modalities- CT, ultrasound, and MRI that our model is capable of producing several possible variants while capturing the frequencies of their occurrences. Comprehensive results show that our proposed approach outperforms existing state-of-the-art ambiguous segmentation networks in terms of accuracy while preserving naturally occurring variation. We also propose a new metric to evaluate the diversity as well as the accuracy of segmentation predictions that aligns with the interest of clinical practice of collective insights.

研究の動機と目的

• 臨床診断におけるレーティング者間ばらつきが inherent であるにもかかわらず、単一のマスクしか出力しない決定的ディープラーニングモデルの限界を是正すること。
• 最良の個々のラベルを模倣するのではなく、1つの確率的フレームワークを用いて、専門放射線科医が提示する複数の妥当なアノテーションの分布をモデル化すること。
• c-VAEベースのアプローチで一般的に必要とされる別個の事前ネットワークを排除するため、確率性を拡散モデルの階層的アーキテクチャに直接埋め込むこと。
• 臨床実践に整合する新たな評価指標（CIスコア）を提案し、正解の専門家アノテーションを基準にした多様性を測定すること。
• CT、超音波、MRIの複数モodal にわたり、曖昧なセグメンテーションにおいて拡散モデルの有効性を示すこと。

提案手法

• 提案されたCIMDフレームワークは、潜在空間におけるセグメンテーションマスクに徐々にノイズを加える前向きの拡散プロセスを逆方向に学習する条件付き拡散確率モデルを用いる。
• 推論時の確率的サンプリングにより、別個の事前ネットワークからの追加推論ステップを必要とせず、1つの入力から複数の妥当なセグメンテーションマスクを自然に生成できる。
• ノイズを複数の階層的特徴レベルに注入することで、セグメンテーションマップの異なるスケールにわたり、多様かつ一貫性のある予測が可能になる。
• 拡散モデルは、1つの入力画像に対して複数のアノテーションを用いてエンドツーエンドで訓練され、妥当なセグメンテーションの同時分布を学習する。
• 画素単位の再構成損失とコントラスト損失を組み合わせた新しい損失関数を採用し、生成マスクの精度と多様性の両方を維持する。
• 既存の拡散ベースのセグメンテーションフレームワークと互換性があり、最小限のアーキテクチャ変更で実装可能である。

実験結果

リサーチクエスチョン

• RQ1 別個の事前ネットワークを必要とせず、1つの拡散モデルが複数の多様で臨床的に妥当なセグメンテーションマスクを生成できるか？
• RQ2 拡散モデルにおける階層的確率性の注入が、曖昧なセグメンテーション出力の多様性と精度に与える影響は何か？
• RQ3 拡散モデルが、異なる医療画像モダリティにおける専門家アノテーションの周波数分布をどれほど正確に学習・再現できるか？
• RQ4 提案手法は、実世界の医療画像データセットにおいて、従来のc-VAEベースおよび決定的セグメンテーションネットワークを上回る精度と多様性を達成できるか？
• RQ5 精度と多様性の両方を捉える新しい評価指標は、Diceスコアのみを用いる標準指標よりも臨床的妥当性をよりよく反映できるか？

主な発見

• CIMDはLIDC-IDRIデータセットで最大Diceスコア（D_max）0.915を達成し、次に優れた手法（Probabilistic U-Net）を大きく上回り、正解分布への一致度において優れた精度を示した。
• LIDC-IDRIデータセットにおいて、CIMDは臨床的多様性（CI）スコア0.759を達成し、専門家アノテーションで観察される多様性と強く整合しており、すべてのベースラインを上回った。
• LIDC-IDRIデータセットにおいて、CIMDは幾何的エネルギー距離（GED）を0.321にまで低減し、DDPM-Prob-Seg（0.417）を上回り、正解との分布的一致性に優れた性能を示した。
• 定性的な結果から、CIMDは超音波やCTスキャンにおける小さなまたは曖昧な病変のような挑戦的ケースでも、多様で高品質なセグメンテーションを生成していることが明らかになった。
• アブレーションスタディの結果、階層的確率性の注入が不可欠であることが確認された。同じ確率的サンプリングプロセスを用いても、その注入を除去すると多様性と精度の両方が低下した。
• 本手法はCT、超音波、MRIの3つのモダリティにわたり一般化可能であり、多様な画像診断状況下でも、強固で臨床的意味のある性能を示した。