Optimizing Data Augmentation through Bayesian Model Selection
Authors: Madi Matymov
Ba-Hien Tran
Michael Kampffmeyer
Markus Heinonen
Maurizio Filippone
Venue: ICLR 2026 Poster (2026)
Keywords: Bayesian Neural Network
Variational Inference
Data Augmentation
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We present OPTIMA, a Bayesian framework that learns augmentation and model parameters jointly, with theory and experiments showing improved accuracy, calibration, and OOD robustness over strong baselines.

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.

Published as a conference paper at ICLR 2026
OPTIMIZING DATA AUGMENTATION THROUGH BAYESIAN MODEL SELECTION
Madi Matymov*
Ba-Hien Tran
Michael Kampffmeyer
KAUST
Huawei Paris Research Center France
UiT The Arctic University of Norway Norwegian Computing Center, Norway
Saudi Arabia
Markus Heinonen
Maurizio Filippone
Aalto University
KAUST Saudi Arabia
Finland
ABSTRACT
Data Augmentation (DA) has become an essential tool improve robustness and generalization of modern machine learning. However, when deciding 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.
this paper, we counter these limitations by proposing novel framework for optimizing DA. In particular, we take a probabilistic view DA, which leads to the interpretation of augmentation parameters as model -lrpaet-paammeters, and the optimization of the marginal likelihood with respect to these parameters as a Bayesian model selection problem. Due to its intractability, derive tractable Evidence Lower BOund (ELBO), which allows 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.
optimize
INTRODUCTION
Data Augmentation (DA) (Van Dyk & Meng, 2001) is an essential element behind the success of modern machine learning (see, e.g., Shorten & Khoshgoftaar, 2019, and references therein). In supervised learning, DA amounts to creating copies of the data in the training set, and perturbing these with sensible transformations that preserve label information. The success of DA is connected with the current trend of employing over preapaaraeteri models based on neural networks, which require large amounts of data to be trained effectively (Alabdulmohsin et al., 2022). It has been shown that DA has strong connections with regularization (Zhang et al., 2017; Dao et al., 2019), and it can provide a better estimation of the expected risk (Shao et al., 2022; Chen et al., 2020; Lyle et al., 2020; Deng al., 2022). Therefore,
expected for DA to enhance generalization.
For given problem, once transformation for DA are chosen, it is then necessary to decide on their parameters. For example, image classification, if choose to apply transformations in the form of rotations, what range of angles should we choose? Careful choices of DA parameters are important to obtain performance improvements. For example, in the case of rotations applied to the popular MNIST dataset, large rotation angles can turn a '9' into a '6', negatively impacting training. In the literature, DA parameters are often suggested after some trial-and-error. Direct optimization DA parameters could also be approached via grid-search or Bayesian optimization by recording performance validation set, but this is very costly due to the need to perform a large number of training runs.
In this paper, we propose novel approach to optimize DA which counters these limitations. In particular, we take a probabilistic view of DA, whereby we treat DA parameters as model (hyper-)
"oorrespnniigg author; madi mat.mmvvkkausstt sdu.sa

---

Published as a conference paper at ICLR 2026
1,00
ECE=0.092
ECE=0,088
ECE=0,017
Accuracy
0.75
Perfect Calibration
No Augmennltionn
Augmentation
0.25
Learoed Augmeniation.
0.0
0.5 Confidence
1.0
0.0
0.5
1.0
0.0
0.5
1.0
Confidence
Confidence
Figure 1: OPTIMA obtains the best calibration. Example of ResNet-18 CIFARIO; see details in Appendix F.I.
parameters. We then consider the optimization of such parameters as a Bayesian model selection problem. Due to the intractability of the Bayesian model selection objective (i.e., the marginal likelihood), we derive a tractable ELBO, which allows us to optimize DA parameters jointly with model parameters, bypassing the need to perform expensive cross-validation or grid search. We provide an extensive theoretical analysis, which indicates robust predictive performance and low Expected Calibration Error (ECE) demonstrated by the experiments (see, c.g., Fig. 1). Our main contributions are as follows:
Methodology: We introduce OPTIMA (OPTImizing Marginalized Augmentations), novel framework to learn DA parameters grounded on Bayesian principles. We then provide tractable variational approximation which allows for the optimization of both model parameters and DA param- eters, yielding practical and fast alternative to manual tuning or expensive black-box optimization of DA parameters.
Theory: We provide comprehensive theoretical analysis, establishing cohesive framework to understand our Bayesian approach to DA, highlighting its principled nature. Our analysis includes: (i) The analysis of the variational approximation's 's quality, guiding DA distribution design ($ 4.1). (ii) derivation of PAC-Bayes generalization guarantees ($ 4.2) and demonstration on how OPTIMA promotes model invariance and smoother decision boundaries ($ 4.3). (iii) A demonstration of im- proved uncertainty quantification and calibration through proper marginalization over DA parameters (8 4.4). (iv) The establishing of empirical Bayes optimality 4.5) for data-driven strategies, complemented by information-theore insights ($ 4.6) into how learned DA enhances inference.
Empirical Validation: We support OPTIMA and the theoretical developments with rigorous empirical validation on various tasks ($ 5), including regression, image classification standard benchmarks (c.g., CIFARIO and IMAGENET). and an additional natural language classification task (SST-5), Across all these settings-ssanningg both continuous geometric transformations and discrete text perturbations our experiments consistently demonstrate that OPTIMA improves generalization, model calibration, and robustness to ditto--diittrbution data compared models trained with fixed or no augmentation strategies.
Overall, work demonstrates how Bayesian principles, specifically through (partial full) variational treatment of both model and augmentation parameters, can be effectively leveraged to develop practical, scalable, and principled framework for optimizing DA, moving beyond expensive trial-and-erod or validation-baas procedures for optimal DA.
2 BACKGROUND AND RELATED WORK
We consider supervised learning tasks, where mappings from inputs x E R. labels y E ROO are learned from N training observations D ((x,) A common approach is to find loss minimizing point estimates, which is equivalent to maximizing a log likelihood log loy ( 10 X), where X and Y denote all inputs and labels, respectively.
Marginal likelihoud and ELBO. In the Bayesian approach we choose prior p(0), and infer the posterior distributions parameters and predictive distribution for new data point x* as:
2

---

Published as a conference paper at ICLR 2026
p(0 D) P 0,
x))
(1)
ply"1x",, - 1 1x"00)p(0 D)d0.
(2)
p (
The denominator of Eq. is the marginal likelihood, representing the data likelihood under the prior:
PY pl p ()p(0 6)d0.
where we made explicit the dependence on continuous hyper-paramettrs &b We can perform model selection by choosing the one with highest log-marginal likelihood, also known as model evidence. The intractability of this objective motivates to employ variational inference to obtain a tractable lower bound to be optimized with respect to parametric surrogate posterior (()).
log P x, Φ) 2 Beco) log P 0, O, Φ)] - [L( || p(O Φ)]
ELBO
Data augmentation in neural Networks. In DA, apply transformations (x) parameterized by y to the inputs at training time. In image classification common transformations include rotations, translations, flips, and color manipulations, while in natural language popular augmentations involve word substitutions and syntactic transformations (Shorten & Khoshgoftaar, 2019; Feng et al., 2021). During training, for each sample in mini-batch, first sample transformation parameter y and then apply T- (x). This approach has proven highly effective in improving generalization in deep learning (Shorten & Khoshgoftuar, 2019).
Augmentations overcount evidence. Naively replicating augmented examples (x,)) y:)) as fully independent effectively multiplies the evidence (3), overcounting the likelihood (Wilson Izmailov, 2020). For single data point y,), this yields a likelihood II ply, TT, (x:), 0), equivalent raising p(y, x,, 0) to the power K. This overcounting artificially shrink posterior uncertainty and degrade calibration, undermining key advantage Bayesian methods.
2.1 RELATED WORKS
Data augmentation optimization. Optimizing data augmentation (DA) parameters has been ap- proached via computationally expensive reinforcement learning (AutoAugment (Cubuk et al., 2019)), made more efficient by population-base based training (Ho et al., 2019). Others formulate it as density matching, often using black-box search like Bayesian optimization (Snock et al., 2012), or differen- tiable policy search (Lim et al., 2019; Cubuk al., 2020; Hataya al., 2020), or as gradient matching (Zheng et al., 2022). Bi-level optimization has also been used, but remains expensive and often relies on strong relaxations (Liu et al., 2021; Li et al., 2020; Hataya et al., 2022: Mounsaveng et al., 2021). These methods typically rely on heuristics and complex search pipelines. More recently, DA has been framed as wrariance-constraan
learning with regularized objectives (Benton et al., 2020) or via non-parametric models solved with costly Markov chain Monte Carlo (MCMC) (Hounie et al., 2023).
Probabilistic perspectives of DA. Probabilistic views of DA have shown perturbed inputs can induce degenerate (Izmailov al., 2021) or tempered likelihoods (Kapoor et al., 2022), informing studies on DA 's role the cold-posterios effect (Wenzel et al., 2020; Bachmann et al., 2022). Nabarro al. (2022) proposed an integral likelihood similar to ours using a Jensen lower bound, and Heinonen et al. (2025) defined an augmented likelihood via label smoothing and input mollification (Tran et al., 2023). Kapoor al. (2022) analyzed augmentations through a Dirichlet likelihood. Related latent- variable formulations also appear in work such as Chen et al. (2020) and Chatzipantazis et al. (2023), which consider probabilistic transformations but do not optimize augmentation parameters within joint Bayesian model. However, these approaches generally use fixed, unoptimized augmentation parameters. In contrast, Wang al. (2023) modeled DA with stochastic output layers and auxiliary variables for MAP optimization via expectation maximization, while Wu & Williamson (2024) applied MixUp (Zhang et al., 2018) for martingale posteriors (Fong et al., 2023). Broader connections link DA to kernel methods for task-specific invariances (Dao et al., 2019), though not directly to Bayesian inference. More directly, van der Wilk et al. (2018) learned invariances via marginal likelihood for Gaussian processes (Williams & Rasmussen, 2006), idea Immer et al. (2022) extended to BNNs (Neal, 1996; Tran et al., 2022) using the Laplace approximation (MacKay, 1992; Daxberger et al., 2021; Immer al., 2021), but without theoretical generalization guarantees.
3

---

