[論文レビュー] Self-supervised Mean Teacher for Semi-supervised Chest X-ray Classification

Fengbei Liu, Yu Tian+4 more|arXiv (Cornell University)|Mar 5, 2021

COVID-19 diagnosis using AI参考文献 39被引用数 9

ひとこと要約

本論文では、自己教師ありの平均教師フレームワークであるS²MTS²を提案する。この手法は、自己教師あり事前学習における無限の正例特徴ペアと微調整段階における一貫性正則化を組み合わせた、半教師ありのチアックスレジオン分類のためのものである。未ラベルデータを事前学習段階で無限の正例特徴ペアとして活用し、微調整段階で一貫性正則化を適用することで、S²MTS²はChestX-ray14、CheXpert、ISIC2018で最先端の性能を達成し、ラベル付きデータが限られている状況でも、従来の手法を大きく上回った。

ABSTRACT

The training of deep learning models generally requires a large amount of annotated data for effective convergence and generalisation. However, obtaining high-quality annotations is a laboursome and expensive process due to the need of expert radiologists for the labelling task. The study of semi-supervised learning in medical image analysis is then of crucial importance given that it is much less expensive to obtain unlabelled images than to acquire images labelled by expert radiologists. Essentially, semi-supervised methods leverage large sets of unlabelled data to enable better training convergence and generalisation than using only the small set of labelled images. In this paper, we propose Self-supervised Mean Teacher for Semi-supervised (S$^2$MTS$^2$) learning that combines self-supervised mean-teacher pre-training with semi-supervised fine-tuning. The main innovation of S$^2$MTS$^2$ is the self-supervised mean-teacher pre-training based on the joint contrastive learning, which uses an infinite number of pairs of positive query and key features to improve the mean-teacher representation. The model is then fine-tuned using the exponential moving average teacher framework trained with semi-supervised learning. We validate S$^2$MTS$^2$ on the multi-label classification problems from Chest X-ray14 and CheXpert, and the multi-class classification from ISIC2018, where we show that it outperforms the previous SOTA semi-supervised learning methods by a large margin.

研究の動機と目的

• 医療画像の専門家によるアノテーションコストを低減するため、チアックスレジオン分類における半教師あり学習（SSL）の改善を目的とする。
• 従来のSSL手法が微調整段階で未ラベルデータを十分に活用していない、または自己教師あり事前学習段階で有限の正例ペアに依存しているという制限を克服することを目的とする。
• 自己教師あり事前学習と一貫性に基づく学生・教師学習を統合した、新たなフレームワークの開発を目的とする。
• ラベル付きデータの割合が異なる複数のベンチマーク（多ラベル：ChestX-ray14、CheXpert；多クラス：ISIC2018）を用いて、手法の有効性を検証することを目的とする。

提案手法

• 本手法は、未ラベル画像から無限の正例特徴ペアを生成するための共同対照学習（JCL）を用い、自己教師あり事前学習段階で表現学習を向上させる。
• 学生・教師モデルは、学生の重みの指数移動平均（EMA）を用いて堅牢な教師ネットワークを構築し、予測の一貫性を強制する。
• 自己教師あり事前学習段階では、未ラベルデータから抽出されたクエリ特徴とキーフィーチャーの無限の正例ペアに対する対照損失を最小化する。
• 微調整段階では、ラベル付きデータに対する交差エントロピー損失と、未ラベルデータにおける学生と教師の予測の間の一貫性損失の二重損失を用いる。
• S4L（自己教師あり事前学習）と一貫性を強制するSSLを組み合わせることで、従来の最先端手法よりも優れた一般化性能を実現する。
• モデルはResNetまたはDenseNetバックボーンを用いてエンドツーエンドで訓練され、最終的な分類器はラベル付きデータと未ラベルデータの両方を用いて微調整される。

実験結果

リサーチクエスチョン

• RQ1 無限の正例特徴ペアを用いた自己教師あり事前学習は、半教師ありチアックスレジオン分類における表現学習を向上させることができるか？
• RQ2 自己教師あり事前学習と平均教師の一貫性正則化を組み合わせることで、従来の最先端SSL手法を上回る性能が得られるか？
• RQ3 自己教師あり学習、JCL、平均教師微調整といった個々のモジュールが、全体の性能にどのように寄与しているか？
• RQ4 完全にラベル付きデータを使用して微調整した場合、本手法は完全に教師ありの設定でも最先端の結果を達成できるか？

主な発見

• ChestX-ray14で20%のラベル付きデータを使用した場合、S²MTS²は平均AUC 81.06%を達成し、従来の最先端手法を0.5～1.5ポイント上回った。
• CheXpertでは、平均AUC 94.71%を達成し、従来の最先端（93.58%）を1.13ポイント上回った。
• ISIC2018で20%のラベル付きデータを使用した場合、S²MTS²はF1スコア62.67%を達成し、前回の最先端（59.10%）を3.57ポイント上回った。
• アブレーションスタディの結果、自己教師あり学習、JCL、平均教師の各モジュールが顕著に寄与しており、完全なモデルはベースライン比で平均AUCを2.7%向上させた。
• 100%のラベル付きデータで訓練した場合でも、S²MTS²はChestX-ray14で平均AUC 71.58%を達成し、教師ありの最先端手法を上回った。
• 本手法は、異なるデータ分割とラベル比率に対してもロバストであり、すべてのベンチマークでAUC、F1、感度の面で一貫した向上を示した。