Learning Sparse Label Couplings for Multilabel Chest X-Ray Diagnosis

TL;DR

This work tackles multilabel chest X-ray diagnosis by explicitly modeling inter-label dependencies through a lightweight Label-Graph Refinement that learns a sparse coupling matrix to refine logits after a strong SE-ResNeXt101 backbone. It combines MIS-based cross-validation, Asymmetric Loss, AMP, EMA, and a compact 1-step message-passing refinement to improve macro-ROC-AUC with minimal compute. The method achieves a macro-AUC of about $0.926$ on held-out data and demonstrates robust, cross-fold consistency, all while remaining hardware-friendly and needing no extra annotations. This approach offers a practical route to more reliable and scalable multilabel CXR classifiers, with potential for extension via context-conditioned couplings and deeper relational reasoning.

Abstract

We study multilabel classification of chest X-rays and present a simple, strong pipeline built on SE-ResNeXt101 $(32 \times 4d)$. The backbone is finetuned for 14 thoracic findings with a sigmoid head, trained using Multilabel Iterative Stratification (MIS) for robust cross-validation splits that preserve label co-occurrence. To address extreme class imbalance and asymmetric error costs, we optimize with Asymmetric Loss, employ mixed-precision (AMP), cosine learning-rate decay with warm-up, gradient clipping, and an exponential moving average (EMA) of weights. We propose a lightweight Label-Graph Refinement module placed after the classifier: given per-label probabilities, it learns a sparse, trainable inter-label coupling matrix that refines logits via a single message-passing step while adding only an L1-regularized parameter head. At inference, we apply horizontal flip test-time augmentation (TTA) and average predictions across MIS folds (a compact deep ensemble). Evaluation uses macro AUC averaging classwise ROC-AUC and skipping single-class labels in a fold to reflect balanced performance across conditions. On our dataset, a strong SE-ResNeXt101 baseline attains competitive macro AUC (e.g., 92.64% in our runs). Adding the Label-Graph Refinement consistently improves validation macro AUC across folds with negligible compute. The resulting method is reproducible, hardware-friendly, and requires no extra annotations, offering a practical route to stronger multilabel CXR classifiers.

Figures (8)

• Figure 1: Vertical architecture/flow for SE-ResNeXt101 (32$\times$4d) with a 14-label sigmoid head and a learned sparse label-coupling refinement applied to the logits.
• Figure 2: Label frequency distribution for the 14 CXR conditions in the training set. Bars show the number of positive cases per label in a multi-label setting.
• Figure 3: Representative chest X-ray images from the training set illustrating heterogeneity in view (PA/AP vs. lateral), patient positioning (upright/portable), presence of support devices, and acquisition contrast/cropping. Such variability motivates MIS-based CV and the use of lightweight augmentations (random resized crop, horizontal flip, $\pm 10^\circ$ rotation) before resizing to $380\times380$ and ImageNet normalization.
• Figure 4: Fold agreement (test predictions). Pairwise agreement among the three fold models aggregated over test cases. Most predictions are unanimous, indicating a stable ensemble.
• Figure 5: Inter-label correlation (ensemble predictions). Pearson correlation across labels. Warm colors denote positive correlations; cool colors denote negative correlations. Clinically related findings cluster together, while No Finding is broadly anti-correlated with pathologies.