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

arXiv: https://arxiv.org/abs/2602.18502 · HTML: https://mlsanity.com/paper/2602.18502/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 51 of its 82 references that are in the map.
- First posted on arXiv 17 Feb 2026.

## Builds on (most cited first, 30 of 51)

- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [Shortcut Learning in Deep Neural Networks](https://mlsanity.com/paper/2004.07780/) (2020)
- [Unsupervised Domain Adaptation by Backpropagation](https://mlsanity.com/paper/1409.7495/) (2014)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [CheXpert: A Large Chest Radiograph Dataset with Uncertainty Labels and Expert Comparison](https://mlsanity.com/paper/1901.07031/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Concept Bottleneck Models](https://mlsanity.com/paper/2007.04612/) (2020)
- [MINE: Mutual Information Neural Estimation](https://mlsanity.com/paper/1801.04062/) (2018)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [Right for the Right Reasons: Training Differentiable Models by Constraining their Explanations](https://mlsanity.com/paper/1703.03717/) (2017)
- [The Variational Fair Autoencoder](https://mlsanity.com/paper/1511.00830/) (2015)
- [Improving Out-of-Distribution Robustness via Selective Augmentation](https://mlsanity.com/paper/2201.00299/) (2022)
- [Causality matters in medical imaging](https://mlsanity.com/paper/1912.08142/) (2019)
- [Simple data balancing achieves competitive worst-group-accuracy](https://mlsanity.com/paper/2110.14503/) (2021)
- [REPAIR: Removing Representation Bias by Dataset Resampling](https://mlsanity.com/paper/1904.07911/) (2019)
- [On Feature Learning in the Presence of Spurious Correlations](https://mlsanity.com/paper/2210.11369/) (2022)
- [Making deep neural networks right for the right scientific reasons by interacting with their explanations](https://mlsanity.com/paper/2001.05371/) (2020)
- [Domain Generalization by Mutual-Information Regularization with Pre-trained Models](https://mlsanity.com/paper/2203.10789/) (2022)
- [Unsupervised domain adaptation in brain lesion segmentation with adversarial networks](https://mlsanity.com/paper/1612.08894/) (2016)
- [Learning Debiased Representation via Disentangled Feature Augmentation](https://mlsanity.com/paper/2107.01372/) (2021)
- [Discover and Cure: Concept-aware Mitigation of Spurious Correlation](https://mlsanity.com/paper/2305.00650/) (2023)
- [MaskTune: Mitigating Spurious Correlations by Forcing to Explore](https://mlsanity.com/paper/2210.00055/) (2022)
- [Examining and Combating Spurious Features under Distribution Shift](https://mlsanity.com/paper/2106.07171/) (2021)
- [FairPrune: Achieving Fairness Through Pruning for Dermatological Disease Diagnosis](https://mlsanity.com/paper/2203.02110/) (2022)
- [Morpho-MNIST: Quantitative Assessment and Diagnostics for Representation Learning](https://mlsanity.com/paper/1809.10780/) (2018)

## Abstract (arXiv)

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.