[論文レビュー] Unpaired Multi-modal Segmentation via Knowledge Distillation

Qi Dou, Quande Liu+2 more|arXiv (Cornell University)|Jan 6, 2020

Domain Adaptation and Few-Shot Learning参考文献 34被引用数 10

ひとこと要約

本稿では、CTとMRI間ですべての畳み込みカーネルを共有しながら、モダリティ固有のバッチ正規化層を用いる、パrameter効率的でペairドでないマルチモodalセグメンテーションフレームワークを提案する。予測分布のKLダイバージェンスに基づく新しい知識蒸留損失を導入し、モダリティ間でセマンティック特徴を整列させる。2Dおよび3Dネットワークを用いた心臓および腹部多臓器セグメンテーションタスクで最先端の性能を達成した。

ABSTRACT

Multi-modal learning is typically performed with network architectures containing modality-specific layers and shared layers, utilizing co-registered images of different modalities. We propose a novel learning scheme for unpaired cross-modality image segmentation, with a highly compact architecture achieving superior segmentation accuracy. In our method, we heavily reuse network parameters, by sharing all convolutional kernels across CT and MRI, and only employ modality-specific internal normalization layers which compute respective statistics. To effectively train such a highly compact model, we introduce a novel loss term inspired by knowledge distillation, by explicitly constraining the KL-divergence of our derived prediction distributions between modalities. We have extensively validated our approach on two multi-class segmentation problems: i) cardiac structure segmentation, and ii) abdominal organ segmentation. Different network settings, i.e., 2D dilated network and 3D U-net, are utilized to investigate our method's general efficacy. Experimental results on both tasks demonstrate that our novel multi-modal learning scheme consistently outperforms single-modal training and previous multi-modal approaches.

研究の動機と目的

• ペアドデータやアライメントが利用できないCTとMRI間のペアドでないマルチモダリティ医療画像セグメンテーションの課題に対処すること。
• ペアドトレーニングデータや複雑なネットワークアーキテクチャを必要とせず、クロスモダリティ知識を活用することでセグメンテーション性能を向上させること。
• すべての畳み込みカーネルをモダリティ間で共有しながら、モダリティ固有の正規化と新しい蒸留損失により分布シフトを最小限に抑える、極めてパrameter効率的なモデルを開発すること。
• 多様なネットワークアーキテクチャとデータ品質条件の下で、実臨床のセグメンテーションタスクに対して本手法を検証すること。

提案手法

• 本手法は、CTとMRI間ですべての畳み込みカーネルを共有し、モデルパラメータを削減し、パrameter効率を向上させる。
• CTとMRIの異なる統計的分布に適応するため、モダリティ固有の内部正規化層（例：バッチ正規化）を用いる。
• 同じ入力から得られる予測分布のKLダイバージェンスを最小化することで、新しい知識蒸留損失を導入する。
• 損失は事前ソフトマックス特徴上で計算され、セマンティック表現を整列させ、視覚的差異がある中でも共有で強固な特徴を学習するよう促進する。
• 本フレームワークは2Dの拡張畳み込みニューラルネットワーク（dilated CNN）および3DのU-Netに適用され、アーキテクチャやタスクにわたる一般化を示している。
• 標準のクロスエントロピー損失と提案されたKD損失を用いて、エンドツーエンドで訓練され、モダリティ固有および共有表現の共同最適化が可能になっている。

実験結果

リサーチクエスチョン

• RQ1 ペアドデータや画像アライメントを必要とせず、ペアドでないCTとMRIデータから効果的に学習できる単一でコンパクトなニューラルネットワークは構築可能か？
• RQ2 2つのモダリティが著しく異なる強度分布と視覚的外観を持つ場合、どのようにしてクロスモダリティ知識を効果的に蒸留できるか？
• RQ3 すべての畳み込み重みを共有しながらモダリティ固有の正規化層を用いることで、別々のエンコーダを用いる手法よりも優れた性能が得られるか？
• RQ4 予測分布整列に基づく知識蒸留損失は、低品質または外れ値を含む画像環境において、ロバスト性とセグメンテーション精度を向上させられるか？
• RQ5 本手法は、2Dと3Dの異なるネットワークアーキテクチャおよび多クラスセグメンテーションタスクに一般化可能か？

主な発見

• 提案手法は、画像アーティファクトを含む外れ値ケースにおいて、腹部多臓器セグメンテーションタスクで81.1のDiceスコアを達成し、すべてのベースラインを上回った。
• 同じ外れ値ケースにおいて、ハウスドルフ距離が3.78にまで低下し、次に優れた手法（4.89）を著しく下回った。境界局在の優位性が示された。
• 一般的なケースにおいて、単一モダリティ学習および最先端のマルチモダリティアプローチ（'Y'型および'X'型ネットワーク含む）を上回り、心臓および腹部セグメンテーションタスクの両方で優れた性能を示した。
• 2Dの拡張CNNおよび3DのU-Netを含む、さまざまなネットワークアーキテクチャにおいて一貫した改善が確認され、一般化能が裏付けられた。
• アブレーションスタディにより、提案された知識蒸留損失が、特に挑戦的な画像環境下で性能に不可欠であることが確認された。
• 本手法は、画像アーティファクトや低品質データに対してもロバストであり、他の手法が著しく劣化する状況でも高い性能を維持した。