MLVICX: Multi-Level Variance-Covariance Exploration for Chest X-ray Self-Supervised Representation Learning

TL;DR

This work tackles the challenge of learning robust chest X-ray representations under limited labels by introducing MLVICX, a self-supervised framework that explicitly regularizes variance and covariance across multiple feature levels. It employs dual Siamese encoders, intermediate representations, and context-bottleneck modules to fuse global and local contextual cues, optimized with VICReg-inspired losses on both global and multi-level embeddings, with the total objective $L_{mlca} = L_G + \lambda L_L$. Empirical results across NIH Chest X-ray-14, Vinbig-CXR, RSNA Pneumonia, and SIIM-ACR Pneumothorax show consistent gains over SOTA SSL methods and domain baselines in both fine-tuning and frozen settings, including strong performance in low-data regimes and segmentation transfer. The approach yields improved interpretability via diagnostic heatmaps and offers practical potential for clinical chest X-ray analysis by enhancing generalization and reducing annotation needs.

Abstract

Self-supervised learning (SSL) is potentially useful in reducing the need for manual annotation and making deep learning models accessible for medical image analysis tasks. By leveraging the representations learned from unlabeled data, self-supervised models perform well on tasks that require little to no fine-tuning. However, for medical images, like chest X-rays, which are characterized by complex anatomical structures and diverse clinical conditions, there arises a need for representation learning techniques that can encode fine-grained details while preserving the broader contextual information. In this context, we introduce MLVICX (Multi-Level Variance-Covariance Exploration for Chest X-ray Self-Supervised Representation Learning), an approach to capture rich representations in the form of embeddings from chest X-ray images. Central to our approach is a novel multi-level variance and covariance exploration strategy that empowers the model to detect diagnostically meaningful patterns while reducing redundancy effectively. By enhancing the variance and covariance of the learned embeddings, MLVICX promotes the retention of critical medical insights by adapting both global and local contextual details. We demonstrate the performance of MLVICX in advancing self-supervised chest X-ray representation learning through comprehensive experiments. The performance enhancements we observe across various downstream tasks highlight the significance of the proposed approach in enhancing the utility of chest X-ray embeddings for precision medical diagnosis and comprehensive image analysis. For pertaining, we used the NIH-Chest X-ray dataset, while for downstream tasks, we utilized NIH-Chest X-ray, Vinbig-CXR, RSNA pneumonia, and SIIM-ACR Pneumothorax datasets. Overall, we observe more than 3% performance gains over SOTA SSL approaches in various downstream tasks.

Figures (2)

• Figure 1: Architecture of the proposed approach. $v$ and $v'$ are two augmented versions of the input $x$, process by weight shared encoders $f_{\theta}$ and $f_{\theta{'}}$ respectively to give $y$ and $y'$. $C_{t_p}$ and $C_{t'_p}$ are the context-bottleneck blocks to refine the intermediate representations $y_p$ and $y'_p$. The refined intermediate representations are aggregated to make compound representations $u$ and $u'$ in each branch. $( || )$ represents the concatenation of multi-level intermediate feature maps, while $\otimes$ represents multiplication.
• Figure 2: Diagnostic heatmaps generated by MLVICX and SSL baselines. The heatmaps represent model interpretations of chest X-ray images fine-tuned with 1% of training samples from the NIH dataset.