[Paper Review] How Can We Tame the Long-Tail of Chest X-ray Datasets?

Arsh Verma|arXiv (Cornell University)|Sep 8, 2023

COVID-19 diagnosis using AIMedicine3 citations

TL;DR

This paper proposes that fine-tuning vision models on chest X-ray (CXR) data before training on a target dataset significantly improves performance on rare, long-tail labels—without requiring complex loss functions or hyperparameter tuning. By using a two-stage pretraining strategy (ImageNet followed by CXR-specific datasets), the method achieves up to a 3 percentage point gain on ultra-rare labels and a 6% absolute mAP improvement via model averaging, demonstrating that initialization quality is a key lever for long-tail generalization in medical imaging.

ABSTRACT

Chest X-rays (CXRs) are a medical imaging modality that is used to infer a large number of abnormalities. While it is hard to define an exhaustive list of these abnormalities, which may co-occur on a chest X-ray, few of them are quite commonly observed and are abundantly represented in CXR datasets used to train deep learning models for automated inference. However, it is challenging for current models to learn independent discriminatory features for labels that are rare but may be of high significance. Prior works focus on the combination of multi-label and long tail problems by introducing novel loss functions or some mechanism of re-sampling or re-weighting the data. Instead, we propose that it is possible to achieve significant performance gains merely by choosing an initialization for a model that is closer to the domain of the target dataset. This method can complement the techniques proposed in existing literature, and can easily be scaled to new labels. Finally, we also examine the veracity of synthetically generated data to augment the tail labels and analyse its contribution to improving model performance.

Motivation & Objective

• To address the challenge of poor performance on rare, long-tail labels in chest X-ray (CXR) classification tasks.
• To investigate whether model initialization via domain-specific pretraining can improve generalization to low-frequency labels without modifying loss functions or data sampling strategies.
• To evaluate the effectiveness of model averaging and synthetic data augmentation in enhancing performance on tail labels.
• To provide a scalable, plug-and-play method for improving model performance on new, rare labels with minimal retraining.

Proposed method

• The method uses a two-stage pretraining strategy: first initializing with ImageNet weights, then fine-tuning on three open-source CXR datasets (NIH CXR, PadChest, CheXpert) to create a domain-adapted initialization.
• The model is then finetuned end-to-end on the MIMIC CXR dataset, which contains a long-tailed distribution of 26 labels, including 5 unique tail labels not present in the pretraining datasets.
• Model averaging is applied by ensembling predictions from models with different backbones (DenseNet161 and ResNeXt101), improving robustness on low-resource labels.
• Synthetic data augmentation is explored by jointly finetuning on MIMIC CXR and a small synthetic dataset (5,000 samples) generated using RoentGen to improve tail label performance.
• All models use binary cross-entropy loss with sigmoid activation and are trained for 20 epochs at 448×448 resolution to preserve fine-grained anatomical details.
• The training pipeline is applied at the patient level, with all available views used per patient, and data splits are created at the patient level to avoid data leakage.

Experimental results

Research questions

• RQ1 Can pretraining on a diverse set of CXR datasets improve model performance on rare, long-tail labels without modifying the loss function or data sampling strategy?
• RQ2 Does a two-stage pretraining approach (ImageNet followed by CXR-specific data) lead to better generalization on tail labels compared to standard ImageNet initialization?
• RQ3 To what extent does model averaging improve performance on low-frequency labels, especially those unique to the target dataset?
• RQ4 Can synthetic data augmentation help mitigate the data scarcity of tail labels, and does it show measurable gains in model performance?

Key findings

• The chained pretraining strategy improved performance on the 5 unique tail labels (Tail-U) by 3 percentage points for DenseNet161, compared to only 0.5 points for ResNeXt101, indicating strong gains on ultra-rare classes.
• Model averaging achieved a 6% absolute increase in mean average precision (mAP) on the Tail-U labels (0.1806 to 0.2430), outperforming all individual models.
• The averaged model achieved the best performance on 19 out of 26 labels on the validation set and led in all category-wise evaluations (head, medium, tail).
• Synthetic data augmentation, though limited in size (5,000 samples), showed a measurable performance increase on tail labels, with the most significant gains observed in the tail category, suggesting promise for future scaling.
• The method is highly scalable: adding new labels requires only a few fine-tuning epochs, making it practical for real-world deployment in clinical settings.
• The results demonstrate that model initialization is a critical, underappreciated factor in long-tail generalization, especially in medical imaging where rare pathologies are clinically significant.