[論文レビュー] MS-KD: Multi-Organ Segmentation with Multiple Binary-Labeled Datasets

Shixiang Feng, Yuhang Zhou+3 more|arXiv (Cornell University)|Aug 5, 2021

Advanced Neural Network Applications参考文献 13被引用数 7

ひとこと要約

本稿では、複数の事前学習済み単一臓器分類ネットワーク（教師）と1つのマルチ臓器学生ネットワークを用いて、マルチ臓器分類モデルを学習する、Multi-Teacher Single-Student Knowledge Distillation（MS-KD）フレームワークを提案する。領域ベースのログティワイズおよび特徴ワイズの監視を適用することで、各教師が自身の臓器領域および全教師が合意する背景領域で学生を指導する。この手法により、ベンチマークデータセットにおける平均DSCで、個々のモデルや最先端手法を最大0.43%上回る優れた性能を達成する。

ABSTRACT

Annotating multiple organs in 3D medical images is time-consuming and costly. Meanwhile, there exist many single-organ datasets with one specific organ annotated. This paper investigates how to learn a multi-organ segmentation model leveraging a set of binary-labeled datasets. A novel Multi-teacher Single-student Knowledge Distillation (MS-KD) framework is proposed, where the teacher models are pre-trained single-organ segmentation networks, and the student model is a multi-organ segmentation network. Considering that each teacher focuses on different organs, a region-based supervision method, consisting of logits-wise supervision and feature-wise supervision, is proposed. Each teacher supervises the student in two regions, the organ region where the teacher is considered as an expert and the background region where all teachers agree. Extensive experiments on three public single-organ datasets and a multi-organ dataset have demonstrated the effectiveness of the proposed MS-KD framework.

研究の動機と目的

• 3Dマルチ臓器分類の高コストなアノテーション問題を、バイナリラベルを備えた既存の単一臓器データセットを活用することで解決すること。
• 異なる単一臓器データセットで学習されたタスク固有の教師モデルから、1つの学生モデルが知識を効果的に学習できる知識蒸留フレームワークを開発すること。
• 条件ベースのアプローチにおける計算効率の低さや、低品質な疑似ラベル依存の既存手法の限界を、複数の教師からのソフトな監視によって克服すること。
• 多様な単一臓器エキスパートの知識を領域ベースの監視によって統合することで、マルチ臓器分類の汎化性能と性能を向上させること。

提案手法

• K個の事前学習済み単一臓器分類ネットワークを教師とし、1つのマルチ臓器分類ネットワークを学生とする。
• 領域ベースの監視を適用：各教師が自身の臓器領域および全教師が合意する背景領域で学生を監視する。
• ログティワイズ監視では、教師の2クラスログティスを次元変換により学生のK+1クラス出力空間にマッピングし、その後に蒸留損失を適用する。
• 特徴ワイズ監視では、分布マッチング損失を用いて、学生の内部特徴分布を教師のものと一致させる。
• 学生はログティワイズ損失と特徴ワイズ損失の組み合わせを用いてエンドツーエンドで学習され、複数のタスク固有の教師からの知識伝達が可能になる。
• 反復的推論や動的ヘッド生成を回避することで、学習時および推論時における計算コストを低く保つ。

実験結果

リサーチクエスチョン

• RQ1 バイナリラベルを備えた複数の単一臓器データセットのみを用いて、1つのマルチ臓器分類モデルを効果的に学習できるか？
• RQ2 タスク固有の条件付けや反復的推論を用いずに、複数のタスク固有の教師モデルから1つの学生モデルへ知識を効果的に伝達できるか？
• RQ3 ログティスと特徴の両方を用いた領域ベースの監視は、標準的な疑似ラベルや一貫性ベースの蒸留よりも優れた性能を発揮するか？
• RQ4 訓練時に単一臓器データのみを用いても、未知のマルチ臓器データセットに十分に一般化できるか？

主な発見

• 3つの単一臓器データセットで学習した場合、MS-KDは個々の単一臓器モデルと比較して平均DSCが0.80%絶対的に向上（91.10% vs. 90.30%）した。
• 特徴ワイズ監視の追加により、平均DSCはさらに0.22%向上し91.32%に達し、中間特徴の蒸留の有効性が示された。
• BTCVマルチ臓器ベンチマークでは、個々のモデルと比較して平均DSCが1.32%高い結果を達成し、優れた一般化能力を示した。
• 追加の推論時間やモデルパラメータを必要としなかったにもかかわらず、Co-training（+0.43%平均DSC）やDoDNetといった最先端手法を上回った。
• 定性的な結果から、MS-KDは特にパンクレアスのような挑戦的な臓器において、より正確で滑らかな境界を生成することがわかった。これは、複数の教師からの滑らかな監視のおかげである。
• 不確実性解析の結果、教師は臓器の境界で最も不確実性が高くなることが判明し、MS-KDはこのソフトな監視を効果的に活用してエッジ検出性能を向上させた。