Mitigating Shortcut Learning via Feature Disentanglement in Medical Imaging: A Benchmark Study

Sarah Müller; Philipp Berens

cs.CVFeb 17, 2026

Mitigating Shortcut Learning via Feature Disentanglement in Medical Imaging: A Benchmark Study

Authors: Sarah Müller, Philipp Berens

Organizations: Hertie Institute for AI in Brain Health, Faculty of Medicine, University of Tübingen, Germany · Tübingen AI Center, University of Tübingen, Germany

Abstract

Although deep learning models in medical imaging often achieve excellent classification performance, they can rely on shortcut learning, exploiting spurious correlations or confounding factors that are not causally related to the target task. This poses risks in clinical settings, where models must generalize across institutions, populations, and acquisition conditions. Feature disentanglement is a promising approach to mitigate shortcut learning by separating task-relevant information from confounder-related features in latent representations. In this study, we systematically evaluated feature disentanglement methods for mitigating shortcuts in medical imaging, including adversarial learning and latent space splitting based on dependence minimization. We assessed classification performance and disentanglement quality using latent space analyses across one artificial and two medical datasets with natural and synthetic confounders. We also examined robustness under varying levels of confounding and compared computational efficiency across methods. We found that shortcut mitigation methods improved classification performance under strong spurious correlations during training. Latent space analyses revealed differences in representation quality not captured by classification metrics, highlighting the strengths and limitations of each method. Model reliance on shortcuts depended on the degree of confounding in the training data. The best-performing models combine data-centric rebalancing with model-centric disentanglement, achieving stronger and more robust shortcut mitigation than rebalancing alone while maintaining similar computational efficiency. The project code is publicly available at https://github.com/berenslab/medical-shortcut-mitigation.

Figures & tables

Figure 1: Overview of shortcut learning and mitigation via feature disentanglement. a Example of a spurious correlation between two binary tasks in the training data that reverses at test time, illustrating a distribution shift. b Causal graph in which a confounder affects both tasks

y_{1}

and

y_{2}

, while the image

X

is generated from both; the predictive association from

y_{2}

y_{1}

constitutes a shortcut. c Subspace disentanglement architecture that splits the latent representation into task-specific subspaces and minimizes their statistical dependence by minimizing a dependence measure

\mathcal{L}_{\text{dep}}

. d Comparison of latent representations showing that disentanglement reduces clustering of the confounding label

y_{2}

in the target-task subspace

z_{1}

Dataset	Primary task $y_{1}$	Confounder $y_{2}$
Morpho-MNIST	digits (0–4 / 5–9)	style (thin / thick)
CheXpert	pleural effusion (no / yes)	sex (female / male)
OCT	drusen (no / yes)	radial notch (no / yes)

Table 1: Summary of the datasets and tasks.

	Morpho-MNIST (digits)			CheXpert (pleural effusion)			OCT (drusen)
Method	Original	Balanced	Inverted	Original	Balanced	Inverted	Original	Balanced	Inverted
Baseline (ERM)	$96\pm 1$	$95\pm 1$	$88\pm 3$	$79\pm 1$	$79\pm 1$	$46\pm 4$	$99\pm 1$	$92\pm 5$	$74\pm 15$
Rebalancing	$98\pm 1$	$97\pm 1$	$97\pm 0$	$88\pm 2$	$88\pm 2$	$84\pm 5$	$99\pm 1$	$96\pm 3$	$90\pm 9$
AdvCl	$98\pm 0$	$97\pm 1$	$96\pm 1$	$83\pm 1$	$83\pm 1$	$64\pm 4$	$99\pm 1$	$97\pm 2$	$93\pm 5$
AdvCl+Rebal	$98\pm 2$	$97\pm 2$	$96\pm 1$	$85\pm 2$	$85\pm 1$	$72\pm 9$	$99\pm 1$	$97\pm 1$	$92\pm 4$
dCor	$97\pm 0$	$96\pm 1$	$92\pm 1$	$79\pm 3$	$79\pm 3$	$49\pm 9$	$99\pm 2$	$95\pm 2$	$88\pm 5$

Table 2: Classification performance of the main task (

z_{1}\rightarrow y_{1}

) measured as AUROC over different datasets and distributions.

