How Data Augmentation affects Optimization for Linear Regression
Year: 2021

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.

heading: 1 Introduction
text: Data augmentation, a popular set of techniques in which data is augmented (i.e. modified) at every optimization step, has become increasingly crucial to training models using gradient-based optimization. However, in modern overparametrized settings where there are many different minimizers of the training loss, the specific minimizer selected by training and the quality of the resulting model can be highly sensitive to choices of augmentation hyperparameters. As a result, practitioners use methods ranging from simple grid search to Bayesian optimization and reinforcement learning [8, 9, 17] 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 in contrast to 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 in [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-based training effective in practice.
The present work takes a step towards bridging this gap. We consider two main questions for a learning rate schedule and data augmentation policy: ∗Equal contribution
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
1. When and at what rate will optimization converge? 2. Assuming optimization converges, what point does it 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) = 1 N ‖Y −WX‖2F . (1.1)
In this setting, we analyze the effect of the data augmentation policy on the optimization trajectory Wt of augmented (stochastic) gradient descent2, which follows the update equation
Wt+1 = Wt − ηt∇WL(W ;Dt) ∣∣ W=Wt . (1.2)
Here, the dataset D = (X,Y ) contains N datapoints arranged into data matrices X ∈ Rn×N and Y ∈ Rp×N whose columns consist of inputs xi ∈ Rn and outputs yi ∈ Rp. In this context, we take a flexible definition data augmentation scheme as any procedure that consists, at every step of optimization, of replacing the dataset D by a randomly augmented variant which we denote by Dt = (Xt, Yt). This framework is flexible enough to handle SGD and commonly used augmentations such as additive noise [14], CutOut [12], SpecAugment [23], Mixup [35], and label-preserving transformations (e.g. color jitter, geometric transformations [26])).
We give a general answer to Questions 1 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 statistics. 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, we find learning rate and augmentation parameters which ensure convergence with rates to the minimum norm optimum for overparametrized linear regression. To sum up, our 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 augmentation-dependent optimum encodes the ultimate effect of augmentation on optimization, and we characterize it in Theorem 3.1. Our results generalize Monro-Robbins 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), a proxy for the popular CutOut and SpecAugment augmentations [12, 23]. In each case, we find that jointly scheduling learning rate and augmentation strength is critical for allowing convergence with rates to the minimum norm optimizer. We find specific schedule choices which guarantee this convergence with rates (Theorems 4.1, 4.2, and 5.1) and validate our results empirically (Figure 4.1). This suggests explicitly adding learning rate schedules to the search space for learned augmentations as in [8, 9], which we leave to future work.

heading: 2 Related Work
text: 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 [1, 11, 4, 27, 22] and the references therein. The last of these, namely [22], finds optimal choices of learning rate and batch size for SGD in the overparametrized linear 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 `2-type regularization on model parameters, features, gradients, and Hessians. The first article of this kind of which we are aware is [3],
2Both GD and SGD fall into our framework. To implement SGD, we take Dt to be a subset of D.
which treats the case of additive Gaussian noise (see §4). More recent work in this direction includes [5, 30, 19, 21]. There are also several articles investigating optimal choices of `2-regularization for linear models (cf e.g. [33, 31, 2]). These articles focus directly on the generalization effects of ridge-regularized minima but not on the dynamics of optimization. We also point the reader to [20], which considers optimal choices for the weight decay coefficient empirically in neural networks and analytically in simple models.
We also refer the reader to a number of recent attempts to characterize the benefits of data augmentation. 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 label-preserving augmentations, 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-NN classification reduces to a kernel method for augmentations transforming each datapoint to a 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.