Published as a conference paper at ICLR 2026
PAC-Bayes generalization bounds. PAC-Bayes bounds (McAllester, 1999; Catoni, 2007; Alquier, 2024) offer theoretical guarantees for Bayesian methods, including in deep learning (Dziugaite & Roy, 2017; Lotfi et al., 2022; Wilson, 2025). However, prior work rarely treats augmentation parameters as latent variables within this framework. We unify these directions by making the augmentation distribution key component of the model's likelihood, deriving novel theoretical results that characterize its benefits.
AUGMENTATION OPTIMIZATION THROUGH BAYESIAN MODEL SELECTION
Augmentation as Marginalization. In this section we treat the optimization of DA parameters as Bayesian model selection. To do so, we start by defining a transformation-augmennte likelihood:
ply 0, () Eptr ab) (T(y|I,(x)x0)
(5)
where (x) is the transformed input under augmentation distribution p(y ()) parameterized by Φ. This formulation treats augmentation as marginalization over transformations rather than data replication. This method averages over transformations to contribute each original example exactly once, as opposed to the overcounting effect in the case of naïve augmentation. As we will see, this yields more calibrated posterior with appropriate uncertainty quantification.
The data likelihood given model parameters and augmentation parameters Φ
N p(D の) = II E ply, (x:)..
(6)
Taking a fully Bayesian treatment, we assign a prior p()) on the augmentation parameters Φ making &b a latent variable alongside O. The joint distribution over all variables is p(D, 0, Φ r)) p(0) p(0)p(())(()) Φ)p(D 0, Φ). Our goal is to approximate the posterior p(0, Φ 'D), which typically intractable. To address this challenge, employ variational inference (Jordan al., 1999).
Augmented Evidence Lower Bound. For variational inference, we introduce variational dis- tribution q0, <) = ()) approximate the posterior p00, db D). The standard ELBO is lower bound on the log marginal likelihood L == log p(D) = log JJJ p(D, 0, o 7) d0 dΦ dry Us- ing Jensen's inequality with @0,, Φ) and with standard manipulations, we obtain the ELBO, which consists of data-fitting term and two regularization terms (L({(0)lIp(() and (L()())lp(())
C M (y(0))(ohr(111
Σ log ply: (x,();,, (L))(()))))) KL(G())IP(O)). (7)
Optimization of the ELBOO. The augmented FI.BO presented in Eq. 7 is optimized by jointly updating the parameters of the variational distributions (0) and ((Φ)) using stochastic gradient-based methods. This involves sampling from these distributions (often via reparameterization) and from the DA distribution p(y () to compute Monte Carlo estimates of the expectation term, and then backpropagating through the objective. A detailed algorithm, specific choices for parameterizing p(>| Φ) and q( (()) for continuous and discrete transformations, and other practical implementation aspects discussed in Appendix D.
THEORETICAL ANALYSIS
We present a comprehensive analysis of the proposed DA approach based on Bayesian model selection, analyzing its properties from multiple perspectives: variational approximation quality, generalization guarantees, invariance properties, and connections to empirical Bayes. Our analysis includes a direct comparison with naive DA, which amounts in treating augmented samples as training samples. This analysis provides a rigorous foundation for OPTIMA while yielding practical insights for implementation.
4

---

Published as a conference paper at ICLR 2026
4.1 VARIATIONAL APPROXIMATION WITH AUGMENTATION
We begin by analyzing the quality of our variational approximation when incorporating DA. Proposition 4.1 (Jensen Gap Bound). The augmentation distribution variance and model sensitivity control the Jensen gap introduced by our lower bound approximation. If f()) = =logp(y (x). ()
L-Lipschitz in y, and 7 ~ p()))) ) is Lub-Gaussian with variance proxy o², then:
log [p(p (x 0)] - (ogply (Tx()),0) 1.202
(8)
Also, this bound is tight when f()) is approximately linear in the high-probability region of p(rlo))
The proof is presented in Appendix B.1. This result has important implications for optimizing the augmentation distribution p(7])):
Corollary 4.2 (Optimal Augmentation Variance). The optimal variance o² for the augmentation distribution balances two competing factors:
1. Increasing o² improves exploration of the augmentation spuce.
2. Decreasing o lightens the variational bound.
For models with high sensitivity to augmentations (large L), smaller variance is preferred to maintain bound tightness.
This corollary provides practical guidance for setting augmentation distribution parameters, suggest- ing that highly sensitive models benefit from more conservative augmentation strategies.
4.2 GENERALIZATION GUARANTEES
To analyze the generalization of our Bayesian popssian- DA. we leverage the PAC-Bayes framework (McAllester, 1999; Catoni, 2007); see Appendix C for a primer. PAC-Bayes theory provides high- probability upper bounds on the true risk (generalization error) of learning algorithm that outputs distribution over hypotheses (a "posterior"). These bounds typically depend on the empirical risk observed on the training data and complexity term, often expressed as the KI. divergence between this posterior and data-independent prior distribution. By extending this framework to setting, we can formally quantify how well the model with learned DA parameters will perform on unseen data. We first present a PAC-Bayes bound for OPTIMA, then provide theorem that explicitly compares OPTIMA to naive DA, demonstrating superior generalization.
Definition 4.3 (True and Empirical Risks). Given the transformation function (x(x,, define:
True risk: R(O, Φ) E(x,y)- [- log Eptris)pyy (x).. 0)].
Our empirical risk: R(0,0) = - EN log Eprroopyy, (x;),
Empirical risk for naive augmentation (K samples per datapoint):
N
K logp(y: Tr (z)), 7k ~ p(r (b)
1
Roavve (0)
Σ
NK
Theorem 4.4 (PAC-Bayes with Augmented Likelihood). For an i.i.d. dataset D {(x;, y;) )
drawn from an unknown distribution P, any prior p(0, Φ). and any posterior ( Φ)) = q (())(() over the hypothesis space 0 x Φ, with probability least 1 - 8 over the draw of D:
KL ())pp( x)) + log 2/N
E R Φ)] & Ey(e [((, Φ)]
(9)
2N
where KL(g(0, (0 Φ)) = (L(g(())lP(())) + KL(g()))lp(O)) if Φ) = p())p(())
To explicitly demonstrate that OPTIMA generalizes better than naive DA. now compare the PAC- BaC-Bayes bounds of both methods, showing that OPTIMA yields a tighter bound due to proper marginalization over transformations. We encourage reader refer Appendix C for further discussion.
5

---

