[論文レビュー] Improving Calibration and Out-of-Distribution Detection in Medical Image Segmentation with Convolutional Neural Networks

Davood Karimi, Ali Gholipour|arXiv (Cornell University)|Apr 12, 2020

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

ひとこと要約

本論文では、畳み込みニューラルネットワーク（CNN）ベースのセグメンテーションにおけるキャリブレーションの向上と分布外（OOD）検出の向上を目的として、異種の医用画像データを用いたマルチタスク学習を提案する。複数の臓器やモダリティのデータを統合して1つのモデルを訓練することで、より良好なキャリブレーションが達成され、CNNの特徴マップに対するスペクトル解析手法を導入した。この手法は、不確実性に基づくOOD検出を上回り、主要な実験で100%の正確性を達成した。

ABSTRACT

Convolutional Neural Networks (CNNs) have shown to be powerful medical image segmentation models. In this study, we address some of the main unresolved issues regarding these models. Specifically, training of these models on small medical image datasets is still challenging, with many studies promoting techniques such as transfer learning. Moreover, these models are infamous for producing over-confident predictions and for failing silently when presented with out-of-distribution (OOD) data at test time. In this paper, we advocate for multi-task learning, i.e., training a single model on several different datasets, spanning several different organs of interest and different imaging modalities. We show that not only a single CNN learns to automatically recognize the context and accurately segment the organ of interest in each context, but also that such a joint model often has more accurate and better-calibrated predictions than dedicated models trained separately on each dataset. Our experiments show that multi-task learning can outperform transfer learning in medical image segmentation tasks. For detecting OOD data, we propose a method based on spectral analysis of CNN feature maps. We show that different datasets, representing different imaging modalities and/or different organs of interest, have distinct spectral signatures, which can be used to identify whether or not a test image is similar to the images used to train a model. We show that this approach is far more accurate than OOD detection based on prediction uncertainty. The methods proposed in this paper contribute significantly to improving the accuracy and reliability of CNN-based medical image segmentation models.

研究の動機と目的

• 小規模なデータセットで訓練されたCNNベースの医用画像セグメンテーションモデルにおける、不良なキャリブレーションと信頼性の低い予測の課題に対処すること。
• テスト時に分布外（OOD）データを検出することで、モデルが静黙的に失敗するのを防ぐことによる、モデルの信頼性の向上。
• 異なる臓器やモダリティを含む異種データセットを共同で学習させることで、専用モデルと比較して一般化性能とキャリブレーションが向上するかを検討すること。
• 従来の不確実性ベースの手法よりも、より正確なOOD検出手法を医用画像セグメンテーション分野に開発すること。

提案手法

• 異なる臓器や画像モダリティをカバーする異種の医用画像データセットの混合物を用いて、モダリティや臓器の追加入力を用いずに1つのCNNモデルを訓練する。
• 最終畳み込み層からの特徴マップのスペクトル解析を用い、活性化のパワースペクトルに基づいてOODスコア（OODM）を計算する。
• OODMにしきい値τを設定し、テスト画像を分布内または分布外に分類する。しきい値の設定には、分布内データを用いてキャリブレーションを行う。
• AUC、正確性、感度、特異度などの指標を用いて、提案手法OODMと予測不確実性に基づく検出手法を比較する。
• 例えば、MRIにおける肝臓や胎児MRIにおける皮質板を含む複数のデータセットを共同で訓練することで、一般化性能とキャリブレーション性能を評価する。
• セグメンテーションの正確性と信頼性のキャリブレーションを評価するために、Dice類似度係数（DSC）とキャリブレーション指標（例：ECE）を用いる。

実験結果

リサーチクエスチョン

• RQ1 異なる臓器やモダリティを含む異種の医用画像データで訓練された1つのCNNモデルは、個々のデータセットで訓練された専用モデルと同等またはそれ以上のセグメンテーション性能を達成できるか？
• RQ2 異種データ上でマルチタスク学習を実施することで、個々のデータセットで訓練した場合と比較して、より良好なキャリブレーションが達成できるか？
• RQ3 CNNの特徴マップのスペクトル解析は、医用画像セグメンテーションにおける分布外（OOD）画像の検出に有効に機能するか？
• RQ4 本研究で提案するスペクトル解析に基づくOOD検出手法は、不確実性ベースのOOD検出手法と比較して、正確性において優れているか？
• RQ5 明示的なモダリティや臓器情報なしに、異なる年齢層や画像プロトコルを含む多様なデータ分布に一般化できるか？

主な発見

• 異なる臓器やモダリティ（例：MRIにおける肝臓、胎児MRIにおける皮質板）を含む異種データセットで訓練された1つのCNNモデルは、個々のデータセットで訓練された専用モデルと同等またはそれ以上のセグメンテーション性能を達成した。
• 異種データで訓練されたモデルは、期待キャリブレーション誤差（ECE）の低さによって、有意に良好なキャリブレーションが達成された。
• 本研究で提案するCNN特徴マップのスペクトル解析に基づくOOD検出手法は、肝臓セグメンテーションタスクにおけるLiver-MRI-SPIRデータからのOOD検出において、100%の正確性、100%の感度、100%の特異度、AUC 1.00を達成した。
• 一方、不確実性ベースのOOD検出手法は、同じ実験で64%の正確性、61%の感度、60%の特異度、AUC 0.65にとどまった。
• 分布内データ（Liver-MRI-DUAL-In および Liver-MRI-DUAL-Out）のOODM値は、しきい値τ=0.0057以下で一貫して低く、分布外データ（Liver-MRI-SPIR）は顕著に高い値を示し、完璧な分離が可能であった。
• テスト時に、Liver-MRI-DUAL-Outデータセットのセグメンテーションが成功した（DSC=0.886）こと、およびそのOODM値がトレーニング時の分布内データと同様の分布を示したことから、分布シフトに対して高いロバスト性が確認された。