heading: 3 Time-varying Monro-Robbins for linear models under augmentation
text: 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 (1.1). Since the augmented dataset Dt at time t is a stochastic function of D, we may view the update rule (1.2) as a form of stochastic optimization for the proxy loss at time t
Lt(W ) := EDt [L(W ;Dt)] (3.1)
which uses an unbiased estimate of the gradient of L(W ;Dt) from a single draw of Dt. 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, we consider scheduling the augmentation, which allows the distribution of Dt to change with t and thus enables optimization to converge to points which are not minimizers of the proxy loss Lt(W ) at any fixed time. Our main results, Theorems 3.1 and 3.2, provide sufficient conditions for jointly scheduling learning rates and general augmentation schemes to guarantee convergence of augmented gradient descent in the linear regression model (1.1). Before stating them, we first give examples of augmentations falling into our framework, which we analyze using our general results in §4 and §5.
• Additive Gaussian noise: For SGD with batch size Bt and noise level σt > 0, this corresponds to Xt = ct(XAt + σt · Gt) and Yt = ctY At, where Gt is a matrix of i.i.d. standard Gaussians, At ∈ RN×Bt has i.i.d. columns with a single non-zero entry equal to 1 chosen uniformly at random and ct = √ N/Bt is a normalizing factor. The proxy loss is
Lt(W ) = L(W ;D) + σ2t ‖W‖2F , (3.2)
which adds an `2 penalty. We analyze this case in §4.
• Random projection: This corresponds to Xt = ΠtX and Yt = Y , where Πt is an orthogonal projection onto a random subspace. For γt = tr(Πt)/n, the proxy loss is
Lt(W ) = 1
N ‖Y − γtWX‖2F +
1
N γt(1− γt)
1 n ‖X‖2F · ‖W‖2F +O(n−1),
adding a data-dependent `2 penalty and applying Stein-type shrinkage on input data. We analyze this in §5.
In addition to these augmentations, the augmentations below also fit into our framework, and Theorems 3.1 and 3.2 apply. However, in these cases, explicitly characterizing the learned minimum beyond the general description given in Theorems 3.1 and 3.2 is more difficult, and we thus leave interpretion of these specializations to future work.
(a) The time t proxy loss Lt is non-degenerate on V‖ with minimal norm minimizer W ∗t . The increment Wt+1 −Wt is in V‖, so only the projection WtQ‖ of Wt onto V‖ changes. (b) The proxy losses Lt and Lt+1 at consecutive times will generally have different minimal norm minimizers W ∗t , W ∗t+1. The increment Ξ∗t = W ∗t+1 − W ∗t measures this change.

heading: 3.1 A general time-varying Monro-Robbins theorem
text: Given an augmentation scheme for the overparameterized linear model (1.1), the time t gradient update at learning rate ηt is
Wt+1 := Wt + 2ηt N · (Yt −WtXt)XTt , (3.3)
where Dt = (Xt, Yt) is the augmented dataset at time t. The minimum norm minimizer of the corresponding proxy loss Lt (see (3.1)) is
W ∗t := E[YtXTt ]E[XtXTt ]+, (3.4)
where E[XtXTt ]+ denotes the Moore-Penrose pseudo-inverse (see Figure 3.1a). In this section we state a rigorous result, Theorem 3.1, giving sufficient conditions on the learning rate ηt and distributions of the augmented matrices Xt, Yt under which augmented gradient descent converges. To state it, note that (3.3) implies that each row of Wt+1 −Wt is contained in the column span of the Hessian XtXTt of the augmented loss and therefore almost surely belongs to the subspace
V‖ := column span of E[XtXTt ] ⊆ Rn, (3.5) as illustrated in Figure 3.1a. The reason is that, in the orthogonal complement to V‖, the augmented loss L(W ;Dt) has zero gradient with probability 1. To ease notation, we assume that V‖ is independent of t. This assumption holds for additive Gaussian noise, random projection, MixUp, SGD, and their combinations. It is not necessary in general, however, and we refer the interested reader to Remark B.2 in the Appendix for how to treat the general case.
Let us denote by Q‖ : Rn → Rn the orthogonal projection onto V‖ (see Figure 3.1a). As we already pointed out, at time t, gradient descent leaves the matrix of projections Wt(Id−Q‖) of each row of Wt onto the orthogonal complement of V‖ unchanged. In contrast, ‖WtQ‖ −W ∗t ‖F decreases at a rate governed by the smallest positive eigenvalue λmin,V‖ ( E [ XtX T t ]) of the Hessian for the proxy
loss Lt, which is obtained by restricting its full Hessian E [ XtX T t ] to V‖. Moreover, whether and at what rate WtQ‖ −W ∗t converges to 0 depends on how quickly the distance Ξ∗t := W ∗ t+1 −W ∗t (3.6)
between proxy loss optima at successive steps tends to zero (see Figure 3.1b). Theorem 3.1 (Special case of Theorem B.1). Suppose that V‖ is independent of t, that the learning rate satisfies ηt → 0, that the proxy optima satisfy
∞∑ t=0 ‖Ξ∗t ‖F <∞, (3.7)
ensuring the existence of a limit W ∗∞ := limt→∞W ∗ t , that
∞∑ t=0 ηtλmin,V‖(E[XtX T t ]) =∞ (3.8)
and finally that ∞∑ t=0 η2tE [ ‖XtXTt − E[XtXTt ]‖2F + ‖YtXTt − E[YtXTt ]‖2F ] <∞. (3.9) Then, for any initialization W0, we have that WtQ‖ converges in probability to W ∗∞.
If the same augmentation is applied with different strength parameters at each step t (e.g. the noise level σ2t for additive Gaussian noise), we may specialize Theorem 3.1 to this augmentation scheme. More precisely, translating conditions (3.7), (3.8), (3.9) into conditions on the learning rate and augmentation strength gives conditions on the schedule for ηt and these strength parameters to ensure convergence. In §4 and §5, we do this for additive Gaussian noise and random projections.
When the augmentation scheme is static in t, Theorem 3.1 reduces to a standard Monro-Robbins theorem [25] for the (static) proxy loss Lt(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 ;Dt) to ensure the total variance of the stochastic gradients is summable. Condition (3.7) is new and enforces that the minimizers W ∗t of the proxy losses Lt(W ) change slowly enough for augmented optimization procedure to keep pace.

