[论文解读] 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²，一种用于半监督胸部X光片分类的自监督均值教师框架，结合了联合对比学习（JCL）与指数移动平均教师训练。通过在预训练阶段利用未标注数据生成无限正样本特征对，在微调阶段应用一致性正则化，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.

研究动机与目标

• 为应对专家标注医学影像的高昂成本，提升胸部X光片分类中的半监督学习（SSL）性能。
• 克服现有SSL方法在微调阶段对未标注数据利用不足，或在自监督预训练中依赖有限正样本对的局限性。
• 开发一种新型框架，整合自监督预训练与基于一致性的师生学习，以提升表征能力与泛化性能。
• 在不同标注数据比例下，于多标签（ChestX-ray14、CheXpert）与多分类（ISIC2018）医学影像基准上验证该方法的有效性。

提出的方法

• 该方法使用联合对比学习（JCL）从未标注图像中生成无限数量的正样本特征对，提升自监督预训练阶段的表征学习能力。
• 通过学生-教师模型训练，利用学生网络权重的指数移动平均（EMA）构建稳健的教师网络，强制预测一致性。
• 自监督预训练阶段最小化从未标注数据中提取的查询特征与键特征之间无限正样本对的对比损失。
• 微调阶段采用双重损失：在标注数据上使用交叉熵损失，在未标注数据上使用学生与教师预测之间的一致性损失。
• 该框架结合了 S4L（自监督预训练）与一致性约束的SSL，使泛化能力优于先前的SOTA方法。
• 模型采用ResNet或DenseNet主干网络端到端训练，最终分类器通过同时使用标注与未标注数据进行微调。

实验结果

研究问题

• RQ1 使用无限正样本特征对进行自监督预训练，能否提升半监督胸部X光片分类中的表征学习能力？
• RQ2 将自监督预训练与均值教师一致性正则化相结合，是否能显著优于现有SOTA SSL方法？
• RQ3 自监督学习、JCL与均值教师微调等各组件对整体性能的贡献如何？
• RQ4 当在100%标注数据上微调时，该方法是否也能在完全监督设置下达到SOTA性能？

主要发现

• 在使用20%标注数据的 ChestX-ray14 上，S²MTS² 实现了81.06%的平均AUC，比之前SOTA方法高出0.5–1.5个百分点。
• 在 CheXpert 上，S²MTS² 实现了94.71%的平均AUC，超越之前SOTA（93.58%）1.13个百分点。
• 在使用20%标注数据的 ISIC2018 上，S²MTS² 实现了62.67%的F1分数，较之前SOTA（59.10%）提升3.57个百分点。
• 消融实验表明，每个组件——自监督学习、JCL与均值教师——均具有显著贡献，完整模型相比基线平均AUC提升2.7%。
• 即使在100%标注数据上训练，S²MTS² 在 ChestX-ray14 上仍实现71.58%的平均AUC，优于监督学习SOTA方法。
• 该方法在不同数据划分与标签比例下均表现出鲁棒性，在所有基准上AUC、F1与灵敏度指标均保持一致提升。