# Changing the Training Data Distribution to Reduce Simplicity Bias Improves In-distribution Generalization

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

## Key facts

- Cited by 3 papers in the map (data through 14 Oct 2026).
- Builds on 67 of its 82 references that are in the map.
- First posted on arXiv 27 Apr 2024.

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

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Neural Architecture Search with Reinforcement Learning](https://mlsanity.com/paper/1611.01578/) (2016)
- [Sharpness-Aware Minimization for Efficiently Improving Generalization](https://mlsanity.com/paper/2010.01412/) (2020)
- [On Large-Batch Training for Deep Learning: Generalization Gap and Sharp Minima](https://mlsanity.com/paper/1609.04836/) (2016)
- [Distributionally Robust Neural Networks for Group Shifts: On the Importance of Regularization for Worst-Case Generalization](https://mlsanity.com/paper/1911.08731/) (2019)
- [Class-Balanced Loss Based on Effective Number of Samples](https://mlsanity.com/paper/1901.05555/) (2019)
- [Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet](https://mlsanity.com/paper/2101.11986/) (2021)
- [Reconciling modern machine learning practice and the bias-variance trade-off](https://mlsanity.com/paper/1812.11118/) (2018)
- [Efficient Neural Architecture Search via Parameter Sharing](https://mlsanity.com/paper/1802.03268/) (2018)
- [Deduplicating Training Data Makes Language Models Better](https://mlsanity.com/paper/2107.06499/) (2021)
- [DivideMix: Learning with Noisy Labels as Semi-supervised Learning](https://mlsanity.com/paper/2002.07394/) (2020)
- [An Empirical Study of Example Forgetting during Deep Neural Network Learning](https://mlsanity.com/paper/1812.05159/) (2018)
- [Deep Double Descent: Where Bigger Models and More Data Hurt](https://mlsanity.com/paper/1912.02292/) (2019)
- [DataComp: In search of the next generation of multimodal datasets](https://mlsanity.com/paper/2304.14108/) (2023)
- [Just Train Twice: Improving Group Robustness without Training Group Information](https://mlsanity.com/paper/2107.09044/) (2021)
- [Progressive Neural Architecture Search](https://mlsanity.com/paper/1712.00559/) (2017)
- [Sharp Minima Can Generalize For Deep Nets](https://mlsanity.com/paper/1703.04933/) (2017)
- [SWAD: Domain Generalization by Seeking Flat Minima](https://mlsanity.com/paper/2102.08604/) (2021)
- [The generalization error of random features regression: Precise asymptotics and double descent curve](https://mlsanity.com/paper/1908.05355/) (2019)
- [Last Layer Re-Training is Sufficient for Robustness to Spurious Correlations](https://mlsanity.com/paper/2204.02937/) (2022)
- [Entropy-SGD: Biasing Gradient Descent Into Wide Valleys](https://mlsanity.com/paper/1611.01838/) (2016)
- [The Pitfalls of Simplicity Bias in Neural Networks](https://mlsanity.com/paper/2006.07710/) (2020)
- [In Search of the Real Inductive Bias: On the Role of Implicit Regularization in Deep Learning](https://mlsanity.com/paper/1412.6614/) (2014)
- [Not All Samples Are Created Equal: Deep Learning with Importance Sampling](https://mlsanity.com/paper/1803.00942/) (2018)
- [Implicit Regularization in Matrix Factorization](https://mlsanity.com/paper/1705.09280/) (2017)
- [Environment Inference for Invariant Learning](https://mlsanity.com/paper/2010.07249/) (2020)
- [Certified Defenses for Data Poisoning Attacks](https://mlsanity.com/paper/1706.03691/) (2017)

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

- [Do We Need All the Synthetic Data? Targeted Image Augmentation via Diffusion Models](https://mlsanity.com/paper/2505.21574/) (2025)
- [A Critical Review of Predominant Bias in Neural Networks](https://mlsanity.com/paper/2502.11031/) (2025)
- [Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMs](https://mlsanity.com/paper/2602.00576/) (2026)

## Abstract (arXiv)

Can we modify the training data distribution to encourage the underlying optimization method toward finding solutions with superior generalization performance on in-distribution data? In this work, we approach this question for the first time by comparing the inductive bias of gradient descent (GD) with that of sharpness-aware minimization (SAM). By studying a two-layer CNN, we rigorously prove that SAM learns different features more uniformly, particularly in early epochs. That is, SAM is less susceptible to simplicity bias compared to GD. We also show that examples containing features that are learned early are separable from the rest based on the model's output. Based on this observation, we propose a method that (i) clusters examples based on the network output early in training, (ii) identifies a cluster of examples with similar network output, and (iii) upsamples the rest of examples only once to alleviate the simplicity bias. We show empirically that USEFUL effectively improves the generalization performance on the original data distribution when training with various gradient methods, including (S)GD and SAM. Notably, we demonstrate that our method can be combined with SAM variants and existing data augmentation strategies to achieve, to the best of our knowledge, state-of-the-art performance for training ResNet18 on CIFAR10, STL10, CINIC10, Tiny-ImageNet; ResNet34 on CIFAR100; and VGG19 and DenseNet121 on CIFAR10. Our code is available at https://github.com/BigML-CS-UCLA/TADA.