heading: 3.2 Convergence rates and scheduling for data augmentation
text: Refining the proof of Theorem 3.1 gives rates of convergence for the projections WtQ‖ of the weights onto V‖ to the limiting optimum W ∗∞. When the quantities in Theorem 3.1 have power law decay, we obtain the following result. Theorem 3.2 (Special case of Theorem B.4). Suppose V‖ is independent of t and the learning rate satisfies ηt → 0. Moreover assume that for some 0 < α < 1 < β1, β2 and γ > α we have
ηtλmin,V‖(E[XtX T t ]) = Ω(t −α), ‖Ξ∗t ‖F = O(t−β1) (3.10) as well as η2tE[‖XtXTt − E[XtXTt ]‖22] = O(t−γ) (3.11) and
η2tE [ ‖E[Wt](XtXTt − E[XtXTt ])− (YtXTt − E[YtXTt ])‖2F ] = O(t−β2). (3.12)
Then, for any initialization W0, we have for any ε > 0 the following convergence in probability:
tmin{β1−1, β2−α 2 }−ε‖WtQ‖ −W ∗∞‖F p→ 0.
It may be surprising that E[Wt] appears in condition (3.12). Note that E[Wt] is the gradient descent trajectory for the time-varying sequence of deterministic proxy losses Lt(W ). To apply Theorem 3.2, one may first study this deterministic problem to show that E[Wt] converges to W ∗∞ at some rate and then use (3.12) to obtain rates of convergence of the true stochastic trajectory Wt to W ∗∞.
In §4 and §5 below, we specialize Theorems 3.1 and 3.2 to 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.

heading: 4 Special Case: Additive Gaussian Noise
text: We now specialize our main results Theorem 3.1 and 3.2 to the commonly used additive noise augmentation [14]. Under gradient descent, this corresponds to taking augmented data matrices

