How Data Augmentation affects Optimization for Linear Regression
Authors: Boris Hanin
Yi Sun
Venue: NeurIPS 2021 Poster (2021)
Keywords: data augmentation
stochastic optimization
convex optimization

Abstract: Though data augmentation has rapidly emerged as a key tool for optimization in modern machine learning, a clear picture of how augmentation schedules affect optimization and interact with optimization hyperparameters such as learning rate is nascent. In the spirit of classical convex optimization and recent work on implicit bias, the present work analyzes the effect of augmentation on optimization in the simple convex setting of linear regression with MSE loss.

We find joint schedules for learning rate and data augmentation scheme under which augmented gradient descent provably converges and characterize the resulting minimum. Our results apply to arbitrary augmentation schemes, revealing complex interactions between learning rates and augmentations even in the convex setting. Our approach interprets augmented (S)GD as a stochastic optimization method for a time-varying sequence of proxy losses. This gives a unified way to analyze learning rate, batch size, and augmentations ranging from additive noise to random projections. From this perspective, our results, which also give rates of convergence, can be viewed as Monro-Robbins type conditions for augmented (S)GD. 

How Data Augmentation affects Optimization for Linear Regression
Boris Hanin"
Yi Sun' Department of Statistics University of Chicago
Department of Operations Research and Financial Engineering
Princeton University bhanindprinceton. edu
yisundstatistics. uchicago. edu
Abstract
Though data augmentation has rapidly emerged key tool for optimization in modern machine learning, a clear picture of how augmentation schedules affect optimization and interact with optimization hyperparameters such learning rate is nascent. In the spirit of classical convex optimization and recent work on implicit bias, the present work analyzes the effect augmentation on optimization the simple convex setting of linear regression with MSE loss.
We find joint schedules for learning rate and data augmentation scheme under which augmented gradient descent provably converges and characterize the resulting minimum. Our results apply to arbitrary augmentation schemes, revealing complex interactions between learning rates and augmentations even in the convex setting. Our approach interprets augmented (S)GD stochastic optimization method for a time-varying sequence of proxy losses. This gives a unified way to analyze learning rate, batch size, and augmentations ranging from additive noise to random projections. From this perspective, our results, which also give rates of convergence, can be viewed Monro-Robbins type conditions for augmented (S)GD.
Introduction
Data augmentation, popular set techniques which data is augmented (i.c. modified) at every optimization step. has become increasingly crucial to training models using gradient-based optimiza- tion. However, in modern overparametrized settings where there are many different minimizers the training loss, the specific minimizer selected by training and the quality of the resulting model can be highly sensitive to choices augmentation hyperparameters. As a result, practitioners use methods ranging from simple grid search Bayesian optimization and reinforcement learning [ 9 [7] to select and schedule augmentations by changing hyperparameters over the course of optimization. Such approaches, while effective, often require extensive compute and lack theoretical grounding.
These empirical practices stand contrast theoretical results from the implicit bias and stochastic
optimization literature. The extensive recent literature on implicit bias [15 29 [32 gives provable guarantees on which minimizer of the training loss is selected by GD and SGD in simple settings, but considers cases without complex scheduling, On the other hand, classical theorems in stochastic optimization, building on the Monro-Robbins theorem (25), give provably optimal learning rate schedules for strongly convex objectives. However, neither line of work addresses the myriad augmentation and hyperparameter choices crucial to gradient-basee training effective in practice.
The present work takes step towards bridging this gap. We consider two main questions for learning rate schedule and data augmentation policy:
Equua contribution
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

1. When and at what rate will optimization converge?
2. Assuming optimization converges, what point does converge to?
To isolate the effect on optimization of jointly scheduling learning rate and data augmentation schemes, we consider these questions in the simple convex model of linear regression with MSE loss:
L(W:D) N
|| W
(1.1)
In this setting, we analyze the effect of the data augmentation policy the optimization trajectory We of augmented (stochastic) gradient descen 2
which follows the update equation
W++ W, "I V ((W; w
(1.2)
Here, the dataset D YX contains N datapoints arranged into data matrices X E R and Y E R N whose columns consist of inputs z; E R" and outputs 34 E RR. In this context, we take flexible definition data augmentation scheme any procedure that consists, at every step of optimization, of replacing the dataset 'D by a randomly augmented variant which we denote by D: (X,. Yt). This framework is flexible enough to handle SGD and commonly used augmentations such as additive noise [1], CutOut [1]], SpecAugment [23], Mixup [35. and label-preserving
transformations (e.g. color jitter, geometric transformations (26))).
We give a general answer to Questions and 2 for arbitrary data augmentation schemes. Our main result (Theorem 3.1 gives sufficient conditions for optimization to converge in terms of the learning rate schedule and simple 2nd and 4th order moment statistics of augmented data matrices. When convergence occurs, we explicitly characterize the resulting optimum in terms of these statisties. We then specialize our results to (S)GD with modern augmentations such as additive noise [14] and random projections (e.g. CutOut [12 and SpecAugment (23)). In these cases, find learning rate and augmentation parameters which ensure convergence with rates to the minimum norm optimum for overparametrized linear regression. To sum up, main contributions are:
1. We analyze arbitrary data augmentation schemes for linear regression with MSE loss, obtaining explicit sufficient conditions on the joint schedule of the data augmentation policy and the learning rate for (stochastic) gradient descent that guarantee convergence with rates in Theorems 3.1 and 3.2 The resulting augmentationn-eeendeen optimum encodes the ultimate effect of augmentation on optimization, and we characterize it in Theorem 3.1 Our results generalize Monro-Robben theorems [25] to situations where the stochastic optimization objective may change at each step.
2. We specialize our results to (stochastic) gradient descent with additive input noise ($4) or random projections of the input ($5). proxy for the popular CutOut and Augment augmentations 112 23 In each case, we find that jointly scheduling learning rate and augmentation strength is critical for allowing convergence with rates the minimum norm optimizer. We find specific schedule choices which guarantee this convergence with rates (Theorems 4.1 4.2 and 55]] and validate our results empirically (Figure 4.1). D. This suggests explicitly adding learning rate schedules to the search space for learned augmentations as in 8 9 which we leave future work.
2 Related Work
In addition to the extensive empirical work on data augmentation cited elsewhere in this article, we briefly catalog other theoretical work on data augmentation and learning rate schedules. The latter were first considered in the seminal work (25). This spawned a vast literature on rates of convergence for GD, SGD, and their variants. We mention only the relatively recent articles [T 27 (2) and the references therein. The last of these, namely (22), finds optimal choices of learning rate and batch size for SGD in the overparametrized líncar setting.
A number of articles have also pointed out in various regimes that data augmentation and more general transformations such as feature dropout correspond in part to f2-type regularization on model parameters, features, gradients, and Hessians. The first article of this kind of which we are aware is [
Both GD and SGD fall into our framework. To implement SGD, take 2, to be subset of D.

---

which treats the case of additive Gaussian noise (see $ More recent work in this direction includes
There are also several articles investigating optimal choices of l Eg-regularization for linear models (cf e.g. [3 31 2. These articles focus directly on the generalization effects of
30 19
ridge-regularized minima but not on the dynamies of optimization. We also point the reader to [00]. which considers optimal choices for the weight decay coefficient empirically in neural networks and analytically in simple models.
We also refer the reader to number of recent attempts to characterize the benefits of data augmenta- tion. In (24], for example, the authors quantify how much augmented data, produced via additive noise, is needed to learn positive margin classifiers. [6], in contrast, focuses on the case of data invariant under the action of a group. Using the group action to generate pabee-pres augmenta- tions, the authors prove that the variance of any function depending only on the trained model will decrease. This applies in particular to estimators for the trainable parameters themselves. [10] shows augmented k--N classification reduces to a kernel method for augmentations transforming each datapoint to finite orbit of possibilities. It also gives a second order expansion for the proxy loss of a kernel method under such augmentations and interprets how each term affects generalization. Finally, the article (34) considers both label preserving and noising augmentations, pointing out the conceptually distinct roles such augmentations play.
3 Time-varying Monro-Robbins for linear models under augmentation
We seek to isolate the impact of data augmentation on optimization in the simple setting of augmented
(stochastic) gradient descent for linear regression with the MSE loss [.T]. Since the augmented
dataset Dr at time is a stochastic function of DD we may view the update rule (122) as a form of stochastic optimization for the proxy loss time
CWW == ED, [C(W; D:)l
(3.1)
which uses an unbiased estimate of the gradient of C(W;D; from single draw of Dr. The connection between data augmentation and this proxy loss was introduced in 3 5). but we now consider it in the context of stochastic optimization. In particular, augmentation, which allows the distribution of D; to change with and thus enables optimization to converge to points which are not minimizers of the proxy loss Le (W) at any fixed time.
consider scheduling the
Our main results, Theorems 3.1 and 3.2 provide sufficient conditions for jointly scheduling learning rates and general augmentation schemes guarantee convergence of augmented gradient descent in the linear regression model (I.I). Before stating them, we first give examples of augmentations falling into our framework, which we analyze using our general results $4 and $
Additive Gaussian noise: For SGD with batch size B, and noise level o, 0, this corresponds to X, = cs A.A,,, + o, Gt) and Y, = =yYA, where G, is a matrix of i.i.d. standard Gaussians, A, N B, has i.i.d. columns with single non-zero entry equal to 1 chosen uniformly random and et /N/B, is normalizing factor. The proxy loss is
(W) C(W; D) of || 1
(3.2)
which adds an le penalty. We analyze this case in $4
Random projection: This corresponds to X, = n and Y, = Y, where II, is an orthogonal projection onto a random subspace. For 7 tr( the proxy loss is
1 Z.(W) = N | wx + N - | |WI + O(n
adding data- dattadeepnndett & penalty and applying Stein-type shrinkage input data. We analyze this in $
In addition to these augmentations, the augmentations below also fit into our framework, and Theorems and beyond the general description given in Theorems and 3.2 more difficult, and we thus leave interpretion of these specializations to future work.
and 3.2 apply. However, in these cases, explicitly characterizing the learned minimum
3

---

{W x
[++1(W) =0]
W
wiQn
We+t
{W R
[[(W) 0)
W2
(W &R Z,(W)
Span [B[X,X]])
ViI Span [[[X Span (B(X)=1X2.11)
(a) The time proxy loss L. non-degenerate on VI with minimal norm minimizer W7. The increment
(b) The proxy losses L. and L(+1 at consecutive times will generally have different minimal norm minimiz- ers WT,, W The increment E; W W? measures this change.
W++1 W, is in Vij so only the projection WiQu of
W, onto Vi changes.
Figure 3.1: Schematic diagrams of augmented optimization in the parameter space Rox,,
transformations: For 2-D image viewed as vector R", geometric transforms (with pixel interpolation) or other aabel-preserving transforms such as color jitter take the form of linear transforms R - R" We may implement such augmentations in our framework by X₂ = A,X and Y, = for some random transform matrix A,-
Mixup: To implement Mixup, we can take X, = X A, and Y; = YA, where A, E R has i.i.d. columns containing two random non-zero entries equal - O and a with mixing coefficient drawn independently from Beta(o,, az) distribution for parameter a,
3.1 A general time-varying Monro-Robbins theorem
Given augmentation scheme for the overparameterizee linear model I.I.. the time t gradient update at learning rate T/r
2T N
W == W
(Y: WxX,)x
(3.3)
where D. (Xi, is the augmented dataset time t. The minimum minimizer of the corresponding proxy loss E, (see (.I is
W == E[XXX]]][[XX
(3.4)
where X[X,X,]+ denotes the Moore-Penrose pseudo-inverss (see Figure 3. In this section we state a rigorous result, Theorem 3.1 giving sufficient conditions on the learning rate nh and distributions of the augmented matrices Xx, Yr under which augmented gradient descent converges. To state it, note that (3.3) implies that each row of We+1 W, is contained in the column span of the Hessian X,X, of the augmented loss and therefore almost surely belongs to the subspace
column span of E[X,X,] C R",
(3.5)
as illustrated in Figure 3. The reason that, in the orthogonal complement VjJ the augmented loss L(W: D.) has zero gradient with probability 1. To ease notation, assume that Vjj is inde- pendent of f. This assumption holds for additive Gaussian noise, random projection, MixUp, SGD, and their combinations. It
necessary in general, however, and refer the interested reader Remark B.2 in the Appendix for how to treat the general case.
Let us denote by Q11 _. + R the orthogonal projection onto Vi (see Figure 3.l.p. As already pointed out, at time f, gradient descent leaves the matrix of projections (Id -Q =1) of each row of
We onto the orthogonal complement of VI unchanged. In contrast, [| W || decreases at
rate governed by the smallest positive cigenvalue Amin, Vi (E [x,xT]] of the Hessian for the proxy

---

loss Lt, which is obtained by restricting its full Hessian E [ to VI. Moreover, whether and at what rate 11 W converges to depends on how quickly the distance
Ei == W+1 W
(3.6)
between proxy loss optima at successive steps tends zero (see Figure 3.1b).
Theorem 3.1 (Special case of Theorem B.I). Suppose that Vii is independent of t, that the learning rate satisfies T)t O, that the proxy optima satisasatissy
Eille 00,
(3.7)
t=0
ensuring the existence of a limit W lim
Wi, that
me Mmin V (B[X,X,]) = 00
(3.8)
and finally that
Σ oZ I E[X,X]]I + + - E[xX
< 80.
(3.9)
Then, for any initialization Wo, we have that WiQj converges in probability to W
If the same augmentation is applied with different strength parameters at each step / (e.g. the noise level of for additive Gaussian noise), we may specialize Theorem 3.1 to this augmentation scheme. More precisely, translating conditions (3.7). (388) 3..9 into conditions on the learning rate and augmentation strength gives conditions on the schedule for 7/ and these strength parameters to ensure convergence. In and and $
we do this for additive Gaussian noise and random projections.
When the augmentation scheme static in t, Theorem 3.1 reduces to standard Monro-Robbins theorem (25 for the (static) proxy loss (W). As in that setting, condition (3.8 enforces that the learning trajectory travels far enough to reach an optimum, and the summand in Condition (3.9) bounds the variance of the gradient of the augmented loss L(W; DD) ensure the total variance of the stochastic gradients is summable. Condition (3.7) is new and enforces that the minimizers W of the proxy losses (W) change slowly enough for augmented optimization procedure to keep pace.
3.2 Convergence rates and scheduling for data augmentation
Refining the proof of Theorem 3.1 gives rates of convergence for the projections WiQj of the weights onto Vii to the limiting optimum W When the quantities Theorem 3.1 have power law decay, we obtain the following result.
Theorem 3.2 (Special case of Theorem B.4p. Suppose Vi is independent of and the learning rate satisfies v + O. Moreover assume that Jor some 0 < a < 1 < 31,32 and >> huve
Ti minn V (E(X,X;)) 52(1 "),
| 01)
(3.10)
as well
T X
X, = O(
(3.11)
and
B [E]..](X,x] - E(X,X,I) (Yxx) B| DMI O(
(3.12)
Then, for any initialization Wo. we have for any E the following convergence in probability:
mmi (at - ", - [ W l1 & o.
may be surprising that W[W appears in condition 3.12 Note that TE [W is the gradient descent trajectory for the varying sequence of deterministic proxy losses (W). To apply Theorem 3.2 may first study this deterministic problem to show that E[W, converges to W at some rate and then use (3.12) to obtain rates of convergence of the true stochastic trajectory W, to W
In and below, we specialize Theorems 3.1 and 3.2 obtain rates of convergence for specific augmentations. Optimizing the learning rate and augmentation parameter schedules in Theorem 3.2 yields power law schedules with convergence rate guarantees in these settings.
5

---

Special Case: Additive Gaussian Noise
We now specialize our main results Theorem 3.1 and 3.2 to the commonly used additive noise augmentation [4. Under gradient descent, this corresponds to taking augmented data matrices
X, X o,G, and Y, =,Y,,
where G2 E R=XN are independent matrices with i..d. standard Gaussian entries, and is strength parameter. Under SGD (with replacement), this corresponds to augmented data matrices
X, (X
G.) and =Y=YYA
where A, E RRNx.. has i.i.d. columns with single non-zero entry equal to chosen uniformly at random and G2 A//BB B, is a normalizing factor. In both cases, the proxy loss is
1 N Y W W N
Co,
(4.1)
which to knowledge was first discovered in
3
Before stating our precise results, we first illustrate how jointly scheduling learning rate and augmen- tation strength impacts GD for overparameterized linear regression, where
N #data points < input dimension
71.
(4.2)
The incquality 4.2 ensures L(W;D) D) has infinitely many minima, of which consider the minimum
norm minimizer
Wiin YXT(XXT)
most desirable. Notice that steps of vanilla gradient descent
2m
W W₂
(Y
-Wx)xx
(4.3)
N
change the rows the weight matrix W, only in the column space Vi = colspan(XXX] S R". Because Vi * R by the overparamettrivation assumption (4.2). minimizing L(W; D) without augmentation cannot change W. the matrix whose rows are the components of the rows of W orthogonal to Víi This means that GD converges to the minimal norm optimizer Wmin only when each row of Wo belongs to VII As this explicit initialization may not be available for more general models, we seek to find augmentation schedules which allow GD or SGD to converge to Wmin without it, in the spirit of recent studies on implicit bias of GD.
4.1 Joint schedules for augmented GD with additive noise to converge to Wmin
We specialize Theorems 3.1 and 3.2 to additive Gaussian noise to show that when the learning rate т)t and noise strength o are jointly scheduled to converge to 0 at appropriate rates, augmented gradient descent can find the minimum norm optimizer Wmin-
Theorem 4.1. Consider any joint schedule of the learning rate ti and noise variance of in which both Te and of tend to O and ag is non-decreasing. If the joins schedule satisfies
DC Σ
mo? 0o
and
Σnh
ho? 00.
(4.4)
then the weights Wt converge in probability to the minimal norm optimum W regardless of the initialization. Moreover, the first condition (4.4) necessary for E[W;] converge to W min-
If we further have T) = =((==) and of ((- with > 0, x + < 1, and 2r + > 1 so that
Ti and of satisfy 4.4 then for small s > 0, we have emin(v. [ W min P 0.
The conditions of 4.4 require that 7) and o, be jointly scheduled correctly to ensure convergence
to Wiim and akin to the Monro-Robbine conditions (25) for convergence of stochastic gradient
methods. We now give an heuristic explanation for why the first condition from (44) is necessary. In this setting, the average trajectory of augmented gradient descent
E[ = E[W,) - V Lo, (W) w=
(4.5)

---

Mean Squared Error (GD)
Mean Squared Error (SGD)
Unauimmented
Unaugmented
104
104
Gauss const
Gauss const
10a
Gauss exp
Gauss exp
Gauss pow
Gauss pow
10a
10a
104
10
10"
2000
4000
6000
8000
10000
10000
20000
30000
40000
50000
Steps
Steps
Orthogonal Weight Norm [W; ||
(GD)
Orthogonal Weight Norm I|W, lle (SGD)
Unaugmented
Unaugmented
Gauss
Gauss const
Gauss exp
Gauss exp
Gauss pow
Gauss pow
10P
109
2000
4000
6000
8000
10000
10000
20000
30000
40000
50000
Steps
Steps
Figure 4.1: MSE and [[W,, ||r for optimization trajectories of GD with additive Gaussian noise augmentation and SGD with additive Gaussian noise augmentation under different augmentation schedules. For both GD and jointly scheduling learning rate and noise variance to have polyno- mial decay is necessary for optimization converge to the minimal norm solution Wmin. min- Gauss const, Gauss exp, and Gauss pow have Gaussian noise augmentation with of = 2, 2e 0.02t 2(1 +[1+, 50 ) -0.33 respectively. All other details are given in $4.3
is given by gradient descent on the ridge regularized losses Co, (W). If = > 0 is constant, then [[W will converge to the unique minimizer W of the ridge regularized loss C.. This point W has zero orthogonal component, but does not minimize the original loss L.
To instead minimize L, must choose schedule satisfying dr O. For the expected optimization trajectory to converge to Wmin for such a schedule, the matrix E[W,, of components of rows of [Wa] orthogonal to Vi must converge to O. The GD steps for this matrix yield
E[ (1 [o
II ==0
(4.6)
Because of approaches 0, this implies the necessary condition noo E oo for E -
This argument illustrates a key intuition behind the conditions 4.4h. The augmentation strength or must decay 0 to allow convergence must be carefully tuned to allow the implicit regularization of the noising augmentation to kill the orthogonal component of W, in expectation. In similar manner. the second expression in 4.4) measures the total variance of the gradients and ensures that only a finite amount of noise is injected into the optimization.
true minimizer of the training loss, but this convergence
Although Theorem 4.1 is stated for additive Gaussian noise, an analogous version holds for arbitrary additive noise with bounded moments. Moreover, optimizing over æ, y, the fastest rate of convergence guaranteed by Theorem 4.1 is obtained by setting т/t l-2/3+e of -1//3 and results in Ott-1/3+)) rate of convergence. It is not evident that this is best possible, however.
7

---

4.2 Joint schedules for augmented SGD with additive noise to converge to Wmin
To conclude our discussion of additive noise augmentations, we present the following result on convergence to W Wmin in the presence of both Gaussian noise and SGD (where datapoints in each batch are selected with replacement).
Theorem 4.2. Suppose of 0 is decreasing, Y)t + O, and we have
Σ not 00 and Σ 1=0
m < 00.
(4.7)
t=
Then the trajectory W, of SGD with additive noise converges probability W min for any initializa- tion. If we further have = (t=²)) and of = (tt with x >
2 y > 0 and r + < 1, then for
any e we have that (min(y,2=)=ee
W min
3 0.
Theorem 4.2 is the analog of Theorem 4.1 for mini-batch SGD and provides example where our framework can handle the composition of two augmentations, namely additive noise and mini- batching. The difference between conditions 4.7 for SGD and (4.4) for GD accounts for the fact that the batch selection changes the scale of the gradient variance at each step. Finally, Theorem 4.2 reveals qualitative difference between SGD with and without additive noise. If 7): has power law decay, the convergence of noiseless SGD (Theorem F. exponential in 6, while Theorem 4.2) gives power law rates.
4.3 Experimental validation
To validate Theorems 4.1 and 4.2 ran augmented GD and SGD with additive Gaussian noise on N = 100 simulated datapoints. Inputs were i.i.d. Gaussian vectors in dimension n 400, and outputs in dim p 1 were generated by random linear map with i.i.d Gaussian coefficients drawn from 1). The learning rate followed a fixed polynomially decaying schedule π) = 100 (batch size) MSE and ||W, |F along single optimization trajectory with different schedules for the variance of used in Gaussian noise augmentation. Complete code to generate this figure is provided in supplement zip in the supplement. ran 30 minutes on standard laptop CPU.
+ ) and the batch size used for SGD was 20. Figure 4.1 shows
For both GD and SGD, Figure 4.1 shows that the optimization trajectory reaches Wmin only when both learning rate and noise variance decay polynomially to zero. Indeed, Figure 4.1 shows that if of zero (blue) or exponentially decaying (green), then while the MSE tends to zero, the orthogonal component Ti, does not tend to zero. Thus these choices of augmentation schedule cause W, to converge to an optimum which does not have minimal norm.
On the other hand, if of remains constant (orange). then while W. tends to zero, the MSE is not minimized. Only by decaying both noise strength and learning rate to sufficiently slow polynomial rates (red) prescribed by Theorem [4. do find both MSE and W.,, tending to 0. meaning that augmented (S)GD finds the minimum norm optimum W min under this choice of parameter scheduling.
5 Special Case: Augmentation with Random Projections
We further illustrate our results by specializing them to class of augmentations which replace each input in batch by its orthogonal projection 11,X onto random subspace. In practice (e.g. when using CutOut 2 or SpecAugment (23]), the subspace is chosen based on a prior about correlations between components of X, but consider the simplified case of uniformly random subspace of R" of given dimension.
At each time step t we fix a dimension h; and fixed --dimensionaal subspace 3, of IR". Define the random subspace Se by
S. : = (..) (Q a S.).
where Qi E (n) a Haar random orthogonal matrix. Thus, S, is uniformly distributed among all k dimeennsionaal subspaces in R At step take the augmentation given by
X, I1X
Y,=Y,
II. Q,
8

---

where II, is the orthogonal projection onto S₁ and hence II, is the orthogonal projection onto St. Denoting by " ke/n the relative dimension of St, direct computation (see Lemma E.1 reveals that the proxy loss (W) equals ; plus
(1 71) X |
71 (1/n 2/n²) N + 1/n - 2/n²) (IIWX
1
|II [W] (5.1) n2
Neglecting terms of order (nn this proxy loss applies Stein-type shrinkage input data by n and adds a data-dependent εg penalty. For 7/8 < 1, the minimizer of the proxy loss
(5.1 is
YXT W %/
21 + 1/n- 2/n² 2/n²
XX
7t 1/n 2/n2
| Id
1 1/n
n
Again, although W does not minimize the original objective for any n 1, the sequence of these proxy optima converges to the minimal norm optimum in the weak regularization limit. Namely, have lim.,-1- W W Specializing our general result Theorem [3.1 to this setting, obtain explicit conditions under which joint schedules of the normalized rank of the projection and the learning rate guarantee convergence the minimum norm optimizer Wmin-
Theorem 5.1. Suppose that т) 00 Te 1 with "t non-decreasing and
Σ (1)= oc and
2E T6) <c.
(5.2)
Then, W. 2 Wmin- Further, is m) = O((=*) and iit = O(( with y > 0, x + < 1, and
2r + y > 1, then for small > we have that emin (v. (v. - W Wmin ll
0.
Comparing the conditions 5.2 of Theorem 5.1 to the conditions 4.4 of Theorem 4.1 we see that
1 - Yi a measure of the strength of the random projection preconditioning. As in that setting, the fastest rates of convergence guaranteed by Theorem 5.1 obtained by setting f ( 2/3++ and V 1/3, yielding a (13+)/ rate of convergence.
6 Discussion and Limitations
We have presented a theoretical framework to rigorously analyze the effect of data augmentation. As be seen in main results, framework applies completely general augmentations and relies only on analyzing the first few moments of the augmented dataset. This allows us to handle augmentations diverse as additive noise and random projections as well as their composition in uniform manner. We have analyzed some representative examples in detail in this work, but many other commonly used augmentations may be handled similarly: label- aabel-preeerving transformations (e.g. color jitter, geometric transformations) and Mixup [35], among many others. Another line of investigation left future work is compare different methods of combining augmentations such as mixing, alternating, or composing, which often improve performance in the empirical literature [16].
Though results provide rigorous baseline compare
complex settings, the restriction of the present work to linear models is a significant constraint. In future work, we hope to extend our general analysis models closer to those used in practice. Most importantly, we intend consider more complex models such as kernels (including the neural tangent kernel) and neural networks by making similar connections stochastic optimization. In an orthogonal direction, our analysis currently focuses on the mean square loss for regression, and we aim to extend it to other losses such as cross-entropy. Finally, our study has thus far been restricted the effect of data augmentation on optimization, and it would be of interest to derive consequences for generalization with more complex models. We hope our framework can provide the theoretical underpinnings for a more principled understanding of the effect and practice of data augmentation.
Broader Impact
Our work provides a new theoretical approach to data augmentation for neural networks. By giving a better understanding of how this common practice affects optimization, we hope that it can lead to

---

more robust and interpretable uses of data augmentation in practice. Because our work is theoretical and generic, we do not envision negative impacts aside from those arising from improving learning algorithms in general.
Acknowledgments and Disclosure of Funding
a pleasure to thank Daniel Park, Ethan Dyer, Edgar Dobriban, and Pokey Rule for number of insightful conversations about data augmentation. B.H. was partially supported by NSF grants DMS- 1855644 and DMSS2133806 and ONR MURT "Theoretical Foundations of Deep Learning." and DMS-2054838.
Y. S. was partially supported by NSF grants
References
[1] F. Bach and E. Moulines. Non-strongly-convex smooth stochastic approximation with con- vergence rate o(1 (1/)).. In Advances in neural information processing systems, pages 773-781, 2013.
[2] L. Bartlett, P. M. Long, G. Lugosi, and A. Tsigler. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 2020.
[3] C. M. Bishop. Training with noise is equivalent to Tikhonov regularization. Neural computation, 1116 1995.
[4] L. Bottou, F. E. Curtis, and Nocedal. Optimization methods for large-scale machine learning. Siam Review, 60(2):223-31 211 2018.
[5] O. Chapelle, J. Weston, L. Bottou, and V. Vapnik. Vicinal risk minimization. In T. K. Leen, G. Dietterich, and V. Tresp, editors, Advances in Neural Information Prucessing Systems 13, pages 416-422 MIT Press, 2001.
[6] S. Chen, E. Dobriban, and J. H. Lee. Invariance reduces variance: Understanding data augmen- tation in deep learning and beyond. stat, 1050:25, 2019.
[7] B. Collins and P. Sniady. Integration with respect to the Haar measure on unitary, orthogonal and symplectic group. Communications in Mathematical Physics, 264(3):773-795, Mar 2006.
[8] F. D. Cubuk, B. Zoph, D. Mane, V. Vasudevan, and Q. V. Le. Autoaugment: Learning augmentation strategies from data. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 113-123, 2019.
[9] E. D. Cubuk, B. Zoph, Shlens, and Q- V. Le. Randaugment: Practical automated data augmentation with reduced search space. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition Workshops, pages 702-703, 2020.
[10] T. Dao, A. Gu, A. Ratner, V. Smith, C. De Sa, and C. Re. A kernel theory of modern data augmentation. volume 97 of Proceedings of Machine Learning Research, pages 1528-1537, Long Beach, California, USA, 09-15 Jun 2019. PMLR.
[11] A. Défossez and F. Bach. Averaged least-mean-sen -queanss Sirs-nariance trade-offs and optimal sampling distributions. In Artificial Intelligence and Statistics, pages 205-213, 2015.
[12] T. DeVries and G. W. Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv: /00:::708.0000 04552, 2017.
[13] P. Goyal, P. Dollár, R. Girshick, P. Noordhuis, L. Wesolowski, A. Kyrola, A. Tulloch, Y. Jia. and K. He. Accurate large minibatch SGD: Training Imagenet hour. arXiv preprint 2017.
[14] Y. Grandvalet and S. Canu. Noise injection for inputs relevance determination. In Advances in intelligent systems, pages 378-382. IOS Press, 1997.
[15] S. Gunasekar, J. Lee, D. Soudry, and Srebro. Characterizing implicit bias in terms of optimization geometry, 2020.
10

---

[16] D. Hendrycks, N. Mu, E. D. Cubuk, B. Zoph. J. Gilmer, and B. Lakshminarayanan. AugMix: simple data processing method to improve robustness and uncertainty. Proceedings of the International Conference on Learning Representations (ICIRR), 2020.
[17] D. Ho, E. Liang, X. Chen, I. Stoica, and P. Abbeel. Population based augmentation: Efficient learning of augmentation policy schedules. In International Conference on Machine Learning, pages 17224,,,, 2019.
[18] D. Huang, J. Niles-Weed, J. A. Tropp, and R. Ward. Matrix concentration for products, 2020.
[19] D. LeJeune, R. Balestriero. H. Javadi, and G. Baraniuk. Implicit rugosity regularization via data augmentation. arXiv preprint arXiv:1905. 1633, 2019.
[20] A. Lewkowyez and G. Gur-Ari. On the training dynamics of deep networks with L2 regulariza- tion, arXiv preprint Niv:20006.08643 2020.
[21] Liu, A. Najmi, and M. Sundararajan. The penalty imposed by ablated data augmentation. arXiv preprint arXiv:2006.04766 2020.
[22] S. Ma, R. Bassily, and M. Belkin. The power of interpolation: Understanding the effectiveness of SGD in modern over- parametrized learning. In International Conference on Machine Learning, pages 3325-3334, PMLR, 2018.
[23] D. Park, W. Chan, Y. Zhang, C.-C. Chiu, B. Zoph, E. D. Cubuk, and Q V. Le. SpecAugment: A simple data augmentation method for automatic speech recognition. Proc. Interspeech 2019, pages 2613-2617, 2019.
[24] Rajput, Z. Feng, 7. Charles, P.-L. Loh, and D. Papailiopoulos. Does data augmentation lead 2019.
to positive margin? arXiv preprint
[25] H. Robbins and S. Monro. A stochastic approximation method. Amn. Math. Statist., 22(3):400- 407, 09 1951.
[26] P. Simard, D. Steinkraus, and J. Platt. Best practices for convolutional neural networks applied visual document analysis. In Seventh International Conference on Document Analysis and Recognition, 2003. Proceedings., volume 2, pages 958-958, 2003.
[27] S. L. Smith, P.-. Kindermans, C. Ying, and Q. V. Le. Don't 't decay the learning rate, increase the batch size. In International Conference on Learning Representations, 2018.
[28] S. L. Smith and Q. Le. A Bayesian perspective on generalization and stochastic gradient descent. In International Conference Learning Representations, 2018.
[29] D. Soudry, E. Hoffer, M. S. Nacson, S. Gunasekar, and N. Srebro. The implicit bias of gradient descent on separable data, 2018
[30] S. Wager, S. Wang, and P. S. Liang. Dropout training as adaptive regularization. In Advances in neural information processing systems, pages 351-359, 2013.
[31] D. Wu and J. Xu. On the optimal weighted Lg regularization in overparameterized linear regression. arXiv preprint arXiv:2006.05800 2020.
[32] J. Wu, D. Zou, V. Braverman, and Q. Gu. Direction matters: On the implicit bias of stochastic gradient descent with moderate learning rate, 2021.
[33] L. Wu, C. Ma, and E. Weinan. How SGD selects the global minima in over pver-parameterized learning: A dynamical stability perspective. In Advances in Neural Information Processing Systems, pages 9279-8288, 2018.
[34] S. Wu, H. R. Zhang, G. Valiant, and C. Ré. On the generalization effects of linear transformations in data augmentation, 2020.
[35] H. Zhang, M. Cisse, N. Dauphin, and D. Lopez-Paz, mixup: Beyond empirical risk mini- mization. arXiv preprint arXiv:1710.09412, 2017.


Reviews:
[{"review_id": "U8w8-baqgvd", "reviewer": "Reviewer_wBTQ", "summary": "This paper studies the joint effect of data augmentation and learning rate scheduling on the optimization trajectory of overparameterized linear model. The authors study the conditions under which (S)GD will converge to the min-norm interpolator, even when the data augmentation distribution changes over time. With a more careful analysis, the authors also study the rate of convergence. Specific augmentations studied here are additive Gaussian noise and random projection.", "questions": "", "limitations": "", "rating": 6, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The proofs are technical and the overall problem studied is interesting. It is important to study how data augmentation affects the optimization trajectory. As the authors have mentioned, my main concerns are that 1) the model discussed here is limited to linear case, how can this case provide us insight for more general classes like kernels? and 2) although the problem formulation is general enough to include many types of augmentations, the detailed results are derived only for additive noise and random projects, which may not be what people mostly use nowadays.\n\nWhat appears more concerning is that it is unclear to me why studying the convergence to the min-norm solution is a good idea. On a high level, augmenting with noise should make the learned minimizer more robust (at least to noise). As the authors have mentioned (and proven before), training with random noise corresponds to training with l2 regularization (eq 3.2). Therefore I'm not sure why the authors consider min-norm solution as the limiting classifier. I'm also curious what if the condition eta_t->0 is relaxed (at least is allowed to decay much slower), the limiting behavior may be more interesting.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "The limitations has been adequately addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "Q6JEaxWyUOP", "reviewer": "Reviewer_sQPi", "summary": "This paper provides a theoretical framework on the convergence of (S)GD when considering learning rate and data augmentation. It presents the consistency and convergence rate results of a general optimization method, and gives some detailed examples of adaptive Gaussian noise and random projection. \n", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: this paper is novel based on my knowledge. It provides a new theoretical framework to justify the convergence of optimization under changing learning rate and (either random or arbitrary) data augmentation. The results not only covers the situations where data augmentation is independent to iteration t, but also include more general results when data augmentation is related to t.\n\nQuality and clarity: the paper is well written and the main text is easy to understand. The claims are well supported by the theoretical results. The detailed examples for the general theorems are clearly written. The authors are careful and mention both strengths and weaknesses (only for linear regression model) of this work. \n\nSome minor issue: \n\n(1) the current numerical experiments are weak to support the claims. However, since this paper is already well supported by theories, please consider moving the experiments to appendix and emphasize some other important things.\n\n(2) the statements from line 120 to 128 are somewhat confusing. Please specify why they cannot be obtained from Theorem 3.1/3.2. If A_t is randomly generated and independent to iteration t, they can be viewed as some new data distribution, so they are supposed to converge in the optimization.\n\nSignificance: this paper provides detailed theories in the optimization of linear model, extending to other complicated algorithms is challenging, but it could potentially provide deeper understanding in data augmentation.\n\n\n\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Some limitations are mentioned in the above.\nThere is no negative social impact.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5"}, "extra_text": {}}, {"review_id": "H9f3qp1HICk", "reviewer": "Reviewer_6v6C", "summary": "This paper rigorously analyzed the convergence of data augmentation in linear regression under the convex setting. Its main result gives sufficient condition on the joint schedule of the learning rate and the augmentation strength under which asymptotic convergence rates are provided. Specializations The above result is then instantiated for additive Gaussian noise, random projections.", "questions": "", "limitations": "", "rating": 3, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Originality: The convergence results of data augmentation algorithms are solid. The analysis techniques seem non-trivial\n\nQuality: Some assumptions seems to strong. How to justify the assumptions (3.10) (3.11) and (3.12) are reasonable? Something for assumptions (3.7) (3.8) and (3.9).\n\nClarity: Overall, the writing is clear to me. \n\nSignificance: First, I'm not convinced by the experimental results. It only runs on simulated data for additive Gaussian noise case. It would be better if the authors can provide the results for random projection case, and more importantly, if some real datasets can be used. I think additional experiments would have made the paper much stronger. Second, this paper only studies optimization of data augmentation. However, I think the most interesting and important thing is about the generalization of data augmentation, since the key goal of using data augmentation is to avoid over-fitting. Besides, it would also be better if the authors could provide some generalization results (such as testing error) in the experiments. ", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No. See \"Significance\" in Main Review section.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "EQVA70H7Jqj", "reviewer": "Reviewer_cnWP", "summary": "This paper rigorously studies the sufficient conditions of the convergence of GD/SGD when data augmentation is applied to learn an overparameterized linear model. Some of the sufficient conditions are also necessary. In particular, the sufficient conditions are formulated in terms of the learning rate, 2nd and 4th moment of augmented data matrices.", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "-\tThis paper presents a solid theory that shed light on the understanding of the hyperparameters of data augmentation and GD/SGD, which affords many interesting insight, like Fig. 3.1. On the other hand, more numerical analysis with synthetic data is needed to further test the theory. In particular,\no\tNo numerical results are presented for the random projection augmentation in Sect. 5\no\tIt would be interesting to see if the sufficient conditions are sharp. If the selection of hyperparameters start to deviate from these conditions, would the error be slowly increasing, or would there be a phase transition from small error to large error?\n-\tDoes the convergence result in Theorem 4.2 hold for any selection of minibatch size S_t that may be time dependent? Does the selection of the minibatch size and sigma_t need to jointly satisfy some constraints?\n-\tIn Fig. 4.1, why does the MSE of Gauss exp schedule performs nearly the best for GD, but only outperforms the Gauss const for SGD? What is it in GD and SGD that leads to such difference?", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations and society impacts are adequately addressed.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}]

Author rebuttal:

Meta-review: As far as I am aware, this is the first work that presents theoretical results for optimization dynamics with data augmentation, and three of the four reviewers found the paper to be of sufficient significance and impact to warrant acceptance. I ultimately found the argument that the paper should be rejected on the grounds that it analyzes optimization dynamics instead of generalization error unconvincing. While I generally agree that the primary motivation for data augmentation is to prevent overfitting, it's not obvious to me that the effects of data augmentation on optimization are well understood or trivial, and it's not even a priori obvious to me that optimization should necessarily converge for all data augmentation schemes outside of basic ones like additive Gaussian noise.
Decision: Accept (Poster)