PaperlayerView exact revision on arXiv

Source-linked AI summary

Delving into Masked Autoencoders for Multi-Label Thorax Disease Classification

Junfei Xiao, Yutong Bai, Alan Yuille, Zongwei Zhou

arXiv:2210.12843v1cs.CV

TL;DR

Medical ViTs have lagged CNNs amid limited data and difficult annotation, prompting the question of their potential with suitable training. The paper pre-trains ViTs with MAE on chest X-rays and customizes pre-training and fine-tuning recipes. The resulting vanilla ViT performs comparably, and sometimes better, than state-of-the-art CNNs across three chest X-ray benchmarks.

  • Problem

    ViTs require more data than CNNs, while medical data are limited and difficult to annotate, and their performance on medical tasks has often lagged state-of-the-art CNNs.

    Source: §1 · p. 1

  • Method

    The paper customizes MAE pre-training and ViT fine-tuning for medical imaging, including large-scale pre-training on unlabeled chest X-rays.

    Source: §1 · pp. 1–2

  • Results

    82.3%, 89.2%, and 99.3% mAUC on NIH ChestX-ray14, Stanford CheXpert, and COVIDx, respectively, as vanilla ViT approaches state-of-the-art CNN performance.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The study finds that vanilla ViT can achieve performance comparable to, and sometimes better than, state-of-the-art CNNs with strong pre-training and fine-tuning recipes.

    Source: §7 · p. 8

    The paper reports no clear evidence that reconstruction capability is positively correlated with transfer-learning performance.

    Source: §4.5 · p. 5

Abstract

from arXiv · show

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.

1. Introduction

The paper asks whether a stand-alone vanilla ViT can match state-of-the-art CNNs in medical imaging, where data are limited and annotations are difficult to obtain. It evaluates customized MAE pre-training and fine-tuning recipes, reporting comparable or better performance across three chest X-ray benchmarks.

1. Introduction

  • Medical data are limited and annotations are difficult to obtain, while Transformers require considerably more data than CNNs.The introduction frames this mismatch as a challenge for applying Transformers directly to medical tasks.Source: §1 · p. 1
  • The paper customizes MAE pre-training and fine-tuning for medical imaging and evaluates the recipe on three chest X-ray datasets.It also releases code and pre-trained ViT-Small, ViT-Base, and CNN models trained on 510K X-ray images.Source: §1 · p. 1
  • The study evaluates ViT pre-training on ImageNet and chest X-rays, highlighting the opportunity for in-domain transfer and self-supervised learning.The comparison uses ImageNet with 14M images and labels and chest X-rays with 0.3M images.Source: §1 · p. 2
  • Medical ViT pre-training uses 266,340 unlabeled X-rays, a 90% masking ratio, and a 0.5∼1.0 random cropping scale.The recipe is designed to learn image representations from chest X-rays efficiently.Source: §1 · p. 2
  • Layer-wise LR decay, RandAug magnitude, and DropPath rate are identified as important fine-tuning hyperparameters for multi-label thorax disease classification.Their importance is evaluated in Table 3.Source: §1 · p. 2
  • 82.3%, 89.2%, and 99.3% mAUC on NIH ChestX-ray14, Stanford CheXpert, and COVIDx, respectively, bring vanilla ViT close to state-of-the-art CNNs on three chest X-ray benchmarks.The paper describes this as among the first efforts to approach vanilla ViT performance to state-of-the-art CNNs on these benchmarks.Source: §1 · p. 2

2. Related Works

Radiography differs from photographic imagery: relevant findings can be subtle and distributed across anatomically consistent images. The paper motivates a fair vanilla ViT–CNN benchmark and contrasts self-supervised approaches suited to medical imaging.

Preliminary

  • Radiographs require models to capture both global and local features because clinically relevant information is dispersed and diseased regions can contain subtle, fine-grained variation.The paper contrasts this with photographic images that often have prominent foreground objects.Source: §2 · p. 2

ViTs or CNNs for medical imaging?

  • ViTs still lag state-of-the-art CNNs in radiography, and prior medical comparisons often lack similar experimental settings or use hybrid architectures.The authors identify no broad fair benchmark accounting for parameters, computation, GPU use, and suitable pre-training schemes.Source: §2 · p. 2

Self-supervised methods in medical imaging

  • Contrastive pre-training may be poorly suited to radiographs because anatomically similar images can create overly similar negative pairs.The authors contrast it with restorative pre-training, which conserves fine-grained textures embedded in image context.Source: §2 · pp. 2–3
  • The study benchmarks ViTs and CNNs directly on public radiography datasets and extends MAE-style pre-training to CNNs at the same medical-data scale.It also presents reusable insights for medical vision research.Source: §2 · p. 3