heading: Xt = X + σtGt and Yt = Y,
text: where Gt ∈ Rn×N are independent matrices with i.i.d. standard Gaussian entries, and σt > 0 is a strength parameter. Under SGD (with replacement), this corresponds to augmented data matrices
Xt = ct(XAt + σt ·Gt) and Yt = ctY At, where At ∈ RN×Bt has i.i.d. columns with a single non-zero entry equal to 1 chosen uniformly at random and ct = √ N/Bt is a normalizing factor. In both cases, the proxy loss is
Lσt(W ) := 1
N ‖Y −WX‖2F + σ2t ‖W‖2F , (4.1)
which to our knowledge was first discovered in [3].
Before stating our precise results, we first illustrate how jointly scheduling learning rate and augmentation strength impacts GD for overparameterized linear regression, where
N = #data points < input dimension = n. (4.2)
The inequality (4.2) ensures L(W ;D) has infinitely many minima, of which we consider the minimum norm minimizer
Wmin := Y X T(XXT)+
most desirable. Notice that steps of vanilla gradient descent
Wt+1 = Wt − 2ηt N · (Y −WtX)XT (4.3)
change the rows of the weight matrix Wt only in the column space V‖ = colspan(XXT) ⊆ Rn. Because V‖ 6= Rn by the overparameterization assumption (4.2), minimizing L(W ;D) without augmentation cannot change Wt,⊥, the matrix whose rows are the components of the rows of Wt orthogonal to V‖. This means that GD converges to the minimal norm optimizer Wmin only when each row of W0 belongs to V‖. 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 σt 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 ηt and noise variance σ2t in which both ηt and σ2t tend to 0 and σt is non-decreasing. If the joint schedule satisfies
∞∑ t=0 ηtσ 2 t =∞ and
∞∑ t=0 η2t σ 2 t <∞, (4.4)
then the weights Wt converge in probability to the minimal norm optimum Wmin regardless of the initialization. Moreover, the first condition in (4.4) is necessary for E[Wt] to converge to Wmin. If we further have ηt = Θ(t−x) and σ2t = Θ(t
−y) with x, y > 0, x+ y < 1, and 2x+ y > 1 so that ηt and σ2t satisfy (4.4), then for small ε > 0, we have t min{y, 12x}−ε‖Wt −Wmin‖F p→ 0.
The conditions of (4.4) require that ηt and σt be jointly scheduled correctly to ensure convergence to Wmin and are akin to the Monro-Robbins conditions [25] for convergence of stochastic gradient methods. We now give an heuristic explanation for why the first condition from (4.4) is necessary. In this setting, the average trajectory of augmented gradient descent
E[Wt+1] = E[Wt]− η∇WLσt(W ) ∣∣ W=E[Wt]
(4.5)
is given by gradient descent on the ridge regularized losses Lσt(W ). If σt ≡ σ > 0 is constant, then E[Wt] will converge to the unique minimizer W ∗σ of the ridge regularized loss Lσ. This point W ∗σ has zero orthogonal component, but does not minimize the original loss L. To instead minimize L, we must choose a schedule satisfying σt → 0. For the expected optimization trajectory to converge to Wmin for such a schedule, the matrix E[Wt,⊥] of components of rows of E [Wt] orthogonal to V‖ must converge to 0. The GD steps for this matrix yield
E[Wt+1,⊥] = (1− ηtσ2t )E[Wt,⊥] = t∏
s=0
(1− ηsσ2s)E[W0,⊥]. (4.6)
Because ηtσ2t approaches 0, this implies the necessary condition ∑∞ t=0 ηtσ 2 t =∞ for E[Wt,⊥]→ 0.
This argument illustrates a key intuition behind the conditions (4.4). The augmentation strength σt must decay to 0 to allow convergence to a true minimizer of the training loss, but this convergence must be carefully tuned to allow the implicit regularization of the noising augmentation to kill the orthogonal component of Wt in expectation. In a 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.
Although Theorem 4.1 is stated for additive Gaussian noise, an analogous version holds for arbitrary additive noise with bounded moments. Moreover, optimizing over x, y, the fastest rate of convergence guaranteed by Theorem 4.1 is obtained by setting ηt = t−2/3+ε, σ2t = t
−1/3 and results in a O(t−1/3+ε) rate of convergence. It is not evident that this is best possible, however.
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 Wmin in the presence of both Gaussian noise and SGD (where datapoints in each batch are selected with replacement). Theorem 4.2. Suppose σ2t → 0 is decreasing, ηt → 0, and we have
∞∑ t=0 ηtσ 2 t =∞ and
∞∑ t=0 η2t <∞. (4.7)
Then the trajectory Wt of SGD with additive noise converges in probability to Wmin for any initialization. If we further have ηt = Θ(t−x) and σ2t = Θ(t
−y) with x > 12 , y > 0 and x+ y < 1, then for any ε > 0 we have that tmin{y, 1 2x}−ε‖Wt −Wmin‖F p→ 0.
Theorem 4.2 is the analog of Theorem 4.1 for mini-batch SGD and provides an example where our framework can handle the composition of two augmentations, namely additive noise and minibatching. 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 a qualitative difference between SGD with and without additive noise. If ηt has power law decay, the convergence of noiseless SGD (Theorem F.1) is exponential in t, while Theorem 4.2 gives power law rates.

heading: 4.3 Experimental validation
text: To validate Theorems 4.1 and 4.2, we 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 a random linear map with i.i.d Gaussian coefficients drawn from N (1, 1). The learning rate followed a fixed polynomially decaying schedule ηt = 0.005 100 · (batch size) · (1 + t20 )−0.66, and the batch size used for SGD was 20. Figure 4.1 shows MSE and ‖Wt,⊥‖F along a single optimization trajectory with different schedules for the variance σ2t used in Gaussian noise augmentation. Complete code to generate this figure is provided in supplement.zip in the supplement. It ran in 30 minutes on a standard laptop CPU.
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 σ2t is zero (blue) or exponentially decaying (green), then while the MSE tends to zero, the orthogonal component Wt,⊥ does not tend to zero. Thus these choices of augmentation schedule cause Wt to converge to an optimum which does not have minimal norm.
On the other hand, if σ2t remains constant (orange), then while Wt,⊥ tends to zero, the MSE is not minimized. Only by decaying both noise strength and learning rate to 0 at sufficiently slow polynomial rates (red) prescribed by Theorem 4.1 do we find both MSE and Wt,⊥ tending to 0, meaning that augmented (S)GD finds the minimum norm optimum Wmin under this choice of parameter scheduling.

