# Optimizing Data Augmentation through Bayesian Model Selection

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 67 references that are in the map.
- First posted on arXiv 27 May 2025.

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

- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Dropout as a Bayesian Approximation: Representing Model Uncertainty in Deep Learning](https://mlsanity.com/paper/1506.02142/) (2015)
- [DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter](https://mlsanity.com/paper/1910.01108/) (2019)
- [CutMix: Regularization Strategy to Train Strong Classifiers with Localizable Features](https://mlsanity.com/paper/1905.04899/) (2019)
- [Benchmarking Neural Network Robustness to Common Corruptions and Perturbations](https://mlsanity.com/paper/1903.12261/) (2019)
- [Categorical Reparameterization with Gumbel-Softmax](https://mlsanity.com/paper/1611.01144/) (2016)
- [RandAugment: Practical automated data augmentation with a reduced search space](https://mlsanity.com/paper/1909.13719/) (2019)
- [Understanding deep learning requires rethinking generalization](https://mlsanity.com/paper/1611.03530/) (2016)
- [AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty](https://mlsanity.com/paper/1912.02781/) (2019)
- [Computing Nonvacuous Generalization Bounds for Deep (Stochastic) Neural Networks with Many More Parameters than Training Data](https://mlsanity.com/paper/1703.11008/) (2017)
- [Bayesian Deep Learning and a Probabilistic Perspective of Generalization](https://mlsanity.com/paper/2002.08791/) (2020)
- [A Survey of Data Augmentation Approaches for NLP](https://mlsanity.com/paper/2105.03075/) (2021)
- [Laplace Redux -- Effortless Bayesian Deep Learning](https://mlsanity.com/paper/2106.14806/) (2021)
- [What Are Bayesian Neural Network Posteriors Really Like?](https://mlsanity.com/paper/2104.14421/) (2021)
- [Fast AutoAugment](https://mlsanity.com/paper/1905.00397/) (2019)
- [User-friendly introduction to PAC-Bayes bounds](https://mlsanity.com/paper/2110.11216/) (2021)
- [Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks](https://mlsanity.com/paper/2002.10118/) (2020)
- [Accounting for Variance in Machine Learning Benchmarks](https://mlsanity.com/paper/2103.03098/) (2021)
- [Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules](https://mlsanity.com/paper/1905.05393/) (2019)
- [Revisiting Neural Scaling Laws in Language and Vision](https://mlsanity.com/paper/2209.06640/) (2022)
- [A Kernel Theory of Modern Data Augmentation](https://mlsanity.com/paper/1803.06084/) (2018)
- [Scalable Marginal Likelihood Estimation for Model Selection in Deep Learning](https://mlsanity.com/paper/2104.04975/) (2021)
- [Faster AutoAugment: Learning Augmentation Strategies using Backpropagation](https://mlsanity.com/paper/1911.06987/) (2019)
- [On the Benefits of Invariance in Neural Networks](https://mlsanity.com/paper/2005.00178/) (2020)
- [PAC-Bayes Compression Bounds So Tight That They Can Explain Generalization](https://mlsanity.com/paper/2211.13609/) (2022)
- [Variational Bayesian Last Layers](https://mlsanity.com/paper/2404.11599/) (2024)
- [Do Bayesian Neural Networks Need To Be Fully Stochastic?](https://mlsanity.com/paper/2211.06291/) (2022)
- [All You Need is a Good Functional Prior for Bayesian Deep Learning](https://mlsanity.com/paper/2011.12829/) (2020)
- [Deep Learning is Not So Mysterious or Different](https://mlsanity.com/paper/2503.02113/) (2025)

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

- [Universal Adaptive Environment Discovery](https://mlsanity.com/paper/2510.12547/) (2025)

## Abstract (arXiv)

Data Augmentation (DA) has become an essential tool to improve robustness and generalization of modern machine learning. However, when deciding on DA strategies it is critical to choose parameters carefully, and this can be a daunting task which is traditionally left to trial-and-error or expensive optimization based on validation performance. In this paper, we counter these limitations by proposing a novel framework for optimizing DA. In particular, we take a probabilistic view of DA, which leads to the interpretation of augmentation parameters as model (hyper)-parameters, and the optimization of the marginal likelihood with respect to these parameters as a Bayesian model selection problem. Due to its intractability, we derive a tractable ELBO, which allows us to optimize augmentation parameters jointly with model parameters. We provide extensive theoretical results on variational approximation quality, generalization guarantees, invariance properties, and connections to empirical Bayes. Through experiments on computer vision and NLP tasks, we show that our approach improves calibration and yields robust performance over fixed or no augmentation. Our work provides a rigorous foundation for optimizing DA through Bayesian principles with significant potential for robust machine learning.