Published as a conference paper at ICLR 2026
Remark. As usual with Monte Carlo estimates, the naive risk Rnaive is consistent approximation of the true marginalization when is large enough, which is the setting assumed in Theorem 4.5. Augmentation). Consider a model
Theorem 4.5 (Generalization Advantage of parameterized by O and let & E parameterize an augmentation distribution p(7) Φ), where r defines transformations (x)..
We consider the following assumptions:
1. The transformation distribution p(a) Φ)) is such that Eptrioply TT(x),0) can be computed or upproximated uccurately.
2. The variational posteriors q(0, +) and q(0) are optimized to minimize their respective bounds.
3. The KI. divergences Kl.(g(0, Φ) [p(O, ()) and Kl]g(()) (l)) are comparable, i.e, the com-
plexity penalties are similar.
Under these assumptions, the PAC-Bayes bound for OPTIMA is righter than that for naive DA:
1x R(0, Φ) - [ A,
(10)
where = Ey(o, #/4 E (x.,).)) 2 O, and (x y,) = log Ep(lo)p(y, (x,)x,),
Eptrie) loy,(y, (x,)) Furthermore, A when ply: Tx(x,),0) varies across 7 indicating a strictly better generalization bound for OPTIMA.
The proofs for Theorem 4.4 and Theorem 4.5 detailed Appendix B.2 and Appendix B.3. This theorem provides several key insights:
(Marginalization Advantage). For fixed Φ. the term Ao(Xi, log P(y, Ty (x:), 0) Epp(rlo) logp(y, (x:),, 0) quantifies the advantage of proper marginalization naïve DA. By Jensen's 's inequality, Dg (x) O. with equality only when ply; (x:), 0) is constant across all y in the support of p()(
Corollary 4.6
Corollary 4.7 (Augmentation p(0, Φ) reflects the invariances induced by the augmentation family. Priors that favor parameters satisfying ply (x), 0) = p(y x, 0) lead to smaller KL. terms and tighter bounds.
Prior). The PAC-Bayes bound is smallest when the prior
These results demonstrate that OPTIMA provides better generalization guarantees than naïve DA and suggests principles for designing priors that complement the strategy.
4.3 INVARIANCE ANALYSIS
We now analyze how OPTIMA promotes invariance to transformations, extending beyond first-order bbase) analysis include higher-order effects. This analysis reveals how the model's sensitivity to input transformations is regularized, encouraging robustness and generalization.
Theorem 4.8 (Higher-Orde Order Invariance). Let fe be viice-differentiable function parameterized by 0, with its Hessian bounded such that 2 < H. For input transformations Tx(x) = x + 8(y), where 5(7) is a perturbation with zero mean, Eptrle) [8] = O, and covariance 1 Eg,
the expected squared difference in the model's output under these transformations is:
Ep(tlo (IT))((,(()) [o(()) =Tr (J,(()) (x))
+ Botioo [6 V fo()) V fo(x)) +(1881))
(11)
where (x) the Jacobian of So at input x, (ol is the Hessian of Sa at x, and ((161) represents higher-order terms that become negligible for small perturbations.
Corollary 4.9 (Input-Space Regularization). The second-order term in Theorem 4.8 acts as regularizer, penalizing high curvature in the model's output with respect to the input. This encourages smoother response surface, promoting robustness transformations and potentially enhancing generalization by reducing sensitivity to irrelevant input variations.

---

Published as a conference paper at ICLR 2026
Corollary 4.10 (Optimal Transformation Covariance). The optimal covariance structure Esp for the augmentation distribution depends on the geometry of the model's response surface. Specifically, E should allocate more variance in directions where the model is approximately invariant (small eigenvalues of J(()) ]x(x)) and less variance in directions of high sensitivity.
Remark 4.11. In our framework, Φ is inferred via ((Φ) by maximizing the augmented ELBO, allowing the DA distribution to adapt to the data and further enhance model robustness.
provides practical guidance for designing augmentation distributions that align with the model's nat- ural invariances, enhancing robustness and generalization. The proof can be found in Appendix B.4.
4.4 MARGINALIZATION Vs. HEURISTIC AUGMENTATION
We now quantify the difference between marginalization approach and naive DA, focusing the impact on posterior uncertainty. The next theorem assumes local Gaussianity of the posterior with full-rank covariance, which might not be the case in practice for over-paramettrized models; however, we believe that the theoretical development gives some useful insights into the behavior of OPTIMA compared naïve DA, and we will attempt more general proof in the future.
Theorem 4.12 (Posterior Shrinkage under Naive Augmentation). Let Prvee (0 D) be the posterior under our marginalization approach and Pnaive (0 (D) be the posterior under naive DA with K augmentations per data point. Under regularity conditions and assuming a locally Gaussian approxi- mation around the MAP estimate 0: Enaive 73 1 Ennee where Enuivee and Lune are full- full-rank covariance matrices of Dnave
K
D) and Pirue (f)) D),, respectively.
The proof is in Appendix B.5, This result has significant implications for uncertainty quantification:
Corollary 4.13 (Uncertainty Propagation). Predictive uncertainty is underestimated by a factor of approximately V under naive augmentation, leading to overconfident predictions, particularly for
inputs.
These results provide a quantitative characterization of the benefits of proper marginalization over naive augmentation, particularly for uncertainty quantification and calibration.
4.5 EMPIRICAL BAYES PERSPECTIVE
Finally, we analyze OPTIMA from an empirical Bayes perspective (Robbins, 1992; Efron, 2012), showing how it naturally leads to optimal augmentation strategies.
Theorem 4.14 (Empirical Bayes Optimality via Augmented ELBO). The augmented ELBO ang (4o-9d (see Eq. 7) holds when q()) = p(0 | (D) and =pá)= D). Consequently, maximizing F.IBOaugg amlBamggg00 96) with respect both q()) and q(Φ) approximates the posterior distri- butions p(0 (D) and p( (D), with the mode or mean of ((Φ) serving as a point estimate analogous an Empirical Bayes solution, regularized by the prior p()).
Corollary 4.15 (Data-Driven Augmentation). The optimization of q(Φ) via the augmented ELBO results in an augmentation distribution p(> () that is specifically tailored the observed data D, with &f ~ q( (). This process effectively selects DA parameters enhancing the ability of the model to explain the data, implicitly performs model selection over the space of augmentation strategies.
Corollary 4.16 (Convergence of Joint Optimization). Under mild regularity conditions (e.g.. continuity and boundedness of the likelihood and prior), the alternating optimization of the variational distributions (0) and q(Φ) converges to a local optimum of the marginal likelihood p(D). This ensures that the learned augmentation distribution p(y () is both data-consistent and aligned with the model's posterior distribution.
These results establish OPTIMA as a principled, data-driven method for learning optimal augmenta- tion strategies within Bayesian framework. The proof is detailed in Appendix B.6.
7

---

Published as a conference paper at ICLR 2026
4.6 INFORMATION THERRETI PERSPECTIVE
We now provide an information-theerrtiic analysis of OPTIMA, offering additional insights into the role of DA in Bayesian inference.
Theorem 4.17 (Information Gain from Augmentation). The expected information gain from DA, measured as the reduction in posterior entropy, is:
AH
-
~ log de Hnosug Hang),
(12)
where Hnoaug and Haug are the Hessians of the negative log- log-likelihood without and with DA, respectively, and p(@ D) uses the marginalized likelihood.
The proof is in Appendix B.7. This information-theoreti perspective provides additional insights:
Corollary 4.18 (Optimal Information Gain). The DA distribution that maximizes information gain while maintaining a fixed KI. divergence from a reference distribution aligns with the eigenvectors of the Fisher information matrix, with variance inversely proportional to the eigenvalues.
Corollary 4.19 (Connection to Information Bottleneck). OPTIMA can be viewed as implement- ing an information bottleneck, where the DA distribution p(r) ()) is optimized to maximize the mutual information between the augmented inputs and the targets, while minimizing the mutual information between the original and augmented inputs.
These information-thhorree results provide complementary perspective OPTIMA, connecting to principles of optimal experimental design and information bottleneck theory.
EXPERIMENTS
5.1 SYNTHETIC REGRESSION EXAMPLE
We begin with a toy regression problem by generating 50 training and 1000 test points from y sin(2r) + 0.5 (35 g + sin with e ~ (0, 0.22) and <V 0.15²). We report results in Fig. In the competing approaches, Fixed Aug augments data by adding Gaussian noise with fixed standard deviation a 0.1. In Naïve Aug. for each training example, average the loss over K = 5 independent augmentations with = 0.1. In OPTIMA, the augmentation shift 7 ~ N(u, o²) has leamable parameters and it has prior A/00, 0.2?).
Although No Aug attains a lower training error, its test error is significantly higher due to overfitting. Conversely. Fixed Aug and Naïve Aug achieve better test performance than augmentation, indicating that input perturbations help regularize the model. Our OPTIMA achieves competitive test MSE. The learned augmentation distribution widens training taking from 0.10 to about 0.18, implying that a broader range of translational perturbations is optimal for this dataset. This dynamic adaption shows the benefit of OPTIMA's ability learn the augmentation distribution, and it is theoretically justified in Corollary 4.2 and Corollary 4.15, which state that OPTIMA tailors the augmentation distribution the observed data. For additional ablations with different intensities on image classification dataset CIFARIO. 10, see Appendix F.2.
Predictions
Training Loss
Test Loss
Convergence of
0.175
0.150
0.125
0.100
-2
2
1000
2000
3000
D
1000
2000
3000
1000
2000
3000
X
Epoch
Epoch
Epoch
Training data
(MSE: Train 0.021 Test: 0.101) - No Aug.
(MSE: Train 0.039
Test. 0.073) Naive Aug.
True function
(MSE. Train 0.6 0.039
-Test: 0.074) Fixed Aug.
(MSE: Train 0.047
Test: 0.072) Learned Aug.
Figure 2: Synthetic regression: (Left) Test predictions compared to the ground-truth function. (Right) Convergence traces for OPTIMA; green dashed line denotes the fixed a = 0.1 used in Fixed Aug.
8

---

Published as a conference paper at ICLR 2026
5.2 IMAGENET AND IMAGENET-C
We next evaluate the robustness of OPTIMA on IMA-
Table IMAGENET and IMAGENET-C with non-Bayesian ResNet-50.
GENET (Deng et al., 2009) and IMAGENET-C, of-distribution (OOD) dataset (Hendrycks & Dietterich, 2019) using a Bayesian ResNet-18 (He et al., 2016), where the final layer is replaced with a BayesianLines treating only the final layer in Bayesian manner-is common and efficient strategy Bayesian deep learning (Harrison et al., 2024). As noted by Sharma et al. (2023), full network stochasticity is often unnecessary; introducing stochasticity in the final layer can be sufficient to capture predictive uncertainty, especially with strong deterministic feature extractors (Kristiadi et al., 2020). This allows the model to represent uncertainty in class probabilities while leveraging the pretrained backbone. With OPTIMA, we optimize augmentation parameters for Mixup (Zhang al., 2018), CutMix (Yun et al., 2019), and AugMix (Hendrycks et al., 2020). Im- portantly our approach is general and can also be applied to standard -ayyssaa)) networks. To illustrate this, we further evaluate OPTIMA on ResNet-50 without Bayesian treatment of its parameters. Implementation details are provided in Appendix E.
out-
Method
Ace (%)
Acc (%)
Clean
Corrupted
module. This partially stochastic
Mixup OPTÍMA Mixup
76.1
40.1
76.8
41.6
Table 2 summarizes the results, confirming that OPTIMA obtains better calibration and robustness for both ID and OOD. Regarding the non- -an-Baaan NNs. our framework also allows us to have better accuracy on the clean and corrupted data (see Table 1), further demonstrating that OPTIMA captures variations in the data better than fixed augmentations. More results can be found in Appendix G.5. These results supports our theoretical analyses - e.g. Theorem 4.5 (improved generalization on test and OOD data) and Theorem 4.12 (enhanced calibration and uncertainty quantification).
Table 2: IMAGENET and IMAGENET-C C with pretrained Bayesian ResNet-50 (last layer) after 10 epochs ("C" -oorrupee,,, "m" mean)
Method
Test Ace CClean Data) (%)
ECE (J)
mCE ()) (unnoraalzee) (C) (%)
MECE (C)
MHOD-AUROC (C)
Fixed Mixup
0.043
0.062
0.820
OPTIMA Mixup
74.97
0,031
61.65
0.045
0.822
Cutmix
74.17
0.036
0.059
OPTIMA Cutmix
74.34
0,034
63,60
0.058
0.820
Hixed Augmix
74.71
0.084
61.45
0.156
0.790
OPTIMA Augmix
75.33
0.083
60.68
0.149
0.793
5.3
COMPUTATIONAL EFFICIENCY AND COMPARISON WITH BAYESIAN OPTIMIZATION
Our method introduces almost no additional computational cost compared to traditional data augmen- tation. The difference lies in our adaptive augmentation strategies, which evolve over iterations rather than remaining fixed. OPTIMA employs Monte Carlo estimates with small sample size like one per data point per iteration and uses the reparameterization trick for efficient, low-variance gradient estimation. To highlight its efficiency, we compare against Bayesian Optimization (BO), strong baseline for augmentation tuning. BO requires costly black-box optimization with many full training runs per hyperparameter setting, whereas OPTIMA's tractable ELBO jointly optimizes augmentation and model parameters within the same training loop- removing the need for separate validation runs.
We evaluate CIFARI0 (Krizhevsky & Hinton, 2009) using pretrained Bayesian ResNet-18 (Bayesian last layer) to optimize augmentation parameters (mean and variance). BO is run for 25 trials of 15 epochs, followed by 50 epochs of training with the optimized parameters, while OPTIMA is trained directly for 50 epochs. For augmentation, we use Mixup and learn the parameter a. We also assess performance on CIFAR 10-C (Hendrycks & Dietterich, 2019) OOD data. As shown in Table 3, OPTIMA achieves higher test accuracy on clean data (with a slight calibration tradc-off) and substantially better accuracy, ECE, and AUROC on OOD data, all in far less time than BO demonstratinn improved calibration and robustness at much lower cost.
Table 3: Comparison between Bayesian optimization and OPTIMA CIFARI0
Method
Test Acc (%)
ECE (%)
mAAcuracy (C)
MECE (C)
OOD AUROC
Time
Bayesian Optimization
93.43
0.010
72.44 78.52
0.127
0.652
x
OPTIMA
95.03
0.047
0.076
0.680

---

