[论文解读] Supervised Transfer Learning at Scale for Medical Imaging

Basil Mustafa, Aaron Loh+13 more|arXiv (Cornell University)|Jan 14, 2021

COVID-19 diagnosis using AI参考文献 49被引用 23

一句话总结

该论文表明，即使在自然图像与医学影像之间存在显著领域差异的情况下，使用大规模自然图像数据集（如JFT和ImageNet-21k）预训练的大型监督微调迁移学习方法，仍能显著提升医学影像任务（包括胸部X光、乳腺X光摄影和皮肤科影像）的性能、数据效率以及分布外泛化能力。

ABSTRACT

Transfer learning is a standard technique to improve performance on tasks with limited data. However, for medical imaging, the value of transfer learning is less clear. This is likely due to the large domain mismatch between the usual natural-image pre-training (e.g. ImageNet) and medical images. However, recent advances in transfer learning have shown substantial improvements from scale. We investigate whether modern methods can change the fortune of transfer learning for medical imaging. For this, we study the class of large-scale pre-trained networks presented by Kolesnikov et al. on three diverse imaging tasks: chest radiography, mammography, and dermatology. We study both transfer performance and critical properties for the deployment in the medical domain, including: out-of-distribution generalization, data-efficiency, sub-group fairness, and uncertainty estimation. Interestingly, we find that for some of these properties transfer from natural to medical images is indeed extremely effective, but only when performed at sufficient scale.

研究动机与目标

• 探究现代大规模预训练是否能够克服自然图像与医学影像之间的领域差异挑战。
• 评估模型规模和预训练数据集规模对医学影像迁移学习性能的影响。
• 评估关键部署特性，如分布外泛化能力、数据效率、子组公平性以及不确定性校准能力。
• 通过分析深层特征重用和权重更新机制，理解性能提升的内在原因。

提出的方法

• 采用基于ResNet架构的Big Transfer（BiT）模型，其在包含100万至3亿张图像的自然图像数据集（如JFT、ImageNet-21k）上进行大规模预训练。
• 在三个多样化的医学影像任务上微调这些大规模预训练模型：胸部X光摄影（CheXpert）、乳腺X光摄影（UK数据集）以及皮肤病学影像（由皮肤科医生标注的图像）。
• 使用标准指标（AUC、准确率）和额外的临床相关指标（分布外泛化能力、数据效率、子组公平性以及期望校准误差ECE）评估性能。
• 应用线性中心化核对齐（CKA）分析预训练模型与微调后模型在不同网络深度上的特征相似性。
• 追踪ResNet模块中参数的移动情况，研究不同预训练规模如何影响微调过程中的权重更新。
• 在特定任务中使用平滑概率标签和焦点损失，以提升模型的校准能力和公平性。

实验结果

研究问题

• RQ1 在存在显著领域差异的情况下，是否能够通过JFT和ImageNet-21k等自然图像数据集进行大规模监督预训练，有效迁移到医学影像任务？
• RQ2 增加模型规模和预训练数据集规模是否能提升医学影像在分布偏移下的数据效率和泛化能力？
• RQ3 大规模迁移学习如何影响子组公平性和模型校准能力——这两项是临床部署中的关键因素？
• RQ4 当在大规模下进行预训练时，深层特征在微调模型中能在多大程度上被有效重用？
• RQ5 预训练规模、模型架构规模与微调过程中的参数更新模式之间存在何种关系？

主要发现

• 在自然图像（JFT、ImageNet-21k）上进行大规模预训练，能显著提升医学影像任务的性能，优于标准ImageNet预训练方法。
• 在更大数据集（如JFT）上预训练的模型，仅使用30%–60%的训练数据即可达到全量数据基线的性能水平，展现出强大的数据效率。
• 随着预训练数据集规模增大，分布外泛化能力显著提升，表明对分布偏移具有更强的鲁棒性。
• 大规模预训练下，子组公平性和模型校准能力得到保持或略有改善，且在不同人口统计子组中未出现公平性指标下降的情况。
• 线性CKA分析显示，预训练于更大数据集的更大模型，能更好地保留并重用来自预训练初始化的高层特征。
• 参数移动分析表明，更大的预训练数据集导致浅层权重更新更大，而深层权重更新更小，表明高层表征被更有效地重用。