[論文レビュー] Self-training with dual uncertainty for semi-supervised medical image segmentation

Zhanhong Qiu, Haitao Gan+3 more|arXiv (Cornell University)|Apr 10, 2023

Radiomics and Machine Learning in Medical ImagingMedicine被引用数 3

ひとこと要約

本稿では、自己学習フレームワークに二重の不確実性推定（サンプルレベルおよびピクセルレベル）を導入することで、半教師あり医療画像セグメンテーションの性能を向上させることを提案する。サンプルレベルの不確実性は、同じ未ラベル付きサンプルに対して複数のチェックポイントにおけるモデル予測の分散を用いて推定され、ピクセルレベルの不確実性は、異なるアップサンプリング手法（例：transpose convolution と deconvolution）を用いた二つのデコーダー出力の差異から算出される。この手法により、難易度の高い未ラベル付きサンプルの選択的再学習と、疑似ラベルの精錬が可能となり、ACDCおよびProstateデータセットにおいて、同一の設定下で最先端の性能を達成した。

ABSTRACT

In the field of semi-supervised medical image segmentation, the shortage of labeled data is the fundamental problem. How to effectively learn image features from unlabeled images to improve segmentation accuracy is the main research direction in this field. Traditional self-training methods can partially solve the problem of insufficient labeled data by generating pseudo labels for iterative training. However, noise generated due to the model's uncertainty during training directly affects the segmentation results. Therefore, we added sample-level and pixel-level uncertainty to stabilize the training process based on the self-training framework. Specifically, we saved several moments of the model during pre-training, and used the difference between their predictions on unlabeled samples as the sample-level uncertainty estimate for that sample. Then, we gradually add unlabeled samples from easy to hard during training. At the same time, we added a decoder with different upsampling methods to the segmentation network and used the difference between the outputs of the two decoders as pixel-level uncertainty. In short, we selectively retrained unlabeled samples and assigned pixel-level uncertainty to pseudo labels to optimize the self-training process. We compared the segmentation results of our model with five semi-supervised approaches on the public 2017 ACDC dataset and 2018 Prostate dataset. Our proposed method achieves better segmentation performance on both datasets under the same settings, demonstrating its effectiveness, robustness, and potential transferability to other medical image segmentation tasks. Keywords: Medical image segmentation, semi-supervised learning, self-training, uncertainty estimation

研究の動機と目的

• セグメンテーションタスクにおける限られたラベル付き医療画像データという重要な課題に対処すること。
• モデルの不確実性に起因するノイズを低減することで、半教師あり学習における自己学習プロセスの安定化を図ること。
• 不確実性推定に基づいて難易度の高い未ラベル付きサンプルを選択的に再学習することで、セグメンテーション精度を向上させること。
• 二重のデコーダーを用いたピクセルレベルの不確実性推定により、疑似ラベルの品質を向上させること。
• 異なる医療画像分野にわたる多様なデータセットにおいて、一般化性と頑健性を示すこと。

提案手法

• サンプルレベルの不確実性は、同じ未ラベル付きサンプルに対して複数のチェックポイントにおけるモデル予測の分散として推定される。
• 未ラベル付きサンプルは、不確実性が低い順（容易なものから困難なものへ）に段階的に学習に導入される。
• ピクセルレベルの不確実性は、異なるアップサンプリング戦略（例：transpose convolution と deconvolution）を用いた二つのデコーダー出力の差異として計算される。
• 疑似ラベルはピクセルレベルの不確実性を用いて精錬され、不確実性が低い領域ほど高い信頼度が割り当てられる。
• モデルは自己学習を繰り返し、誤差伝搬を低減するため、最も不確実性の高いサンプルのみを再学習することで段階的に改善される。
• このフレームワークは、アーキテクチャの大幅な変更なしに、標準的な自己学習パイプラインに不確実性を考慮した疑似ラベル作成を統合している。

実験結果

リサーチクエスチョン

• RQ1 サンプルレベルの不確実性推定は、半教師あり医療画像セグメンテーションにおける疑似ラベルの信頼性を向上させることができるか？
• RQ2 二重のデコーダーから得られるピクセルレベルの不確実性は、疑似ラベルの品質向上とノイズ低減に寄与するか？
• RQ3 二重の不確実性に基づいて難易度の高い未ラベル付きサンプルを選択的に再学習することで、セグメンテーション性能が向上するか？
• RQ4 提案手法は、ベンチマークとして用いられる既存の自己学習アプローチと比較して、どのように性能を示すか？
• RQ5 この手法は、異なる医療画像セグメンテーションタスクにわたって、どの程度一般化可能で頑健であるか？

主な発見

• 提案手法は、ベースライン手法と同一の学習プロトコルを用いた2017年ACDCデータセットにおいて、最先端の性能を達成した。
• 2018年Prostateデータセットでは、五つの既存の半教師ありアプローチを上回るセグメンテーション精度を示した。
• 二重の不確実性推定の統合により、自己学習プロセス中にノイズの多い疑似ラベルから生じる誤差伝搬が顕著に低減された。
• 最も不確実性の高いサンプルのみを再学習することで、低品質な疑似ラベルによる性能劣化を回避し、モデルの頑健性が向上した。
• 本手法は、二つの異なる医療画像分野において一貫した性能向上を示し、優れた一般化可能性を示した。
• 複数のモデルチェックポイントと二重のデコーダーを用いることで、追加の学習コストや推論コストをかけずに信頼性の高い不確実性推定が可能となった。