Published as a conference paper at ICLR 2026
5.4 OPTIMA ON DISCRETE NLP AUGMENTATIONS, SST-5 CASE STUDY
To show that OPTIMA is not restricted to continuous or geometric transformations used in computer vision, we additionally evaluate it NLP classification task where augmentations are inherently discrete. We use the SST-5 benchmark (Socher et al. 2013), a fine-grained 5-class sentiment dataset, and fine-tune& a DistilBERT model (Sanh et al. 2019) for epochs on the full training split.
Discrete augmentation family. We consider token dropout, a stochastic masking transformation used in NLP regularization. For input x ( Tt) with dropout rate Pårop. we sample Nt Bernoulli Pdrop) and replace 24 with [MASK] if y = O, yielding discrete latent variable y. Although non- non-differentiable, OPTIMA optimizes Pdrop via ccoreefunction (REINFORCE) gradient, consistent with Section 3.
Experimental setup. We evaluate OPTIMA on discrete token- -rooout augmentation. The augmen- tation uses dropout probability Parop E [0. Pmax] parameterized as Parop Pmax (ss), where is a learnable scalar. To encode prior preferences for dropout, we place a Gaussian prior directly on pdrop- and OPTIMA jointly learns and the model parameters via our ELBO. We compare OPTIMA against: (i) No Aug: (ii) Fixed Aug, which uses the same initial dropout as OPTIMA, (iii) Fixed Aug (Matched), where Parop is set equal to the value learned by OPTIMA; and (iv) BO-Fixed, which selects Párop through a validation-bast
hyperparameter search. More details provided in Appendix H.
Table 4: SST-5 results for discrete token-dropout augmentation averaged over 5 different seeds.
Method
Accuracy
NLL
ECE
No Aug
0.516 + 0.003
1.240 + 0.010
0.190 + 0.004
Fixed Parop 0.04
0.522 + 0.003
1.180 + 0.006
0.154 + 0.006
Fixed Pdrop 0.0625
0.516 + 0.006
1.162 + 0.007
0.143 + 0.007
OPTIMA with μ 0.1 (Pleamed 0.0625)
0.524 : 0.003
1.161 = 0.007
0.142 = 0.006
BO-Fixed Pdrop = 0.3
0.521 0.004
1.086 + 0.006
0.043 + 0.004
OPTIMA with 0.3 (Pleamed = 0.3)
0.524 + 0.004
1.086 + 0.005
0.046 = 0.002
Results. Table shows that accuracy differences on SST-5 are small, but OPTIMA consistently achieves lower NLL and improved calibration over fixed-augmentation baselines. Notably, it matches the BO-tuned baseline-desye despite BO requiring a full hyperparameter search across multiple runs (approximately 8x more compute)-whil while OPTIMA learns Párop in single training run. This indicates that the gains stem from marginal likelihood optimization rather than dropout tuning. These findings confirm that OPTIMA extends naturally to discrete augmentation spaces and that its theoretical advantages (Sections 4.2 4.33 generalize beyond vision tasks.
6 DISCUSSION AND CONCLUSION
We presented a theoretical and methodological framework for optimizing DA taking inspiration from Bayesian principles, which allow to cast this problem as model selection. We derived variational objective to learn optimal DA strategies from data in a practical way. We also provided extensive theoretical insights on the advantages of our proposed data-driven approach to DA compared to alternatives, revealing improved generalization through PAC-Bayes bounds, enhanced invariance via higher-order regularization, and better calibration through marginalization. Empirical results confirm these theoretical benefits, showing consistent improvements in calibration and predictive performance across various tasks. We believe that OPTIMA A is a key step toward robust and well-calibrated models capable of assisting decision- making in applications where this is of critical importance.
Limitations and future work. While OPTIMA offers significant advantages, it has some limi- tations that suggest directions for future work. Although our main experiments focus on computer vision, OPTIMA itself is not tied to continuous or geometric transformations. Our additional evaluation on a natural language task (SST-5) demonstrates that OPTIMA can also handle discrete, non-geometric transformations by optimizing a latent augmentation distribution in text space. Nev- ertheless, a broader exploration of more expressive or compositional transformations in NLP, time series, or multimodal settings remains important next step. In addition, our theoretical analysis could be strengthened by developing tighter PAC-Bayes bounds and more refined characterizations of the benefits introduced by Bayesian marginalization over augmentation parameters.
10

---

Published as a conference paper at ICLR 2026
REFERENCES
Ibrahim M Alabdulmohsin, Behnam Neyshabur, and Xiaohua Zhai. Revisiting neural scaling laws in language and vision. Advances Neural Information Processing Systems, 2022.
Pierre Alquier. User-friendly introduction to pac-bayes bounds. Foundations and Trends in Machine doi: 10.1561/22000001000. URL https:
Learning, 17(2):174-303, 2024. ISSN 1935-8237.
//doi. orw. 0tt../116//222000000000
Gregor Bachmann, Lorenzo Noci, and Thomas Hofmann. How tempering fixes data augmentation in Bayesian neural networks. In International Conference on Machine Learning, 2022.
Gregory Benton, Mare Finzi, Pavel Izmailov, and Andrew G Wilson. Learning invariances in neural networks from training data. In Advances in Neural Information Processing Systems, 2020.
Charles Blundell, Julien Comebise,, Koray Kavukcuoglu, and Daan Wierstra. Weight uncertainty in neural network. In International Conference on Machine Learning, 2015.
Olivier Catoni. Accounting for variance in machine learning benchmarks. Institute of Mathematical Statistics Lecture Notos-Monograph Series, 2007.
Evangelos Chatzipantazis, Stefanos Pertigkiozoglou, Kostas Danilidis, and Edgar Dobriban. Learn- ing augmentation distributions using transformed risk minimization. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL https: /ttp///ppe///////ee/////ee/eeeet ew net/ forum?
id=IRYENj8XwD
Shuxiao Chen, Edgar Dobriban, and Jane H Lee. A group-theoretic framework for data augmentation. Journal of Machine Learning Research, 2020.
Ekin D. Cubuk, Barret Zoph, Dandelion Mané, Vijay Vasudevan, and Quoe V. Le. AutoAugment: Learning augmentation strategies from data. In Conference Computer Vision and Pattern Recognition, 2019.
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with a reduced search space. In Conference on Vomputer Vsion and Pattern Recognition workshops, 2020.
Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Re. A kernel theory of modern data augmentation. In International Conference Machine Learning, 2019.
Erik Daxberger, Agustinus Kristiadi, Alexander Immer, Runa Eschenhagen, Matthias Bauer, and Philipp Hennig. Laplace redux effortless Bayesian deep learning. In Advances in Neural Information Processing Systems, 2021.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. ImageNet: A large-scale hicrarchical image database. In Conference on Computer Vision and Pattern Recognition, 2009.
Weijian Deng. Stephen Gould, and Liang Zheng. On the strong correlation between model invariance and generalization. Advances Neural Information Processing Systems, 2022.
Gintare Karolina Dziugaite and Daniel M Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Conference on Uncertainty in Artificial Intelligence, 2017.
Bradley Efron. Large-scale inference: Empirical Bayes methods for estimation, testing. and predic- tion. Cambridge University Press, 2012.
Steven Feng, Varun Gangal, Jason Wei, Sarath Chandar, Soroush Vosoughi, Teruko Mitamura, and Eduard Hovy. A survey of data augmentation approaches for NLP. Findings of the Association for Computational Linguistics: ACL-IJCNLP, 2021.
Edwin Fong, Chris Holmes, and Stephen G Walker. Martingale posterior distributions. Journal of the Royal Statistical Society Series B, 15(5):135711391, 2023.

---

Published as a conference paper at ICLR 2026
James Harrison, John Willes, and Jasper Snoek. Variational bayesian last layers. In International Conference on Learning Representations, 2024.
Ryuichiro Hataya, Jan Zdenek, Kazuki Yoshizoe, and Hideki Nakayama. Faster autoaugment: Learning augmentation strategies using backpropagation. In European Conference on Computer Vision, 2020.
Ryuichiro Hataya, Jan Zdenek, Kazuki Yoshizoe, and Hideki Nakayama. Meta approach to data augmentation optimization. In Winter Conference Applications of Computer Vision, 2022.
Kaiming He. Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Conference on Computer Vision and Pattern Recognition, 2016.
Markus Heinonen, Ba-Hien Tran, Michael Kampffmeyer, and Maurizio Filippone. Robust classifica- tion by coupling data mollification with label smoothing. In International Conference Artificial Intelligence and Statistics, 2025.
Dan Hendrycks and Thomas Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference Learning Representations, 2019.
Dan Hendrycks, Norman Mu, Ekin D Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshmi- narayanan. Augmix: simple data processing method to improve robustness and uncertainty. In International Conference Learning Representations, 2020.
Daniel Ho. Eric Liang, Xi Chen, lon Stoica, and Pieter Abbeel. Population based augmentation: Efficient learning of augmentation policy schedules. In International Conference on Machine Learning, 2019.
Ignacio Houmie, Luiz FO Chamon, and Alejandro Ribeiro. Automatic data augmentation via invariance-constrained learning. In International Conference on Machine Learning. 2023.
Alexander Immer, Matthias Bauer, Vincent Fortuin, Gunnar Rätsch, and Khan Mohammad Emtiyaz. Scalable marginal likelihood estimation for model selection in deep learning. In International Conference on Machine Learning, 2021.
Alexander Immer, Tycho van der Ouderaa, Gunnar Rätsch, Vincent Fortuin, and Mark van der Wilk. Invariance learning in deep neural networks with differentiable Laplace approximations, In Advances in Neural Information Processing Systems, 2022.
Pavel Izmailov, Sharad Vikram, Matthew Hoffman, and Andrew Gordon Wilson. What are Bayesian neural network posteriors really like? In International Conference Machine Learning, 2021.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with Gumbel-softtax. In International Conference Learning Representations, 2017.
Michael I Jordan, Zoubin Ghahramani, Tommi S Jaakkola, and Lawrence K Saul. An introduction to variational methods for graphical models. Machine learning, 1999.
Sanyam Kapoor, Wesley Maddox, Pavel Izmailov, and Andrew Gordon Wilson. On uncertainty, tempering, and data augmentation in Bayesian classification. In Advances in Neural Information Processing Systems, 2022.
Agustinus Kristiadi, Matthias Hein, and Philipp Hennig. Being Bayesian, even just a bit, fixes overconfidence in relu networks. In International Conference Machine Learning, 2020.
A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. Master's thesis, Department of Computer Science, University of Toronto, 2009.
Yonggang Li. Guosheng Hu, Yongtao Wang. Timothy Hospedales. Neil M Robertson, and Yongxin Yang. Differentiable automatic data augmentation. In European Conference on Computer Vision, 2020.
Sungbin Lim, Ildoo Kim, Taesup Kim, Chiheon Kim, and Sungwoong Kim. Fast autoaugment. In Advances in Neural Information Processing Systems, 2019.
12

---

Published as a conference paper at ICLR 2026
Aoming Liu, Zehao Huang, Zhiwu Huang, and Naiyan Wang. Direct differentiable augmentation search. In International Conference on Computer Vision, 2021.
Sanae Lotfi, Mare Anton Finzi, Sanyam Kapoor, Andres Potapczynski, Micah Goldblum, and An- drew Gordon Wilson. PAC-Bayes compression bounds so tight that they can explain generalization. In Advances in Neural Information Processing Systems, 2022.
Clare Lyle, Mark van der Wilk, Marta Kwiatkowska, Yarin Gal, and Benjamin Bloem-Reddy. On the benefits of invariance in neural networks. arXiv, 2020.
David JC MacKay. A practical Bayesian framework for backpropagation networks. Neural computa- fion, 1992.
David A McAllester. PAC-Bayesian Model Averaging. In Annual Conference on Learning Theory, 1999.
Saypraseuth Mounsaveng, Issam H. Laradji, Ismail Ben Ayed, David Vázquez, and Marco Peder- soli. Learning data augmentation with online bilevel optimization for image classification. In Proceedings of the IEEE/CVF Winter Conference Applications of Computer Vision (WACV), pp. 1691-1700, 2021.
Seth Nabarro, Stoil Ganev, Adrià Garriga-Alonso, Vincent Fortuin, Mark van der Wilk, and Laurence Aitchison. Data augmentation in Bayesian neural networks and the cold posterior effect. In Conference on Uncertainty in Artificial Intelligence, 2022.
Radford M. Neal. Bayesian Learning for Neural Networks (Lecture Notes in Statistics). Springer, 1996.
Herbert E Robbins. An empirical Bayes approach to statistics. In Breakthroughs in Statistics: Foundations and basic theory, pp. 388-394. Springer, 1992.
Victor Sanh, Lysandre Debut, Julien Chaumond, and Thomas Wolf. Distilbert, a distilled version of 2019.
bert: smaller, faster, cheaper and lighter. arXiv preprint
Han Shao, Omar Montasser, and Avrim Blum. A theory of PAC lcarnability under transformation invariances. Advances Neural Information Processing Systems, 2022.
Mrinank Sharma, Sebastian Farquhar, Eric Nalisnick, and Tom Rainforth. Do Bayesian neural networks need to be fully stochastic? In International Conference Artificial Intelligence and Statistics, 2023.
Connor Shorten and Taghi M Khoshgoftaar. A survey on image data augmentation for deep learning. Journal of Big Data, 2019.
Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical Bayesian optimization of machine learning algorithms. In Advances in Neural Information Processing Systems, 2012.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1631-1642, 2013.
Ba-Hien Tran, Simone Rossi, Dimitrios Milios, and Maurizio Filippone. All you need is good functional prior for Bayesian deep learning. Journal of Machine Learning Research, 2022.
Ba-Hien Tran, Giulio Franzese, Pietro Michiardi, and Maurizio Filippone. tme-line-of-code data mollification improves optimization of likelihood-basee generative models. In Advances in Neural Information Processing Systems, 2023.
Mark van der Wilk, Matthias Bauer, ST John, and James Hensman. Learning invariances using the marginal likelihood. In Advances in Neural Information Processing Systems, 2018.
David A Van Dyk and Xiao-Li Meng. The art of data augmentation. Journal of Computational and Graphical Stafistics, 1051),1-50, 2001.
13

---

Published as a conference paper at ICLR 2026
Yuexi Wang, Nicholas Polson, and Vadim O. Sokolov. Data Augmentation for Bayesian Deep Learning. Bayesian Analysis, 2023.
Florian Wenzel, Kevin Roth, Bastiaan Veeling, Jakub Swiatkowski, Linh Tran, Stephan Mandt, Jasper Snoek, Tim Salimans, Rodolphe Jenatton, and Sebastian Nowozin. How good is the Bayes posterior in deep neural networks really? In International Conference on Machine Learning, 2020.
Christopher KI Williams and Carl Edward Rasmussen. Gaussian processes for machine learning. volume 2. MIT press Cambridge, MA, 2006.
Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in Neural Information Processing Systems, 2020.
Andrew Gordon Wilson. Deep learning is mysterious or different. arXiv, 2025.
Luhuan Wu and Sinead Williamson. Posterior uncertainty quantification in neural networks using data augmentation. In International Conference on Artificial Intelligence and Statistics, 2024.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. Cutmix: Regularization strategy to train strong classifiers with localizable features. International Conference Computer Vision, 2019.
Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals. Understand- ing deep learning requires rethinking generalization. In International Conference on Learning Representations, 2017.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. mixup: Beyond empirical risk minimization. In International Conference on Learning Representations, 2018.
Yu Zheng, Zhi Zhang, Shen Yan, and Mi Zhang. Deep autoaugment. In International Conference Learning Representations, 2022.
14

---

