[論文レビュー] Stochastic Segmentation Networks: Modelling Spatially Correlated Aleatoric Uncertainty

Miguel Monteiro, Loïc Le Folgoc+6 more|arXiv (Cornell University)|Jun 10, 2020

Explainable Artificial Intelligence (XAI)参考文献 48被引用数 79

ひとこと要約

確率的セグメンテーションネットワーク（SSNs）は、ロジット空間での低ランク multivariate normal を用いてセグメンテーションマップの結合分布をモデル化し、複数の空間的一貫性のある仮説と2Dおよび3Dの医用画像に対する効率的推論を実現します。

ABSTRACT

In image segmentation, there is often more than one plausible solution for a given input. In medical imaging, for example, experts will often disagree about the exact location of object boundaries. Estimating this inherent uncertainty and predicting multiple plausible hypotheses is of great interest in many applications, yet this ability is lacking in most current deep learning methods. In this paper, we introduce stochastic segmentation networks (SSNs), an efficient probabilistic method for modelling aleatoric uncertainty with any image segmentation network architecture. In contrast to approaches that produce pixel-wise estimates, SSNs model joint distributions over entire label maps and thus can generate multiple spatially coherent hypotheses for a single image. By using a low-rank multivariate normal distribution over the logit space to model the probability of the label map given the image, we obtain a spatially consistent probability distribution that can be efficiently computed by a neural network without any changes to the underlying architecture. We tested our method on the segmentation of real-world medical data, including lung nodules in 2D CT and brain tumours in 3D multimodal MRI scans. SSNs outperform state-of-the-art for modelling correlated uncertainty in ambiguous images while being much simpler, more flexible, and more efficient.

研究の動機と目的

• 複数のもっともらしい境界が存在する画像分割における固有の不確実性を捉える必要性を動機づける（例：医用画像）.
• ネットワークアーキテクチャを変更せずに結合ラベルマップ分布をモデリングする軽量な確率的フレームワークを提案する。
• 単一のフォワードパスから複数の空間的一致したセグメンテーション仮説を生成可能にする。
• 2D CT肺結節および3D多 modality MRI脳腫瘍セグメンテーションへの適用性を示す。

提案手法

• η|x ~ N(μ(x), Σ(x)) with Σ = P P^T + D を low-rank plus diagonal multivariate normalとしてモデル化する: η|x ~ N(μ(x), Σ(x))
• μ(x)、P、および D をニューラルネットワークの出力としてパラメータ化し、効率的な計算とスケーラビリティを可能にする。
• η に対する不可算積分をロジット空間でのモンテカルロサンプリングで近似する。
• モンテカルロサンプルを通して逆伝播するためにリパラメタライゼーション・トリックを用い、クロスエントロピー様の目的関数で学習する。
• 推論は単一のフォワードパスで複数のロジットマップサンプルを生成する。平均ロジットが最も可能性の高いセグメンテーションを提供する。

実験結果

リサーチクエスチョン

• RQ1 低ランクのロジット空間分布を介して空間的に相関したアレータック不確実性をモデリングすると、一致した複数のセグメンテーション仮説を生み出せるか？
• RQ2 SSN は医用画像タスクで決定論的および既存の確率的セグメンテーション手法に対して予測性能と不確実性の較正を改善するか？
• RQ3 高次元の3Dデータへのスケールと一般的なセグメンテーションバックボーンとの統合はどうなるか？

主な発見

• SSNsは曖昧な医用画像における相関不確実性のモデリングで最先端手法を上回る。
• 低ランク共分散モデルは、肺結節セグメンテーションで対角共分散変異体および決定論的ベースラインよりも予測性能が高く、不確実性の較正が良好である。
• SSNsは2Dおよび3Dで最小限のアーキテクチャ変更と控えめな計算オーバーヘッドで複数のもっともらしいセグメンテーションをサンプリングできる。
• BraTS 2017 3D 脳腫瘍セグメンテーションでは、確率モデルは決定論的モデルと同程度のDiceスコアを達成しつつセグメンテーションの分布を提供する。
• フレームワークは推論後のサンプル操作（例：クラスの有無を調整）や多様性の温度ベース制御を可能にする。