UniChest: Conquer-and-Divide Pre-training for Multi-Source Chest X-Ray Classification

TL;DR

UniChest tackles cross-source heterogeneity in chest X-ray vision-language pre-training by introducing a two-stage Conquer-and-Divide framework. It first learns common cross-source patterns (Conquer) and then partitions source-specific signals into multiple query-network experts (Divide) using a mixture-of-experts and source-contrastive learning, achieving robust zero-shot generalization across diverse datasets. The model employs image and text encoders with a MoE-QN module and optimizes $\mathcal{L}_{\text{BCL}}$, $\mathcal{L}_{\text{BCE}}$, and $\mathcal{L}_{\text{SCL}}$, producing an ensemble prediction $s_i^{\text{All}}$ for diagnosis. Experiments on large multi-source pre-training data and multiple downstream benchmarks demonstrate state-of-the-art or competitive performance, including notable improvements on in-domain and zero-shot tasks, and provide qualitative lesion grounding, highlighting the practical impact of multi-source vision-language pre-training in radiology.

Abstract

Vision-Language Pre-training (VLP) that utilizes the multi-modal information to promote the training efficiency and effectiveness, has achieved great success in vision recognition of natural domains and shown promise in medical imaging diagnosis for the Chest X-Rays (CXRs). However, current works mainly pay attention to the exploration on single dataset of CXRs, which locks the potential of this powerful paradigm on larger hybrid of multi-source CXRs datasets. We identify that although blending samples from the diverse sources offers the advantages to improve the model generalization, it is still challenging to maintain the consistent superiority for the task of each source due to the existing heterogeneity among sources. To handle this dilemma, we design a Conquer-and-Divide pre-training framework, termed as UniChest, aiming to make full use of the collaboration benefit of multiple sources of CXRs while reducing the negative influence of the source heterogeneity. Specially, the ``Conquer" stage in UniChest encourages the model to sufficiently capture multi-source common patterns, and the ``Divide" stage helps squeeze personalized patterns into different small experts (query networks). We conduct thorough experiments on many benchmarks, e.g., ChestX-ray14, CheXpert, Vindr-CXR, Shenzhen, Open-I and SIIM-ACR Pneumothorax, verifying the effectiveness of UniChest over a range of baselines, and release our codes and pre-training models at https://github.com/Elfenreigen/UniChest.

Figures (8)

• Figure 1: Left: The class distribution of the multi-source dataset composed of MIMIC-CXR, ChestX-ray14, CheXpert and VinDr-CXR training sets. The initial five characters of each label serve as the abbreviated x-axis annotation. Right: T-SNE visualization w.r.t. the visual representations of medical images randomly selected from MIMIC-CXR, ChestX-ray14, CheXpert and VinDr-CXR, which characterizes the source heterogeneity.
• Figure 2: Inconsistent improvement is achieved when comparing KAD and KAD-Multi in terms of the AUC metric, when scaling up the training data by multiple sources and evaluating on the VinDr-CXR test set.
• Figure 3: The framework of UniChest, which consists of two training stages. During the "Conquer" stage, two modality encoders first project visual and textual representations into the common space with alignment, then feed them into the first transformer query networks for prediction. The multi-source common patterns are learnt as much as possible at this stage. During the "Divide" stage, we freeze the modality encoders and squeeze the source-specific patterns via the MoE-QN module with the guidance of the enhanced supervised loss and the source contrastive learning.
• Figure 4: Per-category performance of different methods on ChestX-ray14 (left) and Open-I (right). AUC scores of each category are displayed. 1 and 0 are adopted as the maximum and minimal values for each category in the radar chart.
• Figure 5: Per-category performance of 16 seen categories during pre-training in PadChest.