heading: 5 Special Case: Augmentation with Random Projections
text: We further illustrate our results by specializing them to a class of augmentations which replace each input x in a batch by its orthogonal projection ΠtX onto a random subspace. In practice (e.g. when using CutOut [12] or SpecAugment [23]), the subspace is chosen based on a prior about correlations between components of X , but we consider the simplified case of a uniformly random subspace of Rn of given dimension.
At each time step t we fix a dimension kt and a fixed kt-dimensional subspace S̃t of Rn. Define the random subspace St by St := Qt(S̃t) = {Qtx | x ∈ S̃t}, where Qt ∈ O(n) is a Haar random orthogonal matrix. Thus, St is uniformly distributed among all kt-dimensional subspaces in Rn. At step t, we take the augmentation given by
Xt = ΠtX Yt = Y, Πt := QtΠ̃tQ T t ,
where Π̃t is the orthogonal projection onto S̃t and hence Πt is the orthogonal projection onto St.
Denoting by γt = kt/n the relative dimension of St, a direct computation (see Lemma E.1) reveals that the proxy loss Lt(W ) equals L(γtW ;D) plus
1
N γt(1− γt) n ‖X‖2F · ‖W‖2F + γt(1− γt)(1/n− 2/n2) N(1 + 1/n− 2/n2) (‖WX‖ 2 F + 1 n ‖X‖2F · ‖W‖2F ). (5.1)
Neglecting terms of order O(n−1), this proxy loss applies a Stein-type shrinkage on input data by γt and adds a data-dependent `2 penalty. For γt < 1, the minimizer of the proxy loss (5.1) is
W ∗γt = Y X T (γt + 1/n− 2/n2
1 + 1/n− 2/n2 XX T + 1− γt 1 + 1/n− 2/n2 ‖X‖2F n
Id )−1 .
Again, although W ∗γt does not minimize the original objective for any γt < 1, the sequence of these proxy optima converges to the minimal norm optimum in the weak regularization limit. Namely, we have limγt→1−W ∗ γt = Wmin. Specializing our general result Theorem 3.1 to this setting, we obtain explicit conditions under which joint schedules of the normalized rank of the projection and the learning rate guarantee convergence to the minimum norm optimizer Wmin. Theorem 5.1. Suppose that ηt → 0, γt → 1 with γt non-decreasing and
∞∑ t=0 ηt(1− γt) =∞ and ∞∑ t=0 η2t (1− γt) <∞. (5.2)
Then, Wt p→ Wmin. Further, if ηt = Θ(t−x) and γt = 1 − Θ(t−y) with x, y > 0, x + y < 1, and 2x+ y > 1, then for small ε > 0, we have that tmin{y, 1 2x}−ε‖Wt −Wmin‖F p→ 0.
Comparing the conditions (5.2) of Theorem 5.1 to the conditions (4.4) of Theorem 4.1, we see that 1− γt is a measure of the strength of the random projection preconditioning. As in that setting, the fastest rates of convergence guaranteed by Theorem 5.1 are obtained by setting ηt = t−2/3+ε and γt = 1− t−1/3, yielding a O(t−1/3+ε) rate of convergence.

heading: 6 Discussion and Limitations
text: We have presented a theoretical framework to rigorously analyze the effect of data augmentation. As can be seen in our main results, our framework applies to completely general augmentations and relies only on analyzing the first few moments of the augmented dataset. This allows us to handle augmentations as diverse as additive noise and random projections as well as their composition in a uniform manner. We have analyzed some representative examples in detail in this work, but many other commonly used augmentations may be handled similarly: label-preserving transformations (e.g. color jitter, geometric transformations) and Mixup [35], among many others. Another line of investigation left to future work is to compare different methods of combining augmentations such as mixing, alternating, or composing, which often improve performance in the empirical literature [16].
Though our results provide a rigorous baseline to compare to more 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 to models closer to those used in practice. Most importantly, we intend to consider more complex models such as kernels (including the neural tangent kernel) and neural networks by making similar connections to 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 to 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
It is a pleasure to thank Daniel Park, Ethan Dyer, Edgar Dobriban, and Pokey Rule for a number of insightful conversations about data augmentation. B.H. was partially supported by NSF grants DMS-1855684 and DMS-2133806 and ONR MURI “Theoretical Foundations of Deep Learning.” Y. S. was partially supported by NSF grants DMS-1701654/2039183 and DMS-2054838.

Review summaries:
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.

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.

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.