# Learning an Invertible Output Mapping Can Mitigate Simplicity Bias in Neural Networks

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 35 references that are in the map.
- First posted on arXiv 4 Oct 2022.

## Builds on (most cited first, 27 of 27)

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [Domain-Adversarial Training of Neural Networks](https://mlsanity.com/paper/1505.07818/) (2015)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Invariant Risk Minimization](https://mlsanity.com/paper/1907.02893/) (2019)
- [WILDS: A Benchmark of in-the-Wild Distribution Shifts](https://mlsanity.com/paper/2012.07421/) (2020)
- [In Search of Lost Domain Generalization](https://mlsanity.com/paper/2007.01434/) (2020)
- [A Closer Look at Memorization in Deep Networks](https://mlsanity.com/paper/1706.05394/) (2017)
- [Learning to Generalize: Meta-Learning for Domain Generalization](https://mlsanity.com/paper/1710.03463/) (2017)
- [Fine-Tuning can Distort Pretrained Features and Underperform Out-of-Distribution](https://mlsanity.com/paper/2202.10054/) (2022)
- [The Implicit Bias of Gradient Descent on Separable Data](https://mlsanity.com/paper/1710.10345/) (2017)
- [Domain Generalization via Invariant Feature Representation](https://mlsanity.com/paper/1301.2115/) (2013)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [Self-Challenging Improves Cross-Domain Generalization](https://mlsanity.com/paper/2007.02454/) (2020)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [The Pitfalls of Simplicity Bias in Neural Networks](https://mlsanity.com/paper/2006.07710/) (2020)
- [Generalizing Across Domains via Cross-Gradient Training](https://mlsanity.com/paper/1804.10745/) (2018)
- [Gradient Matching for Domain Generalization](https://mlsanity.com/paper/2104.09937/) (2021)
- [Deep learning generalizes because the parameter-function map is biased towards simple functions](https://mlsanity.com/paper/1805.08522/) (2018)
- [Domain Generalization Using a Mixture of Multiple Latent Domains](https://mlsanity.com/paper/1911.07661/) (2019)
- [Ensemble of Averages: Improving Model Selection and Boosting Performance in Domain Generalization](https://mlsanity.com/paper/2110.10832/) (2021)
- [Domain Generalization by Mutual-Information Regularization with Pre-trained Models](https://mlsanity.com/paper/2203.10789/) (2022)
- [Evading the Simplicity Bias: Training a Diverse Set of Models Discovers Solutions with Superior OOD Generalization](https://mlsanity.com/paper/2105.05612/) (2021)
- [Domain-Adjusted Regression or: ERM May Already Learn Features Sufficient for Out-of-Distribution Generalization](https://mlsanity.com/paper/2202.06856/) (2022)
- [Rich Feature Construction for the Optimization-Generalization Dilemma](https://mlsanity.com/paper/2203.15516/) (2022)
- [Domain Generalization using Pretrained Models without Fine-tuning](https://mlsanity.com/paper/2203.04600/) (2022)
- [Controlling Directions Orthogonal to a Classifier](https://mlsanity.com/paper/2201.11259/) (2022)
- [Adaptive Methods for Aggregated Domain Generalization](https://mlsanity.com/paper/2112.04766/) (2021)

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

- [Stochastic Re-weighted Gradient Descent via Distributionally Robust Optimization](https://mlsanity.com/paper/2306.09222/) (2023)
- [Revealing Model Biases: Assessing Deep Neural Networks via Recovered Sample Analysis](https://mlsanity.com/paper/2306.06414/) (2023)
- [Responsible AI (RAI) Games and Ensembles](https://mlsanity.com/paper/2310.18832/) (2023)

## Abstract (arXiv)

Deep Neural Networks are known to be brittle to even minor distribution shifts compared to the training distribution. While one line of work has demonstrated that Simplicity Bias (SB) of DNNs - bias towards learning only the simplest features - is a key reason for this brittleness, another recent line of work has surprisingly found that diverse/ complex features are indeed learned by the backbone, and their brittleness is due to the linear classification head relying primarily on the simplest features. To bridge the gap between these two lines of work, we first hypothesize and verify that while SB may not altogether preclude learning complex features, it amplifies simpler features over complex ones. Namely, simple features are replicated several times in the learned representations while complex features might not be replicated. This phenomenon, we term Feature Replication Hypothesis, coupled with the Implicit Bias of SGD to converge to maximum margin solutions in the feature space, leads the models to rely mostly on the simple features for classification. To mitigate this bias, we propose Feature Reconstruction Regularizer (FRR) to ensure that the learned features can be reconstructed back from the logits. The use of em FRR in linear layer training (FRR-L) encourages the use of more diverse features for classification. We further propose to finetune the full network by freezing the weights of the linear layer trained using FRR-L, to refine the learned features, making them more suitable for classification. Using this simple solution, we demonstrate up to 15% gains in OOD accuracy on the recently introduced semi-synthetic datasets with extreme distribution shifts. Moreover, we demonstrate noteworthy gains over existing SOTA methods on the standard OOD benchmark DomainBed as well.