CXR-CML: Improved zero-shot classification of long-tailed multi-label diseases in Chest X-Rays

TL;DR

This paper introduces a class-weighted, clustering-based method to improve zero-shot classification of long-tailed diseases in Chest X-Rays, significantly enhancing recognition of rare classes in imbalanced datasets.

Contribution

It proposes a novel class-weighting mechanism combined with GMM clustering and metric loss to improve zero-shot multi-label disease classification in chest X-ray images.

Findings

Achieved a 7% average increase in zero-shot AUC scores across 40 classes.

Effectively improves recognition of long-tailed, rare disease classes.

Demonstrates robustness and stability in feature clustering and classification.

Abstract

Chest radiography (CXR) plays a crucial role in the diagnosis of various diseases. However, the inherent class imbalance in the distribution of clinical findings presents a significant challenge for current self-supervised deep learning models. These models often fail to accurately classify long-tailed classes. Current Vision-Language models such as Contrastive Language Image Pre-training (CLIP) models effectively model the manifold distribution of the latent space, enabling high zero-shot classification accuracies. Although CLIP performs well on most of the primary classes in the dataset, our work reveals that its effectiveness decreases significantly for classes with a long-tailed distribution. Our approach employs a class-weighting mechanism that directly aligns with the distribution of classes within the latent space. This method ensures a substantial improvement in overall classification performance, with particular emphasis on enhancing the recognition and accuracy of rarely observed classes. We accomplish this by applying Gaussian Mixture Model (GMM) clustering to the latent space. The subsequent clusters are further refined by Student t-distribution, followed by a metric loss that utilizes the altered embeddings. Our approach facilitates stable and adaptive clustering of the features. This results in a notable average improvement of 7\% points in zero-shot AUC scores across 40 classes in the MIMIC-CXR-JPG dataset from previous SOTA models.