[論文レビュー] Correlation-Aware Mutual Learning for Semi-supervised Medical Image Segmentation

Shengbo Gao, Ziji Zhang+3 more|arXiv (Cornell University)|Jul 12, 2023

Advanced Neural Network Applications被引用数 4

ひとこと要約

本稿では、相関に配慮した相互学習（CAML）という新しい半教師付き医療画像セグメンテーションフレームワークを提案する。この手法は、ラベル付きデータを明示的に活用してラベルなしデータからの特徴学習をガイドすることで、モデル性能を向上させる。CAMLは2つのモジュールを備えた共-teaching方式を採用しており、ラベル付きとラベルなしの特徴間の類似性を特徴レベルでモデル化するための、オムニ相関一貫性（OCC）と、ラベル付きとラベルなしデータ間の密な双方向自己注意機構であるクロスサンプル相互注意（CMA）を含む。このアプローチにより、特にラベルが少ない状況下でも最先端の性能を達成し、ラベル付きデータが5%の状況でDiceスコア87.34%を達成し、20%のラベル付きデータで完全教師あり性能に近づく。

ABSTRACT

Semi-supervised learning has become increasingly popular in medical image segmentation due to its ability to leverage large amounts of unlabeled data to extract additional information. However, most existing semi-supervised segmentation methods only focus on extracting information from unlabeled data, disregarding the potential of labeled data to further improve the performance of the model. In this paper, we propose a novel Correlation Aware Mutual Learning (CAML) framework that leverages labeled data to guide the extraction of information from unlabeled data. Our approach is based on a mutual learning strategy that incorporates two modules: the Cross-sample Mutual Attention Module (CMA) and the Omni-Correlation Consistency Module (OCC). The CMA module establishes dense cross-sample correlations among a group of samples, enabling the transfer of label prior knowledge to unlabeled data. The OCC module constructs omni-correlations between the unlabeled and labeled datasets and regularizes dual models by constraining the omni-correlation matrix of each sub-model to be consistent. Experiments on the Atrial Segmentation Challenge dataset demonstrate that our proposed approach outperforms state-of-the-art methods, highlighting the effectiveness of our framework in medical image segmentation tasks. The codes, pre-trained weights, and data are publicly available.

研究の動機と目的

• 既存の半教師付き医療画像セグメンテーション手法が、通常は初期の監視信号としてのみラベル付きデータを扱うという、その活用不足に対処すること。
• ラベル付きデータとラベルなしデータ間のクロスサンプル相関を明示的にモデル化することで、事前知識の転送を促進し、モデルの汎化性能を向上させること。
• ラベル付きデータとラベルなしデータの両方を共同最適化することで、特徴表現を強化する相互学習フレームワークを開発すること。
• 少数のラベル付き画像しか利用できない状況でも、頑健で一貫性のある性能を達成すること。
• 限られたラベル付きデータから最大限の情報を抽出することで、臨床応用における大規模アノテーションデータセットへの依存を低減すること。

提案手法

• CAMLは、相互偽監視により学習される2つの学生ネットワークを用いた共-teachingフレームワークを採用しており、一貫性と知識蒸留を促進する。
• クロスサンプル相互注意（CMA）モジュールは、複数のサンプル間で密な双方向自己注意を確立し、ラベル付きとラベルなしデータ間で共通の顕著特徴を学習可能にする。
• オムニ相関一貫性（OCC）モジュールは、ラベル付き特徴のメモリバンクを構築し、ラベルなし特徴とこの参照セットとの間でオムニ相関を計算することで、特徴レベルでの類似性をモデル化する。
• OCCは両方のサブモデルのオムニ相関行列の一貫性を強制し、学習プロセスを正則化し、ラベルなしデータにおける統一された特徴分布を促進する。
• このフレームワークは、両モジュールを相互学習ループ内で統合し、両ネットワークを同時に更新しながら、ラベル付きデータからラベルなしデータへの知識伝達を実現する。
• この手法はエンド・トゥ・エンドで学習可能であり、追加の推論コストを発生させないため、臨床応用に適している。

実験結果

リサーチクエスチョン

• RQ1 半教師付きセグメンテーションにおいて、ラベル付きデータがラベルなし医療画像からの表現学習を改善するために効果的に活用可能か？
• RQ2 ラベル付きデータとラベルなしデータ間のクロスサンプル相関をどのようにモデル化すれば、特徴転送と一貫性の向上が達成できるか？
• RQ3 複数のサンプルにわたる相互注意を導入することで、標準的一致性正則化と比較してセグメンテーション性能が向上するか？
• RQ4 ラベル付きとラベルなし特徴間のオムニ相関一貫性は、より頑健で汎用性の高い特徴学習をもたらすか？
• RQ5 極度のラベル不足、例えば5%または10%のラベル付きデータ下で、提案されたCAMLフレームワークはどのように性能を発揮するか？

主な発見

• 5%のラベル付きデータで、CAMLは標準偏差0.05を伴い、Diceスコア87.34%を達成し、ベースラインを顕著に上回り、高い安定性を示した。
• 10%のラベル付きデータで、CAMLはDiceスコア89.62%を達成し、完全教師ありモデルの上限（90.98%）に近づいた。
• 20%のラベル付きデータで、CAMLはDiceスコア90.78%を達成し、完全教師ありの上限との差が最小限に抑えられた。
• アブレーションスタディにより、CMAとOCCモジュールの両方が不可欠であることが確認された。5%のラベル付きデータ下で、両者の組み合わせはDiceスコアを6.42%絶対的に向上させた。
• 本手法は、実行回数にかかわらず標準偏差が低く保たれ、特にラベルが少ない状況下で高い頑健性と汎化性能を示した。
• 定性的な結果から、CAMLはすべての比較手法と比較して、境界領域や小規模構造領域においてより正確で完全なセグメンテーションマスクを生成することがわかった。