Learning Augmentation Distributions using Transformed Risk Minimization
Authors: Evangelos Chatzipantazis
Stefanos Pertigkiozoglou
Kostas Daniilidis
Edgar Dobriban
Venue: Accepted by TMLR ()

Abstract: We propose a new \emph{Transformed Risk Minimization} (TRM) framework as an extension of classical risk minimization. 
In TRM, we optimize not only over predictive models, but also over data transformations; specifically over distributions thereof.
As a key application, we focus on learning augmentations; for instance appropriate rotations of images, to improve classification performance with a given class of predictors. Our TRM method (1) jointly learns transformations and models in a \emph{single training loop}, (2) works with any training algorithm applicable to standard risk minimization, and (3)  handles any transforms, such as discrete and continuous classes of augmentations. To avoid overfitting when implementing empirical transformed risk minimization, we propose a novel regularizer based on PAC-Bayes theory. For learning augmentations of images, we propose a new parametrization of the space of augmentations via a stochastic composition of blocks of geometric transforms. This leads to the new \emph{Stochastic Compositional Augmentation Learning} (SCALE) algorithm. The performance of TRM with SCALE compares favorably to prior methods on CIFAR10/100. Additionally, we show empirically that SCALE can correctly learn certain symmetries in the data distribution (recovering rotations on rotated MNIST) and can also improve calibration of the learned model.

Published in Transactions on Machine Learning Research (07/2023)
Learning Augmentation Distributions using Transformed Risk Minimization
Evangelos Chatzipantazis"
voghat upenn.eeeee.eeeeeeeeeeee
Deportment of Computer and Information Science University of Pennsyloania
Stefanos Pertigkiozolon"
pelefaano@scas.upeee
Department of Computer and Information Science University of Pennsylania
Kostas Daniilidis
kostas@eis. www.ii..speec.... edu
Department of Computer and Information Science University of Pennsyloania
Edgar Dobriban
dobribennuwaartn..uppee edu
Department Statistics of Statistics and Data Science
Universily of Peonsyloania
Reviewed on OpenReview: https: // openreview. net/ form?? id= LRYNN
Abstract
We propose new Transformed Risk Minimization (TRM) framework as extension of classical risk minimization. In TRM, we optimize not only over predictive models, but also over data transformations: specifically over distributions thereof. As a key application, we focus on learning augmentations; for instance appropriate rotations of images, to improve classification performance with given class of predictors. Our TRM method (1) jointly learns transformations and models in single training loop, (2) works with any training algorithm applicable to standard risk minimization, and (3) handles any transforms, such as discrete and continuous classes of augmentations. To avoid overfitting when implementing empirical transformed risk minimization, we propose novel regularizer based on PAC-Bayes theory. For learning augmentations space of augmentations via. a stochastic composition of blocks of geometric transforms. This leads to the new Stochaslic Compositional Augmentation Learning (SCALE) algorithm. The performance of TRM with SCALE compares favorably to prior methods on CIFARIO/IN 100. Additionally, show empirically that SCALE can correctly learn certain symmetrics in the data distribution (recovering rotations rotated MNIST) and can also improve calibration of the learned model.
images, we propose a new parametrization of the
1 Introduction
Adapting lo the structure of data distributions (such as symmetry and transformation invariances) is an important problem in machine learning, Invariances can be built into the learning process by architecture design (such as in convolutional neural nets), or by Traditionally, these approaches require full knowledge of the symmetries of the data distribution. Absent this knowledge, practitioners may mis-specify the set of augmentations, and time-consuming tuning. The current practice of machine learning largely depends on handcrafted aug- mentations, based on extrinsic knowledge about the task to be solved. In contrast, the "best" " augmentations
traanninnntttaannonnnnnggggg the datasct during training.
may need to resort to expensive
"Equal Contribution.

---

Published in Transactions on Machine Learning Research (07/2023)
may also depend on the predictive model, not just on the task; in possibly non-obvious ways. The space of augmentations is vast; and these factors make human fine-tuning hard. Therefore, it could be more efficient to learn augmentations from data. There have been several prior approaches to this problem, but none that satisfy our desired criteria described below.
Our goal is to design an algorithm that learns data transformations jointly with solving downstream task. For instance, we want learn which augmentations rotations or flips are most helpful while learning to classify images. In addition, we want to learn the level of data Pransormationn--matinen angle that will be beneficial for the task. We require that the algorithm is:
rotation
1. Efficient: The goal is to jointly learn augmentations and models in a single truining loop (without a significant increase in training time compared to standard model training).
2. Modular: The method should pair with any learning problem and algorithm for the downstream task. For example, it should be independent of the neural network architecture in an image classification lask and should be incorporated without significant changes into a standard training algorithm for risk minimization.
3. Applicable to any transforms, for instance to both continuous and discrete-valued augmentations. This is challenging, as previous approaches have used specialized approaches (Cubuk et al., 2019a; Benton et al. 2020): for discrete augmentations, one can sometimes enumerate them; for continuous ones one may leverage Lie group theory.
Our approach to this problem (in Section 2 is, intuitively, to expand the usual model space (e.g., neural networks) to include first "stochastic layer" where we optimize over distributions of transformations using novel objective. In the framework of statistical learning theory, the transformations induce a so-called transformed risk: Thus, we formulate our algorithu in a new theoretical framework called Transformed Risk Minimization (TRM). This enables us to leverage the theory of risk minimization and design a unified a algorithm that satisfies all three criteria discussed above. We can summarize our contributions as follows:
We formulate the Transformed Risk Minimization framework (TRM) and connect it to standard risk minimization. We prove a zero-gap theorem, implying that when there is an unknown distributional invariance in the data, TRM can recover both the invariance as well as model that is optimal under the standard risk when the distributional invariance is known.
2. We then propose the new Stochastic Compositional Augmentation Learning (SCALE) algorithm that aims to optimize the transformed risk for finite data. SCALE consists of a new parametrization of the augmentation space via a stochastic composition of blocks and a novel regularizer for training which is derived using PAC-Bayes theory.
We provide experiments with SCALE. Our first experiment shows that on RotMNIST, SCALE learns the correct symmetries and outperforms previous methods; which supports our zero-gap theorem. We also perform experiments on CIFAR 10 and CIFAR 100, showing that SCALE has advantages in both accuracy and calibration compared to prior works, while maintaining the benefit of time efficiency.
2 Transformed Risk Minimization
2.1 Preliminaries Risk Minimization
In conventional supervised statistical learning [Vapnik] 2013). the learner is provided with independently i drawn sample S = i E [n]], where [n] = 2,.
and (X,, Yi) E X x y for all E [n], from an unknown data distribution D. Further, the learner has access to an action space y and a hypothesis space H. C (2**)* containing candidate hypotheses (models) h X' + y" In general, y" might not be equal to y. For example, in a multi-class classification problem, y can be the discrete set of classes, and y can
... n.,
be the set of probability distributions over them.
All objects are measurable with respect to appropriate sigma-algebras, which are used tacitly our work.
2

---

Published in Transactions on Machine Learning Research (07/2023)
(a) Training Data
(b) Distribution with Standard Augm.
(c) Learned Augm.
Distribution with TRM
(d) Learned Model with (e) Learned Model with (f) Learned Model with Standard Augm.
RM
TRM
Figure When only a small amount of training data is available (as shown in (a)), models trained with standard risk minimization (RM) may lead to poor generalization. A common practice to improve the generalization is to perform data augmentation, but this can be harmful when the appropriate distribution of angmentations is not known (as shown in (b), (e) where we apply rotations of 180%). We propose the framework of Transform Risk Minimization (TRM) to learn both the model and the optimal distribution of augmentations. As shown in (c), (f), TRM recovers the appropriate distribution of augmentations (rotations of 45%) and learns an improved model.
Based on the training data S, the learner outputs hypothesis hy whose performance on a test datapoint (z, v) ~ D is evaluated via a loss function y" x y + R. The goal in Risk Minimization (RM) is to minimize the following risk over h E H:
RAM (h) EXX,)>DD(((X),,))))))))))))))()) V))-
(1)
One may aim to control the risk Russ (hs) of the data-dependent predictor hg either in expectation or with high probability with respect to S.
2.2 From fixed augmentations to an augmentation hypothesis space
A common way improve performance is by leveraging domain knowledge in the form of transformations that leave the prediction invariant (c.g., in classification), equivariant trrnsanonnttnn in pro-specified way (e.g., in pose estimation of complex shapes). When the set of transformations admits group structure, it is beneficial to constrain the hypothesis space to rerriiat/eqqvvariant models. However, for more general transformations, the most common practice is data augmentation: adding transformed data to training, If tuned properly, data augmentation can improve performance. However, it assumes that the user fully knows the beneficial transformations. This assumption general not valid; and in practice augmentations require a great dcal of tuning.
To relax this assumption, switch viewpoint and suppose that the user only has certain level of incomplete knowledge about the beneficial transformations. We then aim to learn those transformations jointly with the model. Intuitively, one may think that one should optimize over specific transformations
x + g(z). However, we find it beneficial to relax to distributions over transformations, as this is both more general- includes point masses at specific transformations and more closely mirrors standard stochastic data augmentation. Thus, we consider modification of the classical lcarning paradigm which an additional transformation hypothesis space Q is provided to the learning algorithm, containing candidate transformation distributions Q over sets of transformations GQ- Then, denote the action of g Gq on
3

---