3. Method

The method pre-trains ViTs to reconstruct masked chest X-ray patches, using unlabeled images and a high masking ratio. A lightweight decoder handles reconstruction while the encoder processes only visible patches.

Data

  • The pre-training data comprise three public X-ray datasets, and the procedure uses no dataset-provided annotations.The datasets are NIH ChestX-ray14 (75,312 X-rays), Stanford CheXpert (191,028), and MIMIC-CXR (243,334); images are resized to 256×256 and cropped at scales 0.5∼1.0.Source: §3 · p. 3

Task

  • The model reconstructs masked image patches from visible patches, with mean squared error computed in pixel space over masked patches.The reconstruction task follows the MAE formulation.Source: §3 · p. 3
  • A 90% masking ratio is reported as optimal and accelerates pre-training by 2.5× compared with the original MAE.The figure description also reports that random resized cropping consistently improves MAE pre-training.Source: §3 · p. 3, Figure 1 · p. 3

Model

  • The vanilla ViT encoder processes only visible patches, while a separate, lighter ViT decoder reconstructs the masked patches during pre-training.Positional embeddings preserve patch locations; the decoder has depth 2 and width 512.Source: §3 · p. 3
  • The method uses ViT-S/16 and ViT-B/16 to denote ViT-Small and ViT-Base, respectively, with 16×16 patches; its lightweight decoder processes both visible and masked patches.The decoder has depth 2 and width 512, and its complexity is much smaller than the encoder’s.Source: §3 · p. 3

4. Pre-training: Recipe and Results

The section compares ViT and CNN pre-training across three public datasets and studies how in-domain data, dataset scale, masking, cropping, and fine-tuning choices affect the approach. It reports that X-ray pre-training helps ViTs approach or exceed CNN benchmarks, while reconstruction quality itself is not a reliable proxy for downstream transfer.

4. Pre-training: Recipe and Results

  • ViT-S/16 pre-training used AdamW, 800 epochs, a 1.5e-4 learning rate, batch size 2048, and random resized crop plus horizontal flip.The learning rate was warmed up for 20 epochs and then cosine-annealed.Source: §4 · p. 3

4.1. On the importance of in-domain transfer

  • Table 1 reports that ViT exceeds CNNs pre-trained with state-of-the-art ImageNet schemes, surpasses CNNs pre-trained on the same 0.3M X-rays, and performs comparably or better than literature CNNs.The table compares three public datasets and also reports that scratch training takes 200 versus 75 epochs for fine-tuning, while scratch- or ImageNet-initialized ViT trails CNNs.Source: Table 1 · p. 4
  • On ChestX-ray14, in-domain pre-training raised ViT mAUC from 78.6 to 82.3, while ImageNet-pretrained CNNs scored 82.2 versus 81.2 for in-domain pre-training.ViT-S/16 trained from scratch scored 67.9%, compared with 79.6% after ImageNet pre-training; CNN scores were 80.4% and 82.1%, respectively.Source: §4.1 · pp. 3–4, §4.2 · p. 4

4.2. Learning from 266,340 unlabeled X-rays

  • On ChestX-ray14, ViT-S/16 reached mAUC 79.3%, 81.9%, 82.3%, and 82.3% after pre-training on 75K, 191K, 266K, and 510K X-rays, respectively.The 75K-to-266K improvement was statistically significant (p-value=1.2e-127), but the gain from 266K to 510K was negligible; the study selected 266,340 X-rays for benchmarking.Source: §4.2 · p. 4

4.3. Masking out 90% X-ray content

  • A 90% masking ratio was optimal among tested ratios of 75%–95% and made MAE pre-training 2.5× faster than the original MAE.With 75%-mask pre-training, ViTs generalized better than CNNs when validation inputs had higher masking ratios.Source: §4.3 · p. 4, Figure 2 · p. 5
  • Fine-tuning studies identify layer-wise learning-rate decay, RandAug magnitude, and DropPath rate as important settings, with modest augmentation preferred and stronger regularization needed than for natural images.The ablations use ViT-S/16 on NIH Chest X-ray14 and report 14-class average AUC; the caption lists settings of 0.55, 6, and 0.2, respectively.Source: §4.3 · p. 5, Table 3 · p. 5

4.4. Cropping patches scaled of (0.5∼1.0)

  • Random resized crop consistently benefited X-ray MAE pre-training, and crop scales of 0.5∼1.0 yielded 1.0% higher mAUC than 0.2∼1.0.The authors suggest stronger spatial augmentation can crop informative lesions or organs and introduce noisy annotations.Source: §4.4 · p. 5

