Optimizing Data Augmentation through Bayesian Model Selection

2025-05-27

AI Summary

Data augmentation (DA) parameter selection typically relies on costly validation-based optimization or manual trial-and-error. Method: This paper formulates DA policy selection as a Bayesian model selection problem, treating augmentation parameters as learnable hyperparameters optimized via marginal likelihood maximization. We derive a differentiable, tractable evidence lower bound (ELBO) for end-to-end joint optimization of DA policies and model parameters. Contribution/Results: We theoretically establish the variational approximation quality, generalization bound, and invariance of the ELBO under data transformations, and prove its equivalence to empirical Bayes estimation. Evaluated on computer vision benchmarks, our approach significantly improves model calibration and out-of-distribution robustness over fixed and no-augmentation baselines—introducing the first differentiable, Bayesian inference–based optimization paradigm for data augmentation.

Technology Category

• Machine Learning: Bayesian Learning
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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 Evidence Lower BOund (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 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.

Problem

• Optimizing Data Augmentation parameters via Bayesian model selection
• Replacing trial-and-error DA strategy selection with probabilistic framework
• Enhancing ML robustness and calibration through variational DA optimization

Innovation

• Bayesian model selection optimizes data augmentation
• Tractable ELBO jointly optimizes augmentation parameters
• Probabilistic view enhances generalization and robustness