[Paper Review] Bag of Tricks for Long-Tailed Multi-Label Classification on Chest X-Rays

Hong Feng, Tianjie Dai+3 more|arXiv (Cornell University)|Aug 17, 2023

COVID-19 diagnosis using AIMedicine3 citations

TL;DR

This paper presents a hybrid vision-language framework for long-tailed multi-label chest X-ray classification, integrating visual and textual encoders with transformer-based query networks to model label correlations. By combining data augmentation, loss reweighting, MixUp, test-time augmentation, and exogenous data replenishment, the method achieves 0.349 mAP on the ICCV CVAMD 2023 CXR-LT test set, ranking in the top five.

ABSTRACT

Clinical classification of chest radiography is particularly challenging for standard machine learning algorithms due to its inherent long-tailed and multi-label nature. However, few attempts take into account the coupled challenges posed by both the class imbalance and label co-occurrence, which hinders their value to boost the diagnosis on chest X-rays (CXRs) in the real-world scenarios. Besides, with the prevalence of pretraining techniques, how to incorporate these new paradigms into the current framework lacks of the systematical study. This technical report presents a brief description of our solution in the ICCV CVAMD 2023 CXR-LT Competition. We empirically explored the effectiveness for CXR diagnosis with the integration of several advanced designs about data augmentation, feature extractor, classifier design, loss function reweighting, exogenous data replenishment, etc. In addition, we improve the performance through simple test-time data augmentation and ensemble. Our framework finally achieves 0.349 mAP on the competition test set, ranking in the top five.

Motivation & Objective

• Address the dual challenge of class imbalance and label co-occurrence in clinical chest X-ray diagnosis.
• Improve model performance on rare diseases in long-tailed multi-label datasets.
• Systematically evaluate advanced techniques like data augmentation, loss reweighting, and test-time augmentation in a unified framework.
• Explore the integration of pre-trained text encoders and visual models for better label correlation modeling.
• Achieve state-of-the-art performance on the ICCV CVAMD 2023 CXR-LT competition through ensemble and data replenishment.

Proposed method

• Employ a dual-encoder architecture: ResNet-50 or DenseNet-121 for image features, PubMedBERT or Clinical-T5 for textual label embeddings.
• Use a multi-layer transformer decoder to attend to label embeddings (query) using image features (key/value), enabling cross-modal interaction.
• Apply cross-entropy loss with class-specific reweighting for rare diseases (upweighting factor 2 or 3) to mitigate long-tailed bias.
• Integrate MixUp augmentation with Beta(4) distribution to regularize training and improve generalization.
• Use geometric-mean fusion for test-time augmentation (TTA) to enhance robustness.
• Replenish training data using external datasets (ChestXRay14 and CheXpert) with label space alignment and zero-suppression for unseen labels.

Experimental results

Research questions

• RQ1 How do combined data augmentation and loss reweighting strategies affect performance on long-tailed multi-label chest X-ray classification?
• RQ2 To what extent can pre-trained text encoders improve label correlation modeling in multi-label medical imaging?
• RQ3 Does test-time augmentation and model ensemble improve generalization on rare disease predictions?
• RQ4 How effective is exogenous data replenishment from external CXR datasets in improving performance on rare classes?
• RQ5 Can a vision-language transformer-based query network outperform standard classifier heads in long-tailed multi-label settings?

Key findings

• The combination of separate classifier, loss reweighting (factor 2 or 3), MixUp, and TTA with ResNet-50 and PubMedBERT achieved 0.3280 mAP on the development set, up from 0.3187 without these techniques.
• Class-wise ensemble outperformed model-wise ensemble, achieving 0.348 mAP on the test set.
• The final model, using class-wise ensemble with exogenous data replenishment from ChestXRay14 and CheXpert, achieved 0.349 mAP on the test set, ranking in the top five of the competition.
• Exogenous data replenishment improved performance by aligning labels and leveraging additional training signals, especially for rare diseases.
• Test-time augmentation with geometric-mean fusion contributed to robustness and performance gains, particularly on hard-to-classify cases.
• The use of PubMedBERT and Clinical-T5 as text encoders showed comparable performance, with PubMedBERT yielding slightly better results in final configurations.