[论文解读] Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification

Junfei Xiao, Yutong Bai+2 more|arXiv (Cornell University)|Oct 23, 2022

COVID-19 diagnosis using AI被引用 6

一句话总结

本论文提出了一套基于掩码自编码器（MAE）在266,340张未标注的胸部X光片上进行预训练与微调的强效方法，用于视觉Transformer（ViT），在三个基准数据集上的多标签胸部疾病分类任务中，性能达到或超过当前最优的卷积神经网络（CNN）（如DenseNet-121），在NIH ChestX-ray14、CheXpert和COVIDx数据集上的mAUC分别达到82.3%、89.2%和99.3%。

ABSTRACT

Vision Transformer (ViT) has become one of the most popular neural architectures due to its great scalability, computational efficiency, and compelling performance in many vision tasks. However, ViT has shown inferior performance to Convolutional Neural Network (CNN) on medical tasks due to its data-hungry nature and the lack of annotated medical data. In this paper, we pre-train ViTs on 266,340 chest X-rays using Masked Autoencoders (MAE) which reconstruct missing pixels from a small part of each image. For comparison, CNNs are also pre-trained on the same 266,340 X-rays using advanced self-supervised methods (e.g., MoCo v2). The results show that our pre-trained ViT performs comparably (sometimes better) to the state-of-the-art CNN (DenseNet-121) for multi-label thorax disease classification. This performance is attributed to the strong recipes extracted from our empirical studies for pre-training and fine-tuning ViT. The pre-training recipe signifies that medical reconstruction requires a much smaller proportion of an image (10% vs. 25%) and a more moderate random resized crop range (0.5~1.0 vs. 0.2~1.0) compared with natural imaging. Furthermore, we remark that in-domain transfer learning is preferred whenever possible. The fine-tuning recipe discloses that layer-wise LR decay, RandAug magnitude, and DropPath rate are significant factors to consider. We hope that this study can direct future research on the application of Transformers to a larger variety of medical imaging tasks.

研究动机与目标

• 评估原始视觉Transformer（ViT）在医学影像中的潜力，特别是用于多标签胸部疾病分类。
• 通过在包含266,340张未标注的胸部X光片的大规模数据集上使用掩码自编码器（MAE）进行预训练，解决ViT对数据量的高需求问题。
• 开发一种针对医学图像独特特征（如低对比度解剖结构和微小病灶）量身定制的预训练与微调方案。
• 提供一个公平的基准，将ViT与当前最优的CNN进行比较，包括在相同数据上通过自监督方法预训练的模型。
• 发布代码和预训练模型（ViT-Small、ViT-Base），以支持医学视觉领域的可复现性与未来研究。

提出的方法

• 使用掩码自编码器（MAE）对视觉Transformer进行预训练，其中90%的图像块被掩码，模型从可见块中重建缺失的像素。
• 采用经过修改的数据增强策略，随机裁剪的大小范围为0.5–1.0，显著窄于自然图像预训练中常用的0.2–1.0。
• 采用较高的掩码比例（90%），以更好地适应医学图像重建任务，而非自然图像MAE中使用的75%。
• 使用逐层权重衰减学习率的微调策略，结合RandAug增强（强度为15）和0.1的DropPath率。
• 将ViT的性能与在相同数据和相同评估协议下微调的当前最优CNN（如DenseNet-121）进行比较。
• 使用Grad-CAM分析并比较ViT与CNN的注意力图，以提升可解释性与定位性能。

实验结果

研究问题

• RQ1 当在大规模医学数据集上使用MAE进行预训练时，原始视觉Transformer是否能在多标签胸部疾病分类任务中达到或超过当前最优CNN的性能？
• RQ2 在医学影像（特别是胸部X光片）中应用MAE时，最优的预训练超参数（掩码比例、裁剪范围、数据规模）是什么？
• RQ3 微调超参数（如逐层学习率衰减、RandAug强度、DropPath率）如何影响ViT在医学分类任务中的表现？
• RQ4 在检测微小病灶区域方面，ViT的定位能力与CNN相比如何，其表现通过Grad-CAM和边界框IoU进行评估？
• RQ5 在医学视觉任务中，使用医学领域内数据进行迁移学习（在医学数据上预训练）是否比在ImageNet上预训练更有效？

主要发现

• 在NIH ChestX-ray14数据集上，预训练的ViT取得了82.3%的mAUC，与当前最优的DenseNet-121性能相当。
• 在斯坦福CheXpert数据集上，ViT取得了89.2%的mAUC，优于文献中报道的最佳CNN结果。
• 在COVIDx数据集上，ViT取得了99.3%的mAUC，展现出在具有挑战性、小样本且多样化的数据集上的卓越性能。
• ViT在定位微小病灶（如结节）方面表现更优，而CNN则在定位较大异常（如心脏扩大）方面更占优势。
• 与ImageNet预训练相比，在266,340张胸部X光片上进行领域内预训练显著提升了ViT的性能，凸显了领域特定数据的重要性。
• 逐层学习率衰减、RandAug强度15和0.1的DropPath率被确定为在多标签分类任务中有效微调ViT的关键超参数。