[論文レビュー] Co-training with High-Confidence Pseudo Labels for Semi-supervised Medical Image Segmentation

Zhiqiang Shen, Peng Cao+4 more|arXiv (Cornell University)|Jan 11, 2023

Advanced Neural Network Applications被引用数 7

ひとこと要約

本稿では、不確実性誘導型領域混合（UMIX）を用いて高信頼度の偽ラベルを生成し、共同平均教師（CMT）アーキテクチャによってモデルの不一致を維持することで、co-trainingを強化する半教師あり医療画像セグメンテーションフレームワークである不確実性誘導型共同平均教師（UCMT）を提案する。UCMTは、限定的なラベル付きデータを用いた4つの公的2Dおよび3D医療画像データセットで最先端の性能を達成する。

ABSTRACT

Consistency regularization and pseudo labeling-based semi-supervised methods perform co-training using the pseudo labels from multi-view inputs. However, such co-training models tend to converge early to a consensus, degenerating to the self-training ones, and produce low-confidence pseudo labels from the perturbed inputs during training. To address these issues, we propose an Uncertainty-guided Collaborative Mean-Teacher (UCMT) for semi-supervised semantic segmentation with the high-confidence pseudo labels. Concretely, UCMT consists of two main components: 1) collaborative mean-teacher (CMT) for encouraging model disagreement and performing co-training between the sub-networks, and 2) uncertainty-guided region mix (UMIX) for manipulating the input images according to the uncertainty maps of CMT and facilitating CMT to produce high-confidence pseudo labels. Combining the strengths of UMIX with CMT, UCMT can retain model disagreement and enhance the quality of pseudo labels for the co-training segmentation. Extensive experiments on four public medical image datasets including 2D and 3D modalities demonstrate the superiority of UCMT over the state-of-the-art. Code is available at: https://github.com/Senyh/UCMT.

研究の動機と目的

• co-trainingモデルが半教師あり医療画像セグメンテーションにおいて早期に共通認識に収束し、低信頼度の偽ラベルを生成するという限界に対処すること。
• 訓練中に不確実性を低減しつつ、サブネットワーク間のモデルの不一致を維持することで、偽ラベルの品質を向上させること。
• 限定的なラベル付きデータでより良い一般化を実現するため、モデルの不一致と偽ラベルの信頼度を同時に向上させる統合フレームワークを開発すること。
• co-trainingに基づく半教師あり学習におけるモデルの不一致と偽ラベルの不確実性の相乗的相互作用を調査すること。

提案手法

• UCMTは、2つの学生ネットワーク間のco-trainingを促進するため、3番目のモーメンタム平均化された教師モデルを用いて訓練を正則化し、共通認識の崩壊を防ぐ共同平均教師（CMT）コンponentを導入する。
• 不確実性誘導型領域混合（UMIX）モジュールは、CMTの予測から得られる不確実性マップを用いて、高不確実性領域を特定し、入力画像を部分的に混合することで、高信頼度の偽ラベルの生成を促進する。
• UMIXはCMTの不確実性マップを利用して、曖昧な領域を特定・拡張し、低不確実性領域のみを保持し、高不確実性領域を能動的に精錬する。
• CMTとUMIXを統合した共同訓練スキームを採用し、CMTはUMIXで拡張された入力からの偽ラベルを活用することで、予測の頑健性と多様性を向上させる。
• 本手法は学生ネットワーク間のクロス偽ラベル監視と平均教師正則化を採用し、訓練の安定化と一般化性能の向上を図る。
• UMIXのハイパーパrameter（例：上位k個の領域数、パッチサイズ）の感度分析を行い、さまざまな設定でも安定した性能を示すことがわかった。

実験結果

リサーチクエスチョン

• RQ1 co-trainingに基づく半教師あり医療画像セグメンテーションにおいて、モデルの不一致と偽ラベルの信頼度をどのように同時に最適化できるか。
• RQ2 不確実性に配慮したデータ拡張は、偽ラベルの品質とモデル性能にどのような影響を与えるか。
• RQ3 モーメンタム正則化を施した教師モデルを備えた共同平均教師フレームワークは、自己学習への早期収束を防げるか。
• RQ4 低データ環境下において、不確実性誘導型領域混合は、CutMixのようなランダム拡張手法と比較してどのように優れているか。
• RQ5 co-training設定におけるモデルの不一致、偽ラベルの不確実性、およびセグメンテーション性能の関係は何か。

主な発見

• UCMTは、2Dおよび3Dモodalitiesを含むさまざまな設定下で、ISIC、LA、ACDC、Polypの4つの公的医療画像セグメンテーションデータセットで最先端の性能を達成した。
• 5%のラベル付きデータを用いたISICデータセットでは、UCMTはDiceスコア88.22を達成し、ベースラインおよび先行SOTA手法を上回った。
• アブレーションスタディの結果、CMTとUMIXの併用が最良の性能をもたらし、両方のコンponentを含むUCMT全体モデルは、いずれのコンponentを欠損したバリアントよりも顕著に優れた性能を示した。
• UMIXは、2%のラベル付きデータという低データ環境下でもCutMixを常に上回り、高信頼度の偽ラベルを生成する有効性を示した。
• パラメータ感度分析の結果、UMIXはハイパーパrameterの選択に対して頑健であり、さまざまなパッチサイズや上位k個の領域選択において、最小限の性能低下しか示さなかった。
• 定性的な結果から、UCMTは平均教師やクロス偽ラベル監視ベースラインと比較して、滑らかでより正確なセグメンテーション境界を生成し、細部をよりよく回復していることがわかった。