Published as a conference paper at ICLR 2026
Appendix
TABLE OF CONTENTS
Derivation of the Augmented Evidence Lower Bound
15
B Detailed Proofs
16
B.1 Proof of Proposition 4.1 (Jensen Gap Bound)
16
B.2 Proof of Theorem 4.4 (PAC-Bayes Under Augmentation)
17
B.3 Proof of Theorem 4.5 (Generalization Advantange Bayesian-Optinn Augmen- tation)
17
B.4 Proof of Theorem 4.8 (Higher-order Invariance)
18
B.5 Proof of Theorem 4.12 (Posterior Shrinkage under naive Augmentation)
19
B.6 Proof of Theorem 4.14 (Empirical Bayes Optimality via Augmented ELBBO)
19
B.7 Proof of Theorem 4.17 (Information Gain from Augmentation)
20
A Primer on PAC-Bayes Theory
20
D Algorithm and Implementation
21
D.1 Parameterization of Augmentation Distribution
21
D.2 Practical Considerations
23
E Additional Experimental Details for $ 5.2
23
F Additional Results Different Types of Data Augmentation
23
F.1 Learning Geometric Augmentation for CIFARIO Classification
23
F22 Exploring different intensities of Gaussian translations
24
G Additional Experiment on IMAGENET using Resnet-50
25
G.1 Implementation Details for OPTIMA with Mixup
25
G.2 Implementation Details for OPTIMA with CutMix
26
G.3 Implementation Details for OPTIMA with AugMix (Learnable Severity JSD)
26
G.4 Evaluation and Software/Hardware for all these methods.
27
G.5 Experimental Results
27
H Additional Experimental Details for $ 5.4. Token-dropout implementation details
28
I Broader Impact
28
Reproducibility Statement
28
K The Use of Large Language Models (LLMs)
28
A DERIVATION OF THE AUGMENTED EVIDENCE LOWER BOUND
For variational inference, we introduce a variational distribution q0, Φ) = (() to approximate the posterior p(0,c) & DD). The standard ELBO is lower bound the log marginal likelihood
15

---

