Optimizing Data Augmentation through Bayesian Model Selection

TL;DR

This paper introduces a Bayesian framework for optimizing data augmentation parameters, improving model robustness and calibration by jointly learning augmentation and model parameters through a variational approach.

Contribution

It presents a novel probabilistic method that treats augmentation parameters as hyperparameters and optimizes them via a tractable ELBO, grounded in Bayesian model selection.

Findings

Enhanced model calibration and robustness in vision and NLP tasks

Outperforms fixed or no augmentation strategies in experiments

Provides theoretical guarantees on approximation quality and generalization

Abstract

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.