Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification

TL;DR

This study demonstrates that Vision Transformers pre-trained with Masked Autoencoders on chest X-rays can achieve performance comparable to state-of-the-art CNNs in multi-label thorax disease classification, with tailored pre-training and fine-tuning strategies.

Contribution

It introduces a novel pre-training approach for ViTs using Masked Autoencoders on medical images and provides empirical insights into effective training recipes for medical imaging tasks.

Findings

ViT with MAE pre-training performs comparably to DenseNet-121.

Medical image reconstruction requires smaller image portions and moderate crop ranges.

In-domain transfer learning and specific fine-tuning strategies improve performance.

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.