Published in Transactions on Machine Learning Research (07/2023)
X as g. (instead of g(X) P Now the goal of the learner is to output not just model hg E H. but also distribution Qs E S2 of transformations.
As an example, one transformation distribution Q E Ω might be the uniform distribution over all rotations in the range (-309,300 degrees, i.., Q= Un(-309, 30,, 30") set of all such uniform distributions parametrrizd by their endpoints, e.g., f = (Un(0), 02) 01 < 02. I. f, 0 E [-180", 180"]] The sampled transforms g U= 30") act on their inputs by rotation. For 2D image inputs, they are represented by specific
& The distribution hypothesis space Ω could be the
image I : IR2 - R (2) 2 (2), == [mg-) In this example, we learn the range of the helpful rotations, as well as model that fits the transformed data.
rotation matrices transforming
2.3 Transformed Risk Minimization
A first thought may be try to optimize the standard risk in equation I for the expected prediction over the unknown Q Ω2 i.e., R(h, Q) (xx,))= (xx.y)> [((g)oll((XX)) Y)]- However, this objective is not amenable to standard stochastic optimization due to the difficulty of obtaining unbiased gradients with respect to Q.
In this work we study a notion of transformed risk that bypasses this issue. The goal of the learner in Transformed Risk Minimization (TRM) is to minimize the transformed population risk
RenM (h, Q) == Bxx.))-DEy- Y)]
(2)
simultaneously over augmentation distributions Q E l2 and models H. Figure shows an example of the difference between models learned using the standard risk minimization and our proposed TRM framework.
This objective is inspired by how data augmentation is applied in practice with fixed distribution Q [Baird 1992 Chen et al., 2020). In particular, data augmentation can be viewed averaging the loss over the augmentation distribution. In our setting, this distribution is learned. For losses such that l(,,)) is convex for all y E V, such as typical surrogate losses for classification like the cross-entropy loss, Jensen's inequality shows that R(h, Q) I BrRas(h, Q). Thus a small value of the TRM objective implics that the more difficull-tohhandlee risk of the average model also small. Given a finite dataset, empirical transformed risk minimization (ETRM) would aim to optimize:
n
Ry RMM (h,, Q)
Σ (h(g)((())
over h E Q.Q E Ω. However, as we discuss in Section 3.2 this training objective not good estimate of the population risk (equation 2 as collapses to trivial distributions. In the next section, we discuss how to optimize the population risk by proposing upper bound to equation 2 that is computable from the training data and takes care of this problem.
3
The SCALE Algorithm for Learning Data Augmentations
In this section, building our TRM framework, we design the new Slochastic Compositional Augmentation Learning (SCALE) algorithm, shown in Figure
3.1 Parametrization of Q and H
In practice, we do not optimize over an rast space (H, but over a parametric space (W, O). Considering the complexity of our tasks, we will use rich hypothesis spaces Hussuc as deep networks withh with parameters w for some parameter space W C Rp,pe E i.e., H = +/w; w E W).
To parametrize the augmentation distributions Q in Ω we leverage the rich compositional structure of transformations. Complex transformations can be constructed by compositions of simpler transformations sampled from parametrized base distributions DE We can also include in the composition fixed base
II general, g: + X', h X' +y" We simplify X' X' for clarity.
The symbol U will denote uniform distribution over an appropriate space.
4

---

Published in Transactions on Machine Learning Research (07/202) /023)
Re, F2. (01-0001,00
Prediction: Og (gr)..
Ru (011001111
Horizontal
2-1
by
MUX
Scaling Sy
MuX
qar.
Flip
hw
Training:
BBern(or))
t2 Bern(iz))
Bern((z))
(, (1,00,
Eg- [(hh v)
Flip Probability
Sample
Sample Sy
Reg(O (x))
(055(
-2
-83
Figure 2: SCALE: The computation of gr from random sample g Qo composition of blocks applying randomly sampled transformations to the input. The 7th block in the sequence transforms its input with probability and leaves it unchanged with probability (1 - TT), described in the figure via 2-1 multiplexers "MUX" with probabilistic selectors. For each imput the model samples several gz-s and uses the outputs h(gr) to compute the expected loss during training, and the expected output at test-time inference. During training, we learn the parameters of both the model and the distribution of augmentations by minimizing the expected loss plus the regularization term we derive in Section 3.2
distributions { i- Additionally, we allow choosing which base distributions should be included in the composition independently for each i. This is discrete optimization problem where the optimization variable can take exponential number of values in the number of base distributions. We associate to
each base distribution variable E {0,1} indicating if the distribution is included in the composition. To handle the combinatorial optimization problem, we relax m, to continuous mixture components क. E [0, 1]. E [K] = {l....,K]. K]. Thus, we define Q8 to be the distribution induced by sampling ~ Qo as:
91 92 gk, g: Q i 1=,,..,
(4)
where 0 == ( RK, .1,...,,000 0% and gi ~ Q is:
with probability n; a transformation sampled from a distribution U. This distribution can have learnable parameters ag, or can be fixed (the latter denoted as U:);
2. with probability
n;) the identity transformation leaving the input unchanged.
With this parametrization we can handle compositions of both discrete and continuous distributions of transformations. This is an advantage compared to previous works, as we discuss in the Related Work section.
We will focus on base augmentation distributions Us.,; supported on transformations of the form p:(a), with sampled from the uniform distribution parametrized by a via a ~ U([-a,, ax)), where 0 < a; S A, for some fixed A; > 0 We choose parametrization of our transforms such that P:(0) is the identity map. For instance, (:(a) can be a rotation by degrees with sampled uniformly over the interval 1-30%, 30 We will also consider parameter-free base distributions U,, uniform on some discrete set containing the identity transformation (such as horizontal flips, or cropping). If the first & out of K base distributions are parametric, then the base mixed distributions Q are constructed as
ne )ō + n, -, 0; as).
isk
Q (1 म. )00 + U,0, N 13), k k<isk,
where a; θ, := (0, and m. [0, are the parameters of the mixture we optimize over. Then, θ = 1K x 11 O; and for each E O, cach Qo is constructed as the product measure Qo = II Q
K
NNon-symmetric endpoints handled similarly. We overload notation here. Q6) was supported common subtlety and how can lead tighter bounds the Appendix.
transformations pi(a) the parameters We explain this
5

---

Published in Transactions on Machine Learning Research (07 (07/2023)
on Ii-il-a; a:) x II==+1 + ,0,...,,N; 1}. Finally the augmentation hypothesis space is f2 {Qe : E ).
Distributions of transformations where the parameter is sampled over closed interval [-o;, 0:] are common in practical data augmentation. However, in contrast to their standard usage, here we can learn the endpoints.
The task to learn parameters (w, 0) E (W,O) such that (hw, Q6) E (H, results in a small TRM risk
3.2 Leveraging PAC-Bayes theory for training SCALE
In this section design algorithm aiming to optimize the transformed risk given finite sample. Naively optimizing the empirical TRM criterion (Eq. 3 would typically result in trivial solution for the augmen- tation distribution Qo. namely the Dirac mass around the identity transformation. This is not surprising, since augmenting the initial dataset typically results in worse training error, even though it might result in better generalization; see c.g., Benton ct al. 2020 etc.
Thus, regularization of high importance for identi the useful augmentations. We propose to leverage PAC-Bayes theory (McAllester 1998 Shawe-Taylor & Williamson 1997: Catoni 2007) bound the TRM risk and to obtain novel, thcorctically motivated regularizer.
Due to space limitations we will assume some basic familiarity with PAC-Bayes theory. We provide bound on SCALE building on a PAC-Bayes bound (Alquier et al. 2016) that holds for bounded and some un- bounded losses (c.g., sub-gaussian losses (Germain et al. 2016bb). Since we aim for generality and efficiency in our algorithm, we do not use bounds from (McAllester 2013 derived only for bounded losses or data- dependent priors (Dziugaite & Roy - whhic usually require nested training loops. In our experiments, our regularizer is essential for recovering useful augmentation distributions; even though the bounds them- selves may not provide non-vacuous generalization certificate. After all, we use deep networks for which non-vacuous bounds are uncommon.
Our generalization bound relies on the following reduction: instead of viewing data augmentation as a transformation of the data followed by the application of a hypothesis h E H. we view the entire process as randomized predictive model. For each (h, Q), the transformation distribution Q acts as first "stochastic layer" before applying a deterministic model h. Hence, Eq. 2 becomes the expected risk of a stochastic classifier h h 9, 9 Q. For each input x, the classifier h first draws sample g ~ Q and then predicts (h (z).
Let PH NWw (2), Qu = NWw, I² for some parameters wa, E W be multivariate normal distri- butions with scaled identily covariance matrices. Let Qa be the probability measure constructed in Sec- tion 3.1 and, similarly to Qs, let P be the product measure PA I P(o on A = I- A.] x
II. = 10, (0,.. ... N: 1}, where
380 (1 111--901 B)U[-A, A.),
isk
P()
360 + -LU-OO N₂ 1H,, k k<isK
for a fixed 3 E 10, In practice, we choose & to be a small positive number 6 Consider the prior and posterior probability measures PH PA and QH x Qo on W x A. Then, the following bound holds for the TRM risk:
Theorem 3.1 (SCALE generalization bound). If the map w + (xhx( is L-Lipschitz continuous
for any (X,Y) E X y and bounded in (a, b] then, with probability least over the draw of D it holds simultaneously over all E W, 0 {ai...., an, ......, a; E (0, E [k] and E [0, 1],2 [K] that:
S (X),Y)))
.......Ty. mK} E where
1
RTRM (hw: Qo) S RYRM (he Q0) + Vn
22 | wp wo 12 Reg(() + en
6g 0 is necossary that Q absotutcly continuous with respect to PA and KL(Q well-defined.
In the Appendix we provide general bound for any loss satisfying so-called Hoeffdking assumption Alquier al. 2016

---

Published in Transactions on Machine Learning Research (07 2023)
where
K
Reg(0) Σ
Ki mi) +
Σ K((Ti).
(5)
i=k+1
Here, for S k,
1 - log
x,A,,3
Ki(a,, mh) log
B
(1
Ti aT) 33
KL (L(B(TT))IB(LL - B)) r, (Ul-ai, aillu A.))
and for 1 <ish K,
B (111 N mi
B (11- 1 N₂
(1-))) .))
K2(T,)) = KI.
(1 -
and c. s., = (1) vanishes as n1 + 0c.
The proof is provided in Appendix A.2
The steps are as follows:
PAC-Bayes for joint bound: Using our alternative viewpoint discussed above, we leverage PAC- Bayes theory to derive generalization bounds on the expected risk jointly over models and transfor- mations.
Derandomization: However, in TRM we are interested in deterministic models. Thus, we apply the Lipschitz continuity assumption to provide bound on the relative distance between the expected risk of the joint model and of deterministic model used in TRM.
3. Use of Augmentation Space: Next, we derive the explicit form of the regularizer Reg(O) in Eq. resorting to the Radon-Nikodym derivative a/;; since neither the augmentation prior nor the augmentation posterior Q are absolutely continnous with respect to the Lebesgue measure.
dQ
Optimization of bounds: Finally, we optimize the bound over its free parameters, which leads to the form of Cn.
Observe that if the loss 1(:, v) is convex for all V, then this bound also upper bounds the risk RAM of the aggregale predictors x > Bg-Qu Phm (gx))
3.3 Building the SCALE Algorithm
We propose to optimize the upper bound from Thm. 3.1 to learn simultaneously the parameters of the hypothesis space H = (hwi W} and the parameters 0 ( ai,,..,, ag) E θ of the distribution Qo. We absorb the weight decay term in the optimizer and write the training objective L(w, 0) as:
(1me)
K
L(w, 0) = Bg v1 Qa)
Σ
Eg Q(xi
(
hw II
gf
U/s
Aveg Reg(O = { at)
(6)
where Reg(0= = (m, at)) is the regularization term suggested by Eq. 5 and a = [ai....,, .x). π [ ok] This regularizer minimized when each म, + - x 1 and each a; - A,, increasing the support of the uniform distributions Ur,,,- Thus the regularizer promotes frequent augmentations (large m;) and a broad range of augmentations (large a;). This leads a trade-off with the training loss, as desired. During the experiments, use B = 0.01.
At test time, we aim to use the learned transformation distribution Qo and model h to compute the expected output (x) Ey-Qo [h and we approximate the expectation by Monte Carlo sampling. In the experiments we refer to the use of Anp as test-time augmentation.
End to end training requires the computation of the gradients with respect to the local variables a,, m that appear in the distributions Q in the iterated expectations of equation 0 The general case of this setting has been adrressed in Titsias & Lávaro-Gredillal (215)). Tn case, the specific form of the TRM loss allows

---

Published in Transactions on Machine Learning Research (07/2023)
Algorithm SCALE: Training for learning augmentations using TRM
Input: Data = {2,, yt 3 j {hw] H
Initialize K: o0 = E, 0 = (T Initialize model parameters w = wp-
for epoch=1 to Total Epochs do for batch=(batch_x> batch_y in do for (x. y) in batch do
Sample M fid transforms أاو Qu, j ...., M, according equation 4
Add all (g(#),,)) 3) to augmented batch end for
Compute batch _otpput = h (augmented batch_
Compute batch loss L(w,00 0) = (batch y, batch output) according to equation
Compute L, VoL)
Set parameters w 5- w Am V I, a - clamp(a- - da Val, min = 0, max = A.)
clamp(T A L, min O, max 1)
end for end for
us to compute unbiased gradient estimators V V with respect to a, without the need to resort to generic estimators such as REINFORCE [Williams 1992). In particular, the form of Q permits us to compute the gradients with respect to the n, 's in closed form while we use the reparametrization trick for the ag similar Benton al. (2020). In Appendix A.3 we derive the unbiased gradient estimators for both parametric and parameter-free base distributions. Algorithm shows the training process for optimizing over the parameters of the augmentation distribution and of the network.
4 TRM under Distributional Invariance
In this section, we provide a theoretical result in the setting where the data have some unknown symmetry in the form of a distributional invariance. By drawing a connection between the risks RRM and RyRM from Eq. Eq. 2 prove that TRM recovers the unknown invariance, well model that is optimal under the oracle i.e., the standard risk RAM that knows the invariance.
We denote the respective optimal risks by RRM and RT and define the set GR UQen GQ of all transforms considered. We also define the hypothesis space of transforms composed with models H. Ho Ga (x + (x)) H. g E Ga].
The first proposition provides lower bound on the TRM risk when the hypothesis space H is "large enough" that H. Ga C H. that it is closed under Q-expectations (so for any Q Ω h H implies + Eg~oh(gz) E H). For detailed proofs we refer to Appendix A.I.
Proposition 4.1 (Lower Bound). For any H and Q E Qt we have BTRM (h, (h, Q) 2 RisM under either one of the following two sufficient conditions:
H Gn S H;;00 or
2. if ((.y)) conver for all y Y, and H is closed under Q-expeclations for all Ω.
Condition is not restrictive, because model spaces of interest are typically very large and expressive. For example if Ω contains distributions on 2D affine transformations, then even the simple model class H. of all linear classifiers satisfies the condition.
The next proposition shows how an assumption about the type of augmentations within Go can lead to an upper bound on the risk RT
8

---

Published in Transactions on Machine Learning Research (07/2023)
Rotation
Rotation
Scalingg
ScalingX
Scaling
ScalingY
045
Parameters
Shearing
Parameters
2
Shearing
180d Rot
180d Rot
Crop
Crop
HFllp
l
22
Trained ratMnist
Trained on Mnist
272
202
0.
epochs
Redians rotation
Figure Parameters ny during training using SCALE on Left: MNIST, Center: rotMNIST. Right:
Accuracy of models trained on MNIST and roLMNIST, using MNIST images rotated by various angles.
Proposition 4.2 (Upper Bound). Let Gine = ig E Gn (X,Y) =d (gX, Y)X be the set of distributional invariances 9. If there is a distribution Q E Ωl such that its support GQ included the set of distributional invariances, i.e., GQ C Gine, then RTRM S RRM--
When the conditions for both the upper and lower bounds hold, then the optimal classifier under the TRM risk also optimal for the standard risk RRM- We let haM any optimal classifier under the standard risk RRM. Further, a classifier is G-invariant if for all g E G and all E X, h(gz) h(z).
Corollary 4.3 (Zero gap under distributional invariance).
1. Assume H Go C H and that there is Q E Ω such that GQ S Gine. Then, Ri = Riss and any pair (hRM Q) minimizes the TRM risk.
If in addition the conditions in Proposition 4.1 part 2, hold, and there is a compact group G S Gine such that Unif(G) 5 then there erists G-invariant TRM-optimal classifier lyTMM-
This corollary is consistent with the experimental results in Section 5.2 for ROIMNIST an approximately rotationally invariant dataset- and serves as a proof of concept for the functionality our SCALE algorithm.
5
Experiments
5.1 Base Transformations
We parametrize Qe using a set of base transformations commonly used in image processing. Specifically, in both experiments presented in Sections 5.2 5.3 we learn a distribution Qs where each sample ~ Qo is constructed as g 9192.......... 9K. with 7
The first four transformations g1,.... : correspond to rotation, X-axis scaling, Y-axis scaling and X-axis shearing. Each transformation is applied with probability Ti, and its parameter is sampled uniformly from [-a,, a.). The transformations Is, 9o correspond to discrete rotations by 180 degrees, and horizontal flips, respectively. They are applied with probabilities Ts, TT, respectively, and once they are applicd, they perform the discrete rotation or horizontal flip with probability 0.5. Finally 97 is a random cropping of the image, applied with probability म7. (For more details see Appendix A.5
5.2 Experiments on MNIST, rotMNIST: Discovering Distributional Invariances
Using the parametrization from Section 5.1
we train on MNIST [LeCun et al. 1998 and rotated MNIST
(rotMNIST). We train a simple CNN similar to the model used in Benton et al. (2020). Figure 3 shows how
the probabilities n, of the transformation distributions change during the training.
For MNIST, the parameters म. giving the probabilitics of rotation, scaling, shearing and random cropping increase and converge to values above 0.6, while the -5- for 180 rotation and horizontal flips converge to
8The theorem holds mutatis mutandis for locally compact groups too, extend Q contain more general (not just. probability) measures, egg, the rivaiain Haar measure Ug; and replace expectations with integrals.

---

Published in Transactions on Machine Learning Research (07/202) 2022)
CIFAR 10
CIFAR 100
SCALE (Ours)
SCALE
FasienAA
FastenvA
Past-AA
Fast-AA
Augering
ECE (%6)
ECE (%6
Rank-1 BNN
Rank-1 BNN
SNGP
SNSP
Baseline
Baseline
0.5-
0.0
3.25
3.50
3.75
4.25
16.5
17.0
16.5
19.0
10.5
21.0
Classification Error (%)
Classification Error (%)
Figure 4: Comparison of the ECE and Classification Error achioved by the various methods on CIFAR10/100 In the methods where lest-time augment was performed the standard deviation in the results is indicated by the width and the height of the ellipses (corresponding to two standard deviations). For the methods drawn with a triangle the standard deviation not provided.
zero. For the rotation, the learned distribution corresponds to angles uniformly sampled from [-0.31, 0.31) radians, which is consistent with the common practice of augmenting with small rotations.
When training rotMNIST, we observe that the parameters Ti...., TT for rotation (continuous and dis- crete), scaling and shearing converge to values above 0.6, while the parameters Te, т7 for horizontal flips and random crops go close to zero. This indicates that on rotMNIST, SCALE selects augmentations that are useful on the original dataset (MNIST) and also all rotations that correspond to the invariance of the dataset. Additionally, the right plot of Figure 3 shows that contrary to MNIST on rotMNIST SCALE results an invariant classifier. These results are consistent with Corollary 4.3 On rotMNIST we achieve test accuracy of 99.1%, which compares favorably Augerino, that achieves test accuracy of 98.9% on the same network. According to Benton et al. (2020) only Weiler & Cesa (0199) have reported better result using a network equivariant to rotations.
5.3 Learning Data Augmentations on CIFAR 10/ 100
Using SCALE and the parametrization shown Section 5.1 we learn useful data augmentations on the CIFAR 10 and CIFAR 100 datasets (Krizhevsky 2009). For both CIFAR 10 and CIFAR 100 we used a Wide ResNet 28-10 Zagoruyko & Komodakis 2016) (for training details see Appendix A.4.
We compare the results for our method and for prior works on learning data augmentations that are discussed
in Section G For Trivial-Augment (TA) (Müller & Hutter 2021 and DADA (Li et al. 2020 we evaluate
using the same augmentation space as our method (presented in Section 5.1 and Appendix A.5) First. we evaluate the trained models on the original test images, without applying any test-time augmentation. Additionally, for each method, we perform test-time augmentation by sampling augmented copics from the augmentation policy learned by the corresponding method and taking the average prediction. the accuracy for CIFAR 10 and CIFAR 100, both using test-time augmentation with 4 and 8 transforms, and without test-time augmentation. Finally Table 2 shows an ablation study on the use of our proposed parametrization and regularizer.
[
On both CIFAR 10 and CIFAR 100, SCALE converges to random rotations, scalings, shearings, horizontal flips and crops, while does not pick 180 discrete rotations. SCALE outperforms Augerino and DADA on both CIFAR 10 and CIFAR 100. Also, SCALE performs similarly AutooAuunuug Lim et al. (2019) (Fast-AA) and Faster-AutoAugmeent Hataya et al. (2020) (Faster-AA) that first search for augmentation policy and then train model-an on CIFAR 10 and outperforms both methods CIFAR 100.
5.4 Learning Augmentations Improve Calibration
In addition to accuracy, many applications require a model to be well calibrated, which means that it provides an accurate uncertainty estimate for its predictions. In classification setting we can use the output of the softmax layer indication of the model's confidence that the input belongs to a specific
10

---

Published in Transactions on Machine Learning Research (07/2023)
Table 1: Test accuracy of Wide ResNet 28-10 trained on CIFAR 10/100, using several method for learning data augmentations. We show the test accuracy of the models when we do not perform test-time augmen- tation (No TestAug.) and when we perform test-time augmentation (TestAug.) with 4, samples. The values presented are the average over five trials. For all methods, the standard deviation is close to 0.1% (please see the appendix for table containing the standard deviations) We also report the additional time that each method spends in policy search in GPU Hours relative to the standard training (No Augm. column)
No AUGM.
AUGERINO +FLIPS
FAST-AA
AASTER-AA
TA
DADA (GEOMETRIC)
SCALE
(Ours)
AUG
95.7%
97.2%
97.1%
97.2%
96.6%
CIFAR10
TESTAUG. (N=4)
95.7%
97.0%
96.6%
96.8%
96.5%
96.7%
TESTAUG. (N=8)
95.8%
97.2%
96.8%
96.9%
96.8%
96.9%
No AUGTAUG.
73.1%
79.4%
82.5 %
81.9%
83.1%
79.4%
82.7%
CIFAR100
TESTAUG. (N=4)
79.2% 2%
82.5 %
82.0%
82.9%
79.8%
83.0%
TESTAUG. (N=8)
79.3%
82.9 %
42.4%
83.6%
80. 1%
83.3%
GPU HOURS (CIFAR 10/100)
4-81
4-31
1H
4-2.5H
4-1.0
Table 2:
Ablation study on the proposed regularization and parametrization. First we evaluate model using the Augerino parametrization and our proposed PAC-Bayes regularizer (without the addition of flips).
Second we evaluate a model using the SCALE parametrization and the regularizer used in Augerino (L2 regularizer). Finally we show our method with both our proposed regularizer and parametrization
AUGERINO (No FLIPS)
SCALE
SCALE
PB REGULARIZER
+ AUG REGULARIZER
(Ours)
No PESTAUG:
94.2%
96.6%
96.7%
CIFAR10
TESTAUG. (N=4)
94.3%
96.7%
96.7%
TESTAUG (N=8)
94.4%
96.8%
96.9%
class. In particular, consider a model X + RD that assigns score hy (say, softmax output) to cach class y. As discussed in [Cox. 1958 DeGroot & Fienberg 1983; Gneiting al. 2007: Guo et al. 2017) given random datapoint (X, Y) drawn from the data distribution D, model h is perfectly calibrated if IP PiD (Y y/ hy (X) p) p. Vp 1].
A commonly used metric for quantifying the calibration of a model is the Expected Calibration Error (ECE) (see e.g., [Hrrell] 2015 Nacini et al. 2015)): the samples are partitioned into M >> 0 equally spaced bins ac- cording to the confidence (x(x) for the predicted category & = argmax, (xy( (hgmax,(h,)r))). Let Bm be the set containing the indices of datapoints that belong the m-th bin, for m E [M]. For the m-th bin we define the mcan ac- curacy acc(B,, (m.) |
1 EieB hg. (x;). Bmm ΣieB.. 1 yi) and the mean confidence conf (Bm) [Bm.
We can then define the Expected Calibration Error as: ECE ΣM (acc(Bm)) conf (Bm)) )1 |Bm 1/n./.
We compare basclinc deterministic model trained with the standard augmentations presented in Zagoruyko & Komodakis (2016) and evaluated without test-time augmentation to models trained with SCALE, Augerino, Fast-AA AA and Faster-AA and evaluated with test-time augmentation using 4 augmented datapoints Gawlikowski et al. (2021). Additionally we compare with methods designed to improve the calibration of a model that are presented in Nado et al. (2021). Specifically we compared with Batch Ensemble Wen et al. 22020), Rank-1 BNN Dusenberry et al. (020)) and SNGP Ensemble Liu et al. (2020). For the latter three mothods we use the implementation and results provided in Nado et al. (2021). For both CIFAR 10 and CIFAR 100, all methods the same WideResNet. 28-10 network architecture.
Figure compares the calibration and classification errors achieved by the various methods on CIFAR 10/100, SCALE outperforms the bascling and Augerino, and obtains classification error and ECE comparable to more computationally expensive methods that learn augmentation policies (Faster-AA, Fast-AA), and also to methods that aim to improve the calibration models. It is important to note that SCALE mainly aims to minimize the classification error of model, and not its calibration error. Models that only optimize the classification performance need not be well calibrated. Thus, perhaps surprisingly, we show that SCALE results in models not only with low classification error (on CIFAR 100 it achieves the lowest error) but also
11

---

Published in Transactions on Machine Learning Research (07/2023)
Baselie ICIPAR 10)
Augerine
Frat-A4 (CINAS
Faster-M (CIPAR 10I
(CTMAR
BC5-2.5%
BCE
ECF-0.6N
ECE-0.9%
ECE=O.877
Cala, Errorwel 0%
Cls..
Errorma.59%
Confidence
SCE- 7.0%
E.2-6.2%
EC1-2.1%
EGE-2.9%
ECE-2.1%
Accuraty
Conficence
Confidence
Confidence
0
Figure 5: Reliability diagrams of various methods for learning data angmentations on CIFAR 10/100.
low with calibration
close to the error achieved by methods that mainly aim to improve calibration. We also present reliability diagrams that show, for each confidence bin, the mean accuracy achieved by the corresponding mothod. FiguroB]shows the reliability diagrams of the bascline deterministic model and the models trained with SCALE, Augerino, Augt-mnt and Faster- Aaater-AttAAugmen which are evaluated using test-time augmentation with augmented transforms.
5.5 Learning Data Augmentations on Language Navigation Task
We showcase the use of our method on a recently proposed task of language navigation, where the set of best augmentations is not generally known. Specifically this task, given a semantic map and a natural language instruction, we aim to infer a sequence of waypoints that define a path that follows the given instruction. For our experiment we follow the definition of the lask proposed in Krantz et al. (2020), where the maps are extracted from scenes of the Matterport3D dataset (Chang et al., 2017) and the instructions are given from the R2R. datasct (Anderson et al. 2018).
We apply SCALE over the path prediction model shown in Figure 6b which was proposed in Georgakis et al. (2022). We can measure the performance of the predicted waypoints by comparing them to the waypoints of the ground truth path. We consider a predicted waypoint as correct if it is located at most one meter away from the ground truth. We focus on learning transformations applied on the imput semantic map using the parametrization described Scction 5.1 In the figure Ga we show the performance achieved by our method compared to the performance of the baseline model as function of the regularization parameter Xreg- We observe that for values of Areg ranging from 5 10-4 to 3.5 10-3 our method learns augmentations that result in increased performance compared to the performance achieved by the baseline.
6 Related Work
There is vast amount of related work; we can only discuss the most closely related ones.
VRM. We view ETRM generalization of Vicinal Risk Minimization (VRM) (Chapelle et al. 2001 In VRM the Dirac masses around each training datapoint are replaced with vicinal functions, usually uniform or Gaussian distributions centered at each point. It has been proven empirically (Cao & Rockett 2015) and theoretically (Zhang et al.) 2018) that VRM enjoys some benefits over ERM, but is highly dependent on the right choice of the vicinal lunctions. Instead fixed vicinal functions, we learn transformation distributions. In this sense, our learning algorithm adapts the vicinal functions the task.
Invariance. Data augmentation can be viewed as an implicit restriction to models invariant to predefined transformations. When the invariance is known, and the transformations have a group structure, exact way to guarantec the invariance is to "bake it in" the network design. This can be traced back to (Fukushima, 1980 LeCun el al.) 1989) with the development of CNNs, which use convolutions to build (ranslation equivariant layers. Architectures that extend the idea to discrete (Cohen Welling (0166) or continuous (Weiler et al. (0188) rotational invariance and more general Lie Groups (Finzi et al.) 20200 have been
proposed. More recently they have been generalized beyond Euclidean domains (Esteves et al. 20188 Maron
al. 2019 Cohen et al. 2019). When the data symmetry unknown, Anselmi al. (019)) assume that the
12

---

Published in Transactions on Machine Learning Research (07 (07/202) 2203))
Naviqation
aren fetwerk malle boster was dear to pork
Baxalin
Theoretical a (1ivn)
UIRT
Path Allertion
00.5
(a)
(b)
Figure 6: (a) Percentage of correct waypoints (PCW) predicted by the language navigation model when we train only the original model (Baseline) and when we incorporate SCALE into the training of the model with different regularization parameters Areg. (b) A high level overview of the language navigation model used in Georgakis et al. (2022), ; which is given a semantic map and an instruction and outputs ten waypoint predictions (red dots) that aim to match the ten waypoints of the ground truth path (blue dots).
dataset is partitioned into full orbits of an unknown group and propose symmetry-adaptedd regularization to find it.
Data augmentation (DA). The concept of data augmentation in machine learning goes back al least
to Baird (1992). More recently, generative networks have been used to generate augmented data (Mirza
& Osindero 2014 Antoniou et al. 2018 Robey et al., 2020). There is a variety of work on designing
augmentations (Zhong et al. 2017 Lopes et al. 2019 Perlin 1985 Zhang et al., 2017). Vedaldi al. (2011)
propose method for learning equivariant SVM regressors. Novotny et al. (2018) introduce augmentations
the fly to learn equivariant descriptors in a self-supervised setting. Deeper theoretical understanding the relative efficiency gains of DA has been obtained by Chen et al. (2020), who develop group-theoretic framework, and Lyle et al. (2019), who compare feature averaging with data augmentation on the grounds of generalization using PAC-Bayes theory under the assumption that (X, Y) =a (gX,Y) Y). We derive the bound jointly over (Q,4) h) without this assumption.
Ratner et al. (2017) use a generative adversarial approach to learn sequences of predefined valid transfor- mations. Similarly AutoAugment (Cubuk et al. 2019a) learns augmentation policics using reinforcement learning to search over a predefined space of transformations. Many works aim to reduce the high compu- tational complexity of the AutoAugment policy search: RandAugment (Cubuk et al., 2019b) simplifies the search space to two hyperparameters, Auso--Auto (Lim et al. 2019) proposes a more efficient search strategy based on density matching and Faster- Autto--ettt (Hataya et al., 2020) introduces a gradicnt- based search method that minimizes the distance between the original and the augmented dataset. Similarly to our method, [in et al. (2019) propose to formulate the augmentations as parametrized distributions. They suggest bi-level optimization algorithm to train the model simultancously with the augmentation parameters using policy gradients. While faster than AutoAugment this approach still introduces a very big overhead over standard training. [Hataya et al. (2022) also propose bi-level opt problem to directly optimize the validation accuracy without proxy tasks or smaller augmentation spaces by approximating the inverse Hessian via a truncated von Neumann series, an approximation that was first studied in Lorraine et al. 2019g. similar automatic augmentation algorithm was proposed in Raghu et al. (2022) for electro- cardiograms. Finally Müller & Hutter (2021) propose Trivial-Anmentt, a simple tuning- -frnng-free augmentation policy, and argue it achieves comparable or better performance than the above methods presented.
van der Wilk et al. (0018) cast the problem Bayesian model selection and learns data transformations jointly with downstream task by maximizing the marginal likelihood. This is applicd to Guassian processes by imposing an invariance constraint on the kernel. The method is restricted to differentiable transformalions and not directly applicable to deep neural networks due to the inefficiency to scale the marginal likelihood. Extending this work to invariant deep kernel learning (Schöoell learningSchwöbel et al. (2022) propose to compute the marginal likelihood only in the last layer of deep network and show some performance gains but also report scal- ability issues. Lastly, [mmer et al. (2022) extend this formalism to deep networks by proposing a Taplace
13

---

Published in Transactions on Machine Learning Research (07/2023)
approximation to the marginal likelihood. While they demonstrate some performance gain the computation that involves gradients of the log determinant of the Hessian of the model introduces big overhead over standard training.
Augerino [Benton et al. 2020 optimizes over a continuous parametrization of augmentation distributions, and thus excludes discrete transformations (e.g., horizontal flips). Augerino uses the negative norm of the parameters of the learned augmentation distribution as regularization term. This choice of regularizer may be more unstable, because it leads to minimizing concave function. We compare favorably with Augerino in both accuracy and calibration in the experimental section. Similarly to our method, DADA (Li et al, (020)) relaxes the discrete data augmentation policy selection to differentiable optimization problem. They rely on the Gunbel-Softmax reparametrization trick introduced in Maddison et al. (2017); Jang et al. (2017) and apply variance reduction to their proposed gradient estimator. We compare favorably against DADA in all CIFAR 10/100 and SVHN.
7 Conclusion
this paper, we formulated the Transformed Risk Minimization framework, and used to design SCALE, a new algorithm for learning augmentation distributions. This relies on new parametrization of the aug- mentation space via a composition of stochastic blocks, and leverages PAC-Bayes theory to derive novel training objectívc. We provided experimental results that corroborate
thcorctical analysis, and showed that SCALE can lead to advantages over prior approaches in both accuracy and calibration and can be in- corporated easily in large-scale problems like language navigation and benefit them in terms of performance.
14

---

Published in Transactions on Machine Learning Research (07 (07/2023)
References
Pierre Alquier, James Ridgway, and Nicolas Chopin. On the properties of variational approximations of gibbs posteriors. Journal of Machine Learning Research, 17(236):1-41. 2016.
Peter Anderson, Qi Wu, Damien Teney, Jake Bruce, Mark Johnson, Niko Sünderhauf, Ian Reid, Stephen Gould, and Anton Van Den Hengel. Vision-and-language navigation: Interpreting visually-grounded navi- gation instructions in real environments. In Proceedings of the IEEE Conference on Compuler Vision and Pattern Recognition, pp. 3674-3683 3683, 2018.
Fabio Anselmi, Georgios Evangelopoulos, Lorenzo Rosasco, and Tomaso Poggio. Symmetry-adapted repre- sentation learning. Pattern Recognition, 86:201-208, 2019.
Anthreas Antoniou, Amos Storkey, and Harrison Edwards. Data augmentation generative adversarial net- works, 2018
Henry S Baird. Document image defect models. In Structured Document Image Analysis, pp. 546-556.
Springer, 1992.
Gregory Benton, Mare Finzi, Pavel Tzmailov, and Andrew G Wilson. Learning invariances in neural networks from training data. In Advonces in Neural Information Processing Systems, volume 33, pp. 17605 17616, 2020.
Yilong Cao and Peter I. Rockett. The use of vicinal-risk minimization for training decision trees. Applied Soft Computing, 31:185 195, 2015.
Olivier Catoni. Pac-bayesian supervised classification: The thermodynamies of statistical learning. IMS Lecture Notes Monograph Series, pp. 1-163, 2007.
Angel Chang, Angela Dai, Thomas Funkhouser, Maciej Halber, Matthias Niessner, Manolis Savva, Shuran Song, Andy Zeng, and Yinda Zhang. Matterport3d: Learning from rgb-d data in indoor environments. International Conference on 3D Vision (3DV), 2017.
Olivier Chapelle, Jason Weston, Léon Bottou, and Vladimir Vapnik. Vicinal risk minimization. In Advances in Neural Information Processing Systems, volume 13. MIT Press, 2001.
Shuxiao Chen, Edgar Dobriban, and Jane Lee. A group-theoretic framework for data augmentation. Advances in Neurol Information Processing Systems, JMLR, 33, 2020.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International conference on machine learning, pp. 2990-2099, 2016.
Taco Cohen, Maurice Weiler, Berkay Kicanaoglu, and Max Welling, Gauge equivariant convolutiona ional net- works and the icosahedral CNN. In Proceedings of the 36th International Conference Machine Learning, volume 97, pp. 1321 1330, 2019.
David R Cox. Two further applications of model for binary regression. Biometrika, 45(3/4):562 565, 1958.
Ekin D. Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoe V. Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR), June 2019a.
Ekin Cubuk, Barret Zoph, Jonathon Shlens, and Quoc Le. Randaugment: Practical data augmentation with no separate search. arXiv preprint arXiv: 13719, 2019b.
Morris H DeGroot and Stephen E Fienberg. The comparison and evaluation of forecasters. Journal of the Royal Statistical Sociely: Series D (The Statistician), 32(1-2):12 22, 1983.
15

---

Published in Transactions on Machine Learning Research (07/2023)
Michael Dusenberry, Ghassen Jerfel, Yeming Wen, Yian Ma, Jasper Snock, Katherine Heller, Balaji Lak- shminarayanan, and Dustin Tran. Efficient and scalable Bayesian neural nets with rank-1 factors. Tn Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of mlr.
Machine Learning Research, pp. 2782-2792. PMLR, 13-18 Jul 2020. URL 0hees//111//bbssmberr220.....1
Gintare Karolina Dziugaite and Daniel M. Roy. Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data. In Proceedings of the 33rd Annual Conference on Uncertainty in Artificial Intelligence (UAT), 2017.
Gintare Karolina Dziugaite and Daniel M Roy. Data-dependent pac-bayes priors via. differential privacy.
In Advances in Neural Information Processing Systems, volume 31. Curran Associates, Inc., 2018. URL
pdf
Carlos Esteves, Christine Allen-Blanchette, Ameesh Makadia, and Kostas Daniilidis. Learning so(3) equiv- ariant representations with spherical cnns. In The European Conference on Computer Vision (ECCV), September 2018.
Marc Finzi, Samuel Stanton, Pavel Izmailov, and Andrew Gordon Wilson. Generalizing convolutional neural networks for equivariance to lie groups on arbitrary continuous data. In Proceedings of the 37th Inler- national Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pp. 3165-3176. PMLR, 13-18 Jul 2020.
Donald Alexander Stuart Fraser. Structural probability and a generalization. Biometrika, 53(1-2):1-9, 1966.
Kunihiko Fukushima. Neocognitron: A self-organizing neural network model for a mechanism of pattern recognition unaffected by shift in position. Biological cybernetios, 36(4):193-200 202, 1980.
Jakob Gawlikowski, Cedrique Rovile Njieutchen Tassi, Mohsin Ali, Jongscok Lee, Matthias Humt, Jianxiang Feng, Anna Kruspe, Rudolph Triebel, Peter Jung, Ribana Roscher, et al. A survey of uncertainty in deep neural networks. arXiv preprint orXiv:2107.0380 2021.
Georgios Georgakis, Karl Schmeckpeper, Karan Wanchoo, Soham Dan, Eleni Miltsakaki, Dan Roth, and Kostas Daniilidis. Cross-modal map learning for vision and language navigation. In IEFE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Pascal Germain, Francis Bach, Alexandre Lacoste, and Simon Lacoste-Julien. Pac-bayesian theory meets bayesian inference. In D. Lee, M. Sugiyama, U. Luxburg. I. Guyon, and R. Garnett (ods.), Advances in Neural Information Processing Systems, volume 29. Curran Associates, Inc., 2016a. URL https: //proceddings.o newrips. http://////../11000000000031/110
pdf
Pascal Germain, Francis Bach, Alexandre Lacoste, and Simon Lacoste-Julien. Pac-bayesian theory meets bayesian inference. In Advances in Neural Information Processing Systems, volume 29. Curran Associates,
Inc., 2016b.
Tilmann Gneiting, Fadoua Balabdaoui, and Adrian E Raftery. Probabilistic forecasts, calibration and sharp- ness. Journal of the Royal Statistical Society: Series B (Statistical Methodology), 29(2):243322)) 2007.
Chuan Guo, Geoff Pleiss, Yu Sun, and Kilian Q. Weinberger. On calibration of modern neural networks. In Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1321- PMLR, 06-11 Aug 2017.
Frank E Harrell. Regression modeling strategies: with applications to linear models, logistic and ordinal regression, and survical analysis, volume 3. Springer, 2015.
Ryuichiro Hataya, Jan Zdenek, Kazuki Yoshizoe, and Hidcki Nakayama. Faster autoaugment: Learning augmentation strategies using backpropagation. In European Conference on Computer Vision, pp. 1-16. Springer, 2020.
16

---

Published in Transactions on Machine Learning Research (07/2023)
Ryuichiro Hataya, Jan Zdenek, Kazuki Yoshizoe, and Hideki Nakayama. Meta approach to data augmen- tation optimization. In Proceedings 2022 IREE/CVF Winter Conference on Applications of Compuler Vision, WACV 2022, Proceedings 2022 LEEE/CVF Winter Conference on Applications of Computer Vision, WACV 2022, pp. 3535- 3544, United States, 2022. Institute Electrical and Electronics Engineers 458.5511......2...3 2022.0039.
Inc. doi:
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 770-778, 2016.
A. Immer, T. F. A. van der Ouderaa, G. Rätsch, V. Fortuin, and M. van der Wilk. Invariance learning in deep neural networks with differentiable laplace approximations. In Advances in Neu- val Information Processing Systems 35 (NeurIPS 2022), volume 35, pp. www.es//appee/2222//aaaa ciales, Inc., December 2022. URL https::/ newrips www..nwsw....pppp.....e.e. 55000f998cc04600077777/aaaaaaaa2
Curran Asso-
Max Jaderborg, Karen Simonyan, Andrew Zisserman, et al. Spatial transformer networks. In Advances in Neurul Information Processing Systems, pp.
2015.
Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Truck Proceedings. OpenReview.nee nel, 2017.
Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. In 2nd International Conference on Learning Representations, ICLR 2014, 2014.
Jacob Krantz, Erik Wijmans, Arjun Majumdar, Dhruv Batra, and Stefan Lee. Beyond the nav-graph: languuae navigation in continuous environments. In European Conference on Computer Vision (ECCV), 2020.
Alex Krishevsky. Learning multiple layers of features from tiny images. 2009.
John Langford and Rich Caruana. (Not) bounding the true error. In Advances in Neural Information Processing Systems, volume 14. MIT Press, 2002. URL httpsss/pproceedngg..errrrrr www.eip...cpppppeeeeeeeeeeeeetee Paper.par .dff
fi hil/9/6799880884c66ccc4aattfaac2
Yann LeCun, Bernhard Boser, John $ Denker, Donnie Henderson, Richard E. Howard, Wayne Hubbard, and Lawrence D Jackel. Backpropagation applied to handwritten zip code recognition. Neural computation, 1 (4):541-551, 1989.
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient-based learning applied to docu- ment recognition. Proceedings of the IEEE, 86(11):2278 2324, 1998.
Yonggang Li, Guosheng Hu, Yongtao Wang. Timothy M. Hospedales, Neil Martin Robertson, and Yongxin Yang. DADA: differentiable automatic data augmentation. 2020.
Sungbin Lim, Ildoo Kim, Taesup Kim, Chiheon Kim, and Sungwoong Kim. Fast autoaugment. In Advances in Neural Information Processing Systems, volume 32, pp. 6665-6675, 2019.
C. Lin, M. Guo, C. Li, X. Yuan, W. Wu, J. Yan, D. Lin, and W. Ouyang. Online seper-paramette learning for auto-augmentation strategy. 2019 IEFE/CVE International Conference on Computer Vision (ICCV), pp. 1578-65877 Los Alamitos, CA, USA, nov 2019. IEEE Computer Society. doi: 10.1109/ICCV:2019. 00668. URL httph:: //tts: www.ooatttessssstttttttttoooom
20099.00668
17

---

Published in Transactions on Machine Learning Research (07/2023)
Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshmi- narayanan. Simple and principled uncertainty estimation with deterministic deep learning via. dis- tance awareness. In Advances in Neural Information Processing Systems, volume 33, pp. 7498 7512. Curran Associates, Inc.. 2020. URL https://proceednss.ooo hww.sss.errrpcccoo httpr//rp..ccppppp////////1e11.. 4438837422244477222220000a666333 Papeee Paper.pdff
Raphael Gontijo Lopes, Dong Yin, Ben Poole, Justin Gilmer, and Ekin D Cubuk. Improving robustness without sacrificing accuracy with patch gaussian augmentation. orXiv preprint arXiv: 02611, 2019.
Jonathan Lorraine, Paul Vicol, and David Kristjanson Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International Conference on Artificial Intelligence and Stotistics, 2019.
llya Loshchilov and Frank Hutter. SGDR: stochastic gradient descent with warm restarts. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24- 26-26, 2017, Conference Track Proceedings. OpenReview net, 2017. URL https: net/forum ?id= Skg89xxX
Clare Lyle, Marta Kwiatkowksa, and Yarin Gal. An analysis of the effect of invariance generalization in neural worworks. In International conference on machine learning Workshop on Understanding and Improving Generalization in Deep Learning, 2019.
David JC MacKay. Information theory, inference and learning algorithms. Cambridge university press, 2003.
Chris J. Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. In International Conference on Learning Representations, 2017. URL https:
//openreview.o Soopenreries..et/ffruumtii wtt.://oororiidd-111.......
Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. In 7th International Conference on Learning Representations, TCLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
David McAllester. A pac-bayesian tutorial with a dropout bound, 2013.
David A. McAllester. Some pac-bayesian theorems. In Proceedings of the Fleventh Annual Conference on Computational Learning Theory, COLT" 98, pp. 230 234, New York, NY, USA, 1998. Association for Computing Machinery. ISBN 1581130570.
David A. McAllester. Pac-bayesian model averaging. In Proceedings of the Twelfth Annual Conference on Computational Learning Theory, COLT '99, pp. 164-170, New York, NY, USA, 1999. Association for Computing Machinery. ISBN 1581131674. doi: 10.1145/307400.30742 URL https:///olt org/ 1145/307400.3074.
Mehdi Mirza and Simon Osindero. Conditional generative adversarial nets. arXi preprint arXiv:1411.1784, 2014.
Samuel G. Müller and Frank Hutter. Trivialaugment: Tuning-free yet stale-of-the-arl data augmentation. In Proceedings of the LEEE/CVF International Conference on Computer Vision (ICCV), pp. 784-782, October 2021.
Zachary Nado, Noil Band, Mark Collier, Josip Djolonga, Michael Dusenberry, Sebastian Farquhar, Angelos Filos, Marton Havasi, Rodolphe Jenatton, Ghassen Jerfel, Jeremiah Liu, Zelda Mariet, Jeremy Nixon, Shreyas Padhy, Jie Ren, Tim Rudner, Yeming Wen, Florian Wenzel, Kevin Murphy, D. Sculley, Balaji Lakshminarayanan, Jasper Snoek, Yarin Gal, and Dustin Tran. Uncertainty Baselines: Benchmarks for uncortainty robustness in deep lcarning. ar preprint orXiv:2106. 04015, 2021.
Mahdi Pakdaman Nacini, Gregory Cooper, and Milos Hauskrocht. Obtaining well calibrated probabilities using bayesian bimning. In Ninth AAAI Conference on Artificial Intelligence, 2015.
Yurii Nesterov. A method of solving convex programming problem with convergence rate (11k32). Doklady Soviet Mathematics, 269:543 547, 1983.
18

---

Published in Transactions on Machine Learning Research (07/2023)
Yuval Netzer, Tiefie Wang, Adam Coates, A. Bissacco, Bo Wu, and A. Ng. Reading digits in natural images with unsupervised feature learning. 2011.
David Novotny, Samuel Albanie, Diane Larlus, and Andrea Vedaldi. Self-supervised learning of geometrically stable features through probabilistic introspection. In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognation, pp. 3637-3645, 2018.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang, Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjic Bai, and Soumith Chintala. Pytorch: An imperative style, high-performance deep learning library. In Advonces in Newral Information Processing Systems 32, pp. 8024 8035. Curran Associates, Inc., 2019.
Ken Perlin. An image synthesizer. ACM Siggraph Computer Graphics, 19(3):287-296, 1985.
Aniruddh Raghu, Divya Shanmugam, Eugene Pomerantsev, John Guttag, and Collin M Stultz. Data augmentation for rocardioossg In Gerardo Flores, George H Chen, Tom Pollard, Joyce C Ho, and Tristan Naumann (eds.), Proceedings of the Conference on Health, Inference, and Learning, vol- ume 174 of Proceedings of Machine Learning Research, pp. 282-310. PMLR, 07-08 Apr 2022. URL
https://proceediin mlr. press/v174/raghuu22 html
Alexander J Ratner, Henry Ehrenberg, Zeshan Hussain, Jared Dunnmon, and Christopher Ré. Learning compose domain-specific transformations for data augmentation. In Advances in neural information processing systems, pp. 3236 3246, 2017.
Alexander Robey, Hamed Hassani, and George J. Pappas. Model-based robust deep learning: Generalizing natural, out-of-distribution data. arXiv preprint wxXi:22055.02443 2020.
Pola Schwöbel, Martin Jørgensen, Scbastian W. Ober, and Mark der Wilk. Last layer marginal likelihood for invariance learning. In Gustau Campss Valls, Francisco J. R. Ruiz, and Isabel Valera (eds.), International Conference on Artificial Intelligence and Statistics, AISTATS 2022, 28-30 March 2022, Virtual Event,
volume 151 of Proceedings of Machine Learning Research, pp- 3542-3555, PMLR, 2022. URL https: html
//proceedings.
mml.eeddiinsmmm.......m....e..em mlr. pptps//111//ccsooossaeooo
John Shawe- Taylor and Robert C. Williamson. A pac analysis of a bayesian estimator. In Proceedings of the Tenth Annual Conference on Compulational Learning Theory, COLT '97, pp. 2-9, New York, NY, USA, 1997. Association for Computing Machinery. ISBN 0897918916. doi: 10.1145/267460.267466 URL
https://doi. org/ 10. /11111////666../1 267466
Michalis Titsias and Miguel Lázaro-Gredill Local expectation gradients for black box variational inference. In C. Cortes, N. Lawrence, D. Lee, M. Sugiyama, and R.. Garnett (eds.), Advances in Newral Information Processing Systems, volume 28. Curran Associates, Inc., 2015. URL https://procceednngs neurips. cc/ paper http:///pppr/001faaaaaaaaaaaaaoa
.ddf
Mark van der Wilk, Matthias Bauer, ST John, and James Hensman. Learning invariances using the marginal likelihood. In Advances in Neural Information Processing Systems, volume 31, pp. 9938-999999 (988-994,, 2018.
Vladimir Vapnik. The nature of statistical learning theory. Springer science & businoss media, 2013.
Andrea Vedaldi, Matthew Blaschko, and Andrew Zisserman. Learning equivariant structured output svm regressors. In 2011 International Conference on Computer Vision, pp. 959 969-966. TEEE, 2011.
M.J. Wainwright and M.I. Jordan. Graphical Models, Exponential Families, and Variational Inference. Foundations and Trends in Machine Learning, 1(1-2):1-305, 2008.
Maurice Weiler and Gabriele Cesa. General equ)eequivrrian stecrable enns. In Advances in Neural Infor- mation Processing Systems, volume 32. Curran Associates, Inc., 2019.
19

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Maurice Weiler, Fred A. Hamprecht, and Martin Storath. Learning steerable filters for rotation equivariant cnns. In 2018 TEEE Conference Computer Vision and Pallern Recognation, CVPR 2018, Sall. Lake City, UT, USA, June 18-22, 2018, pp. 849 858. IEEE Computer Society, 2018.
Yeming Wen, Dustin Tran, and Jimmy Ba. Batchensemble: alternative approach to efficient. ensemble and lifelong learning. In International Conference Learning Representations, 2020. URL https: //openreview.net//e55e wwwt:/orummidd-ikifiiaa.
Ronald J. Williams. Simple statistical gradient-following algorithms for connectionist reinforcement learning. Mach. Learn., 8(3-4):33)))) 4):229 22:2222556, May 1992.
Sergey Zagoruyko and Nikos Komodakis. Wide residual networks. In Proceedingss of the Brilish Machine Vision Conference (BMVC), pp. 17.12. BMVA Press, September 2016.
Chao Zhang, Min-Hsiu Hsieh, and Dacheng Tao. Generalization bounds for vicinal risk minimization prin- ciple. CoRR, abs/1811 105415, 2018.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz. Paz. Mixup: Beyond empirical risk minimization. arXiv preprint arXiv:1710. arXi::1710000412. 2017.
Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. Random erasing data augmentation. Viviv preprint ar Xiv: 1708. 04896, 2017.
20

---

Published in Transactions on Machine Learning Research (07 2023)
A Appendix
A.1 TRM Propositions
For technical reasons, in the definitions of optimal risks, we switch from minimization to taking infima. This makes the definitions more broadly applicable. When the minima are well-defined, we recover the definitions from the main body. We also assume that all functions we consider are measurable, including after taking the appropriate infimum over the hypothesis space the definitions of optimal risks.
In this section provide the proofs of the propositions in Section extensively in the proofs:
The following definitions will be used
RM risk and its optimal value:
R(h) == Bxx,))-(((((xxx)))
RM == inf R(h). hCH
TRM risk and its optimal value:
R(h, Q) = Fxx.))
RT = inf NEN.QER
R(h, Q)-
Fized g-risk for a fixed transformation g:
gy (xxx)>-l((((xxx)))
Set of distributional incarionces gF
Ginv (g E Gn (X,Y) =a ()X
(7)
Proof of Proposition 4.1 (Lower Bound). First we prove the following lemma:
Lemma A.1. For every distribution Q,
inf E (l 2 Eg- inf Ry NER
Proof of Lemma
A.1 Using the definition the infimum and the monotonicity of expectations
(y) 2 inf heH
(h, Vh E H
E E R (h. (h. 2 Ey-Q inf (h, wh e H = hCH
inf ( By-Q inf Ry(h, NEW
NEH.
Next, for any fixed g we have:
(h, R(h o) = inf
Ry(h,g) inf R(h),
(8)
kef
heHog
where H :og:= (h glh H}.
Using the lemma and observation in equation g now complete the proof of the proposition.
Proof of Proposition
4.1 We prove the two claims in turn.
1. It suffices to show that Ripas 2 RiaM For this, we have
RT inf R(h, Q) inf inf R(h, Q) kEH,QEn
bEH
inf inf IE (Ly-o Y)] inf inf (g-.R QER
Qen heR
21

---

Published in Transactions on Machine Learning Research (07/2023)
where in the first line used that the joint infimum is the same the sequential infimum. In the second line we used the Fubini-Tooe theorem to change the order of expectations, along with the definition of fixed g-risk. By Lemma A.I the last expression is lower bounded by
inf B inf (h 9) inf Bg-Q inf R(h) 2 inf Sq-QRRa BRM KEH
QER
Qen
kgHog
QEQ
For the first equation, used the observation in equation [ that for any fixed g E G2 we have H. S H GR C H from the assumption of the proposition.
and in the following inequality we used
2. We have that if h E H then (x) == Eg-ch(gr) E H for any Q E Ω. Since the loss 1: 3" x y + convex for all E y we have from Jensen's inequality that for all h E H, Q E Qt:
R(h, Q) = (66 (
B Y)] R(hq)
(9)
M inf R(h) RRM hEH
Discussion. The condition H. supp(S?) H means that for each transformation g G used by any distribution in St and each h H. can write h(gr) ((x) for some h' E H. Thus, this means that minimizing the average risk over H. supp([)) is type of constraint compared to minimizing the risk over all hypotheses in H.
Data augmentation is not usually viewed as imposing a constraint. The goal is to learn an invariant predictor, but the motivation for data augmentation that we do not know how to compute one, so we use a heuristic way to expand our dataset. Thus, in this sense, averaging the risk over transforms can be viewed as form of constraint.
Proof of Proposition 4.2 (Upper Bound).
Proof. Suppose Qine is such that Qinv E S2 and GQ... S Ginv- Then,
R(h, Qinv) = Exy.yr--BB--.... [(((gg)), Y)] = B (xxy)--
= =Eg-Qu.. Bx Y)] = R(h),
(10)
where we used the Fubini-Tonelli theorem to change the order of expectations, and then used that (X,Y) =d (gX,Y) for all g E Go... to conclude that ((xx,))-[(((((gg), Y)] = (oxx))=D[((h(x))) Y)]. Therefore,
R
inf NER,QER
R(h, Q) < inf R(h, Qine) inf R(h) RRM. bER
This finishes the proof.
The upper bound uses that we include the "correct" invariances in a distribution Q Q2 the risk does not increase. This seems to be generally justified setting for TRM. The data distribution is invariant under some transforms in G, so we use this knowledge to minimize the empirical transformed risk, better estimator of the test error than the simple empirical risk. From arguments in Chen et al. (2020). it follows that for any fixed Q supported on Gine, the empirical TR has the same mean, but a smaller variance than the empirical risk. Clearly, the same still holds uniformly over all such Q E Ω, which provides justification that the TRM framework provides uniformly more accurate estimator of the population error than RM.
Discussion.
Proof of Corollary 4.3 (Zero Gap).
22

---

Published in Transactions on Machine Learning Research (07 (072))
Proof.
1. From Proposition 4.1 (Part 1), using the assumption H. Go S H, we obtain the lower bound RTRM- 2 RRM- Since there is a distribution Qinv E 12 such that GQ C Ginr. from Proposition we obtain the upper bound RTRM S RRM. Thus, RTRM = -RRM:
Assume now that hMM is minimizer of the risk R(h) over H. Then from equation 10, R(HRMM Qine) R() = RRM = RTRM- So, the pair RM, Qino) is minimizer of the TRM risk.
2. Set U = U(G). From part above, if hRM is minimizer of the standard risk R(h) over H. then (hiM, U) is minimizer of the TRM risk R(h, Q) over H. x L2 since G S Gine and U E Q from the assumptions of the theorem.
From equation we have that RRM RTRM RM U) 2 R(hRM,U), where (x) Egou RAs (gx) E H. Honce, RRM, is also a minimizer the RM risk over H and the pair (AR U) a minimizer of the TRM risk over H. x Ω.
Now observe that RRM is G-invariant, since for all in the group G have:
1 = Eggwir (d gz) I 1R ( go)dd (g')
| Eea
hRM (āx) (U0g) (0g I gea
RRM (ū
Ez-u RM (gx) (x).
Above we changed variables to / = g'g, used that the group is "closed', i.e., for any fixed E G, -g-1E EGG = g E Gg = G, and that dU (āg dU' (ğ), since the Haar measure is right-invariant.
Discussion. This "zero gap" result follows by putting together the upper and lower bounds. requires that we include the correct invariances in distribution Q, but that the transformations in this distribution do not go outside the hypothesis space H. This is an ideal setting, in the sense that use invariances as "inductive biases" to obtain better estimator of the population error, but the hypothesis space is rich enough that all the transforms use can be realized by actual hypotheses.
In addition, the result shows that if the loss is convex, then the optimal classifier is invariant to G, but of course H needs to contain this classifier. In that case, there are many optimization objectives that have that hypothesis as a minimizer, in particular RM and TRM. However, TRM has the advantages described before: (1) reduced variance, and (2) fast inference.
A.2 Proof of PAC-Bayes bound (Theorem 3.1
We will break the proof into three parts. We start by discussing the preliminary theorems from the PAC- Bayes literature that we will use to derive our bound. We then use the assumptions of Lipschitz continuity from the theorem to prove some key lemmas. Finally, we derive the bound on TRM and compute the KL divergence terms for the parametrization of SCALE.
A.2.1 Preliminaries
Consider the hypothesis space parametrized on some space W C RP for some e N. Also consider RRM the standard risk from equation [ and its empirical counterpart RRM (ha Ein (X,,, Y,), and denote = (X,, y D". Assume π prior probability measure on W. Then, we have the following definition due to Alquier et al. (006))
23

---

Published in Transactions on Machine Learning Research (07 (272))
Definition A.2 (Hocffding assumption).
The prior satisfies the Hoeffding assumption if there crists
function and a nonemply open interval I C R2 such that for any X AET and for any w E W:
B (X.))D exp (
R & exp(f(),n)).
Using this assumption, Alquier et al 2016 prove the following thcorem:
Theorem A.3 (Theorem 4.1 (Alquier et al. 2016)). Given any distribution ı over W satisfying the Hoeffding
assumplion, and for any å E (0, IJ, À > O, with probability at least 1 - over the drow of S - (X;, Yi) D", we have simultoneously for all probability measures p & x on W that:
+ ( (KL(Ol(T) + log (7)+ I(,
Evn RRM (hw) S
In particular, the Hooftding assumption above holds for sub- Gaussian losses. The following definition is due to Germain al. (2016a).
Definition A.4. A loss function 1 3" x + IR is sub- Gabb-aassin with variance factor s² under a prior and a data distribution D if V RRM(hhu) 1(h (X), ) is a s² ss-- Gaussian random voriable, i.e.,
x
In Bwwa B(xYY exp (AV) S 2 VEER.
The following lemma is derived from Theorem A.3 above:
Lemma A.5 (Germain et al. (2016a)). Given any distribution т over W, given E (0, >> 0, if the loss is sub- Gaussian with variance factor s², with probability at least 1 - over the drow of S Yi), ~ D", we have simullaaneossly for all probability mensures & π on W that:
Ew-p (lww) S Eu.- + 1 (KLLelIm) + log (t) +
(
X232
2n
In particular, the authors study the case that the loss bounded in [a, bj. More specifically, Hoeffding bound implies that 1 y" x [a, b] is a sub- Gaussian loss with variance factor s² (b a)² and the bound
becomes:
vn (KI(p]mm) + log (t) (b (b-a)2)
2 )
Ewspp Reas (hw) S RRM +
Improved Lemma: This bound can be improved slightly, since we can derive directly from Theorem
A.3 In particular for any fixed w the random variable 1(h (X), Y) is bounded in [a, b] and we know that RaM (hw) (xx (X), Then, from Hoelfding's lemma:
(A2 a)2) [²(6-a)²))
E(X,)) exp A (RRM(hw) 1(h (X), & exp
E
8
Using the above result we find for every w E W and X E lk:
A ES-D" exp A (BRM(hw) (RM.n(hw)) = ESNDD exp n
n
Σ
(RRM(hw)) - (X.,, Y:))
I|
(RRM((hw)) ((hwo( Xi))
i.idd
E B... exp ( (A (hw) (X;), Yi))
exp
11
(b a)² = exp (12 (1²(b-a)) a)
exp
8n
8n
So, the Hoeffding assumption satisfied with f(),n)
(b-a)
and the PAC-Bayes bound becomes:
8n
[TT(pl==) + log (1) +
2 (b a)²
Egrop RRM(hha ) & Bwa RRM (bw) +
(11)
8
24

---

Published in Transactions on Machine Learning Research (07 (07/202) /023)
A.2.2 Lipschitz Lemmas
We prove some key lemmas below. We denote 1(w; X,)) (X)m (), RR BRA((Wo) and R (w) (hm.)
Lemma A. If the map Y() is L-Lisschitz continuous for any (X,Y) E X' x y; then both R and R (.) are L-Lipschitz continuous.
Proof. For any (w, w')= W2
[ (w') - It' (w)l = (lx') -((x X,
S (Exx,))=>l''' ;X,)) -l(w': X, Y) s B(x,y)=DL|lu' -wll =l|lo' w/
For a the analogous argument holds by replacing ExX with the empirical expectation 1 n -E==-
Lemma A.7. ((: X) is L--iipshillz continuous for any (X, E X x then for any distribulion p W with Bu~p[w] wo E W, it holds that:
(wo) - Em (w)) S IN IrCov
Similorry,, for R R
(wo)) - Ew Rh..oRRRMMmn (w)) S L trCov
Proof. Since I(:; X,Y is L-Lipschitz for all (X,Y) E X x V, from Lemma A.6 we know that R'nss (-) is I-Lipschitz. Thus, for wo E W and for any w E W we have:
RRss (wa) Rr (w)) S T.1lu0o w||
(w) T.lluo w|l S RR (wo) S T.|lwo- - will + RRMr(W); = Ewnptthu (w) L.wupplly WD w|l S R (wo) S Weo (w) + LEw wl 1RRM (wo) Ewwp R'R (w) I LEw<pl
w|l-
From the Cauchy-Schward inequality, EEwrypl lwwp w|ll S /Eoool Iwo wl.. Using that Ew-p[w] wp we have: Ewvp w/ - Ew--pl [w]]l trCovw~pllue The desired result follows. The proof is similar for RRM ]
Denote (W, 0) Ry (Bw Qo) and RTRM, 0) = RT Qo). Also define the fized-g loss (w: X,Y) = (g
(gX), Y) 1(wg8,Y)) and the fized-g risk Rg(w)) = RnM (hro g) for all g E Go- Lemma A.8. If 11:;11,,11 Y) is L-Liscchitz continuous for any (X, E X' x y and a distribution Qu on W with EwrnQ. [wo] = w E W, then holds that:
w + RTRM (w, 0) and w + RTRMM (w, 0) are L-Lisschitz continuous for all e e.
2. |RtrM (w, Ew RT (w', 0)1 S L LLC/tooooaaauulu (w v(w,0) 0) E W x θ and same for RTRM
3. In particular, Qm N(w, 0²1) we find:
[ 0) En'nqu (TRM(w' 0) < LyDo, Wwx
and the analogous inequality also holds for RYnM,
Proof. First observe that since I(:; X,Y) is L-Lipschitz for all (X,Y) X ) and gX X, we have that
(-) is L-Lipschitz for all g E Gq. Then, we observe that Ry (w) = R (w). By applying Lemma A.6 to
(-) we find that Rg (-) T-Lipschitz for all g E Ga- Then, applying Lemma A.7 to lg and Q. we find:
(w) Ew'nq.. Ry(w')) [ Ew'nq. (w')l() s LV trCovw (w]].
(12)
25

---

Published in Transactions on Machine Learning Research (07/2023)
Next observe that R (w, 0) = Ry( thus for all w, w' W,
RR 8) (URM 8)1 = -Eg Ro Ry(w") - Qy Rg(w)] =
Eg-200 (Ro(w') - Rn(w))] - 2g R2 (w') - (w)l S Eg- - w| - L| - w|l.
This shows that RTTM (w, 0) is L-Lipschitz, for all 0 E O With similar proof we can also show that w + BTRM, 0) is L-Lipschitz continuous for all 0 E e.
Also, using Eq. 12 and Qu = =N(w, 0²T);
RTRM (w, 0) Bw'nqu B (w',w000 ,0u, - 1B Rig (w) Bw Qy Ro (w')l = Q (w((w) - By Rg(w')) S Q (wx) Ew Rg (2'))) s L trCov [ = L Lypo
The argument for RTRM, is similar.
For some set A, consider a parametrization such that each E A corresponds to transformation ge Go and each pair (w, a) E W x A indexes a predictor I + (hw o (z). We define the fized-parame parametrized-g risk as
Rp(w a) == (R.(()) = Exx, o go (X),
Y) = RRM (Nw o 9a).
Similarly, Rp a) RRM, (Nw ga). Thus, instead of transformations acting on X via Ia_,, we view them as an extension of hypothesis space containing models of the form h 9a. We can now use the bound in Theorem A.3 on the space W x A for Rp(w.a).
A.2.3 BAC-Bayes Bound
The proof continues with the following steps: First, by applying the PAC-Bayes bound on Ry (w,, a) on product posterior measures over W x A, the TRM risk appears naturally, but in the form of an expectation over the models. We then use Lipschitz continuity and the Lemma A.S that we proved to derive the desired bound on TRM. Since the bound holds simultancously for all posteriors over W x A, we further constrain it appropriately. We conclude the proof by computing the KL divergence terms for the transformation and the models. At the end of the argument, we discuss optimization over the free parameters.
Proof of Theorem 3.1
Step 1 (Bound on the expected TRM loss): Consider the prior P3 : x PA with PA = IT P( from Theorem 3.1 Px = =N(wu,, for some fixed wo E W, and posteriors Qu,o QH x Qo on W A with Qo= I Qin as specified by the theorem, with Qи Nww,, (ww, Since QH,, PR, it only remains to prove that Qa < PA, v0 O, which we do later in Step Assuming the Hoeffding assumption holds with (A,n) (bounded losses being example), we apply the bound from Theorem A.3 to d E 10, 0 and Pp to Rp(w,a) a). We find that with probability at least 1 - å over the sampling of S [(x, YOH ~ D" we have simultaneously over any probability measure Qu,o that:
Eninqu Ba-qq Rp(w', a) S Qu Eanq, R ( a) + (KI(Qu,))lP))) log ())+ + f().
=
EwinQu R 0) S Ew RT (w. 0) [KLQQu, + log (1) + f(x.n)).
(13)
Step 2 (Bound on TRM for Lipschitz losses): Now using that X() is L-Lipschitz for any (X, Y) e
X' V/ as well as Lemma A.8 we can derive bound on TRM:
Lem
RTRM(W, 0)
Bnd Qu. RTRM(W. 0) + L IrCovar=qull [w]]]
Eq. B Q RIRM, (w',0) ,0) (KLQQw + log (t) f(2.)) + L ttrCoo [w]]]
Lem
RTnM, (w,0)+ [T..Q mo log (t) s(a,n)) +2 22T///C (oo
26

---

Published in Transactions on Machine Learning Research (07 (07/2023)
Now, considering Qu -N(w, we have:
RTRM (w. 0) S (T,M,)(W;)) (I((Q.,)lI))) + log (5) + s(2.n)) + 21./po.
Since the bound holds for any o > O, we can optimize over this parameter. Moreover, as the bound holds simultaneously for all Qu.o = N(w, o²1), when we optimize the bound, o is permitted to depend on the optimization paramcters. We will discuss optimization over o, after finding the form of the KI-divergence terms.
Step 3 (KL divergence terms): Now will compute the KL(Qw,ellPp) term. We have
K KL(Qw.ellPp) KL(QH. Q 2=1
K
K
x II Pa) = (L + Σ (L(Q
==1
1=1
We start with lo Since Qu N(w, o² and Pæ = Nuwo 22) is known that
1 KL(Qu,ulIPA) (PH) 2p log -
1 s2 1100 woll?
o² s2P
Then, we need derive the expressions for each (1)())))P()).... If i S then the base distributions are parametric, and we denole the KL term by K, m.) while for l < < K the base distributions are parameter-free, and denote the KT. term K₃(z)).
We start with the more challenging K, m.) term. We will omit the index to keep the notation uncluttered. Since none of the Q()): P(s) are absolutely continuous with respect to the Lebesgue measure, we need to resort to the general definition of KI-diivrgence using the Radon-Nikodym derivative.
Let (R, L) be our measurable space. We will use the following definitions in the proof:
1. Dirac measure oa, where (F) 1 E] for all measurable sets E.
2. Lebesgue mcasure À the real line.
Uniform measure U,-A, A], where [[-A,A[(E)) XEn]-A.A) for all measurable sets E. 24
4. Prior mcasure P 36, + (1 3)U(- A], where > 0, BE (0.
5. Posterior measure Q (1 7) TU[-a, a), where E A)
Both P and Q are o-finite measures, since they are probability measures. Recall that a measure is absolutely continuous with respoct to another mcasure v, written as & v, for all measurable sets such that v(A) = O, we also have (uA) 0. Here dy U[-A, A), and Q are absolutely continuous with respect to Indeed, if there some measurable set F. such that P(F) = 0, then
36. (1-3) =
= 50 (E) = U[-A, =
since 3 E (0, Moreover, Q(E) (1- x))(EE) xUI-a a 0, where we also used that [-a, al S [-A, A] which implies Ul-a,a] < A[.
However, P and Q are not absolutely continuous with respect to the Lebesgue measure, since > 0 while A((0)) = 0. However, 6, and U[-a,ai al are absolutely continuous with respect to P, and thus we have:
dQ d((1 - ()6 + = = -1-x) do,
dU[-a
dp +
dP
dP
dP
We can readily verify that the Radon-Nikodym derivatives have the form
dōg (z)
du a a)
A 1/r [-a, (0)).
1/2 0).
(x)
dP
3
dP
a
27

---

Published in Transactions on Machine Learning Research (07 (07/2023)
where the equalities hold P-almost everywhere. Indeed, check that these functions satisfy the required conditions, and then use the uniqueness of the Radon-Nikodym derivative up to P-null sets. Next we compute the Kullback-Leibler divergence ()
/ log
dQ dP
dQ d[( [[[1--]60- + at]-a, all
KL(Q|P)=
dQ
log
JR
dP
3
dQ
dQ dP
= (1 T)
log d
dg +
log
dul a
dQ dP
log dQ dP
= (1 - π) log
+ 2x
1s a- 0))
du -
R
Now, using our formulas for
dQ db
and
dU dU|- -a.a
find
dP ap:
dP
1 - π
1
TA
1 - π 3
A
(1 T)) log
3 2n
+
log
(1
π)) log
+ n log
+ log
2
(1 3)a
1 3
= =L B)) =
where B(T) denotes the Bernoulli distribution with parameter
For (i),),, omitting again the index i to keep the notation clean, we find
Q = (1- + n U/O0 N 1}} +
π
π
N
N
N
B 1 B
8
= 360 + (1 (1- -
B +
N
N
N
Then,
N-1
g(z) p(x)
(1 -- F
+ (N - 1) N/N log x/N
KL(Q1IP) Σ (xx) log
(1 - π + N
log
(3 + 1-9
(1-B)/N
z=0
N
(1 (N - 1)5T
(1 (N (N-1))
+ (N 1)TT
(N - 1)T/NN
log
N
log (N 1) 33/N
N
(N (-11)(1-B)
N
KL B (11-
1 NN
a |IB (1-
1 N (1 (1-8) (1-3))
Step 4 (Putting it all together): We write the whole bound for completeness. Given 5 E [0, A > 0 and
prior measure Px x II Pie it holds with probability at least 1-8 over the draw of S = {(X,, Y:) D and simultancously over w E W and A E O, where 0 = -(aa..... ak, TI...., Rx} and a, E (0, and n, E 10, 1]. that:
RIRM(W, 0) S R (w, +)+
1
1
2p log 8 p+
1 5 |Iw woll== + o² p + Reg((f) + log (t) + f(h.)) + 2LV20 (14)
2
where (eg(O) is defined in equation in Theorem 3.1
TRM bound for a countable W: In the above analysis, Lipschitz continuity used to move from a probabilistic version of TRM to a deterministic one. However, Lipschitz continuity of the loss not needed is when W countable. In that case, by choosing QH = 8 where 6, is the discrete Dirac mass on ha the KL divergence term is equal to: log
+ >/ (I(Q@(lIP(()) and the bound is:
K
KL(QmllP(o) + +((mn)
RT < R (w, +)+
log Pre
(15)
Optimization over the "hyperparameters" a, A;
28

---

Published in Transactions on Machine Learning Research (07 (07/2023)
1. Optimization over 0: As we discussed above, since the bound holds for any > 0, we optimize over this variable before optimising with respect to w. From the right hand side of Eq. 14 we see that the terms involving are:
P A
p logo a + 2 + 21./po.
This is convex function of a, attaining its minimum at o da = (v X (V
vp
The bound would hold simultaneously for all > 0 if we considered posteriors of the form HH,, N(w, r). Then, a-which also appears in the "training error' term-would be permitted to depend on the parameters during optimization.
2. Optimization with respect to X: The bounds we derived hold for any À > O. However, they do not hold uniformly over A. In practice, this implies that we have to set some fixed in the beginning before optimizing over the parameters of the posterior. In particular, it is not permitted for À to depend on these paramcters. In some cases, we can optimize the right hand side of the bound over X, set À to this value, and then optimize over the parameters. This is the case for example in Eq. 11 when K ( does not depend on p. For example if p(w) = 6w
countable space W, 1 V 1n 7[w]8- In our case, even for this simple then (T(p) log त(w in which case X = 5-a scenario discussed in equation
15 optimization over À is still not possible, since the KI-divergence also depends on the parameters 0 of our distribution of transformations Qa.
There are various techniques to address the dependence on x. Typically, union bound over a discrete grid A is used to obtain the PAC-Bayes bound uniformly over X on this grid. This loosens the initial bound by a logarithmic term log |A] in the cardinality of the grid. Langford & Caruana (2022) choose a grid to loosen the bound only by an iterated logarithm factor log log [A] in the cardinality of the grid. Moreover, they relax the upper bound even further to transform the discrete grid into continuous grid. Optimizing over a discrete grid requires multiple training loops, which is not efficiently feasible in our application. Moreover, gradient updates jointly on the parameters of the model, the transformations, and the regularization parameter À may result in unstable training, which moreover can also depend strongly on the initialization.
Another method is to use a bound that is not dependent on any parameter X, such as the bound in (McAllester) 1999)- However, that has been derived only for bounded losses. That bound suggests an objective similar to the one are currently using.
Experimentally, we found out that À v works well in practice. Then, using the optimal o* the bound becomes (e.g., for subgaussian losses with variance factor s²);
vF VT 2T.
RTRM (w, 0) S RTTM, (w, 0)+
1
1
w wo 12 + o² 22 + Reg(O) + log (+)
s²
Vn
2
2p 2plog
p z2
2L
RIRM,
0)
1
w|w woll Reg(0)) cn
(16)
Vn 25
where On = 1
Vn (plog 2Laym
a 0. log 1 + 4 + {) 3
Remark A.9. As mentioned in Dziugoite 8 Roy 2017), bounding the risk RRM(W) with respect to the parameters w E W instead of the elements hw E H. might result in looser bounds due to overparametrization. Much tighter bounds would be obtained if we considered as the quotient space W/~, where the equivalence relation is defined as w w' = hw = had and re-derived the bound directly on this quotient space for Rus (bw). However, finding these equivalences in deep network may in general be intractable. In our selup,
29

---

Published in Transactions on Machine Learning Research (07 2003))
would further tighten the bounds if derived them for Q supported on transformations p E GQ namely Ras (hw: Qo) molher than on the parameters of the tnanformattions E A, namely (T, a).
We could do that in two steps. First, for each base distribution Q()) on A(s), consider GQta as the space Ata/~~ where a ~ d' + p(a) = =p(d). Since the base distributions are usually simple, we could actually establish an isomorphism in this onse by cureful construction of A(o) (e.g., rotations not exceeding 2x). As a second step, consider the distributions Q on the space A, and view GQ as A/~, where for p(a) := Pi (a1) .... isomorphism can be established in this cose too, depending on the structure of the group. For ezample, when the assumptions on the base transformations above hold, and the total transformation space is inner semi- direct product of a larger group, i.e., GQ = GQ ...G..,) the space of compositions can be identified with the Cartesian product.
.............. (x(ak) E GQ have ~ d' = p(a) = =(aa). There are practical scenarios where
Further, we could also consider all posteriors Q S2 instead of the posteriors indered by θ i.e., O 0' - Qo = Qg-- In our example with mixtures of base distributions, Q and Q are actually isomorphic given the conditions above and if not all parometer-free bases are point masses at identity. Finally, one more improvement could be to derive the bounds for posteriors H. GQ, namely RRM (hw 9a) instead of the product space H.S. x Go, i.e., the map (hw,gul + RAM (hw o 9a).
A.3 Unbiased Gradient Estimates
In Eq. G we define the loss for SCALE as:
1 n
L(w 0) Σ Qg (Qro() v:)] + Regg
with 0= (mi,..., K .1,..., .R). We note again that, as stated in 3.2 from the K base transformation distributions we assume that only the first & contain learnable parameters a,, i E [k], respectively.
Computing the gradient of the regularization term ReggRe straightforward, so we will focus on the gradient computation for the first term of the loss. For a fixed s E (1,..., K}, we can write the first term of the loss as:
E ~ Qu) ... Quo ... Egn Quxy [Mh 96. (xx.),
1
(1 (1-Tss=
Bgg Ban ... gg [1
Is- 19+t1 ... ve)
i=1
+
E Q1) ... EgguUs.. ... Egx QUK) (91 ... ................................ ..... gxz.),
(17)
The partial derivative with respect to m is
E Qai
Byx=Ques [h ...1... 1---19-111 ... (x.x))
......g.. 9s ....... + Em ... Eg, Egg Q(x) (g ...
......................... IKTz), v:)))))))))))
(18)
To compute an unbiased estimator of this gradient we first sample M iid copies from the distribution of g 01 ... 9--1- ..- 0K, where each 9x with x is sampled from Q(x), and u, is sampled from U....-
(i) We denote by 9x and u. the 7th sample from the distribution of 9x and us respectively. This leads to the following unbiased Monte Carlo estimator of the partial derivative with respect to shown in equation 18
M
M Σ 5=1
).
ml) ...99K
(..)).1 yi
(ol
...99"000 9x x.), yb
(19)
nM
30

---

Published in Transactions on Machine Learning Research (07/2023)
Now suppose that {I....,k) k) the s-th transformation corresponds to parametric distribution Q(s) with learnable parameter ds. For the partial derivative with respect to Os,, assume that the sample u, ~ Us.., can be written as u, = pa(a) with a a,E, ~ U[-1, and with being a differentiable transformation. We assume that all 9x with X < s are differentiable with respect to their input, i.c., as functions from X to X. Importantly, this covers "discrete" transformations such as flips, which are differentiable transforms of the input object.
Under these assumptions, using the reparametrization trick (Kingma & Welling 2014). which is essentially the same idea as functional/struutural inference, e.g.: Fraser 1966). find the partial derivative with respect to 0g::
T1 Eg Qus - Eg Q (1h -s.... .....gxx.). vi))
aas
1=1
T R Bea(-1.1)
... ((.. ... (x3r),
E QgyeQQK)
(20)
n
da.
Similarly to before, an unbiased estimator of the gradient from equation 20 can be computed by taking M samples from the distribution of each 9x Q(x) with X and M samples from the distribution
~ U[- This leads to the following unbiased Monte Carlo estimator of the partial derivative with respect to 0gg:
M
al(h (o(a)el)) (..66)).)... 96 x;),
T
(21)
nM
da.
Efficient gradient estimation: During training, for each datapoint we take M samples from the distri- bution of gr I[ 9.z. For each 9x: we first sample tx Bern( and ux Uxay if x S k uy ~ Ux if x > k. Then we set 9x = 1tx 0)T + 1 (tx = 1) This means that among the M samples from the distribution of gr, the expected number of transforms Ix sampled from the transformation distribution Uxiox or Ux M. To obtain an efficient training procedure, we develop estimators for the derivatives with respect to all parameters of Qa using only these M samples.
For {l,...,K] K),, we modify the estimator from equation (when n, (0, 1)) follows:
19 for the partial derivative with respect to R
G) Σラ= 1(15) =1) 1 (No ... u. ...9K x.),
M
n
E/ 1( =)
1=1
ΣM 1( =0 (1)
(1) ()) ... ( x2)x).11 vb
...9--190H
(22)
Σ/ 1((h)) =0)
For s E {1.....,}} we modify the partial derivative with respect to as as follows:
M
(0(ha(()))) po(a, (1) 22).
1(19) =)
(23)
nM
da
==
One can readily verify that these estimators remain consistent as the number M of samples tends to infinity.
The estimator from equation equation 22 can be problematic when is very close to zero or unity, which results in 1(4, 1) or 1(1, 0) being zero with high probability. When this happens, we set the corre- sponding term in the estimator (i.e., the fraction whose denominator is zero), to zero. In addition, to avoid this scenario, especially early on in the training process, before convergence, we- somewatthat heuristically- constrain π, to the interval [c, 1 -c for some E (0, 0.5). as we will describe in Section A.4 During training, reduce linearly, so that at the end of training is effectively zero.
31

---

Published in Transactions on Machine Learning Research (07/ 2223)
A.4 Training Details using SCALE
When we train using SCALE, optimize both the network parameters and the parameters of the distri- bution Qo in a single training process. We have generally observed that the algorithm is not too sensitive to the hyperparameters, and a relatively straightforward and quick tuning process (say, a small grid search over a few values) is enough to find good defaults. The experiments were executed using GeForce RTX 2080 Ti GPUs.
First we present the hyperparameters that are kept fixed in all datasets. We found these hyperparameterss by grid search to ensure stable convergence behavior across multiple datasests. We initialize all a, to 0.1 and all my to 1/K, where recall that K is the number of transformations composed. Additionally, to avoid having H, converge to zero or unity at the beginning of the training, we constraint all т; to the interval [c, 1 c-. We initialize c = 0.4/K and reduce it linearly (by a constant after each epoch) so that at the end of the training it equals zero. Since we have a pre-determined number of epochs for each dataset (discussed below), this determines the amount to decrease after cach epoch. For the computation of the gradient presented in A.3 we observe that with M 2 4, the training stays nearly unaffected by the number of Monte Carlo samples M. We M 4 to reduce the required computation. For while for CIFAR 10/100 we use a batch size of 64.
we use batch size of 128,
Finally we use the following network-specific hyperparameters:
MNIST/roLMNIST: We use the same network architecture used in Benton et al. (2020 which consists of five convolutional layers and one fully connected layer. To optimize the network's param- eters we use the Adam optimizer (Kingına & Ba 2015) with a learning rate of 0.02 and with cosine annealing Loshchilov & Hutter (2017). For the regularization term, in the rotMNIST and MNIST experiments Areg is set to 0.006. To learn the augmentation parameters Tú, after each epoch, we linearly reduce the learning rate from 0.001 to zero in order to make the training more robust to the choice of the hyperparameter Areg. We train for 200 epochs by jointly optimizing the parameters of Qs and the parameters of the network, then we train for another 100 epochs by optimizing only the parameters of the network.
CIFAR 10/100: We use a WideResNet 28-10 (Zgoruyko & Komodakis 2016). We optimize the parameters of the network with an SGD optimizer with learning rate of 0.1, cosine anncaling, weight decay of 0.0001 and Nesterov momentum [Nesterov] 1983) with value 0.9. For the CIFAR 10 experiment, Areg is set to 0.01, for CIFAR 100 Areg is set to 0.02. To learn the augmentation parameters Ty, after each epoch, we linearly reduce the learning rate from 0.001 to zero. We train for 200 epochs by jointly optimizing both the parameters Qo and the network parameters.
For all reported results (when no test-time augmentation is applied), we present the average test accuracy over 5 trials over all elements of randomness the training (random initialization, random sampling of transforms, etc). The standard deviation of the test accuracy is 0.12% for the RotMNIST experiments, 0.09% for the CIFAR 10 experiments, and 0.1% for CIFAR 100.
Additionally, when we apply test-time augmentation, there is an additional variance in the model accuracy due to the randomly sampled augmentations during inference. Table g shows the mean accuracy and the standard deviation in the results due to the random choice of test-time augmentations. Again, we can see that the variance of the test accuracy is small enough that it does not affect the quantitative comparison between the models. Finally Figure shows the final paramoters Te, a learned using SCALE on CIFAR10/100.
A.5 Implementation Details for Parametrization
In the parametrization presented in Section 5.1 we include the following transformations: rotation (both a continuous version and a discrete rotation by 180"), X-axis scaling, Y-xxis scaling, X-axis shearing, horizontal flip, and random cropping. We implement the continuous rotation, scaling, and shearing using the differentiable image sampling method proposed in Spatial Transformers Networks by Jaderberg et al. 20155. We implement each of the above transformations follows:
32

---

Published in Transactions on Machine Learning Research (07/2023)
Test (Mean accuracy/ Standard deviation) of a Wide ResNet 28-10 trained on CIFAR 10/100, using various methods for learning data augmentations. All methods are evaluated using test-time augmentation with = 4, 8 samples. The standard deviation is computed using five trials over the random choice of test-time augmentations.
Table 3:
AUGERINO
FAST-AA
FASTER-AA
DADA GEOMETRIC
SCALE
+FLIPS
(Ours)
TESTAUG. (N=4)
95.7% 0.1%
97.0% 0.08%
96.6%
0.09%
96.5%/0.1%
96.7%/ 0.09%
CIFAR1O
TESTAUG. (N=8)
95.8% 0.1%
97.2% 0.04%
96.8%
0.07%
96.8% 0.11%
96.9%/ 0.07%
TESTAUG. (N-4)
79.2%
0.2%
82.5 % 0.1%
82.0% / 0.09%
79.8% 0.1%
83.0%
0.12%
CIFAR100
TESTAUG. (N 8)
79.3% 0.2%
82.9%
0.08%
82.4% 0.1%
80.1% 0.1%
83.3% 0.1%
(ateggofeclppoocaallllleee/e/sce Sealing
Categorical probabilities (n) CIFAR100 Sealing
Initial Learned
Initial Learned
Sheongd
Shearte
0.6
0.6
1144
134
callngs
HFlp
Hilp
Cn9
Crop
Alaha parameters CIFARIC
Alp ha parameters CIFARIOU
Rotcttoo
Rotation
Initial Leanred
Initial Learned
Ro Caf
6.22
Bny
Scaling
ScalingX Scaling
aaing
Shearing
Shearing
Figure 7: Initial values of ni, as and the final parameters learned after training with SCALE on CIFAR10 and CIFAR100
Rotation: For the continuous rotation, sample the angle a (in radians) by sampling U[-1,1] and taking dj ays. Then perform the rotation using the following rotation matrix:
cos(a1) sin(a1)
sin (ar) 0
(as
0
(24)
0
1
33

---

Published in Transactions on Machine Learning Research (07/2023)
Discrete rotation by 180%: This is performed using the same rotation matrix from equation 24
the general continuous rotation, with angle replacing on in equation 24 The angle is sampled uniformly from {0, 7).
X-axis scaling: Similar to the rotation, we sample = a2E by sampling E ~ U[-1,1]. Then the scaling is performed using the transformation matrix:
E 0
0
1
1
Y-axis scaling: We sample as = age by sampling E ~ U[-1,1]. Then the scaling is performed using the transformation matrix:
1 0
0 c
0
0
1
X-axis shcaring: We sample a4e by sampling using the transformation matrix:
U[-1,1]. Then the shearing is performed
1 04 0
0
0
0
1
Horizontal flip: To perform the horizontal flip, sample dg uniformly from 13. If as -1, then the order of the pixels on the X-axis is flipped, otherwise the input image stays unchanged. To combine the flips more efficiently with the rotations, scalings, and shearings, we use the trans- formation matrix
as
0 0
0
1
0
0
0
1
Random Crops: For the random cropping we use the implementation from the Torchvision package
within PoTorch (Paske et al. 2019) with padding equal to four.
A.6 SVHN Experiment: Precomputed policies are not transferable
Precomputed augmentation policies that are produced by computationally expensive search mothods similar Fast-AA are available only for a limited set of the most. frequently used datasets. Additionally, there is no guarantee that we can succesfully transfer these policies new datasets. As result, for any new dataset, user needs to search for best policy, which highlights the importance of methods with low computational cost.
In this section we perform experiments on the SVHN dataset (Ntzer et al. 2011 showcasing that pre- computed policies are not perfectly transferable even between similar datasets. SVHN contains three sets of samples: the "train" set containing 73,257 difficult training samples, the "extra" set containing 531,131 less difficult training samples and the "test" set containing 26,032 test samples. For the precomputed policy, we use a Fast-AA policy computed for the whole SVHN dataset ("train" set and "extra" set), to train a network smaller training set that contains only the more difficult training samples ("train" set).
Table compares the accuracies achieved using SCALE, Augerino, DADA (with the same augmentation space) and the precomputed Fast-AA policy to train on the "train" set of SVHN. We show that SCALE outperforms both the precomputed Fast-AA policy and Augerino and DADA.
2016 trained with learning rate 0.001, cosine annealing and the Adam optimizer (Kingma & Bal 2015). For the SCALE algorithm we used A initialized at 0.02 and for Augerino used Areg 0..5.
For this experiment we used a ResNet34 network (He
34

---

Published in Transactions on Machine Learning Research (07/2023)
Table 4: Test Accuracy on SVHN when the model is trained on the "train" set using Augerino, the Fast-AA policy (precomputed for the whole SVHN), DADA and SCALE. We show the average accuracy over 5 trials. For all methods the standard deviation is less than 0.12%
AUGERINO
FAST-AA (PRECOMPUTED POLICY) GEOMETRIC
DADA
SCALE
(Ouns) 97.4%
SVHN
97.2%
97.1%
96.6%
A.7 A Bayesian Perspective of ETRM
In this section, we will see how variational inference (MacKay 2003 Wainwright & Jordan, 2008) on a specific graphical model suggests optimizing a similar objective as the one we arrived to in Theorem 3.1 using PAC-Bayes bounds. By optimizing this objective we find the model hn E and distribution Q that minimize tradeoff between the ETRM loss (equation 3) and negative entropy.
First build the corresponding graphical model. To simplify the discussion assume that all random variables have a density. A reasonable analogue in the setting of graphical models to the transformation hypothesis space Q in TRM is a corresponding prior on the transformations before socing any data, we can assume p(g) is the uniform density on Gạ: this choice can be viewed as a non-informative prior that maximizes the prior entropy. Instead of the standard observation model conditional censity-p(ylla)
X + X. In particular,
supervised learning, our observation model now is p(ylg. During training, we assume that the functional form of this conditional distribution is known, but depends on unknown parameters For example, in binary classification Pw (ylg..2 x) Bern(y; (gz)), where the form of hw is known. Training the graphical model amounts to learning the unknown non-random parameters w. We train the graphical model using a dataset S = (m,, v:) 1- We denote Sy = { 1: S2 {z;) We further assumne that the a,s are independent. (æ; .IL xz), while the yos are conditionally independent (y, 11 M3 Vzlg). Figure depicts one such graphical model.
W
yi
Npst
Yteet
Figure 8: Probabilistic Graphical Model analogue. Circles represent random variables and arrows represent conditional probabilities. Filled circles are observed variables. Dots represent parameters.
We do not need to model the density of 2,s since during training we maximize the conditional likelihood (marginalizing over g), solving max, log log (S, (.)). Note that for + j. we do not know that yı II. yj (4, when is not observed (i.e., we may have ys not D-separated in the graphical model). So, the above expression not separable in the samples and thus not casily optimized. Explicitly writing the expectation over does not seem make the optimization more convenient, since the expression becomes:
max mog log (S,1(S = log Eg) II (ys
(25)
This is not idcally suited for stochastic optimization, bocause we cannot casily construct unbiased estimators of the gradient with respect to the parameters w. Additionally, the product of distributions may not be numerically stable as increases.
35

---

Published in Transactions on Machine Learning Research (07/2023)
Since exact training is not convenient, we can train the graphical model using variational inference by maxi- mizing a lower bound on this objective. First we need to approximate the posterior P(g]S). To be consistent with the TRM analogy, we choose the variational posterior family to be the same as the transformation hypothesis space Q2 In particular, since in TRM we sample independent transforms, this family contains distributions Q E f2 for which Q(gly, = Q(1).
This might seem restrictive family first glance, however, since g transformation acting on I, the transformed sample gr is still dependent on the initial sample z. Only the distribution of the transformation Q(g) is sample-independent For example, this allows for uniform distribution over rotations in some fixed range to be in the variational family, assigning the same probability to two different rotated vectors given the original vector, i.e., pl (Rox)z>) pl Raa [z'), where Rg is the rotation matrix with angle 0. On the other hand, a uniform distribution over rotations in some range that is dependent on the sample x cannot be in the variational family.
Variational inference suggests finding the best approximate posterior in the variational family as measured by the Kullback-Leiin Leibler divergence between the variational posteriors Q(g) Q and the true posterior p(g]S). Finding both the optimal w and the optimal Q E Ω involves maximizing the following objective jointly over w and Q:
[(w, Q) log Pw. (SylS,) KL(Q())lP(((sS)
log P.. (S,],,,) (KL(Q())))() (g~(()) log (S,1S)) pno
By cancelling terms, this further equals
Eg) log Pw. Ps (Sy1S.. 9) KL((())))()) = Qg) log Pu (v (. KL(Q()))P())
2 E (G)) log g) KL(Q()))())
(26)
We used the independence IL for all E [na] and yi TT Vylg. Z4, T. for i A from the graphical model. The objective is clearly a lower bound for equation equation 25 as the Kullback-Teibh divergence is non-negative and can be zero only if p(g]S) E
Using the prior on g defined above as a uniform density on Gq (which has larger support than the variational posterior, thus the Kullback-Leibler term is well-defined), minimizing Equation equation 26 over w, Q is equivalent to solving:
Σ Egg (w. + HQQ).
(27)
QEB
where H(Q) == -Eg-ooloog Qg)] is the differential entropy of Q.
The equation above is similar to the one from Theorem 3.1
1. The first term is the ETRM objective with loss - log P.. (vlg, x). For example, if the observation model is vlg. ~ Bern(hw then - log (ws mx) - [y; log (gr;)) (1 (1-20 w:) log(1 - h (ga:)) which is the binary cross entropy loss, commonly used for binary classification.
2. The second term is regularizer that promotes the selection of distributions with high entropy.
At test time, the Bayesian approach is to find Pw (Btest (reat, S), which results in averaging over the posterior of the transformations as follows:
(ytst (restt S) =E, (eppl())) (pm (ytest Ig,, Ftest))
(28)
36

---

Published in Transactions on Machine Learning Research (07/2023)
where used that TT Trest[S and Gteat IL Sig, Ftest from the graphical model. Hence, if use the variational posterior Q" found by optimizing equation 27 to approximate p(g]S), and the model h found by optimizing the same objective, we can construct the estimate:
P(ytest Ttest, S) = By-qq [p(Utext (gItest))]
which is a direct analogue of using test time augmentation.
If the global variable were substituted by local variables for each Ii then amortized variational inference in the resulting now graphical model would also result in an objective similar to the one from Theorem 3.1 However, the prediction step would be slightly different.
37


Reviews:
[{"review_id": "ZZBPQazazx", "reviewer": "Reviewer_QAqa", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "--- The strengths of this paper are:\n​\n1. The paper is clearly written and easy to follow.\n\n2. The proposed TRM framework extends the standard RM framework to factor in data transformations. It provides a practical solution to the long-standing problem of choosing proper data augmentation strategies with incomplete knowledge only.\n\n3. The SCALE algorithm is simple to implement and incorporate into standard training algorithms. Experiment results show that it is comparable in terms of classification performance to those methods that are more computationally expensive. In addition, the models trained by SCALE also tend to be well-calibrated.\n​\n--- The weaknesses (mostly nitpicking) are:\n​\n1. The regularizer involves some hyperparameter (e.g., $\\lambda_{\\text{reg}}$ that requires extra tuning for each dataset, which can be time-consuming.\n\n2. Lack of explanation of the necessity of linearly decay of regularizer $\\lambda_{\\text{reg}}$ Does it imply that the joint training does not converge when using a fixed $\\lambda_{\\text{reg}}$?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "​\nThis paper first proposes a new risk minimization (RM) objective, i.e., transformed risk minimization (TRM), that accounts for both the data distribution and the transformation distribution. The authors then derive an upper bound for the population risk in terms of the empirical risk and regularization terms. They also provide a practical algorithm, stochastic compositional augmentation learning (SCALE), for joint optimization over the model and data augmentation parameter space. They show theoretically that TRM can recover the distributional invariance and the corresponding optimal model under certain conditions. Empirically, SCALE is comparable to state-of-that-art methods and recovers the rotational invariance on rotMNIST dataset.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. There are a few works that use gradient-based methods and jointly learn the model and data augmentation in one training loop [1] [2]. How is SCALE compared with them in terms of classification performance and training time cost (GPU hours)?\n\n2. The paper does not include experiments on larger datasets such as SVHN and ImageNet. It is good to see whether SCALE still has good generalization when the dataset becomes larger.\n​\n[1] Lin, Chen, et al. \"Online hyper-parameter learning for auto-augmentation strategy.\" Proceedings of the IEEE/CVF International Conference on Computer Vision. 2019.\n​\n[2] Li, Yonggang, et al. ‘Differentiable Automatic Data Augmentation’. *Computer Vision -- ECCV 2020*, edited by Andrea Vedaldi et al., Springer International Publishing, 2020, pp. 580–595.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "OglUB6xTjR", "reviewer": "Reviewer_tzE7", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths**\n\n- I think that the general way of posing the problem as Transformed Risk Minimisation, akin to vicinal risk minimisation, is elegant and clean. As I detail in the \"Requested changes\" section, I think it should be more clearly mentioned that this is essentially identical to previous approaches to learning data augmentations, but I think that none of these approaches wrote things in such a clean and simple way.\n\n- The parametrization of augmentations, combined with unbiased gradient estimates, is quite clever, and could be useful.\n\n- A big challenge when learning data augmentations by risk minimisation is regularization, and the PAC-based approach of the authors is a fresh perspective on these issues\n\n**Weaknesses**\n\n- The authors did miss a few recent works that tackle the same problem in similar ways. The idea of dealing with discrete augmentations with mixtures was previously used by Hataya et al. (2020) and in a more applied setting  by Raghu et al. (2022). These works relied on continuous relaxations of discrete random variables, whereas the authors here use a nice trick to derive unbiased estimates (Appendix A.3), so there is some novelty, but it is hard to assess it clearly.\n\n- The PAC-based regularization is interesting, but not really well-motivated. Regularisation is clearly needed, as the authors argue, but why go in that direction? Doing a fair comparison (in the form of an ablation study, with grids of hyper-parameters) between this new regularization and simpler ones (e.g. the L2 regularization used by Benton et al.) would enable the community to assess the usefulness of this new regularization. There are also more complex approaches to avoir learning singular augmentation by overfitting, e.g. the line of work initiated by Lorraine et al. (2020), and followed by Hataya et al. (2020) and Raghu et al (2022) ; or in a Bayesian context, Schwöbel et al. (2022), Immer et al. (2022)\n \n- It feels difficult to gain some clear scientific insights from the experiments. What are the new bricks designed by the authors that are most useful, is it the regularization, the way to compute gradients, the parametrization of augmentations?\n\n\n**Additional references**\n\n\nImmer et al., Invariance Learning in Deep Neural Networks with Differentiable Laplace Approximations, NeurIPS 2022\n\nLorraine et al., Optimizing Millions of Hyperparameters by Implicit Differentiation, AISTATS 2020\n\nHataya et al., Meta Approach to Data Augmentation Optimization, arXiv:2006.07965, 2020\n\nSchwöbel et al., Last Layer Marginal Likelihood for Invariance Learning, AISTATS 2022\n\nRaghu et al., Data Augmentation for Electrocardiograms, CHIL 2022\n\n\n\n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The authors tackle learning of data augmentation distributions.\n\nThey cast this as a risk minimisation problem, where the data augmentation is an additional model parameter; the goal is therefore to learn a data augmentation that generalizes well, i.e. that minimises the so-called *transformed population risk*. To this end, they minimize a PAC-Bayes-inspired regularization of the empirical risk of the augmented classifier. They also provide some simple yet insightful theoretical results on transformed risk minimization.\n\nTo optimise efficiently the augmentation distribution, they rely on a clever decomposition combined with the reparametrisation trick.\n\nThey show that, in practice, their approach can learn efficiently useful augmentations.\n\n", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "\n**Thinks that could improve the paper:**\n\n- Many previous approaches can be recast as TRM (in the sense that they try to minimise the transformed population risk), with sole specific choices of loss functions, and regularizer, for instance Augerino, Loraine et al. (2020), etc. It is perhaps a bit of a stretch to call TRM \"a new theoretical framework\" as the authors do on page 2 because of this. It would be interesting to list and comment such examples.  This would illustrate nicely the usefulness of the framework, and make the theory of Section 4 more broadly applicable.\n\n- Doing more ablation-like experiments (see the \"Strengths and Weaknesses\" section). I think it would me more interesting to known which contribution are most useful in general, and not necessarily that the results are better than others.\n\n- I feel like putting the related work section after the experiments harms the narration, I think putting it after the introduction would be much clearer\n\n**Small clarifications:**\n\n- I find it a bit confusing that the authors say \"we aim to optimize [the empirical transformed risk]\" just before Equation (3), whereas, as they mention later, actually optimizing it would lead to degenerate augmentations. I feel like I would be clearer to say that the ultimate goal is to optimize the transformed population risk, and that to that end, they optimise a regularized version of the empirical transformed risk.\n\n- For MNIST, you mention that the learned range of the rotations was [-0.31,0.31], what about rotMNIST? It would be nice to do curves similar to those of Fig. 3 (left and center) but for the range of the rotations and not just for $\\pi_i$-s, perhaps in the appendix.\n\n\n\n", "broader_impact_concerns": "I do not have particular concerns there.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "YhMRRZXQ6m", "reviewer": "Reviewer_idNZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Unfortunately I don't think this idea contributes much to the task of generalization. First of all, the parameterized family of augmentations would have to capture the actual desired invariances in order for this to be helpful (e.g., using rotations on rotated MNIST), but this is a chicken-and-egg problem. We don't know what the correct family of distributions are---if we did, we could just apply it directly. Showing that this does something meaningful on toy data doesn't surprise me and I don't expect it to generalize.\n\nI understand the idea of \"learning\" an augmentation distribution automatically, but this isn't a particularly inspired or creative idea, and it really doesn't seem too different from just trying out several augmentations and choosing the best one via a holdout set. Furthermore, it could even be _more_ amenable to overfitting---there's no real reason to expect this approach to improve generalization, and in fact I think this approach can actually _increase_ model complexity. I had this thought before seeing the experiments section, which confirm that existing augmentation methods do better quite consistently.\n\nFinally, isn't this idea basically what AutoAugment already does? It learns a policy for augmentations based on the data and training process of the model, so it's already doing the same thing? I just don't see what this work does that is new or interesting, even as a proof-of-concept, not to mention the fact that empirically it doesn't seem to actually work.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes to learn a distribution over input transformations which will hopefully improve generalization. They make this tractable by making an approximation to the true \"ideal\" objective by pulling the expectation wrt the augmentations outside, and then parameterize a family of distributions of augmentations to optimize over using SGD.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "No", "audience_explanation": "", "requested_changes": "Why is this idea different from existing methods which learn augmentations, like AutoAugment?\n\nDo you have any experiments on non-toy settings where this method does better than (or even matches) existing methods? This last point would not be necessary in my opinion if the work were doing something fundamentally new, but I don't believe it is.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Please include the following changes in the final revision:
1. Make it clear that you are not claiming that the gradient estimators you describe are novel and cite the relevant prior work. The reparameterization of the uniform distribution for continuous augmentation parameters has already been used in Augerino, and the mixture-based gradient estimator for the discrete case is an application of Local Expectation Gradients [1].
2. Explain the exact setup and motivation for annealing the regularization term in the loss. Currently the paper seems to suggest that you are linearly reducing the weight of the regularizer in the loss to zero during training, but in a response to a reviewer, you said that the linear decay also applied to the other term in the loss function and thus is equivalent to using linear learning rate decay. If this is indeed the case, why is this necessary if you are already using (cosine) learning rate decay as described in the paper?
 
[1] Local Expectation Gradients for Black Box Variational Inference, Michalis K. Titsias, Miguel Lázaro-Gredilla, NeurIPS 2015
Decision: Accept (poster)