[論文レビュー] SAUNet: Shape Attentive U-Net for Interpretable Medical Image Segmentation

Jesse Sun, Fatemeh Darbehani+2 more|arXiv (Cornell University)|Jan 21, 2020

Advanced Neural Network Applications参考文献 39被引用数 22

ひとこと要約

SAUNetは、従来のテクスチャベースの特徴量に加えて構造的特徴量を学習する並列の形状ストリームを統合することで、医療画像セグメンテーションの解釈可能性とロバスト性を向上させる形状に配慮したU-Netアーキテクチャを提案する。このモデルはSUN09およびAC17心臓MRIデータセットで最先端の性能を達成するとともに、推論時にマルチスケールのサリエンシー・マップを生成し、後処理による勾配ベースの解釈可能性手法の必要性を排除する。

ABSTRACT

Medical image segmentation is a difficult but important task for many clinical operations such as cardiac bi-ventricular volume estimation. More recently, there has been a shift to utilizing deep learning and fully convolutional neural networks (CNNs) to perform image segmentation that has yielded state-of-the-art results in many public benchmark datasets. Despite the progress of deep learning in medical image segmentation, standard CNNs are still not fully adopted in clinical settings as they lack robustness and interpretability. Shapes are generally more meaningful features than solely textures of images, which are features regular CNNs learn, causing a lack of robustness. Likewise, previous works surrounding model interpretability have been focused on post hoc gradient-based saliency methods. However, gradient-based saliency methods typically require additional computations post hoc and have been shown to be unreliable for interpretability. Thus, we present a new architecture called Shape Attentive U-Net (SAUNet) which focuses on model interpretability and robustness. The proposed architecture attempts to address these limitations by the use of a secondary shape stream that captures rich shape-dependent information in parallel with the regular texture stream. Furthermore, we suggest multi-resolution saliency maps can be learned using our dual-attention decoder module which allows for multi-level interpretability and mitigates the need for additional computations post hoc. Our method also achieves state-of-the-art results on the two large public cardiac MRI image segmentation datasets of SUN09 and AC17.

研究の動機と目的

• ディープラーニングベースの医療画像セグメンテーションモデルにおける解釈可能性とロバスト性の欠如に対処すること。
• テクスチャベースの特徴量よりもロバストな形状特徴量を学習することで、モデルの一般化性能を向上させること。
• SmoothGradのような計算コストの高い後処理によるサリエンシー手法に依存しないように、マルチスケールの解釈可能性をモデルに直接統合すること。
• 大規模な公開の心臓MRIセグメンテーションベンチマーク（SUN09およびAC17）で最先端の性能を達成すること。
• 診断応用における透明性と信頼性をサポートする臨床的採用可能なセグメンテーションフレームワークの開発

提案手法

• 主なテクスチャストリームと補助のゲート付き形状ストリームを備えた二重ストリームエンコーダーデコーダー構造を導入し、形状依存の特徴量を並列に学習する。
• 推論中にマルチスケールの空間的および形状アテンションマップを生成する二重アテンションデコーダー・モジュールを採用し、組み込み型の解釈可能性を実現する。
• セグメンテーションの精度と形状特徴量の学習を同時に最適化する二重タスク損失関数を採用し、モデルのロバスト性を向上させる。
• エンコーダーとデコーダーのパス間のスキップ接続を適用し、空間的コンテキストを保持するとともに、スケール間での特徴量統合を可能にする。
• 追加の前向き・後向きパスを必要とするSmoothGradなどの手法とは異なり、異なるデコーダー段階でのアテンションマップを用いてサリエンシー・マップを生成する。
• SUN09およびAC17データセットの学習にRAdam最適化手法を用い、初期学習率5e-4、バッチサイズ4を設定する。

実験結果

リサーチクエスチョン

• RQ1 明示的に形状特徴量を学習する二重ストリームU-Netアーキテクチャは、標準のCNNと比較してセグメンテーションのロバスト性を向上させることができるか？
• RQ2 形状ストリームを統合することで、学習データ（SUN09）から異なるテストセット（AC17）に一般化する際の性能低下が軽減されるか？
• RQ3 組み込み型のアテンションマップは、SmoothGradのような後処理による勾配ベースの手法に依存せずにマルチスケールの解釈可能性を提供できるか？
• RQ4 形状に配慮した特徴量は、両心室の心臓MRIセグメンテーションタスクにおけるセグメンテーション精度をどの程度向上させるか？
• RQ5 提案手法は、内在的な解釈可能性を維持しながら最先端の性能を達成できるか？

主な発見

• 形状ストリームを備えたモデルは、SUN09の学習セットで90.84%のmIoUを達成し、形状ストリームなしのベースライン（90.20%）を上回った。
• AC17の検証セットでは、形状ストリームを備えたモデルが81.73%のmIoUを達成し、ベースライン（79.33%）と比較して1.76%の性能低下が小さくなった。
• 形状ストリームの導入により、SUN09からAC17への移行時のmIoU低下が1.76%小さくなったことから、一般化性能の向上が裏付けられ、より高いロバスト性が示された。
• 提案手法は、AC17検証セット全体に対して全マルチスケールのサリエンシー・マップを20秒で生成したのに対し、SmoothGradでは24分を要した。これにより、顕著な計算効率の向上が示された。
• 形状ストリームからのアテンションマップは、左室および右室の解剖学的形状に明確に局在化しており、モデルが意味のある構造的特徴量を学習していることが確認された。
• 二重アテンションデコーダーにより階層的な解釈可能性が実現され、高分解能ブロックは微細な詳細に注目し、低分解能ブロックはグローバルな形状の事前知識を捉えていた。