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

Decoupling Disease, Covariates, and Individual Variability: A Unified Disentanglement Framework for Medical Image Classification
Sep 22, 2026Shengjie Zhang, Jinglin Zhang, Zhuangzhuang Jiang +10Medical Image Classification Disentangled Representation Learning
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
Discovery and Spatial Characterisation of Multiple Shortcut Groups for Auditing Vision Model Bias
Aug 14, 2026Akshit Achara, Vishnunarayan Manickam, Thomas Day +3Shortcut Learning Algorithmic Bias
School of Biomedical Engineering and Imaging Sciences, King’s College London, UK
Prevalence calibration as shortcut mitigation
Sep 7, 2026Mohamed Amine Kina, Eike PetersenPost-Hoc Calibration Medical Image Classification
Universität Bremen, Bremen, Germany · Fraunhofer Institute for Digital Medicine MEVIS, Bremen, Germany

Sep 22, 2026J/K move · Enter open · S save

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

Shengjie Zhang, Jinglin Zhang, Zhuangzhuang Jiang, Ziqi Yu+9

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

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

Prevalence calibration as shortcut mitigation