# Chroma-VAE: Mitigating Shortcut Learning with Generative Classifiers

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

## Key facts

- Cited by 17 papers in the map (data through 14 Oct 2026).
- Builds on 43 of its 60 references that are in the map.
- First posted on arXiv 28 Nov 2022.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization](https://mlsanity.com/paper/1610.02391/) (2016)
- [Intriguing properties of neural networks](https://mlsanity.com/paper/1312.6199/) (2013)
- [Auto-Encoding Variational Bayes](https://mlsanity.com/paper/1312.6114/) (2013)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (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)
- [ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness](https://mlsanity.com/paper/1811.12231/) (2018)
- [CheXpert: A Large Chest Radiograph Dataset with Uncertainty Labels and Expert Comparison](https://mlsanity.com/paper/1901.07031/) (2019)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Generalized Cross Entropy Loss for Training Deep Neural Networks with Noisy Labels](https://mlsanity.com/paper/1805.07836/) (2018)
- [ChestX-ray8: Hospital-scale Chest X-ray Database and Benchmarks on Weakly-Supervised Classification and Localization of Common Thorax Diseases](https://mlsanity.com/paper/1705.02315/) (2017)
- [Right for the Wrong Reasons: Diagnosing Syntactic Heuristics in Natural Language Inference](https://mlsanity.com/paper/1902.01007/) (2019)
- [Deep Learning and the Information Bottleneck Principle](https://mlsanity.com/paper/1503.02406/) (2015)
- [Evasion Attacks against Machine Learning at Test Time](https://mlsanity.com/paper/1708.06131/) (2017)
- [Annotation Artifacts in Natural Language Inference Data](https://mlsanity.com/paper/1803.02324/) (2018)
- [Recognition in Terra Incognita](https://mlsanity.com/paper/1807.04975/) (2018)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Opening the Black Box of Deep Neural Networks via Information](https://mlsanity.com/paper/1703.00810/) (2017)
- [Semi-Supervised Learning with Deep Generative Models](https://mlsanity.com/paper/1406.5298/) (2014)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Confounding variables can degrade generalization performance of radiological deep learning models](https://mlsanity.com/paper/1807.00431/) (2018)
- [Fairness Without Demographics in Repeated Loss Minimization](https://mlsanity.com/paper/1806.08010/) (2018)
- [Noise or Signal: The Role of Image Backgrounds in Object Recognition](https://mlsanity.com/paper/2006.09994/) (2020)
- [An Investigation of Why Overparameterization Exacerbates Spurious Correlations](https://mlsanity.com/paper/2005.04345/) (2020)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Hidden Stratification Causes Clinically Meaningful Failures in Machine Learning for Medical Imaging](https://mlsanity.com/paper/1909.12475/) (2019)
- [No Subclass Left Behind: Fine-Grained Robustness in Coarse-Grained Classification Problems](https://mlsanity.com/paper/2011.12945/) (2020)
- [Do CIFAR-10 Classifiers Generalize to CIFAR-10?](https://mlsanity.com/paper/1806.00451/) (2018)

## Built on by (most cited first, 17 of 17)

- [Simple and Fast Group Robustness by Automatic Feature Reweighting](https://mlsanity.com/paper/2306.11074/) (2023)
- [The Clever Hans Mirage: A Comprehensive Survey on Spurious Correlations in Machine Learning](https://mlsanity.com/paper/2402.12715/) (2024)
- [Learning Data Representations with Joint Diffusion Models](https://mlsanity.com/paper/2301.13622/) (2023)
- [Decompose-and-Compose: A Compositional Approach to Mitigating Spurious Correlation](https://mlsanity.com/paper/2402.18919/) (2024)
- [Mitigating Simplicity Bias in Deep Learning for Improved OOD Generalization and Robustness](https://mlsanity.com/paper/2310.06161/) (2023)
- [Shortcut Detection with Variational Autoencoders](https://mlsanity.com/paper/2302.04246/) (2023)
- [Mitigating Reward Hacking in RLHF via Bayesian Non-negative Reward Modeling](https://mlsanity.com/paper/2602.10623/) (2026)
- [Debiasing Multimodal Models via Causal Information Minimization](https://mlsanity.com/paper/2311.16941/) (2023)
- [Task Confusion and Catastrophic Forgetting in Class-Incremental Learning: A Mathematical Framework for Discriminative and Generative Modelings](https://mlsanity.com/paper/2410.20768/) (2024)
- [Gaze-Guided Learning: Avoiding Shortcut Bias in Visual Classification](https://mlsanity.com/paper/2504.05583/) (2025)
- [Bridging Distribution Shift and AI Safety: Conceptual and Methodological Synergies](https://mlsanity.com/paper/2505.22829/) (2025)
- [Superclass-Guided Representation Disentanglement for Spurious Correlation Mitigation](https://mlsanity.com/paper/2508.08570/) (2025)
- [Exploiting Supervised Poison Vulnerability to Strengthen Self-Supervised Defense](https://mlsanity.com/paper/2409.08509/) (2024)
- [Improving Adversarial Robustness in Android Malware Detection by Reducing the Impact of Spurious Correlations](https://mlsanity.com/paper/2408.16025/) (2024)
- [Shortcut Invariance: Targeted Jacobian Regularization in Disentangled Latent Space](https://mlsanity.com/paper/2511.19525/) (2025)
- [Mitigating Shortcut Learning via Feature Disentanglement in Medical Imaging: A Benchmark Study](https://mlsanity.com/paper/2602.18502/) (2026)
- [Deciphering Shortcut Learning from an Evolutionary Game Theory Perspective](https://mlsanity.com/paper/2605.02658/) (2026)

## Abstract (arXiv)

Deep neural networks are susceptible to shortcut learning, using simple features to achieve low training loss without discovering essential semantic structure. Contrary to prior belief, we show that generative models alone are not sufficient to prevent shortcut learning, despite an incentive to recover a more comprehensive representation of the data than discriminative approaches. However, we observe that shortcuts are preferentially encoded with minimal information, a fact that generative models can exploit to mitigate shortcut learning. In particular, we propose Chroma-VAE, a two-pronged approach where a VAE classifier is initially trained to isolate the shortcut in a small latent subspace, allowing a secondary classifier to be trained on the complementary, shortcut-free latent subspace. In addition to demonstrating the efficacy of Chroma-VAE on benchmark and real-world shortcut learning tasks, our work highlights the potential for manipulating the latent space of generative classifiers to isolate or interpret specific correlations.