[論文レビュー] Double-Uncertainty Weighted Method for Semi-supervised Learning

Yixin Wang, Yao Zhang+5 more|arXiv (Cornell University)|Oct 19, 2020

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

ひとこと要約

本稿では、教師-生徒フレームワークを用いた半教師付き医療画像セグメンテーションのための二重不確実性重み付け手法を提案する。ベイジアンディープラーニングの枠組みにおいて、モンテカルロドロップアウトを用いてセグメンテーションの不確実性と新規特徴量の不確実性を推定することで、一貫性損失を動的に重み付けし、教師モデルの信頼性と生徒の学習を向上させる。2つの公開データセットにおいて、最小限のラベル付きデータで最先端の性能を達成した。

ABSTRACT

Though deep learning has achieved advanced performance recently, it remains a challenging task in the field of medical imaging, as obtaining reliable labeled training data is time-consuming and expensive. In this paper, we propose a double-uncertainty weighted method for semi-supervised segmentation based on the teacher-student model. The teacher model provides guidance for the student model by penalizing their inconsistent prediction on both labeled and unlabeled data. We train the teacher model using Bayesian deep learning to obtain double-uncertainty, i.e. segmentation uncertainty and feature uncertainty. It is the first to extend segmentation uncertainty estimation to feature uncertainty, which reveals the capability to capture information among channels. A learnable uncertainty consistency loss is designed for the unsupervised learning process in an interactive manner between prediction and uncertainty. With no ground-truth for supervision, it can still incentivize more accurate teacher's predictions and facilitate the model to reduce uncertain estimations. Furthermore, our proposed double-uncertainty serves as a weight on each inconsistency penalty to balance and harmonize supervised and unsupervised training processes. We validate the proposed feature uncertainty and loss function through qualitative and quantitative analyses. Experimental results show that our method outperforms the state-of-the-art uncertainty-based semi-supervised methods on two public medical datasets.

研究の動機と目的

• ラベル付き医療画像データの不足という課題に対処すること。これは取得に費用と時間がかかるため。
• トレーニングプロセスに不確実性推定を組み込むことで、教師-生徒型半教師付き学習の信頼性を向上させること。
• 畳み込みカーネルチャネルにわたる不確実性を捉える新しい特徴量不確実性指標を提案・検証すること。
• モデルの予測と連携する学習可能な不確実性一貫性損失を設計し、教師出力を真のラベルなしで精緻化すること。
• 不確実性を動的重み付けメカニズムとして用いることで、教師付きと教師なしのトレーニングのバランスをとること。

提案手法

• 教師モデルはモンテカルロドロップアウトを用いてベイジアンニューラルネットワークを近似し、セグメンテーションの不確実性と特徴量の不確実性の両方を推定可能にする。
• 特徴量の不確実性は、特に第4のアップサンプルレイヤーの畳み込み層のチャネルにわたって計算され、各チャネルの特徴抽出の信頼性を評価する。
• 学習可能な不確実性一貫性損失を導入し、不確実性推定値で重み付けされた一貫性損失を用いて、生徒と教師の出力の不一致を最小化することで、反復的に教師の予測を精緻化する。
• 総損失関数は、ラベル付きデータに対する教師付き交差エントロピー損失と、ラベル付きおよびラベルなしデータに対する不確実性重み付き一貫性損失の組み合わせから構成される。
• 不確実性推定値が一貫性損失における適応的重みとして機能し、トレーニング中に高不確実性の予測の影響を低減する。
• 本手法は、2つの公開医療画像ベンチマーク（2018 Atrial Segmentation Challenge および MICCAI 2019 KiTS Challenge）で評価された。

実験結果

リサーチクエスチョン

• RQ1 畳み込みカーネルチャネルにわたる不確実性という、特徴量不確実性の新しい形態が、半教師付きセグメンテーションにおけるモデルの解釈可能性と性能を向上させることができるか？
• RQ2 モデルの予測と連携する学習可能な不確実性一貫性損失が、真のラベルなしで教師モデルの精度を向上させることができるか？
• RQ3 不確実性に基づく一貫性損失の重み付けが、半教師付き学習における教師付きと教師なしのトレーニングのバランスにどのように影響するか？
• RQ4 ラベル付きデータが限られる状況で、本手法が既存の不確実性ベースの半教師付きアプローチを上回るか？
• RQ5 不確実性推定が、信頼性の低い予測をフィルタリングし、医療画像セグメンテーションにおける一般化性能の向上にどの程度寄与するか？

主な発見

• KiTS 2019腎臓セグメンテーションデータセットにおいて、本手法はわずか10枚のラベル付きスキャンでのみ、Diceスコア0.9206を達成し、すべてのベースラインを上回った。
• 32枚のラベル付きスキャンを用いた場合、Diceスコアは0.9365に達し、最先端の手法に対して一貫した改善を示した。
• 提案された不確実性一貫性損失により、トレーニングの進行に伴い教師モデルの予測不確実性が低下したことが、図2(c)の平均セグメンテーション不確実性値の減少から示された。
• 特徴量不確実性の分析から、チャネルごとの不確実性値がトレーニングに伴い減少したことが判明し、特徴抽出の信頼性が向上したことが示された。
• 非常に少ないラベル付きサンプル（例：4枚）でも、本手法は性能が著しく向上し、Diceスコアで1位を記録した。
• 不確実性重み付けメカニズムにより、UA-MTのようなベースライン手法の性能も向上し、Diceスコアで0.0087の改善が得られた。