Published as a conference paper at ICLR 2026
log pDD) = log [J] p(D, 0, Φ, 7) d0 dΦ dry Using Jensen's inequality with q(0, ().. we have:
log P(D, 8, Φ)
logp(D) 2 E8 (0,
where p(D,0, p(D, dy. dy. (13)
q00 Φ)
BLBO(g)
Applying Jensen's inequality further to the log of the likelihood term,
N
N
log Φ) - =log Ep(r o)
I ply. (x.), 0) 2 Epivl +)
Σ
log py, [T 0) (14)
Substituting this into the ELBO in Eq. 13, with P(D,0,0) - p(D 0, ()))()))(( we can obtain the augmented ELBO as follows:
N log ply; (x;)x;;, 0) + log p(0)p(p)
ELBO 2 (e.0))
Eptr1
Σ
(15)
(())(()))
N
= But logrly, (x:(),,,
(L(((())l((()
(l(()))ll(
(16)
parameter prior
sugmentation prior
The augmented ELBO consists of three terms: data-fitting term that averages over the varia- tional distributions (0), ( and the augmentation distribution p(> | Φ); regularization term (L(q(0))lp) that penalizes divergence from the prior over model parameters; and another regular- ization term (L(g( (((()) that aligns the augmentation parameters with their prior.
B DETAILED PROOFS
This section provides expanded proofs for the theoretical results presented in the main paper.
B.1 PROOF OF PROPOSITION 4.1 (JENSEN GAP BOUND)
Proof. Let f(r) p(y (x), 0). For sub-Gaussian 7 with mean E(y] and variance proxy σ², we use standard moment generating function bounds:
log Eyle(() = log B. [(W
(17)
=f(}) + log Byy (e(
(18)
Since f is L-Lipschitz, we have (f(r) f(4) s L|ly - 1l|1 Using the sub- Gaussan property:
- s Ey - s
120
(19)
Therefore:
1.20 log Ey le [(>]] s f(m) +
(20)
2
Since Ey[f(r)] = f(a) when Byly 11] = 0, the gap is:
Gap log -B
(21)
(22)
For tightness, when f(y) is approximately linear the bighhbrobability region of p(>(@), the bound approaches equality.
16

---

Published as a conference paper at ICLR 2026
B.2 PROOF OF THEOREM 4.4 (PAC-BAYES UNDER AUGMENTATION)
Proof. Define the loss function as:
E(0, Φ, (x. y))
log LogEptre p( (p(y) (x). 0))
(23)
The empirical risk is:
N
1 N Σ e(0, Φb (x:,y.)).
R(0 .)
(24)
Since y(); are ii.d., and the expectation y p(r ()) is computed independently for each sample, the terms E(O, Φ, (x,.y,) are independent for fixed 0, Φ. Applying the standard PAC-Bayes theorem over the joint space (0, <b) (McAllester, 1999; Catoni, 2007; Alquier, 2024):
V (L(g(0.0)ll(() +) Φ)) + log NVN
Ee Φ)) [)] () R(0, 6))
(25)
2N
This completes the proof.
B.3 PROOF OF THEOREM 4.5 (GENERALIZATION ADVANTANGE OF BAYESIAN-OPTIMIZED AUGMENTATION)
Proof. Step 1: PAC-Bayes Bounds
For OPTIMA, the PAC-Bayes bound is:
(0 [R(,, Φ)] ( A(0, + KL(g(0, (L(g(O,))ll((,
@)l(p(O, x)) + log 2/N
(26)
2N
For naïve augmentation, the bound is:
(L(g(0)lp(()) log N
S Eg(o) [Ratvee (0) +
(27)
2NV
Step 2: Relationship Between Empirical Risks
By Jensen's inequality, for each data point (x..y,)
log Epproe)Pyy (x..)x)),T
(x:)).) 0) 2 Ep(r(o) logp(y, (x:)),) 0).
(28)
Thus.
log Eptrispyyy Tx(x)), < - Ep(rio) logply, Tx,)),,,
(29)
For large K, the naïve empirical risk approximates:
Riaaee (0) R
Σ
Eptr logp(y, (x(xx)), 0).
(30)
Therefore,
17

---

Published as a conference paper at ICLR 2026
N Σ
R(O, Φ)
log P(t (xx()),, 0) S R (0);
(31)
N
with equality only if y;) 0 for all i, i.e. when ply, (x.), 0) constant across Ty.
Step 3: Bound Comparison
Assuming (L(g(0,)))pp x)) (L(((0))lp())) the difference in bounds driven by the empiri- cal risks:
() <)] < E Ey +) complexity term,
(32)
E E [ S Ey( Ramive (0) + complexity term.
(33)
Since R((0 ②) S Raaive ( and the complexity terms are similar, our bound is tighter. Specifically,
N
0) [f(0, ( [)] = Eggo, Rnanv (0)) Eyco,,)
A A (xx)))))
(34)
leading to:
(R(0, x)] & Ey(or RRaive (0) A + complexity term, Eg(0 [R(O,
(35)
where A = ()) 1 EM (x) 2 0.
Thus, OPTIMA bound is lower by A, proving better generalization. When p(y, (x,). 0) varies across 7, A > making our bound strictly tighter.
B.4 PROOF OF THEOREM 4.8 (HIGHER-ORDER INVARIANCE)
Proof. Using a second-order Taylor expansion of fe around x:
(x) fo(x) (x(x) +
2 8 fxx)) 0(11)13).
(36)
fo (x)T((x)) - fx(x) = (x)) 1 8' f 0(16113).
(37)
Squaring this difference and taking the expectation over p(7(Ø):
Ep(rlo) [II)o(T((x)) [o(x) Eptr Φ) [1|J{(x)013]
(38)
2 [ 27
+ Ep(rio)
(39)
+ Ep(r
(Jz(x)))
(40)
(116115).
(41)
The cross-term Ep(r(0) (J)(x)6) (6)))) involves odd powers of which vanish since
B 0. Thus: Ep(rio) [x(x)) Eperio) [x'
(42)
+ 4 Ep(ylos (6" (xx))) 2] + +(11811))
1
(43)
18

---

Published as a conference paper at ICLR 2026
Using properties of quadratic forms for zero-mean random variables:
Bi()))) [8 (x(x)* Jy(x) Tr (Jx(x))
(44)
Jo B)(r)0) [o0)) Tr (V -
( 2Tr (VV
(45)
For small perturbations, the dominant term is Tr (V (x ( and higher moments contribute to ((]])])). Therefore:
Bp(r))) [l]f fo(x)) (J ( [V
(46)
This approximation holds for small || completing the proof.
B.5 PROOF OF THEOREM 4.12 (POSTERIOR SHRINKAGE UNDER NAIVE AUGMENTATION)
Proof. Under locally Gaussian approximation with full-rank covariance, the posterior covariance is approximately the inverse of the Hessian of the negative log posterior at the MAP estimate. For the true posterior, marginalizing over Φ:
p(D p(D|0)= 0) 0(D)0, ()p(o) dØ.
Σ-1 20 Line
22 Pogo (0 0= = 2 ((0)
(47)
42 log
(o ()p())
1=
1=1
For the naive posterior:
N K
Σ-1 x3 22 log Pnaive (0 -
(0)
log logply, [Ty, (x:),
naive
(48)
Assuming 7x ~ (e.g., using point estimate of Φ), this approximates:
2 log
2 log E P( [ply, (x:):)), 0))lo- K. E-1
(49)
Therefore, Emuive RJ B
1 Emme-
B.6 PROOF OF THROREM 4.14 (EMPIRICAL BAYES OPTIMALITY VIA AUGMENTED ELBBO)
Proof. The proof leverages variational inference principles. Start with the log marginal likelihood:
P(D) = =log / /
p p(D, 0...0.0.1 d0
(50)
Introduce the variational distribution q(,, Φ) = ((0)(())
log (D. Φ)
log (DD) = ())0,
+ (((( & D)).
(51)
q0, Φ)
Since the KL divergence is non-negative, we obtain the lower bound:
log p(D, 0, Φ))
log 2 By(0,))
(52)
9 Φ)
19

---

Published as a conference paper at ICLR 2026
Now, factor in the augmentation variable 7. The joint likelihood is:
pD Φ) Eplr +)
I ply, (x(xx),
(53)
Applying Jensen's inequality:
log
Epty +)
Σ log ply, 1
(54)
Substitute into the lower bound:
- log p(0)p(()
log (D) 2 (0)
Eptr (b)
Σ
log ply, (Tx(x,),))
(55)
g(0)g(()
Rewrite:
N log p(y; (x())), 0) - (l(((())lp(()) Kl(G()))P(O)).
= (0) Ry(o) ()
(56)
Thus, arrive at the augmented ELBO. When (0) ((0 and ((Φ) p( D), the bound becomes tight, confirming the result.
B.7 PROOF OF THEOREM 4.17 (INFORMATION GAIN FROM AUGMENTATION)
Proof. Under a Gaussian approximation to the posterior, the entropy proportional the log determinant of the covariance matrix. Using the results from Theorem Theorem 4.12, we have:
AH x log det - log det det((2aug))
(57)
where Enoaug H-1 noougg Sang Hal and Hang incorporates the effect of augmentation. Thus:
AH log det (Haug log det(1 + Hnowag (Haug Hnoaug))
(58)
Approximating the effect of augmentation as an effective increase in Fisher information, we obtaín the stated result.
C A PRIMER PAC-BAYES THEORY
The Probably Approximately Correct (PAC)- -ayes framework, pioneered by (McAllester, 1999) and further developed by (Catoni, 2007) among others, provides a powerful tool for deriving generalization bounds for Bayesian-inspires learning algorithms. Unlike traditional PAC learning which often focuses on a single hypothesis, PAC-Bayes theory considers a distribution over hypotheses.
The central idea is to bound the true risk (expected loss on unseen data) of a posterior distribution Q over hypothesis class H. This bound is typically expressed in terms of the empirical risk (average loss on the training data) under Q, and a complexity term that measures how much Q deviates from data-independent prior distribution P over H. The guarantee holds with high probability (at least 1 8) over the random draw of the training dataset.
Core Idea
Key Components
Hypothesis Class (H)): The set of all possible models (e.g., sets of parameters 0).
Prior Distribution (P): A distribution over H chosen before observing any training data. reflects initial beliefs about good hypotheses.
Posterior Distribution (Q): A distribution over H. that is typically learned from the training data D. In PAC-Bayes, can be any distribution, necessarily true Bayesian posterior.
20

---

Published as a conference paper at ICLR 2026
Loss Function (E(h, 2)): Measures the enror of hypothesis h E H. on data point 2 (x,y).
True Risk (R(Q)): The expected loss of hypothesis drawn from Q on the true (unknown)
data distribution: R(Q) e Pez-Dm. (1(h, =)]].
Empirical Risk (R(Q)): The average loss of a hypothesis drawn from Q on the N training
samples: (() En-Ql E: ((h,
Kullback-Leible (KI.) Divergence (KL( Measures the "distance" or "information
gain" from the prior / to the posterior Q. It serves as complexity penalty: if Q is very
different from the penalty is high.
A Common Form of PAC-Bayes Bound A typical PAC-Bayes generalization bound (e.g.. McAllester's 1999 bound or variations) states that for any 5 E (0, 1), with probability at least 1 8 over the draw of i.i.d. training set D of size N, for all posterior distributions Q:
R(Q) S R(Q) KL(QIIP) + In( +C
(59)
2N
where C is constant that can depend the range of the loss or other factors (e.g., In(2v/) used in our paper, which is a common variant for empirical Bernstein bounds).
Interpretation and Significance
The bound guarantees that the true risk is unlikely to be much larger than the empirical risk, plus a term that penalizes the complexity of Q relative to P.
It highlights trade-off:: to achieve good generalization, learning algorithm should find posterior Q that both fits the training data well (low R(Q)) and does not deviate too much from the prior (low P)).
The bounds are often tighter than uniform convergence bounds for complex hypothesis classes like neural networks, especially when good prior is available.
They provide theoretical justification for regularization techniques and guide the design of learning algorithms.
Relevance to This Paper In our work ($ 4.2), we adapt this framework to derive generalization
bounds for augmented likelihood approach. Here, the "hypothesis" space effectively includes both the model parameters and the augmentation (hyper)parameters Φb. The priors p(0) and p(Φ) and the variational posteriors (() and (()) play the roles of and Q. Our Theorem Theorem 4.4 provides such bound, and Theorem Theorem 4.5 uses PAC-Bayes reasoning to show the theoretical advantage of marginalized approach over naïve data replication. The KL terms in our augmented ELBO (Eq. 7) naturally appear as complexity measures in these PAC-Bayee bounds.
D ALGORITHM AND IMPLEMENTATION
We now present practical algorithm for implementing our Bayesian-optimized data augmentation approach. The algorithm employs stochastic gradient-based optimization of both model parameters and augmentation distribution parameters.
D.I PARAMETERIZATION OF AUGMENTATION DISTRIBUTION
For continuous transformation parameters, we typically use Gaussian distribution for p(> Φ):
p (V Ex),
(60)
where Φ (1i, E.p. For q()), might Gaussian:
q( (Φ)) (V Mgg Eggh
(61)
21

---

Published as a conference paper at ICLR 2026
Algorithm 1 Augmented Variational Inference with Learned Augmentation
Input: Dataset = {xx; y,)] transformation family (-)()
2: Initialize: Variational distributions g(0) and g(()
3: while not converged do
4: Sample minibatch
from 2>
5: Sample model parameters 0 ~ (()) (or use reparametcrization)
6: Sample augmentation parameters ((Φ) (or use reparameterizztion))
7:
for each (x, in the minibatch do
8:
Sample augmentation parameters Y; p(7())
9:
Apply transformation x = Tyy (x:)
Compute logeliiiddood log log,(y, 11: end for
Estimate ELBO:
N
B
ELBO aug
log olg 0) KL(g(0))p(() ((0)) KL(O())lPp(O))
B
=11
13:
Update variational parameters in g()) using gradient of ELBO aug
14:
Update variational parameters in g( using gradient of ELBOau aug
15:
end while
16: Output: Optimized variational distributions (()) and ((
Algorithm 2 Partial Variational Inference with Learned Augmentation
Input: Dataset {xx, / transformation family T- (-)
2: Initialize: Model parameters @ and distribution g(Φ)
3: while not converged do
4: Sample minibatch {(x, He from D
5: Sample augmentation parameters & ~ ((())
6: for each (x,, yz) in the minibatch do
7:
Sample 7; ~ p(y(Ø)
8:
Apply transformation x TTy (x:)
9: Compute log- anglliiel log puggryy,x.. 0)
10:
end for Estimate ELBO:
11:
N
ELBO
log logply,12/ KL(g())llp(b))
Update 0 using gradient of ELBÖ aug
13:
Update ((Φ)) using gradient of ELBO
14: end while
15: Output: Optimized parameters 0 and distribution ((Φ)
learning My and Eg. This allows for reparameterization during sampling:
Φ = p4 + 11/2,, e~~
(62)
followed by:
= up + 1/22 2 e'~N(0,))
(63)
For discrete transformations, we can use a categorical distribution:
p(>( = Cat(>/TT),
(64)
where Ty represents the probabilities, and use the Gumbel-Softmax trick (Jang et al., 2017) for differentiable sampling.
22

---

Published as a conference paper at ICLR 2026
D.2 PRACTICAL CONSIDERATIONS
Adaptive Variance Scheduling. Based on Corollary 4.2, we can implement an adaptive schedule for the augmentation variance within (Φ)), adjusting the variance of over training to balance exploration and bound tightness.
Marginalization Advantage Monitoring. Following Corollary 4.6, we can monitor the marginal- ization advantage term Do(2,,81 34) during training to assess the benefit of OPTIMA over naive augmentation.
Curvature-Awer Aware Augmentation. Inspired by Corollary 4.10, we can adapt the augmentation distribution based on the model's sensitivity to different transformations, allocating more variance to directions where the model is approximately invariant.
Computational Efficiency. For large models, use Monte Carlo estimates with small number of samples (c.g., one per data point per iteration) to approximate the expectations in the ELBO. The reparameterization trick ensures low-variance gradient estimates.
E ADDITIONAL EXPERIMENTAL DETAILS FOR 5.2
We ResNet-50 architecture with Bayesian linear layer at the end (for non- non- -on-asssiann case, just use ResNet-500 without any replacements). We apply standard preprocessing for IMAGENET. We use the Adam optimizer with learning rate of 1 x 10 10-5 for model parameters. In OPTIMA, learning a parameter in Beta distribution for Mixup, in uniform distribution for Cutmix, and Dirichlet, depth and Beta distribution parameters jointly for Augmix augmentations. For these augmentations, we use lognormal distribution as a prior because of the simplicity. The augmentation parameters have separate learning rate (1 x 10-3) to facilitate faster exploration. We regularize the augmentation parameters with a KL weight of SkLaug = 1, balancing data-fit fit and prior alignment. For all methods, we include a small KL weight Bua not 10-1 on the model parameters to maintain a Bayesian prior but it works with any weight on the model parameters This acts as Bayesian regularizer the final layer weights, preventing overfitting within that layer and ensuring consistency with the variational Bayesian framework (Blundell et al., 2015). Training proceeds for 30 epochs with a batch size of 256.
Evaluation Metrics. We compute the Expected Calibration Error (ECE) by dividing predictions into 10 bins based on confidence and measuring the difference between average confidence and accuracy in each bin:
10
[B&l n
ECE
Σ 1=1
(Bc)BB)) - conf( B:));
(65)
where B; is the set of examples in bin i, n is the total number of examples, acc(B:) is the accuracy bin and conf (B:) the average confidence in bin
For diitt--diitributtin detection, use the AUROC metric, which measures the area under the ROC
curve when using predictive entropy as the detection score:
H[p(y[=)] Σ ply loggyy = c
(66)
Higher entropy indicates higher uncertainty, which should correlate with out-of-distribution examples.
F ADDITIONAL RESULTS ON DIFFERENT TYPES OF DATA AUGMENTATION
F.I LEARNING GEOMETRIC AUGMENTATION FOR CIFARIO CLASSIFICATION
23

---

Published as a conference paper at ICLR 2026
Train Accuracy (%)
Test Accuracy (%)
Absolute Rotation (rad)
Absolute Translations
100
Aug.
0.06
Iwl
Itll It
90
Pized Aug:
Learned Aug.
0.04
80
70
60
0.00
10
30
20
Epoch
Epoch
Epoch
Epoch
Figure 3: (Left Two) Convergences of training and test accuracy on CIFARIO. OPTIMA generelizes better than the other approaches. (Right Two) Evolutions of the data augmentation parameters.
Setups. We use ResNet-18 architecture with Bayesian linear layer at the end. We apply standard preprocessing: nor- malization with mean (0.4914, 0.4822, 0.4465) and standard deviation (0.2023, 0.1994, 0.2010). We use the Adam opti- mizer with a learning rate of 10 for model parameters. In OPTIMA, we are learning 7 {w, lxz ( jointly, where w rotation (radians), and tz and by are horizontal and vertical shifts. The augmentation parameters have separate learning rate (1 x 10-3) to facilitate faster exploration. We regularize the augmentation parameters with KL weight of BKLaug balancing data-fit fit and prior alignment. For all methods, we include a small KL weight Sut not 0.1 on the model parameters
Table 5: CIFARIO classification.
Fixed Aug uses fixed rotation 0.1 and translations of 0.1.
Method
Acc (%)
ECE (L)
No Aug
80.90
0.092
Fixed Aug
80.73
0.088
OPTIMA
81.35
0.017
maintain Bayesian prior but it works with any weight the model parameters. This acts as Bayesian regularizer on the final layer weights, preventing overfitting within that layer and ensuring consistency with the variational Bayesian framework (Blundell et al., 2015). Training proceeds for 30 epochs with a batch size of 128.
Results. We assess calibration using 100 Monte Carlo (MC) samples. Fig. presents reliability diagrams for No Aug, Fixed Aug, and OPTIMA, revealing that the learned augmentation strategy yields the lowest calibration error (ECE), with the reliability curve closely aligning with perfect calibration. Table summarizes the final ECE values, confirming that OPTIMA leads to more accurate confidence estimates than fixed or augmentation. Moreover, Fig. (second panel) shows test accuracy over time: the learned augmentation generalizes better, while No Aug and Fixed Aug exhibit overfitting and poorer generalization.
F.2 EXPLORING DIFFERENT INTENSITIES OF GAUSSIAN TRANSLATIONS
We use the same implementation details as in Appendix F.I, except that we choose Gaussian Translation an augmentation parameter and validate on different values of K in naïve augmentation, and different prior variance a in OPTIMA. As an OOD data, we choose the SVHN dataset, since this mismatch makes it a widely adopted benchmark for OOD testing of classifiers trained on CIFARIO.
Table Effect of marginalization vs. naive augmentation with different numbers of augmentations example CIFAR 10 using Pretrained Bayesian R8 (last layer) and Gaussian Translation after 30 epochs. Test accuracy and BCE are on CIFAR 10, and OOD AUROC on the SVHN dataset.
Method
Test Acc (%)
ECE 1
OOD AUROC
Aug
94.09
0.0381
0.9069
naive Aug (K=2)
95.03
0.0298
0.9425
naive Aug (K=5)
95.21
0.0327
0.9383
naive Aug (K=10)
93.75
0.0424
0.9560
OPTIMA (011..1)
93.30
0.0192
0.9446
OPTIMA (a (o=0..5) 0.5)
93.87
0.0165
0.9576
OPTIMA (a=
90.25
0.0175
0.9647
Results. Table shows that OPTIMA allows us to get much better calibration than naive and augmentation cases. Because of the overcounting problem in naïve case, obviously consumes
24

---

Learnable Mixup Augmenter. In the OPTIMA Mixup variant, the Mixup hyperparameter (controlling the Beta distribution Beta a, a) from which the mixing coefficient À is sampled) is made learnable. We parameterize a Normal distribution over logit(a) with learnable mean pea and learnable log standard deviation log de.. where to logit(a). The initial value for Pıa is set to logit (0.2). corresponding to an initial Cinit = 0.2. The initial log déa. is set to log(0.1), promoting a small initial variance for the learned distribution over logit(a). A prior distribution p(logit(a)) is defined as N (llogit(amnn) oZ), where the prior standard deviation on = 2.0. The KL divergence between the learned variational posterior g(logit(a)) (ea, o2a) and this prior is added to the training objective, Sampled À values clamped to the Published as a conference paper at ICLR 2026
around K times more time than our approach demostrating that get good generalization and better robustness for OOD in short time with our approach.
G ADDITIONAL EXPERIMENT ON IMAGENET USING RESNET-50
Here we use OPTIMA for Imagenet in order to learn Mixup, Cutmix and Augmix augmentations.
G.1 IMPLEMENTATION DETAILS FOR OPTIMA WITH MIXUP
We evaluate our OPTIMA framework with the Learnable Mixup augmentation on the IMAGENET (Deng al., 2009) dataset for image classification. Performance is assessed on the standard IMA- GENET validation set, and robustness is measured on the IMAGENET-C (Hendrycks & Dietterich, 2019) benchmark.
Model Architecture and Preprocessing. We employ a ResNer-50 ar- chitecture (He
al. 2016). initialized with pretrained weights from IMAGENETIK_VE
torchvision models........o..mmm.mmmmm...... nected layer is replaced with a new linear layer mapping to the 1000 IMAGENET classes. For input preprocessing during training, images are transformed by RandomRes zedCrop to 224 x 224 pixels followed by a wanddmmorrizotalFFFae resized to 256 pixels their shorter edge and then center-cropped to 224 x 224. All images are subsequently converted to tensors and normalized using the standard IMAGENET mean MimageNet (0.485, 0.456, 0.406) and standard deviation I'mageNet (0.229, 0.224, 0.225).
The final fully con-
Validation and test images are
weighted by the hyperparameter beta_ satta__aummnttrree range [10 -6 - 10-6] for numerical stability.
Training Configuration. Models were trained for 10 epochs! using the AdamW optimizer.
The base learning rate for the ResNet-50 parameters was set to 1 x 10-4, The learnable parameters of the Mixup augmenter (Mea, log ota) utilized learning rate of 1 x 10-3 (10 times the base learning rate). A cosine annealing learning rate scheduler with warm restarts aesiiennneallngkaraeeeeee..
was employed, with parameters 'To 10f epochs, "Tmult = 2', and "etamin set to 1/100 of the initial learning rate. The weight decay for net- work parameters
wow_eettwwwrrrroooo was 0.01. The coefficient for the KL divergence term of the augmentation parameters with a global batch size of 256 distributed across 4 NVIDIA A100 GPUs using Distributed Data Parallel (DDP). We used a precision of '"16" (interpreted as 16-bit native mixed precision) and wwtplppprriciimmeeeeeeeie
(eeeeneerrrrr)e was 1.0. Training was performed
Gradient clipping was applied with a maximum norm of 1.0. The number of data loader workers was set to 8 per GPU process.
set "torch.set__ seaattt3_ 1130222mamall_11111111111
Baselines.
We compare OPTIMA Mixup against:
Fixed Mixup: Standard Mixup augmentation with a fixed a = 0.2. The training setup (optimizer, scheduler, epochs, batch size) was identical to that of OPTIMA Mixup, excluding elements specific to learnable augmentation parameters.
'Whiee longer training (e.g., 90-990001 100 epochs) is standard for IMAGENET, these experiments were conducted for 10 epochs to demonstrate the behavior of the leamable augmentation parameters and compare against fixed augmentation under identical short-run conditions.
25

---

Published as a conference paper at ICLR 2026
Pretrained ResNet-50 (No Augmentation Eval): The ResNet-50 model with weights from t torchvision sion. models ReaNet ReaNeet50_Weights IMAGENETIK v2, evaluated directly on the validation and IMAGENET-C sets without any fine-tuning under experimental setup. This serves as a standard reference.
G.2 IMPLEMENTATION DETAILS FOR OPTIMA WITH CUTMIX
For evaluating OPTIMA with CutMix, we follow similar experimental setup on the IMAGENET (Deng et al., 2009) dataset, with robustness assessed on IMAGENET-C (Hendrycks & Dietterich, 2019).
Model Architecture and Preprocessing. We use the ResNet-50 architecture (He et al., 2016) pretrained with torchvis on mowels........ ResNet Weight IMAGENETIK_ V2, replacing the final classifier layer for the 1000 IMAGENET classes. Input preprocessing during training includes RandomRest zedCrop to 224 x 224 and RandomHor zontalF lipip. Validation and test images resized (256 shorter edge) and center-cropped 224 x 224. Standard IMaGeNET normalization applied.
Learnable CutMix Augmenter. In CutMix (Yun al., 2019), a patch from one image is pasted onto another. and labels are mixed proportionally to the area of the patches. The mixing ratio X (determining the area of the first image to keep, and thus (1-X) the area of the patch from the second image) is typically sampled from a Beta(ex, ar) distribution. For our OPTIMA CutMix, this parameter of the Beta distribution is made learnable. We parameterize a Normal distribution over log(a) with learnable mean Mlog o and learnable log standard deviation log olog The initial value for Mlog is set to log(1.0), corresponding to an initial Ciii = 1.0 (a common default for CutMix). The initial log o is set to log(0.1). prior distribution p(log(a)) is defined as A (log(amn ) o with prior standard deviation op 2.0. The KI. divergence between the learned variational posterior for lo and this prior is added to the training loss, weighted by beta augmenter__ reg. The sampled values are clamped to [10-4, 100.0] before being used in the Beta distribution. The resulting mixing coefficient Atiall (coefficient for the first image's label) determined by the actual area of the pasted patch after clipping to image boundaries.
Training Configuration. Models were trained for Nepochs epochs (e.g., 15) using the AdamW opti- mizer. The base learning rate for network parameters was x 10-4, while the learnable CutMix param- eters (Hiog o log llg used learning rate of x 10-3 A ieeosinennnaall wnniin www.starrrttrrrrr learning rate scheduler was used (TT_0=10² or '15', 'T_mult=',, 'eta_min" '=1110 of ini- tial LR). Network weight decay (beta_networkk__
was 0.01, and the KL coefficient (beta_auu augmenter reg) was 1.0. Training used a global batch size of 256 on 4 NVIDIA A100 GPUs with DDP, '116" precision, 'torch.set float32 2 matmul praciisiorriiionttte ("medium') and gra- dient clipping at 1.0. Data loader workers were set to 8 per GPU.
G.3 IMPLEMENTATION DETAILS FOR OPTIMA WITH AUGMIX (LEARNABLE SEVERITY + JSD)
We evaluate OPTIMA by learning component of the AugMix (Hendrycks al., 2020) augmenta- tion strategy, specifically its overall aug_severi tty, while also employing the Jensen-Shanme Divergence (JSD) consistency loss. Experiments conducted IMAGENET (Deng al., 2009) and IMAGENET-C (Hendrycks & Dietterich, 2019).
Model Architecture and Preprocessing. The model ResNet-50 (He al., 2016) initialized with parchvision models. Re _ett0__Weehhhstteeeee IMAGENETIK V2, with the final classifier layer adapted for 1000 classes. During training, input PII. images undergo RandomRest zedCrop These PII. images are then passed to our module. Validation and test images use stan-
to 224 x 224 and HandomHorizontalll. Learnable www.aaatwieeeeertyyyyy.raeoeeeea dard resizing, center cropping, and ToTensor conversion, followed by IMAGENET normalization.
Learnable AugMix Severity + JSD Augmenter. implemented an Module.
The Telleeernbbeeeallls.ssoetttstsnn
26

---

Published as a conference paper at ICLR 2026
For each imput PIL image, three views are generated: the original, and two independently augmented versions using AugMix. Each AugMix version a convex combination (m ~ Beta(1, 1, 1)) of the original image and a mixture of K augmix mixturt ___d__ttth (default 3) augmentation chains. Fach chain consists of D (default random 1-3, controlled by augmi x__mx__mm- www.iituueeee.mm xttue_depte)) basic operations (e.g., rotate, shear, color jitter) sampled randomly. Mixing weights for chains wi are from Dirichlet 1). All PIL operations and the final conversion to tensors (for each of the three views) happen within this augmenter module, ensuring output tensors are on the correct device. We utilize predefined list of tensor-basee based augmentation operations where possible improve performance over PlL-only operations.
AugMix Core:
Learnable Severity: The overall intensity of the basic augmentations, _is_severity (typically a value between 0-10), is made learnable. We parameterize a Normal distribu- tion over (asgaaugsveert
with learnable mean Mlog and learnable log standard deviation log Olg s. The initial Plog corresponds to an initial anitial_aug_sete tial log olg = log(0.1). The prior for org(aug_severigy seeeerrtt_st__eeeeeeeee
of 3.0, and ini- N (Ioo(inittial aug severity, Pss) euaug (default 1.0). A KL divergence term, regularizes these learned severity parameters. The sam- pled severity is clamped to [0.1, 10.0].
with 0p5 weighted by beta_aagmeeeeeeeeeeeeeeee www.aaameeteerreee
JSD Loss: The three output image tensors (original, AugMix view I, AugMix view 2) are passed through the network. A JSD consistency loss is calculated between their softmax predictions, weighted by beta_jjd (default 12.0), and added to the primary cross-entropy loss (calculated on the original view).
The data loader for training uses a custom collate function to provide a list of PIL images to the augmenter.
Training Configuration. Training was conducted for Nepochs epochs (in our case 6 epochs because of the computational complexity related tensor and PII. tranformations) with the Adam optimizer. The base learning rate for network parameters was 10-1. The learnable severity parameters (Hiog s, log Slg s) used a learning rate of 1 x 10 -3 A CosineAnneal wwweeeaiiniaaaaaaaaaeeeea scheduler was used ('T 0T_0=10' or '15', T_mult=2', 'eta_min [oa_min"=//100 of initial LR). Network weight decay (beta_ seeet_neework_ree
was 0.01. The KL coefficient for severity parameters (beta_aue www.aaagggeere.. reg) 1.0. The JSD loss coefficient (beta_ j_sd was 12.0. Training used a global batch size of 128 (reduced due to processing three views) on 4 NVIDIA A A100 GPUs with DDP, '16' precision, 'Morch.se__flou33 [aeee_/oool___mmmmma_am mnalmallpprecssiii meddiimm)) and gradient clipping at 1.0. Data loader workers were 8 per GPU.
G.4 EVALUATION AND SOFTWARE/HARDWARE FOR ALL THESE METHODS.
Models are evaluated on the standard IMAGENET validation set for top-1 accuracy and cross-entropy loss. Robustness is assessed on the IMAGENET-C benchmark, reporting the normalized mean Corruption Error (mCE normalized by AlexNet baseline) across all corruptions and severities. For IMAGENET-C, images are processed using the same validation transforms as for the clean IMAGENET validation set. All final reported evaluations are performed on single GPU using 32-bit poating-point precision.
Experiments were conducted using PyTorch version 2.0.1 and PyTorch Lightning version 2.0. Training and evaluation four utilized NVIDIA A100 (80GB) GPUs.
G.5 EXPERIMENTAL RESULTS
Table The result of IMAGENET and IMAGENET-C using ResNet-50 for each augmentations. We evaluate the average test error for each corruption type
glon
kig
PP
NN Auy
B0.8
T6-7
71
13
35
81
34
57
62
S2
32
6)
33
35
QITTARA (0oge/10
71.11
55
33
33
33
33
30
62
51
31
48
33
30
OPTTINA Cames movilc
2A62
T0V
66
687
428
716
826
757
754
717
788
96
SE1
613
816
705
682
OMP specter
78.09
$6.1
5137
SCS
6U/
SUR
SKT1
SESA
9W/6
54.84
44
3019
45'7
S4.5
5:15
27

---

Published as a conference paper at ICLR 2026
that OPTIMA allows us get better test aceuracy clean data with non- non-Bayeen B---ssaysian ResNet-50. OPTIMA Mixup, Cutmix and Augmix are beating the baseline results within (5-15 epochs). The mCE of OPTIMA Augmix is lower than the benchmark. This can be explained by the very few training epochs (6 epochs) which we could run due to the computational complexity of this experiment; with computational resources it takes around 15 hours for one full epoch.
Results. In Table 7,
H ADDITIONAL EXPERIMENTAL DETAILS FOR $ 5.4. TOKEN-DROPOUT IMPLEMENTATION DETAILS
Parameterization.
The augmentation module applies token dropout with probability
Pdrop Pmax 0 (8),
where s is a trainable scalar (initialized at = - in our implementation) and ar-) is the logistic function. The constant Pmux sets an upper bound on the amount of dropout; we use Puax = 0.5.
Prior. We place a Gaussian prior directly on Pdrops
p(Prrpp) exp
(Pdrop 23²
implemented as a quadratic penalty in the ELBO objective. We use μ {0.1, 0.3) represent weak and strong prior preferences for token dropout, and a = 0.1.
OPTIMA initialization. All methods begin from the same initial dropout rate Pärop Pmax (so) se 0.04, ensuring mild initial augmentation.
Baselines. We compare the following: (i) No Aug (Parpp = 0); (ii) Fixed Aug, using the same initialization OPTIMA; (iii) Fixed Aug (Matched), using OPTIMA's learned dropout; (iv) BO-Fixed, selecting Parop via validation NLL over a grid in [0, 0.3).
Optimizers. DistilBERT is trained with learning rate x 10-5 and is trained with learning rate 5 x 10-2, Further hyperparameters are unchanged from standard HuggingFace defaults.
I BROADER IMPACT
OPTIMA has significant potential for improving the reliability of Bayesian deep learning in high- stakes applications. such as medical imaging, autonomous driving, and scientific discovery. The ability learn optimal augmentation strategies from data also reduces the need for manual tuning, making Bayesian methods more accessible to practitioners across domains.
J REPRODUCIBILIITY STATEMENT
All experiments are fully described in the submission, including dataset details, hyperparameters, and training procedures. The accompanying code is provided to ensure that our results can be independently reproduced.
K THE USE OF LARGE LANGUAGE MODELS (LLMs)
We used large language models (LLMs) solely for non-substantive assistance, including grammar refinement and summarizing relevant literature, All research ideas, analyses, and conclusions are the authors' own.
28


Reviews:
[{"review_id": "tHZUyPuPgS", "reviewer": "Reviewer_5Upk", "summary": "This work proposes a Bayesian optimization strategy to optimize data augmentation parameters; while the marginal likelihood is intractable, the work proposes a tractable ELBO which variational approximation to the joint $q(\\theta, \\phi) = q(\\theta)q(\\phi)$ and introduces regularization terms. The authors include several analyses of this framework, namely: (1) a characterization of the quality of the variational inference approximation in Section 4.1, (2) a PAC-Bayes generalization bound in Section 4.2, (3) a second-order bound on the expected difference in the model’s output under certain transformations (which are perturbations added to  with mean 0) in Section 4.3, (4) a comparison of naive DA compared to the true marginalization (not the variational inference approximation) in terms of covariance of the true posteriors being smaller under naive DA by a factor of K (the number of transformations applied to each data point) in Section 4.4, (5) interpret the solutions which maximize the augmented ELBO as empirical Bayes solutions for $p(\\theta | D)$ and $p(\\phi | D)$ in Section 4.5, and (6) show a reduction in posterior distribution entropy under the true marginalization (not the variational inference approximation) compared to having no augmentation in Section 4.6. The authors also employ their method on synthetic data, ImagetNet, ImageNet-C, and CIFAR-10 to demonstrate their theoretical findings.\n\nNote: I have reviewed an earlier version of this manuscript, and see that the Appendix F.3 has now been moved to the main body to compare the computational efficiency this work to Bayesian optimization. There also appears to have been effort to better relate the experiments to theoretical results, namely top of page 9. Some other changes which had been promised but not applied are:\n- Relation to Chen et al. 2020 (and maybe other works, like Chatzipanzis et. al, 2021) in Section 2. This may fit under probabilistic perspectives of DA.\n- Some mathematical precision comments by a couple reviewers, e.g., simply defining terms or noting the overloading of notation would suffice. I would argue that these imprecisions should be clarified explicitly in the manuscript or addressed.\n    - For example, the proof of Theorem 4.5 needs K (the number of augmentations/replications) to be large enough that the naive empirical risk is close to the true one. I still think this is a bit imprecise, but the minimal changes that the authors can do is state “for K large enough” in the statement of Theorem 4.5.", "questions": "- Why do the authors choose to report their results for Cutmix, etc. reported in Table 2 and have not done the same comparisons, for instance, the setup in Appendix F.2 which uses Gaussian Translation. I think it would confuse some readers why ResNet is achieving under 80% accuracy on these datasets.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- The theoretical results are extensive and address multiple aspects of this framework. The work appears to give practical insight on what relevant quantities and choice of distribution over data augmentation parameters affect various aspects of model performance within this framework of data augmentation via Bayesian optimization. The corollaries written throughout the text after theorems give insight and practical advice, and seem to cover many different lens through which to understand how this framework (alternating between the variational inference approximation and the true posterior under this framework) compares to naive DA and having no augmentation.", "weaknesses": "- The noted mathematical imprecision comments from previous reviews.\n- Some experimental results can be confusing to some readers. In Section 5.1, Figure 2, the “test loss” figure is unclear and does not appear to clearly show what the text is saying. Also, perhaps a reference to Appendix F.2 or including it in the main body would be appropriate and explain why ResNet is getting less than 80% accuracy in Table 2.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CCvfuAU7xc", "reviewer": "Reviewer_oKY4", "summary": "The authors propose OPTIMA, a Bayesian framework that optimizes data augmentation parameters by maximizing the marginal likelihood instead of relying on validation-based hperparameter tuning. \n\nThey validate their approach with theoretical analysis based on information-theoretic principles, to show that the method is probabilisticly sound. \n\nThey evaluated on ImageNet, ImageNet-C, and CIFAR-10 and shows that OPTIMA achieves better calibration and competitive accuracy.", "questions": "How does OPTIMA fundamentally differ from existing bilevel optimization approaches for data augmentation, such as Mounsaveng et al. (WACV 2021) [1]?\n\nCan the proposed framework be extended to more complex tasks like segmentation or detection, and what challenges would come up in doing so?\n\nHow were  the hyperparameters chosen on datasets without a validation set (e.g., CIFAR-10)?\n\n[1] Mounsaveng et al., Learning Data Augmentation With Online Bilevel Optimization for Image Classification, WACV 2021.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The authors address the interesting, useful task of data augmentation as a technique that helps improve model performance across different downstream tasks. \n\nThe authors provide solid theoretical foundations for their proposed method, OPTIMA, highlighting the types of invariances it promotes and the uncertainty quantification it enables. \n\nThey also present experiments on standard datasets such as ImageNet, showing that OPTIMA slightly outperforms methods that do not use data augmentation", "weaknesses": "Methods like bilevel optimization [1] also aim to optimize data augmentation functions, and it would be nice to show how this paper differentiates OPTIMA from these related approaches in terms of how good they perform augmentation.\n\nThe field of image classification is already quite saturated; it would be more compelling to see results on harder tasks such as segmentation or object detection.\n\nWhile data augmentation optimization is valuable, it doesn’t provide fundamentally new information, unlike larger supervised models (e.g., transformers) that may implicitly learn augmentation effects by training on a large corpus in a semi-supervised way.\n\nThe reported improvements on CIFAR-10 and ImageNet appear insignifcant, and given how saturated these benchmarks are, it’s unclear whether the gains are statistically meaningful.\n\nThe paper lacks standard deviations or error bars so it is hard to know if any results are actually significant.\n\nThe implementation details are unclear, how were the hyperparameters selected/tuned without a validation set (like in CIFAR10)?\n\n[1] Mounsaveng et al., Learning Data Augmentation With Online Bilevel Optimization for Image Classification, WACV 2021.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tez6jNcbzR", "reviewer": "Reviewer_stCk", "summary": "This paper proposes the OPTIMA framework, which optimizes data augmentation parameters by using Bayesian methods. The core idea is to treat DA parameters as latent variables and introduce variational inference methods to simultaneously optimize both model parameters and DA parameters. An extensive theoretical analysis is conducted, including generalization guarantees and invariance properties. Experiments demonstrate that OPTIMA enhances confidence and robustness in regression and classification tasks.", "questions": "Q1: How does OPTIMA perform in non-computer-vision domains, such as time series prediction or natural language processing?\n\nQ2: How high a dimension of data augmentation parameters can the proposed method handle?\n\nQ3: The comparison algorithms selected in this paper were mainly published around 2020. Considering the rapid development in the field of automatic data augmentation, many more representative methods have emerged in recent years. It is recommended that the authors supplement the comparison experiments with cutting-edge methods proposed in the last year or two (e.g., 2022-2024). This will enable a more comprehensive and fair evaluation of the proposed methods' advancement and effectiveness, thereby significantly enhancing the persuasiveness of this paper's contributions.\n\nQ4: What are the innovations compared to the traditional Bayesian selection optimization algorithms based on variational inference?\n\nQ5: The proposed method may get stuck in local optima. Have convergence guarantees or initialization strategies been considered?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "-The paper proposes a Bayesian framework for the joint optimization of data augmentation parameters and model parameters.\n\n-The theoretical analysis is comprehensive, including approximation bounds, generalization guarantees, and invariance derivations.", "weaknesses": "-The paper primarily focuses on computer vision tasks (such as image classification) and does not extend to other modalities (such as natural language or time-series data), which limits its generalizability claims.\n\n-The compared algorithms are outdated.\n\n-The impact of the dimensionality of data augmentation parameters on the proposed method is not discussed.\n\n-Some theorems rely on mathematical assumptions that are difficult to apply in practice.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "m6T6mxiDDK", "reviewer": "Reviewer_8dA2", "summary": "The paper introduces a Bayesian model selection framework for online optimization of the data augmentation (DA) hyperparameters, provides a tractable ELBO for variational approximation of the marginalized likelihood, performs extensive analysis of the theoretical characteristics of the augmented model, uncertainty quantification, calibration, decision boundary smoothness and generalization. The theoretical results are backed by several experiments with synthetic and image datasets optimizing the DA parameters of simple and also state of the art baselines in DA and improving the prediction and calibration with almost no additional computational cost.", "questions": "Would the proposed method also work for when transformations are not continuous as in natural language e.g. based on word substitutions and syntactic transformations? \n\nCan the authors provide the intuition of THM4.5? In the the naïve DA the variable $gamma$ is integrated out in $R(\\theta)$ with a uniform empirical distribution as opposed to the latent distribution $p(\\gamma) = \\int p(\\gamma|\\phi)p(\\phi|\\theta)d\\phi$ or its variational approximation $p(\\gamma) = \\int p(\\gamma|\\phi)q(\\phi)d\\phi$?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "The paper is well written and easy to follow. The authors have laid out the background and  the motivation and provide good structured explanation of the proposed method and its theoretical characteristics. Several experiments are provided to compare the Bayesian DA's performance to other baselines and highlight the theoretical results.\n\nThe authors use DA hyperparameters as latent mode parameters and use a Bayesian framework with ELBO optimization to efficiently find the best DA parameters according to the data invariance and model sensitivity. In the theoretical analysis, Bayesian DA method is compared to naïve DA in terms of the PAC-Bayes generalization bounds, regularization and smoothening effect of Bayesian DA and improvement in calibration of the predictions. \n\nThe analysis also sheds light on the connections of DA with information bottleneck which I find particularly novel.", "weaknesses": "It might be easier to follow the paper if a brief explanation of the theoretical results of the Theorems are provided before the theorems, as there are multiple theorems and moving from one technical result to the next requires some context in between. \n\nMinor comment: \nTypo: line 107 wrt -> with respect to\nTypo: line 119 nav̈e -> naive\nTypo: line 229 extra s", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)