4.5. Quality assessment of image reconstruction

  • ViTs and CNNs reconstructed overall anatomy but not fine texture such as shoulder bones, and the paper reports no clear evidence that reconstruction ability correlates positively with transfer performance.The authors assess learned representations by generalizability and transferability on downstream tasks rather than reconstruction quality alone.Source: §4.5 · p. 5

5. Fine-tuning: Recipe and Results

Across CheXpert, COVIDx, and ChestX-ray14, the fine-tuned ViT achieves performance comparable to leading CNNs, with some disease-level and resolution-specific advantages. The section also describes a fine-tuning recipe whose key choices are reused across the three radiography tasks.

5. Fine-tuning: Recipe and Results

  • Layer-wise LR decay, RandAug magnitude, and DropPath rate are identified as crucial fine-tuning choices and reused across the three datasets.Models are fine-tuned for 75 epochs; more aggressive mixup and cutmix augmentation yielded no improvement and could produce noisy labels.Source: §5 · pp. 5–6
  • The benchmarks use dataset-specific evaluation: five-class mAUC on CheXpert, accuracy and COVID-19 sensitivity on COVIDx, and 14-class mean AUC on ChestX-ray14.COVIDx testing covers three classes; ChestX-ray14 follows the official split.Source: §5.1 · p. 6, §5.2 · p. 6, §5.3 · p. 7

5.1. Stanford CheXpert

  • ViT-S achieves 89.2% mAUC on CheXpert, close to the best performance of 89.4%, and leads on Atelectasis (83.5%) and Edema (94.0%).The benchmark reports mean AUC across five classes on the official validation set.Source: §5.1 · p. 6, Table 4 · p. 6

5.2. COVIDx

  • At 224×224 resolution, ViT-S beats the two compared COVIDx models, reaching 95.2% accuracy and 94.5% COVID-19 sensitivity.The reported evaluation uses the three-class testing set.Source: §5.2 · pp. 6–7, Table 5 · p. 6

5.3. NIH ChestX-ray14

  • 82.3% mean AUC over 14 diseases places vanilla ViT-S close to the previous best CNN’s 82.6% on ChestX-ray14; ViT-S has the best classification performance on 6 of 14 diseases.The prior best Transformer result was 81.0%.Source: §5.3 · p. 7, Table 6 · p. 7

6. Discussion

The discussion examines MAE reconstruction as a possible way to reveal anomalies and compares CNN and ViT disease localization using Grad-CAM. It reports different localization strengths by disease size and cautions that class activation maps are more suitable for visualizing CNN-type models.

Can MAE detect anomalies from an image?

  • MAE reconstruction sometimes replaces anomalies with normal patterns, so the difference map between the reconstruction and original X-ray can reveal and localize them.The authors describe this as a possible approach that requires neither image-level nor pixel-level annotations.Source: §6 · p. 7

Weakly-supervised disease localization by ViT and CNN

  • On 787 ChestX-ray14 cases with bounding boxes for eight diseases, CNN localizes larger diseases better while ViT is more robust on smaller diseases such as nodules.Localization uses thresholded Grad-CAM heatmaps, the largest connected component, and box regression; detection is evaluated with AP at IoU thresholds of 25% and 50%.Source: §6 · pp. 7–8, Table 7 · p. 8
  • Despite comparable classification AUC (82.1% for CNN versus 82.3% for ViT), CNN significantly exceeds ViT in localization ability.Their Grad-CAM attention maps differ: CNN attention is relatively larger and more concentrated; the authors consider class activation maps more suitable for visualizing CNN-type models.Source: §6 · p. 8

7. Conclusion and Future Work

The conclusion presents recipes and benchmarking as a route to competitive vanilla ViT performance on chest X-rays, with code and pretrained models available. Future work would expand the data and modalities and explore paired radiology reports.

7. Conclusion and Future Work

  • Vanilla ViT achieves performance comparable to, and sometimes better than, state-of-the-art CNNs after the authors devise medical-domain pre-training and fine-tuning recipes.The conclusion highlights computational efficiency, customized augmentation, larger data scale, and learning-parameter optimization.Source: §7 · p. 8
  • The paper states that code and pre-trained models are available.Source: §7 · p. 8
  • The benchmark compares vanilla ViT with literature-reported state-of-the-art CNNs, ImageNet-pretrained CNNs, and CNNs pretrained on the same medical data.The comparison accounts for parameters, computation, GPU usage, and pre-training schemes.Source: §7 · p. 8
  • Future work includes assembling more public X-rays, extending ViT to 3D modalities such as CT and MRI, and pre-training with paired reports and images.The authors note that 3D applications would require considerable computation and more pre-training data.Source: §7 · p. 8

Related papers

Loading 2210.12843v1…