Causal representation learning with observational grouping for chest X-rays

In brief

• Deep neural networks trained on chest radiographs frequently exploit spurious demographic shortcuts (such as patient sex, race, or imaging view) rather than true pathology.
• These shortcut correlations degrade diagnostic reliability and cause performance disparities across patient subpopulations.
• Observational grouping enforces causal invariance in the learned latent space, guaranteeing that disease representations remain robust across diverse acquisition environments.

This retrospective explainer is dated to the September 2025 FAIMI workshop at MICCAI. It discusses Causal Representation Learning with Observational Grouping for CXR Classification by Rajat Rasal, Avinash Kori and Ben Glocker.

When an emergency medicine doctor examines a chest radiograph (CXR), their diagnosis of pneumonia or pleural effusion is based on pathological indicators: alveolar consolidations, blunting of the costophrenic angles, or air bronchograms. The patient’s biological sex, race, or whether the scan was taken with a portable bedside machine vs an upright scanner should not change the physical definition of lung fluid.

Yet deep neural networks trained on large clinical repositories (such as MIMIC-CXR or CheXpert) often take shortcuts. Because certain demographic groups or hospital wards have higher disease prevalence due to socioeconomic healthcare disparities, classifiers learn to detect patient demographics or scanner metadata as proxy indicators for disease.

When deployed in a new clinic with different demographics, these models suffer severe performance drops and perpetuate systemic diagnostic biases.

This paper tackles this problem using Causal Representation Learning with Observational Grouping, establishing an end-to-end framework that guarantees representation invariance across demographic groups.

Shortcuts as confounding in the causal graph

Consider the generative process behind medical imaging:

• Let $Y$ be the true disease status.
• Let $E$ be an environmental or demographic attribute (e.g. biological sex, self-reported race, or scanner view AP vs PA).
• Let $X$ be the resulting chest radiograph pixels.

In hospital datasets, $E$ and $Y$ are often statistically correlated ($E \not\perp!!\perp Y$) due to selection bias and healthcare access patterns. When a standard convolutional network or Vision Transformer minimizes empirical risk, it extracts all features that correlate with $Y$, blending genuine pathological features $Z_{\text{causal}}$ with spurious demographic signatures $Z_{\text{shortcut}}$:

\[Z = \phi_\theta(X) = \left[Z_{\text{causal}}, Z_{\text{shortcut}}\right].\]

If a model uses $Z_{\text{shortcut}}$ to predict $Y$, its predictions are brittle: intervening on the environment or changing the hospital population causes the correlation to break down.

Enforcing invariance via observational grouping

To isolate $Z_{\text{causal}}$ from $Z_{\text{shortcut}}$, the authors group the observational training data into distinct environment slices:

\[G_k = \{ (x, y) \mid E = k \}.\]

An invariant feature extractor $\phi_\theta$ processes samples from different groups simultaneously using shared weights. The latent representations are then structured via an explicit Latent Space Organisation Loss $\mathcal{L}_{\text{INV}}$.

Figure from the paper. Architecture overview from Rasal et al. Pairs of images $x_{k'} \sim G_{k'}$ and $x_k \sim G_k$ from different demographic or scanner groups pass through an Invariant Feature Extractor $\phi_\theta$ with shared weights. The feature classifier $\psi$ optimizes binary cross-entropy $\mathcal{L}_{\text{BCE}}$, while the Latent Space Organisation regulariser $\mathcal{L}_{\text{INV}}$ aligns the latent manifold geometry $\mathcal{Z}' \to \mathcal{Z}$. Source: Rajat Rasal, Avinash Kori and Ben Glocker. Extracted from the source manuscript; click to enlarge.

The overall training objective combines diagnostic classification with geometric manifold alignment:

\[\mathcal{L}_{\text{Total}} = \mathcal{L}_{\text{BCE}}\left(\psi(\phi_\theta(x)), y\right) + \lambda \, \mathcal{L}_{\text{INV}}\left(\phi_\theta(x_k), \phi_\theta(x_{k'})\right),\]

where $\mathcal{L}_{\text{INV}}$ penalizes discrepancies in feature distributions between pairs of environments:

\[\mathcal{L}_{\text{INV}} = \mathbb{E}_{k \neq k'} \left[ \mathcal{D}_{\text{align}}\left( p(\phi_\theta(X) \mid Y, E=k), \, p(\phi_\theta(X) \mid Y, E=k') \right) \right].\]

By conditioning the alignment on true disease label $Y$, the regulariser preserves disease discriminability while stripping away group-dependent variations.

Clinical robustness and fairness

The paper evaluates the method on multi-label chest radiograph benchmarks, testing generalization under severe demographic and view shifts:

1. Demographic Parity: The performance gap between sub-populations (evaluated across race, biological sex, and age brackets) drops significantly compared to standard empirical risk minimization.
2. Scanner View Robustness: When tested on images acquired in Anteroposterior (AP) view after training predominantly on Posteroanterior (PA) view, models trained with observational grouping maintain high Area Under the ROC Curve (AUC > 0.84), whereas standard baselines degrade substantially.
3. Pathological Interpretability: Saliency checks confirm that the invariant network focuses its visual attention on actual lung parenchyma and pleural boundaries, completely ignoring extraneous hardware, lead markers, and breast tissue.

Paper and sources

• Rajat Rasal, Avinash Kori and Ben Glocker. Causal Representation Learning with Observational Grouping for CXR Classification. FAIMI @ MICCAI 2025 (Springer LNCS).