[論文レビュー] Local contrastive loss with pseudo-label based self-training for semi-supervised medical image segmentation

Krishna Chaitanya, Ertunç Erdil+2 more|arXiv (Cornell University)|Dec 17, 2021

Advanced Neural Network Applications被引用数 7

ひとこと要約

本論文は、ラベルなしデータとラベルありデータの両方を活用する、擬似ラベルに基づく自己学習を用いた新しいローカルコントラスト損失を提案する。この手法では、ランダムな増幅に依存するのではなく、意味的擬似ラベルに基づいてポジティブペアを定義することで、ラベルありおよびラベルなしデータを有効に活用する。本手法は、1つまたは2つの3Dラベル付きボリュームのみを用いて心臓および前立腺MRIデータセットで最先端の性能を達成し、既存の半教師あり学習およびコントラスト学習手法を著しく上回る。

ABSTRACT

Supervised deep learning-based methods yield accurate results for medical image segmentation. However, they require large labeled datasets for this, and obtaining them is a laborious task that requires clinical expertise. Semi/self-supervised learning-based approaches address this limitation by exploiting unlabeled data along with limited annotated data. Recent self-supervised learning methods use contrastive loss to learn good global level representations from unlabeled images and achieve high performance in classification tasks on popular natural image datasets like ImageNet. In pixel-level prediction tasks such as segmentation, it is crucial to also learn good local level representations along with global representations to achieve better accuracy. However, the impact of the existing local contrastive loss-based methods remains limited for learning good local representations because similar and dissimilar local regions are defined based on random augmentations and spatial proximity; not based on the semantic label of local regions due to lack of large-scale expert annotations in the semi/self-supervised setting. In this paper, we propose a local contrastive loss to learn good pixel level features useful for segmentation by exploiting semantic label information obtained from pseudo-labels of unlabeled images alongside limited annotated images. In particular, we define the proposed loss to encourage similar representations for the pixels that have the same pseudo-label/ label while being dissimilar to the representation of pixels with different pseudo-label/label in the dataset. We perform pseudo-label based self-training and train the network by jointly optimizing the proposed contrastive loss on both labeled and unlabeled sets and segmentation loss on only the limited labeled set. We evaluated on three public cardiac and prostate datasets, and obtain high segmentation performance.

研究の動機と目的

• ラベル付きデータセットの規模を大きくすることに依存しないようにすることで、医療画像セグメンテーションにおける高コストなアノテーション問題に取り組む。
• 擬似ラベルからの意味的ラベル情報を取り入れることで、半教師ありセグメンテーションにおけるローカルレベルの表現学習を向上させる。
• 従来のローカルコントラスト手法が無意味な代替ラベルや空間的近接性に依存するという限界を克服する。
• ラベル付きデータに対するセグメンテーション損失と、ラベル付きおよびラベルなしデータに対するコントラスト損失を併用する共同学習フレームワークを開発する。
• 繰り返し行われる擬似ラベルの精練と一貫性正則化を通じて、モデルの汎化性能とロバスト性を向上させる。

提案手法

• 同じ擬似ラベルを持つ画素をポジティブペア、異なる擬似ラベルを持つ画素をネガティブペアとみなすローカルコントラスト損失を定義する。
• 学生ネットワークを用いてラベルなし画像の擬似ラベルを生成し、学習中に繰り返し精練する。
• ラベル付きデータセットとラベルなしデータセットの両方で提案されたコントラスト損失を最適化し、ラベル付きデータセットでのみ標準的なセグメンテーション損失を最適化することで自己学習を実行する。
• 同じ画像の増幅ビュー間の特徴量を比較するためのインライン内および異なる画像間の特徴量を比較するためのインライン間の表現マッチング方式を用いる。
• 変換されたビュー間の重複度に基づいて擬似ラベルをフィルタリングすることで一貫性正則化を適用し、予測品質を向上させる。
• ランダムな重み初期化または自己教師あり事前学習から得た重みを用いてモデルを初期化する。

実験結果

リサーチクエスチョン

• RQ1 ランダムな増幅や空間的近接性と比較して、意味的擬似ラベルは半教師あり医療画像セグメンテーションにおけるローカル表現学習を改善できるか？
• RQ2 ラベル付きおよびラベルなしデータの両方でコントラスト損失を共同最適化することで、標準的な自己学習よりも優れたセグメンテーション性能が得られるか？
• RQ3 真のアノテーションが存在しない状況下で、擬似ラベルの品質がコントラスト損失の性能に与える影響は何か？
• RQ4 インライン内対比とインライン間対比の異なる表現マッチング戦略が、モデルの収束性および精度に与える影響は何か？
• RQ5 ラベルなしデータで事前学習を実施することで、提案されたコントラスト自己学習フレームワークの性能がさらに向上するか？

主な発見

• 本手法は、1つのラベル付き3Dボリュームのみを用いてACDCデータセットで平均DSC 0.881を達成し、最先端の半教師あり手法を上回った。
• 8つのラベル付きボリュームを用いた場合、インライン内マッチングではDSC 0.885、インライン間マッチングではDSC 0.897を達成し、堅牢性とスケーラビリティが確認された。
• Dice重複度が0.7または0.8の閾値を用いた擬似ラベルのフィルタリングは、すべての擬似ラベル（閾値 = 0）を用いる場合よりも高い性能をもたらさなかった。これは、この設定では高信頼度のフィルタリングが有益でないことを示している。
• 2つのラベル付きボリュームを用いた場合、インライン間マッチングはインライン内マッチングを上回った。これは、異なる被験者間のアライメントが一般化性能を向上させることを示唆している。
• 特に1〜2つのラベル付きボリュームという低ショット設定において、同時期のコントラスト学習および半教師あり学習ベースラインと比較して顕著な性能向上を達成した。
• ラベルなしデータでの事前学習と本手法のコントラスト自己学習フレームワークの組み合わせによりさらなる性能向上が得られた。これは、事前学習と自己教師あり学習の相乗効果を確認するものである。