Dataset	Baseline	Rebal	AdvCl	AdvCl+ Rebal	dCor	dCor+ Rebal	MINE	MINE+ Rebal	MMD	MMD+ Rebal
Morpho-MNIST	$z_{1}$ $z_{2}$ $y_{1}$ 87 52 $y_{2}$ 80 100	$z_{1}$ $z_{2}$ $y_{1}$ 86 50 $y_{2}$ 84 100	$z_{1}$ $y_{1}$ 91 $y_{2}$ 59	$z_{1}$ $y_{1}$ 91 $y_{2}$ 55	$z_{1}$ $z_{2}$ $y_{1}$ 90 52 $y_{2}$ 64 99	$z_{1}$ $z_{2}$ $y_{1}$ 93 50 $y_{2}$ 52 100	$z_{1}$ $z_{2}$ $y_{1}$ 90 50 $y_{2}$ 48 100	$z_{1}$ $z_{2}$ $y_{1}$ 90 50 $y_{2}$ 52 100	$z_{1}$ $z_{2}$ $y_{1}$ 89 54 $y_{2}$ 72 99	$z_{1}$ $z_{2}$ $y_{1}$ 84 58 $y_{2}$ 98 100
CheXpert	$z_{1}$ $z_{2}$ $y_{1}$ 68 63 $y_{2}$ 85 89	$z_{1}$ $z_{2}$ $y_{1}$ 76 57 $y_{2}$ 67 90	$z_{1}$ $y_{1}$ 74 $y_{2}$ 69	$z_{1}$ $y_{1}$ 77 $y_{2}$ 61	$z_{1}$ $z_{2}$ $y_{1}$ 71 56 $y_{2}$ 76 90	$z_{1}$ $z_{2}$ $y_{1}$ 80 54 $y_{2}$ 54 86	$z_{1}$ $z_{2}$ $y_{1}$ 80 50 $y_{2}$ 52 90	$z_{1}$ $z_{2}$ $y_{1}$ 80 51 $y_{2}$ 50 90	$z_{1}$ $z_{2}$ $y_{1}$ 71 61 $y_{2}$ 82 93	$z_{1}$ $z_{2}$ $y_{1}$ 73 60 $y_{2}$ 60 77
OCT	$z_{1}$ $z_{2}$ $y_{1}$ 84 61 $y_{2}$ 84 95	$z_{1}$ $z_{2}$ $y_{1}$ 90 54 $y_{2}$ 80 97	$z_{1}$ $y_{1}$ 91 $y_{2}$ 76	$z_{1}$ $y_{1}$ 92 $y_{2}$ 65	$z_{1}$ $z_{2}$ $y_{1}$ 88 52 $y_{2}$ 77 96	$z_{1}$ $z_{2}$ $y_{1}$ 94 51 $y_{2}$ 55 95	$z_{1}$ $z_{2}$ $y_{1}$ 91 50 $y_{2}$ 53 95	$z_{1}$ $z_{2}$ $y_{1}$ 94 51 $y_{2}$ 54 96	$z_{1}$ $z_{2}$ $y_{1}$ 87 76 $y_{2}$ 88 96	$z_{1}$ $z_{2}$ $y_{1}$ 85 80 $y_{2}$ 96 96

Table 3: Disentanglement performance in subspace-label confusion matrices. Performance is measured as

k

NN classification accuracy (

k=30

) across all subspace-label combinations, datasets, and methods. Darker shades indicate higher accuracy; off-diagonal values near 50% reflect optimal disentanglement with minimal confounder leakage.

Figure 4: Qualitative scatter plots showing the two-dimensional subspace $z_{1}$ of the first fold, highlighted by the binary labels $y_{2}$ . Feature disentanglement was successful if there is no structure visible with respect to $y_{2}$ .

Figure 5: Relative AUROC improvement over the Baseline on the inverted test distribution (

p(y_{1}=1\mid y_{2}=1)=5\%

) across different conditional prevalence levels used during training. Positive values indicate gains relative to the Baseline; variants combined with Rebalancing are shown with dashed lines.

	Morpho-MNIST		CheXpert		OCT
Method	Epochs	Time	Epochs	Time	Epochs	Time
Baseline	97	1	9	6	11	5
Rebalancing	33	1	2	2	6	6
AdvCl	961	10	24	19	24	11
AdvCl+Rebal	209	3	6	8	13	12
dCor	536	6	19	14	24	11

Table 4: Comparison of number of epochs and compute time in minutes until convergence.

Figure 6: Disentanglement performance (diagonal dominance) of different methods as a function of convergence time in minutes.

Explore similar work

Sep 22, 2026

Decoupling Disease, Covariates, and Individual Variability: A Unified Disentanglement Framework for Medical Image Classification

