[論文レビュー] Leveraging Labeling Representations in Uncertainty-based Semi-supervised Segmentation

Sukesh Adiga V, José Dolz+1 more|arXiv (Cornell University)|Mar 10, 2022

Advanced Neural Network Applications被引用数 8

ひとこと要約

本論文は、スパarsely annotatedデータからラベル表現を学習する事前学習済みノイズ除去オートエンコーダ（DAE）を用いて、半教師あり3次元左房画像セグメンテーションのための新規な不確実性推定手法を提案する。DAEを介してモデル予測を妥当なセグメンテーションマスクにマッピングすることで、1回の推論で画素単位の不確実性を計算し、計算コストを低減するとともに、最先端手法を上回るセグメンテーション精度を達成した。10%ラベルデータで86.58%のDiceスコアと11.82mmのHDを達成した。

ABSTRACT

Semi-supervised segmentation tackles the scarcity of annotations by leveraging unlabeled data with a small amount of labeled data. A prominent way to utilize the unlabeled data is by consistency training which commonly uses a teacher-student network, where a teacher guides a student segmentation. The predictions of unlabeled data are not reliable, therefore, uncertainty-aware methods have been proposed to gradually learn from meaningful and reliable predictions. Uncertainty estimation, however, relies on multiple inferences from model predictions that need to be computed for each training step, which is computationally expensive. This work proposes a novel method to estimate the pixel-level uncertainty by leveraging the labeling representation of segmentation masks. On the one hand, a labeling representation is learnt to represent the available segmentation masks. The learnt labeling representation is used to map the prediction of the segmentation into a set of plausible masks. Such a reconstructed segmentation mask aids in estimating the pixel-level uncertainty guiding the segmentation network. The proposed method estimates the uncertainty with a single inference from the labeling representation, thereby reducing the total computation. We evaluate our method on the 3D segmentation of left atrium in MRI, and we show that our uncertainty estimates from our labeling representation improve the segmentation accuracy over state-of-the-art methods.

研究の動機と目的

• モンテカルロドロップアウトやアンサンブルのような複数回の推論を要する手法に起因する、半教師あり医療画像セグメンテーションにおける不確実性推定の高コストな計算を緩和すること。
• 追加のトレーニングタスクを必要とせず、アラインされた画像を必要としない、学習済みラベル表現を形状事前分布として活用することで、セグメンテーションの信頼性を向上させること。
• 収束性や安定性の問題を抱える共トレーニングや敵対的トレーニングのような複雑なトレーニングスキームへの依存度を低減すること。
• 事前学習済みDAEからの1回の推論によって、学生モデルをガイドするよりロバストで効率的な不確実性推定メカニズムを構築すること。
• ラベル表現から導出される不確実性が、エントロピーに基づくまたはアンサンブルに基づく不確実性推定よりも優れたセグメンテーション性能をもたらすことを実証すること。

提案手法

• 利用可能なセグメンテーションマスクから、妥当な解剖学的形状を捉えることができるラベル表現を学習するために、事前学習済みノイズ除去オートエンコーダ（DAE）が用いられる。
• DAEは、教師モデルのセグメンテーション予測を再構成された妥当なマスクにマッピングし、効果的に形状事前分布として機能する。
• 画素単位の不確実性は、元の予測とそのDAE再構成との間のL2差分として計算され、1回の推論による不確実性推定を可能にする。
• 不確実性マップは、平均教師フレームワークに正則化項として統合され、学生モデルが信頼性の高い予測領域に注目するようにガイドされる。
• 教師モデルの重みは指数移動平均（EMA）を用いて更新され、教師モデルと学生モデルの両方のトレーニングに、教師付き損失と一貫性損失を組み合わせて用いる。
• 本手法は、2018年アトリウムセグメンテーションチャレンジの3次元左房MRIデータを用いて評価され、最先端のベースラインと比較された。

実験結果

リサーチクエスチョン

• RQ1 スパarsely annotatedデータから学習されたラベル表現が、計算コストを低減しつつ、半教師ありセグメンテーションにおける不確実性推定を改善できるか？
• RQ2 DAEによって再構成されたマスクからの不確実性が、エントロピーに基づくまたはアンサンブルに基づく不確実性推定よりも優れたセグメンテーション性能をもたらすか？
• RQ3 低ラベルレジーム下で、本手法は最先端の半教師ありセグメンテーション手法と比較して、Diceスコアおよびハウスドルフ距離の観点で優れているか？
• RQ4 不確実性重みγおよび一貫性重みβといった異なるハイパーパrameter設定において、本手法の不確実性推定はロバストか？
• RQ5 アラインされた画像や形状予測や敵対的損失といった追加のトレーニングタスクを必要としない状況でも、ラベル表現に基づく不確実性推定は効果的に適用可能か？

主な発見

• 本手法は、10%ラベルデータ設定で86.58%のDiceスコアと11.82mmのハウスドルフ距離（HD）を達成し、最良のベースライン（SASSnet）を2.88%のDiceスコアおよび5.08mmのHDで上回った。
• 20%ラベルデータ設定では、88.60%のDiceスコアと7.61mmのHDを達成し、最良のベースライン（URPC）を0.02%のDiceスコアおよび5.49mmのHDで上回った。
• アブレーションスタディの結果、しきい値またはエントロピー戦略を用いたラベル表現に基づく不確実性推定は、UAMTを上回る性能を示した。フルメソッドが最も優れた結果を達成した。
• 本手法は、1トレーニングステップあたり1回のDAE推論で十分であり、1ステップあたり8回の推論を要するUAMTに比べ、計算コストを大幅に低減した。
• γおよびβのハイパーパramータに関するアブレーションでは、本手法が異なる設定において一貫してベースラインを上回ることが確認され、γ=1が最適な性能を示した。
• 可視化比較では、本手法がSASSnet、DTC、UAMT、URPCに比べ、より滑らかで正確なセグメンテーションを生成し、穴が少なく、境界に沿った適合性が優れていることが示された。