Accurately isolating disease-related features from confounding covariates (e.g., age, gender, site) and individual variations remains a fundamental challenge in medical image classification. Traditional regression-based approaches may ignore non-linear relations between image features and true covariates. To overcome this issue, we present a generalized Medical Imaging Disentanglement Learning (MedIDL) framework. MedIDL maps image features into three mutually orthogonal latent spaces through specialized disentanglement heads: a disease classification head guided by a supervised loss, a covariate-alignment head constrained by cross-subject similarity matching, and a Gaussian head absorbing individual variations. We evaluated our framework across 7 datasets encompassing diverse imaging modalities. MedIDL outperforms state-of-the-art supervised and self-supervised classification methods in accuracy across all datasets. Association analyses demonstrate that MedIDL successfully isolates target-specific latent representations. Gradient-based interpretability mappings localize pathognomonic patterns aligning with established clinical literature.

Shengjie Zhang, Jinglin Zhang, Zhuangzhuang Jiang +10

Department of Radiology & Faculty of Medical Imaging Technology, Ruijin Hospital, College of Health Science and Technology, Shanghai Jiao Tong University School of Medicine, Shanghai 200025, China · Clinical Neuroscience Center, Ruijin Hospital, Shanghai Jiao Tong University School of Medicine, Shanghai 200025, China · School of Data Science, Fudan University, Shanghai 200433, China +3

Medical Image Classification Disentangled Representation Learning Explainable Medical Image Analysis

Ask PDF

Aug 14, 2026

Discovery and Spatial Characterisation of Multiple Shortcut Groups for Auditing Vision Model Bias

Deep learning models trained on datasets with spurious correlations can achieve high average accuracy whilst relying on shortcut features that do not generalise out of distribution. Whilst out-of-distribution testing highlights subgroup performance disparities arising from shortcut learning, it does not localise the regions within images that are associated with it. Existing research mostly uses attribution maps from interpretability methods to understand the spatial nature of spurious correlations. For example, conditional alignment methods separate task-relevant evidence from evidence tied to spurious correlations by comparing attribution maps from a task model, a sensitive attribute model, and a bias-reduced reference model. This yields shortcut-aligned and task-aligned contribution maps for each image. However, existing methods aggregate these maps across the dataset, potentially masking recurring spatial shortcut patterns that occur only in subsets of images. We address this limitation by grouping per-image shortcut and task contribution maps into recurring spatial patterns using K-means and non-negative matrix factorisation, and visualising the resulting shortcut groups through contribution maps and representative examples. Across CelebA, CheXpert, Waterbirds, Camelyon17, and ISIC2019, and across ResNet and ViT models, the discovered shortcut groups reveal both shared and distinct spatial patterns of shortcut and task contribution, with varying subgroup composition and error rates, enabling targeted inspection of image subsets with higher error rates. We perform input occlusion and internal test-time interventions to show that masking or suppressing task contribution regions substantially degrades the model classification performance and propose a combined shortcut suppression and task amplification feature intervention approach which generally reduces performance disparities.

Akshit Achara, Vishnunarayan Manickam, Thomas Day +3

School of Biomedical Engineering and Imaging Sciences, King’s College London, UK

Shortcut Learning Algorithmic Bias Explainable Image Classification

Ask PDF

Sep 7, 2026

Prevalence calibration as shortcut mitigation

Shortcut learning denotes the widespread situation in which a classifier exploits spurious correlations rather than diagnostic features. Existing mitigation strategies mostly aim to learn shortcut-invariant representations; their empirical success is limited and they cannot be applied to classifiers using frozen foundation model encoders. We propose to reframe shortcut learning as fundamentally a calibration problem: unconstrained learning implicitly calibrates each shortcut group to its training set disease prevalence, rendering the resulting classifier necessarily over-confident in one group and under-confident in the other. Building on this insight, we prevalence-equalize calibration between shortcut groups through two encoder-agnostic methods, an in-processing regularizer and a post-hoc prevalence-equalized recalibration step. Across chest-drain-pneumothorax benchmarks on CheXpert and SIIM-ACR, spanning fine-tuned CNNs and frozen foundation-model backbones, both methods substantially outperform all baselines. Post-hoc recalibration of a standard ERM-trained DenseNet raises misaligned-group AUROC from 0.23 to 0.73, indicating that shortcut reliance degrades the classification head rather than the underlying representation. Besides two new state-of-the-art shortcut mitigation approaches, our findings more fundamentally connect shortcut learning to calibration theory and algorithmic fairness.

Mohamed Amine Kina, Eike Petersen

Universität Bremen, Bremen, Germany · Fraunhofer Institute for Digital Medicine MEVIS, Bremen, Germany

Post-Hoc Calibration Medical Image Classification Group Fairness Shortcut Learning

Ask PDF