Year: 2020

Abstract:
We present a theoretical framework recasting data augmentation as stochastic optimization for a sequence of time-varying proxy losses. This provides a unified language for understanding techniques commonly thought of as data augmentation, including synthetic noise and label-preserving transformations, as well as more traditional ideas in stochastic optimization such as learning rate and batch size scheduling. We then specialize our framework to study arbitrary augmentations in the context of a simple model (overparameterized linear regression). We extend in this setting the classical Monro-Robbins theorem to include augmentation and obtain rates of convergence, giving conditions on the learning rate and augmentation schedule under which augmented gradient descent converges. Special cases give provably good schedules for augmentation with additive noise, minibatch SGD, and minibatch SGD with noise.

heading: 1 INTRODUCTION
text: Implementing gradient-based optimization in practice requires many choices. These include setting hyperparameters such as learning rate and batch size as well as specifying a data augmentation scheme, a popular set of techniques in which data is augmented (i.e. modified) at every step of optimization. Trained model quality is highly sensitive to these choices. In practice they are made using methods ranging from a simple grid search to Bayesian optimization and reinforcement learning (Cubuk et al., 2019; 2020; Ho et al., 2019). Such approaches, while effective, are often ad-hoc and computationally expensive due to the need to handle scheduling, in which optimization hyperparameters and augmentation choices and strengths are chosen to change over the course of optimization.
These empirical results stand in contrast to theoretically grounded approaches to stochastic optimization which provide both provable guarantees and reliable intuitions. The most extensive work in this direction builds on the seminal article (Robbins & Monro, 1951), which gives provably optimal learning rate schedules for stochastic optimization of strongly convex objectives. While rigorous, these approaches are typically are not sufficiently flexible to address the myriad augmentation types and hyperparameter choices beyond learning rates necessary in practice.
This article is a step towards bridging this gap. We provide in §3 a rigorous framework for reinterpreting gradient descent with arbitrary data augmentation as stochastic gradient descent on a time-varying sequence of objectives. This provides a unified language to study traditional stochastic optimization methods such as minibatch SGD together with widely used augmentations such as additive noise (Grandvalet & Canu, 1997), CutOut (DeVries & Taylor, 2017), Mixup (Zhang et al., 2017) and label-preserving transformations (e.g. color jitter, geometric transformations (Simard et al., 2003)). It also opens the door to studying how to schedule and evaluate arbitrary augmentations, an important topic given the recent interest in learned augmentation Cubuk et al. (2019).
Quantitative results in our framework are difficult to obtain in full generality due to the complex interaction between models and augmentations. To illustrate the utility of our approach and better understand specific augmentations, we present in §3 and §5 results about arbitrary augmentations for overparameterized linear regression and specialize to additive noise and minibatch SGD in §4 and §6. While our results apply directly only to simple quadratic losses, they treat very general augmentations. Treating more complex models is left to future work. Our main contributions are:
• In Theorem 5.1, we give sufficient conditions under which gradient descent under any augmentation scheme converges in the setting of overparameterized linear regression. Our
result extends classical results of Monro-Robbins type and covers schedules for both learning rate and data augmentation scheme. • We complement the asymptotic results of Theorem 5.1 with quantitative rates of conver-
gence furnished in Theorem 5.2. These rates depend only on the first few moments of the augmented data distribution, underscoring the flexibility of our framework. • In §4, we analyze additive input noise, a popular augmentation strategy for increasing
model robustness. We recover the known fact that it is equivalent to stochastic optimization with `2-regularization and find criteria in Theorem 4.1 for jointly scheduling the learning rate and noise level to provably recover the minimal norm solution. • In §6, we analyze minibatch SGD, recovering known results about rates of convergence for
SGD (Theorem 6.1) and novel results about SGD with noise (Theorem 6.2).

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 Robbins & Monro (1951). This spawned a vast literature on rates of convergence for GD, SGD, and their variants. We mention only the relatively recent articles Bach & Moulines (2013); Défossez & Bach (2015); Bottou et al. (2018); Smith et al. (2018); Ma et al. (2018) and the references therein. The last of these, namely Ma et al. (2018), 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 Bishop (1995), which treats the case of additive Gaussian noise (see §4). More recent work in this direction includes Chapelle et al. (2001); Wager et al. (2013); LeJeune et al. (2019); Liu et al. (2020). There are also several articles investigating optimal choices of `2-regularization for linear models (cf e.g. Wu et al. (2018); Wu & Xu (2020); Bartlett et al. (2020)). 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 Lewkowycz & Gur-Ari (2020), 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 Rajput et al. (2019), for example, the authors quantify how much augmented data, produced via additive noise, is needed to learn positive margin classifiers. Chen et al. (2019), 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. Dao et al. (2019) 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 Wu et al. (2020) considers both label preserving and noising augmentations, pointing out the conceptually distinct roles such augmentations play.

heading: 3 DATA AUGMENTATION AS STOCHASTIC OPTIMIZATION
text: A common task in modern machine learning is the optimization of an empirical risk
L(W ;D) = 1 |D| ∑ (xj ,yj)∈D `(f(xj ;W ), yj), (3.1)
where f(x;W ) is a parameterized model for a dataset D of input-response pairs (x, y) and ` is a per-sample loss. Optimizing W by vanilla gradient descent on L corresponds to the update equation
Wt+1 = Wt − ηt∇WL(Wt;D). In this context, we define a data augmentation scheme to be any procedure that consists, at every step of optimization, of replacing the dataset D by a randomly augmented variant, which we will denote
by Dt. Typically, Dt is related to D in some way, but our framework does not explicitly constrain the form of this relationship. Instead, certain conditions on this relationship will be required for our main results Theorems 5.1 and 5.2 to give useful results for a specific augmentation scheme. A data augmentation scheme therefore corresponds to the augmented update equation
Wt+1 = Wt − ηt∇WL(Wt;Dt). (3.2) SinceDt is a stochastic function ofD, it is natural to view the augmented update rule (3.2) as a form of stochastic optimization for the proxy loss at time t
Lt(W ) := EDt [L(W ;Dt)] . (3.3) The update (3.2) corresponds precisely to stochastic optimization for the time-varying objective Lt(W ) in which the unbiased estimate of its gradient is obtained by evaluating the gradient of L(W ;Dt) on a single sampleDt drawn from the augmentation distribution. The connection between data augmentation and this proxy loss was introduced for Gaussian noise in Bishop (1995) and in general in Chapelle et al. (2001), but we now consider it in the context of stochastic optimization.
Despite being mathematically straightforward, reformulating data augmentation as stochastic optimization provides a unified language for questions about learning rate schedules and general augmentation schemes including SGD. In general, such questions can be challenging to answer, and even evaluating the proxy loss Lt(W ) may require significant ingenuity. While we will return to more sophisticated models in future work, we henceforth analyze general augmentations in the simple context of overparameterized linear regression. Though there are many ways to perform linear regression, we restrict to augmented gradient descent both to gain intuition about specific augmentations and to understand the effect of augmentation on optimization. We therefore consider optimizing the entries of a weight matrix W ∈ Rp×n by gradient descent on
L(W ;D) = 1 |D| ∑ (x,y)∈D ||y −Wx||2F = 1 N ||Y −WX||2F , (3.4)
where our dataset D is summarized by data matrices X ∈ Rn×N and Y ∈ Rp×N , whose N < n columns consist of inputs xi ∈ Rn and associated labels yi ∈ Rp. Following this notation, a data augmentation scheme is specified by prescribing at each time step an augmented dataset Dt consisting of modified data matrices Xt, Yt, whose columns we denote by xi,t ∈ Rn and yi,t ∈ Rp. Here, the number of columns in Xt and Yt (i.e. the number of datapoints in Dt) may vary. We now give examples of some commonly used augmentations our framework can address.
• Additive Gaussian noise: This is implemented by setting Xt = X + σt · G and Yt = Y for σt > 0 and G a matrix of i.i.d. standard Gaussians. We analyze this in §4.
• Mini-batch SGD: To implement mini-batch SGD with batch size Bt, we can take Xt = XAt and Yt = Y At where At ∈ RN×Bt has i.i.d. columns containing a single non-zero entry equal to 1 chosen uniformly at random. We analyze this in detail in §6.
• Random projection: This is implemented by 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 ) = ‖Y − γtWX‖2F + γt(1− γt)n−1 Tr(XXT)‖W‖2F +O(n−1), which adds a data-dependent `2 penalty and applies a Stein shrinkage on input data. • Label-preserving transformations: For a 2-D image viewed as a vector x ∈ Rn, geomet-
ric transforms (with pixel interpolation) or other label-preserving transforms such as color jitter take the form of linear transforms Rn → Rn. We may implement such augmentations in our framework by Xt = AtX and Yt = Y for some random transform matrix At. • Mixup: To implement Mixup, we can takeXt = XAt and Yt = Y At, whereAt ∈ RN×Bt
has i.i.d. columns containing with two random non-zero entries equal to 1− ct and ct with mixing coefficient ct drawn from a Beta(αt, αt) distribution for a parameter αt.
Our main technical results, Theorems 5.1 and 5.2, give sufficient conditions for a learning rate schedule ηt and a schedule for the statistics of Xt, Yt under which optimization with augmented gradient descent will provably converge. We state these general results in §5. Before doing so, we seek to demonstrate both the utility of our framework and the flavor of our results by focusing on the simple but already informative case of additive Gaussian noise.

heading: 4 AUGMENTATION WITH ADDITIVE GAUSSIAN NOISE
text: An common augmentation in practice injects input noise as a regularizer (Grandvalet & Canu, 1997):
Dt = {(xi,t, yi,t), i = 1, . . . , N}, xi,t = xi + σtgi,t, yi,t = yi,
where gi,t are i.i.d. standard Gaussian vectors and σt is a strength parameter. This section studies such augmentations using our framework. A direct computation reveals that the proxy loss
Lt(W ) = Lσt(W ) := L(W ;D) + σ2t ||W || 2 F
corresponding to additive Gaussian noise adds an `2-penalty to the original loss L. This is simple but useful intuition. It also raises the question: what is the optimal relation between the learning rate ηt and the augmentation strength σt (i.e. the `2-penalty)?
To get a sense of what optimal might mean in this context, observe first that if σt = 0, then directly differentiating the loss L yields the following update rule:
Wt+1 = Wt + 2ηt N · (Y −WtX)XT. (4.1)
The increment Wt+1 −Wt is therefore contained in the column span
V‖ := column span of XXT ⊆ Rn (4.2)
of the model Hessian XXT. Overparameterization implies V‖ 6= Rn. The component Wt,⊥ of Wt that is in the orthogonal complement of V‖ thus remains frozen to its initialized value. Geometrically, this means that there are some directions, namely those in the orthogonal complement to V‖, which gradient descent “cannot see.” Optimization with appropriate step sizes then yields
lim t→∞
Wt = W0,⊥ +Wmin, Wmin := Y X T(XXT)+,
where Wmin is the minimum norm solution of Y = WX . The original motivation for introducing the `2-regularized losses Lσ is that they provide a mechanism to eliminate the component W0,⊥ for all initializations, not just the special choice W0 = 0, and they can be used to regularize non-linear models as well. Indeed, for σ > 0, the loss Lσ is strictly convex and has a unique minimum
W ∗σ := Y X T ( XXT + σ2N · Idn×n )−1 ,
which tends to the minimal norm solution in the weak regularization limit limσ→0W ∗σ = Wmin. Geometrically, this is reflected in the fact that `2-penalty yields non-trivial gradient updates
Wt+1,⊥ = Wt,⊥ − ηtσ2Wt,⊥ = (Id− ηtσ2)Wt,⊥ = t∏
s=1
(Id− ηsσ2)W0,⊥, (4.3)
which drive this perpendicular component ofWt to zero provided ∑∞ t=1 ηt =∞. However, for each positive value of σ, the `2-penalty also modifies the gradient descent updates for Wt,‖, ultimately causing Wt to converge to W ∗σ , which is not a minimizer of the original loss L. This downside of ridge regression motivates jointly scheduling the step size ηt and the noise strength σt. We hope that driving σt to 0 at an appropriate rate can guarantee convergence of Wt to Wmin. Namely, we want to retain the regularizing effects of `2-noise to forceWt,⊥ to zero while mitigating its adverse effects which prevent W ∗σ from minimizing L. It turns out that this is indeed possible: Theorem 4.1 (Special case of Theorem 5.1). Suppose σ2t , ηt → 0 with σ2t non-increasing and
∞∑ t=0 ηtσ 2 t =∞ and
∞∑ t=0 η2t σ 2 t <∞. (4.4)

heading: Then, Wt
text: p→ Wmin. Further, if ηt = Θ(t−x) and σ2t = Θ(t−y) with x, y > 0, x + y < 1, and 2x+ y > 1, then for any ∈ (0,min{y, x/2}), we have that
tmin{y, 1 2x}− ‖Wt −Wmin‖F p→ 0.
Let us give a few comments on Theorem 4.1. First, although it is stated for additive Gaussian noise, an analogous version holds for arbitrary additive noise with bounded moments, with the only change being a constant multiplicative factor in the second condition of (4.4).
Second, that convergence in probability Wt p→ Wmin follows from (4.4) is analogous to a MonroRobbins type theorem (Robbins & Monro, 1951). Indeed, inspecting (4.3), we see that the first condition in (4.4) guarantees that the effective learning rate ηtσ2t in the orthogonal complement to V‖ is sufficiently large that the corresponding component Wt,⊥ of Wt tends to 0, allowing the result of optimization to be independent of the initial condition W0. Further, the second condition in (4.4) guarantees that the variance of the gradients, which at time t scales like η2t σ 2 t is summable. As in the usual Monro-Robbins setup, this means that only a finite amount of noise is injected into the optimization. Further, (4.4) is a direct specialization of (5.5) and (5.6) from Theorem 5.1.
Third, by optimizing over x, y, we see that 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 the best possible rate, however.
Finally, although we leave systemic study of augmentation in non-linear models to future work, our framework can be applied beyond linear models and quadratic losses. To see this, as noted for kernels in Dao et al. (2019), augmenting inputs of nonlinear feature models correspond to applying different augmentations on the outputs of the feature map. To give a concrete example, consider additive noise for small σt. For any sufficiently smooth function g, Taylor expansion reveals
E [g(x+ σtG)] = g(x) + σ2t 2 ∆g(x) +O(σ4t ),
where ∆ = ∑ i ∂ 2 i is the Laplacian and G is a standard Gaussian vector. For a general loss of the form (3.1) we have
Lt(W ) = L(W ;D) + σ2t 2|D| ∑
(x,y∈D)
Tr [ (∇xf)T (Hf `)∇xf ] + (∇f `)T ∆xf +O(σ4t ),
where we have written Hf ` for the Hessian of some convex per-sample loss ` with respect to f and ∇x,∇f for the gradients with respect to x, f , respectively. This is consistent with the similar expansion done in the kernel setting by Dao et al. (2019, Section 4). If σt is small, then the proxy loss Lt will differ significantly from the unaugmented loss L only near the end of training, when we expect ∇f ` to be small and Hf ` to be positive semi-definite. Hence, we find heuristically that, neglecting higher order terms in σt, additive noise with small σt corresponds to an `2-regularizer
Tr [ σ2t 2 (∇xf)T (HfL)∇xf ] =: σ2t 2 ||∇xf ||2HfL
for the gradients of f with respect to the natural inner product determined by the Hessian of the loss. This is intuitive since penalizing the gradients of f is the same as requiring that f is approximately constant in a neighborhood of every datapoint. However, although the input noise was originally isotropic, the `2-penalty is aligned with the loss Hessian and hence need not be.

heading: 5 TIME-VARYING MONRO-ROBBINS FOR LINEAR MODELS UNDER AUGMENTATION
text: In this section, we state two general results, Theorems 5.1 and 5.2, which provide sufficient conditions for jointly scheduling learning rates and general augmentation schemes to guarantee convergence of augmented gradient descent in the overparameterized linear model (3.4).

heading: 5.1 A GENERAL TIME-VARYING MONRO-ROBBINS THEOREM
text: Given an augmentation scheme for the model (3.4), the time t gradient update at learning rate ηt is
Wt+1 := Wt + 2ηt N · (Yt −WtXt)XTt , (5.1)
where Dt = (Xt, Yt) is the augmented dataset at time t. The minimum norm minimizer of the corresponding proxy loss Lt (see 3.3) is
W ∗t := E[YtXTt ]E[XtXTt ]+, (5.2)
where E[XtXTt ]+ denotes the Moore-Penrose pseudo-inverse. In this section we state a rigorous result, Theorem 5.1, giving sufficient conditions on the learning rate ηt and distributions of the augmented matrices Xt, Yt under which augmented gradient descent converges. In analogy with the case of Gaussian noise, (5.1) shows Wt+1 −Wt is contained in the column span of the Hessian XtX T t of the augmented loss and almost surely belongs to the subspace
V‖ := column span of E[XtXTt ] ⊆ Rn. (5.3)
To ease notation, we assume that V‖ is independent of t. This assumption is valid for additive Gaussian noise, random projection, MixUp, SGD, and their combinations. We explain in Remark B.2 how to generalize Theorems 5.1 and 5.2 to the case where V‖ varies with t.
Let us denote by Q‖ : Rn → Rn the orthogonal projection onto V‖. At time t, gradient descent leaves the projection Wt(Id−Q‖) 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 ]) := λmin ( Q‖E [ XtX T t ] Q‖ )
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 must depend on how quickly
Ξ∗t := W ∗ t+1 −W ∗t (5.4)
tends to zero. Indeed, ||Ξ∗t ||F is the distance between proxy loss optima at different times and hence must tend to zero if ||WtQ‖ −W ∗t ||F converges to zero. Theorem 5.1. Suppose that V‖ is independent of t, that the learning rate satisfies ηt → 0, that the proxy optima satisfy
∞∑ t=0 ‖Ξ∗t ‖F <∞, (5.5)
ensuring the existence of a limit W ∗∞ := limt→∞W ∗ t , and that
∞∑ t=0 ηtλmin,V‖(E[XtX T t ]) =∞. (5.6)
If either ∞∑ t=0 η2tE [ ‖XtXTt − E[XtXTt ]‖2F + ‖YtXTt − E[YtXTt ]‖2F ] <∞ (5.7)
or the more refined condition ∞∑ t=0 η2tE [ ‖XtXTt − E[XtXTt ]‖2F + ‖E[Wt](XtXTt − E[XtXTt ])− (YtXTt − E[YtXTt ])‖2F ] <∞ (5.8) hold, then for any initialization W0 we have WtQ‖ p→W ∗∞.
The conditions of Theorem 5.1 can be applied to the choice of joint schedule for the learning rate and augmentation scheme applied to gradient descent. If the same augmentation is applied with different strength parameters at each step t such as σt for Gaussian noise, they impose conditions on the choice of joint schedule for ηt and these strength parameters. In the example of Theorem 4.1 for Gaussian noise, the condition that σ2t is non-increasing implies (5.5), the first condition of (4.4) implies (5.6), and the second condition of (4.4) implies (5.7).
In addition to the conditions Theorem 5.1 imposes on Dt, the proxy optima W ∗t and their limit W ∗∞ are determined by the distribution of Dt. Therefore, for W ∗∞ in Theorem 5.1 to be a desirable set of parameters for the original dataset D, the augmented dataset Dt must have some relation to D.
When the augmentation procedure is static in t, Theorem 5.1 reduces to a standard Monro-Robbins theorem Robbins & Monro (1951) for the (static) proxy loss Lt(W ). As in that setting, condition (5.6) enforces that the learning trajectory travels far enough to reach an optimum. Condition (5.7) implies the weaker condition (5.8); the second summand in (5.8) is the variance of the gradient of the augmented loss L(W ;Dt), meaning (5.8) implies the total variance of the stochastic gradients is summable. Condition (5.5) is new; it enforces that the minimizers W ∗t of the proxy losses Lt(W ) change slowly enough that the augmented optimization procedure can keep pace.
Though it may be surprising that E[Wt] appears in this condition, it may be interpreted as the gradient descent trajectory for the deterministic sequence of proxy losses Lt(W ). Accounting for the dependence on E[Wt] allows us to give more precise rates using the variance of the stochastic gradient in (5.8); we include both (5.7) and (5.8) to allow a user of our results to separately analyze E[Wt] to obtain stronger conclusions.

heading: 5.2 CONVERGENCE RATES AND SCHEDULING FOR DATA AUGMENTATION
text: A more precise analysis of the the proof of Theorem 5.1 allows us to obtain rates of convergence for the projections WtQ‖ of the weights onto V‖ to the limiting optimum W ∗∞. In particular, when the quantities in Theorem 5.1 have power law decay, we obtain the following result. Theorem 5.2 (informal - Special case of Theorem B.4). If V‖ is independent of t, the learning rate satisfies ηt → 0, and for some 0 < α < 1 < β1, β2 and γ > α we have
ηtλmin,V‖(E[XtX T t ]) = Ω(t −α), ‖Ξ∗t ‖F = O(t−β1) (5.9) and η2tE[‖XtXTt − E[XtXTt ]‖22] = O(t−γ) (5.10) and
η2tE [ ‖E[Wt](XtXTt − E[XtXTt ])− (YtXTt − E[YtXTt ])‖2F ] = O(t−β2), (5.11)
then for any initialization W0, we have for any > 0 that
tmin{β1−1, β2−α 2 }− ‖WtQ‖ −W ∗∞‖F p→ 0.
Theorem 5.2 measures rates in terms of optimization steps t, but a different measurement of time called the intrinsic time of the optimization will be more suitable for measuring the behavior of optimization quantities. This was introduced for SGD in Smith & Le (2018); Smith et al. (2018), and we now generalize it to our broader setting. For gradient descent on a loss L, the intrinsic time is a quantity which increments by ηλmin(H) for a optimization step with learning rate η at a point where L has Hessian H . When specialized to our setting, it is given by
τ(t) := t−1∑ s=0 2ηs N λmin,V‖(E[XsX T s ]). (5.12)
Notice that intrinsic time of augmented optimization for the sequence of proxy losses Ls appears in Theorems 5.1 and 5.2, which require via condition (5.6) that the intrinsic time tends to infinity as the number of optimization steps grows.
Intrinsic time will be a sensible variable in which to measure the behavior of quantities such as the fluctuations of the optimization path f(t) := E[‖(Wt − E[Wt])Q‖‖2F ]. In the proofs of Theorems 5.1 and 5.2, we show that the fluctuations satisfy an inequality of the form
f(t+ 1) ≤ f(t)(1− a(t))2 + b(t) (5.13)
for a(t) := 2ηt 1N λmin,V‖(E[XtX T t ]) and b(t) := Var[||ηt∇WL(Wt)||F ] so that τ(t) = ∑t−1 s=0 a(s). Iterating the recursion (5.13) shows that
f(t) ≤ f(0) t−1∏ s=0 (1− a(s))2 + t−1∑ s=0 b(s) t−1∏ r=s+1 (1− a(r))2
≤ e−2τ(t)f(0) + t−1∑ s=0 b(s) a(s) e2τ(s+1)−2τ(t)(τ(s+ 1)− τ(s)).
For τ := τ(t) and changes of variable A(τ), B(τ), and F (τ) such that A(τ(t)) = a(t), B(τ(t)) = b(t), and F (τ(t)) = f(t), we find by replacing a right Riemann sum by an integral that
F (τ) - e−2τ [ F (0) + ∫ τ 0 B(σ) A(σ) e2σdσ ] . (5.14)
In order for the result of optimization to be independent of the starting point, by (5.14) we must have τ → ∞ to remove the dependence on F (0); this provides one explanation for the appearance of τ in condition (5.6). Further, (5.14) implies that the fluctuations at an intrinsic time are bounded by an integral against the function B(σ)A(σ) which depends only on the ratio of A(σ) and B(σ). In the case of minibatch SGD, we compute this ratio in (6.2) and recover the commonly used “linear scaling” rule for learning rate.
In Section 6, we specialize Theorem 5.2 to obtain rates of convergence for specific augmentations. Optimizing the learning rate and augmentation parameter schedules in Theorem 5.2 allows us to derive power law schedules with convergence rate guarantees in these settings.

heading: 6 IMPLICATIONS FOR MINI-BATCH STOCHASTIC GRADIENT DESCENT (SGD)
text: We now apply our framework to study mini-batch stochastic gradient descent (SGD) with the potential presence of additive noise. Though data augmentation commonly refers to techniques aside from SGD, we will see that our framework handles it uniformly with other augmentations.

heading: 6.1 MINI-BATCH SGD
text: In mini-batch stochastic gradient descent, Dt is obtained by choosing a random subset Bt of D of prescribed batch size Bt = |Bt|. Each datapoint in Bt is chosen uniformly with replacement from D, and the resulting data matrices Xt and Yt are scaled so that Lt(W ) = L(W ;D). Concretely, this means that for the normalizing factor ct := √ N/Bt we have
Xt = ctXAt and Yt = ctY At,
whereAt ∈ RN×Bt has i.i.d. columnsAt,i with a single non-zero entry equal to 1 chosen uniformly at random. In this setting the minimum norm optimum for each t are the same and given by
W ∗t = W ∗ ∞ = Y X T(XXT)+,
which coincides with the minimum norm optimum for the unaugmented loss. Our main result for standard SGD is the following theorem, whose proof is given in Appendix D.1. Theorem 6.1. If the learning rate satisfies ηt → 0 and
∞∑ t=0 ηt =∞, (6.1)
then for any initialization W0, we have WtQ‖ p→ W ∗∞. If further we have that ηt = Θ(t−x) with 0 < x < 1, then for some C > 0 we have
eCt 1−x ‖WtQ‖ −W ∗∞‖F p→ 0.
Theorem 6.1 recovers the exponential convergence rate for SGD, which has been extensively studied through both empirical and theoretical means (Bottou et al., 2018; Ma et al., 2018). Because 1 ≤ Bt ≤ N for all t, it does not affect the asymptotic results in Theorem 6.1. In practice, however, the number of optimization steps t is often small enough that BtN is of order t
−α for some α > 0, meaning the choice of Bt can affect rates in this non-asymptotic regime. Though we do not attempt to push our generic analysis to this granularity, this is done in Ma et al. (2018) to derive optimal batch sizes and learning rates in the overparametrized setting.
Our proof of Theorem 6.1 shows the intrinsic time is τ(t) = ∑t−1 s=0 2ηs 1 N λmin,V‖(XX T) and the ratio b(t)a(t) in (5.14) is by (D.4) bounded uniformly for a constant C > 0 by
b(t) a(t) ≤ C · ηt Bt . (6.2)
Thus, keeping b(t)a(t) fixed as a function of τ suggests the “linear scaling” ηt ∝ Bt used empirically in Goyal et al. (2017) and proposed via an heuristic SDE limit in Smith et al. (2018).

heading: 6.2 MINI-BATCH SGD WITH ADDITIVE SYNTHETIC NOISE
text: In addition to handling synthetic noise and SGD separately, our results and framework also cover the hybrid case of mini-batch SGD with batch size Bt and additive noise at level σt. Here,
Xt = ct(XAt + σtGt) and Yt = ctY At,
where ct and At are as in Section 6.1 and Gt ∈ Rn×Bt has i.i.d. Gaussian entries. The proxy loss is
Lt(W ) := 1 N E [ ‖ctY At − ctWXAt − ctσtWGt‖2F ] = 1 N ‖Y −WX‖2F + σ2t ‖W‖2F ,
with ridge minimizer W ∗t = Y X T(XXT + σ2tN · Idn×n)−1. Like with synthetic noise but unlike noiseless SGD, the optima W ∗t converge to the minimal norm interpolant Wmin = Y X T(XXT)+. Theorem 6.2. Suppose σ2t → 0 is decreasing, ηt → 0, and for any C > 0 we have ∞∑ t=0 (ηtσ 2 t − Cη2t ) =∞ and ∞∑ t=0 η2t σ 2 t <∞. (6.3)
Then we have Wt p→ Wmin. If we further have ηt = Θ(t−x) and σ2t = Θ(t−y) with x, y > 0 and 0 < x+ y < 1 < 2x+ y, we have for any > 0 that
tmin{y, 1 2x}− ‖Wt −Wmin‖F p→ 0.
Theorem 6.2 provides an example where our framework can handle the composition of two augmentations, namely additive noise and SGD. It reveals a qualitative difference between SGD with and without additive noise. For polynomially decaying ηt the convergence of noiseless SGD in Theorem 6.1 is exponential in t, while the bound from Theorem 6.2 is polynomial in t. This is unavoidable. Indeed, for components of Wt orthogonal to colspan(X), convergence requires that ∑∞ t=0 ηtσ 2 t =∞ (see (4.3)). This occurs only if σt has power law decay, causing the ||W ∗t −Wmin||F to have at most power law decay as well. Finally, the Monro-Robbins conditions (6.3) are more restrictive than the analogous conditions in the pure noise setting (see (4.4)), as the latter allow for large ηt schedules in which ∑∞ t=0 η 2 t diverges but ∑∞ t=0 η 2 t σ 2 t does not.

heading: 7 DISCUSSION
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 mini-batch SGD 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), random projections (DeVries & Taylor, 2017; Park et al., 2019), and Mixup (Zhang et al., 2017), 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 (Hendrycks et al., 2020).
Though our results provide a rigorous baseline to compare to more complex settings, the restriction of the present work to linear models is of course 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 the cross-entropy loss. 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.

heading: A ANALYTIC LEMMAS
text: In this section, we present several basic lemmas concerning convergence for certain matrix-valued recursions that will be needed to establish our main results. For clarity, we first collect some matrix notations used in this section and throughout the paper.

heading: A.1 MATRIX NOTATIONS
text: Let M ∈ Rm×n be a matrix. We denote its Frobenius norm by ‖M‖F and its spectral norm by ‖M‖2. Ifm = n so thatM is square, we denote by diag(M) the diagonal matrix with diag(M)ii = Mii. For matrices A,B,C of the appropriate shapes, define
A ◦ (B ⊗ C) := BAC (A.1)
and Var(A) := E[AT ⊗A]− E[AT]⊗ E[A]. (A.2)
Notice in particular that Tr[Id ◦Var(A)] = E[‖A− E[A]‖2F ].

heading: A.2 ONE- AND TWO-SIDED DECAY
text: Definition A.1. Let At ∈ Rn×n be a sequence of independent random non-negative definite matrices with
sup t ||At|| ≤ 2 almost surely,
let Bt ∈ Rp×n be a sequence of arbitrary matrices, and let Ct ∈ Rn×n be a sequence of nonnegative definite matrices. We say that the sequence of matrices Xt ∈ Rp×n has one-sided decay of type ({At}, {Bt}) if it satisfies
Xt+1 = Xt(Id−E[At]) +Bt. (A.3) We say that a sequence of non-negative definite matrices Zt ∈ Rn×n has two-sided decay of type ({At}, {Ct}) if it satisfies
Zt+1 = E[(Id−At)Zt(Id−At)] + Ct. (A.4)
Intuitively, if a sequence of matrices Xt (resp. Zt) satisfies one decay of type ({At}, {Bt}) (resp. two-sided decay of type ({At}, {Ct})), then in those directions u ∈ Rn for which ||Atu|| does not decay too quickly in t we expect that Xt (resp. Zt) will converge to 0 provided Bt (resp. Ct) are not too large. More formally, let us define
V‖ := ∞⋂ t=0 ker [ ∞∏ s=t (Id−E[As]) ] = { u ∈ Rn ∣∣∣∣ limT→∞ T∏ s=t (Id−E[As])u = 0, ∀t ≥ 1 } ,
and let Q‖ be the orthogonal projection onto V‖. It is on the space V‖ that that we expect Xt, Zt to tend to zero if they satisfy one or two-side decay, and the precise results follows.

heading: A.3 LEMMAS ON CONVERGENCE FOR MATRICES WITH ONE AND TWO-SIDED DECAY
text: We state here several results that underpin the proofs of our main results. We begin by giving in Lemmas A.2 and A.3 two slight variations of the same simple argument that matrices with one or two-sided decay converge to zero. Lemma A.2. If a sequence {Xt} has one-sided decay of type ({At}, {Bt}) with
∞∑ t=0 ‖Bt‖F <∞, (A.5)
then limt→∞XtQ‖ = 0. Proof. For any > 0, choose T1 so that ∑∞ t=T1 ‖Bt‖F < 2 and T2 so that for t > T2 we have∥∥∥∥∥( t∏ s=T1 (Id−E[As]) ) Q‖ ∥∥∥∥∥ 2 < 2 1 ‖X0‖F + ∑T1−1 s=0 ‖Bs‖F .
By (A.3), we find that
Xt+1 = X0 t∏ s=0 (Id−E[As]) + t∑ s=0 Bs t∏ r=s+1 (Id−E[Ar]),
which implies for t > T2 that
‖Xt+1Q‖‖F ≤ ‖X0‖F ∥∥∥∥∥( t∏
s=0
(Id−E[As]) ) Q‖ ∥∥∥∥∥ 2 + t∑ s=0 ‖Bs‖F ∥∥∥∥∥( t∏ r=s+1 (Id−E[Ar]) ) Q‖ ∥∥∥∥∥ 2 .
(A.6)
Our assumption that ||At|| ≤ 2 almost surely implies that for any T ≤ t∥∥∥∥∥( t∏
s=0
(Id−E[As]) ) Q‖ ∥∥∥∥∥ 2 ≤ ∥∥∥∥∥( T∏ s=0 (Id−E[As]) ) Q‖ ∥∥∥∥∥ 2
since each term in the product is non-negative-definite. Thus, we find
‖Xt+1Q‖‖F ≤ [ ‖X0‖F +
T1−1∑ s=0 ‖Bs‖F ]∥∥∥∥∥( t∏
s=T1
(Id−E[As]) ) Q‖ ∥∥∥∥∥ 2 + t∑ s=T1 ‖Bs‖F < .
Taking t→∞ and then → 0 implies that limt→∞XtQ‖ = 0, as desired.
Lemma A.3. If a sequence {Zt} has two-sided decay of type ({At}, {Ct}) with
lim T→∞
E ∥∥∥∥∥( T∏ s=t (Id−As) ) Q‖ ∥∥∥∥∥ 2
2  = 0 for all t ≥ 0 (A.7) and
∞∑ t=0 Tr(Ct) <∞, (A.8)
then limt→∞QT‖ZtQ‖ = 0.
Proof. The proof is essentially identical to that of Lemma A.2. That is, for > 0, choose T1 so that∑∞ t=T1 Tr(Ct) < 2 and choose T2 by (A.7) so that for t > T2 we have
E ∥∥∥∥∥( t∏
s=T1
(Id−As) ) Q‖ ∥∥∥∥∥ 2
2
 < 2
1 Tr(Z0) + ∑T1−1 s=0 Tr(Cs) .
Conjugating (A.4) by Q‖, we have that
QT‖Zt+1Q‖ = E [ QT‖ ( t∏ s=0 (Id−As) )T Z0 ( t∏ s=0 (Id−As) ) Q‖ ]
+ t∑ s=0 E [ QT‖ ( t∏ r=s+1 (Id−Ar) )T Cs ( t∏ r=s+1 (Id−Ar) ) Q‖ ] .
Our assumption that ||At|| ≤ 2 almost surely implies that for any T ≤ t∥∥∥∥∥( t∏
s=0
(Id−As) ) Q ∥∥∥∥∥ 2 ≤ ∥∥∥∥∥( T∏ s=0 (Id−As) ) Q ∥∥∥∥∥ 2 .
For t > T2, this implies by taking trace of both sides that
Tr(QT‖Zt+1Q‖) ≤ Tr(Z0)E ∥∥∥∥∥( t∏
s=0
(Id−As) ) Q‖ ∥∥∥∥∥ 2
2
+ t∑ s=0 Tr(Cs)E ∥∥∥∥∥( t∏ r=s+1 (Id−Ar) ) Q‖ ∥∥∥∥∥ 2
2  (A.9)
≤ [ Tr(Z0) +
T1−1∑ s=0 Tr(Cs)
] E ∥∥∥∥∥( t∏
s=T1
(Id−As) ) Q‖ ∥∥∥∥∥ 2
2
+ t∑ s=T1 Tr(Cs)
< ,
which implies that limt→∞QT‖ZtQ‖ = 0.
The preceding Lemmas will be used to provide sufficient conditions for augmented gradient descent to converge as in Theorem B.1 below. Since we are also interested in obtaining rates of convergence, we record here two quantitative refinements of the Lemmas above that will be used in the proof of Theorem B.4.
Lemma A.4. Suppose {Xt} has one-sided decay of type ({At}, {Bt}). Assume also that for some X ≥ 0 and C > 0, we have
log ∥∥∥∥∥( t∏
r=s
(Id−E[Ar]) ) Q‖ ∥∥∥∥∥ 2 < X − C ∫ t+1 s r−αdr
and ‖Bt‖F = O(t−β) for some 0 < α < 1 < β. Then, ‖XtQ‖‖F = O(tα−β). Proof. Denote γs,t := ∫ t s r−αdr. By (A.6), we have for some constants C1, C2 > 0 that
‖Xt+1Q‖‖F < C1e−Cγ1,t+1 + C2eX t∑
s=1
(1 + s)−βe−Cγs+1,t+1 . (A.10)
The first term on the right hand side is exponentially decaying in t since γ1,t+1 grows polynomially in t. To bound the second term, observe that the function
f(s) := Cγs+1,t+1 − β log(s+ 1)
satisfies
f ′(s) ≥ 0 ⇔ C(s+ 1)−α − β 1 + s
≥ 0 ⇔ s ≥ ( β
C
)1/(1−α) =: K.
Hence, the summands are monotonically increasing for s greater than a fixed constant K depending only on α, β, C. Note that
K∑ s=1 (1 + s)−βe−Cγs+1,t+1 ≤ Ke−CγK+1,t+1 ≤ Ke−C ′t1−α
for some C ′ depending only on α and K, and hence sum is exponentially decaying in t. Further, using an integral comparison, we find
t∑ s=K+1 (1 + s)−βe−Cγs+1,t+1 ≤ ∫ t K (1 + s)−βe− C 1−α ((t+1) 1−α−(s+1)1−α)ds. (A.11)
Changing variables using u = (1 + s)1−α/(1− α), the last integral has the form
e−Cgt(1− α)−ξ ∫ gt gK u−ξeCudu, gx := (1 + x)1−α 1− α , ξ := β − α 1− α . (A.12)
Integrating by parts, we have∫ gt gK u−ξeudu = C−1ξ ∫ gt gK u−ξ−1eCudu+ (u−ξeCu)|gtgK
Further, since on the range gK ≤ u ≤ gt the integrand is increasing, we have
e−Cgtξ ∫ gt gK u−ξ−1eCudu ≤ ξg−ξt .
Hence, e−Cgt times the integral in (A.12) is bounded above by
O(g−ξt ) + e −Cgt(u−ξeCu)|gtgK = O(g −ξ t ).
Using (A.11) and substituting the previous line into (A.12) yields the estimate
t∑ s=K+1 (1 + s)−βe−Cγs+1,t+1 ≤ (1 + t)−β+α,
which completes the proof.
Lemma A.5. Suppose {Zt} has two-sided decay of type ({At}, {Ct}). Assume also that for some X ≥ 0 and C > 0, we have
logE ∥∥∥∥∥( t∏
r=s
(Id−Ar) ) Q‖ ∥∥∥∥∥ 2
2
 < X − C ∫ t+1 s r−αdr
as well as Tr(Ct) = O(t−β) for some 0 < α < 1 < β. Then Tr(QT‖ ZtQ‖) = O(t α−β).
Proof. This argument is identical to the proof of Lemma A.4. Indeed, using (A.9) we have that
Tr ( QT‖ ZtQ‖ ) ≤ C1e−Cγ1,t+1 + C2eX t∑ s=1 (1 + s)−βe−Cγs+1,t+1 .
The right hand side of this inequality coincides with the expression on the right hand side of (A.10), which we already bounded by O(tβ−α) in the proof of Lemma A.4.
In what follows, we will use a concentration result for products of matrices from Huang et al. (2020). Let Y1, . . . , Yn ∈ RN×N be independent random matrices. Suppose that
‖E[Yi]‖2 ≤ ai and E [ ‖Yi − E[Yi]‖22 ] ≤ b2i a2i
for some a1, . . . , an and b1, . . . , bn. We will use the following result, which is a specialization of (Huang et al., 2020, Theorem 5.1) for p = q = 2. Theorem A.6 ((Huang et al., 2020, Theorem 5.1)). For Z0 ∈ RN×n, the product Zn = YnYn−1 · · ·Y1Z0 satisfies
E [ ‖Zn‖22 ] ≤ e ∑n i=1 b 2 i n∏ i=1 a2i · ‖Z0‖22
E [ ‖Zn − E[Zn]‖22 ] ≤ ( e ∑n i=1 b 2 i − 1 ) a2i · ‖Z0‖22.
Finally, we collect two simple analytic lemmas for later use. Lemma A.7. For any matrix M ∈ Rm×n, we have that
E[‖M‖22] ≥ ‖E[M ]‖22.
Proof. We find by Cauchy-Schwartz and the convexity of the spectral norm that
E[‖M‖22] ≥ E[‖M‖2]2 ≥ ‖E[M ]‖22. Lemma A.8. For bounded at ≥ 0, if we have ∑∞ t=0 at =∞, then for any C > 0 we have
∞∑ t=0 ate −C ∑t s=0 as <∞.
Proof. Define bt := ∑t s=0 as so that
S := ∞∑ t=0 ate −C ∑t s=0 as = ∞∑ t=0 (bt − bt−1)e−Cbt ≤ ∫ ∞ 0 e−Cxdx <∞,
where we use ∫∞ 0 e−Cxdx to upper bound its right Riemann sum.

heading: B ANALYSIS OF DATA AUGMENTATION AS STOCHASTIC OPTIMIZATION
text: In this section, we prove generalizations of our main theoretical results Theorems 5.1 and 5.2 giving Monro-Robbins type conditions for convergence and rates for augmented gradient descent in the linear setting.

heading: B.1 MONRO-ROBBINS TYPE RESULTS
text: To state our general Monro-Robbins type convergence results, let us briefly recall the notation. We consider overparameterized linear regression with loss
L(W ;D) = 1 N ||WX − Y ||2F ,
where the dataset D of size N consists of data matrices X,Y that each have N columns xi ∈ Rn, yi ∈ Rp with n > N. We optimize L(W ;D) by augmented gradient descent, which means that at each time t we replace D = (X,Y ) by a random dataset Dt = (Xt, Yt). We then take a step
Wt+1 = Wt − ηt∇WL(Wt;Dt) of gradient descent on the resulting randomly augmented lossL(W ;Dt) with learning rate ηt. Recall that we set V‖ := column span of E[XtXTt ] and denoted by Q‖ the orthogonal projection onto V‖. As noted in §5, on V‖ the proxy loss
Lt = E [L(W ;Dt)] is strictly convex and has a unique minimum, which is
W ∗t = E [ YtX T t ] (Q||E [ XtX T t ] Q||) −1.
The change from one step of augmented GD to the next in these proxy optima is captured by Ξ∗t := W ∗ t+1 −W ∗t .
With this notation, we are ready to state Theorems B.1, which gives two different sets of timevarying Monro-Robbins type conditions under which the optimization trajectory Wt converges for large t. In Theorem B.4, we refine the analysis to additionally give rates of convergence. 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 <∞, (B.1)
ensuring the existence of a limit W ∗∞ := limt→∞W ∗ t and that
∞∑ t=0 ηtλmin,V‖(E[XtX T t ]) =∞. (B.2)
Then if either ∞∑ t=0 η2tE [ ‖XtXTt − E[XtXTt ]‖2F + ‖YtXTt − E[YtXTt ]‖2F ] <∞ (B.3) or ∞∑ t=0 η2tE [ ‖XtXTt − E[XtXTt ]‖2F
+ ∥∥∥E[Wt](XtXTt − E[XtXTt ])− (YtXTt − E[YtXTt ])∥∥∥2
F
] <∞ (B.4)
hold, then for any initialization W0, we have WtQ‖ p→W ∗∞.
Remark B.2. In the general case, the column span V|| of E[XtXTt ] may vary with t. This means that some directions in Rn may only have non-zero overlap with colspan(E[XtXTt ]) for some positive but finite collection of values of t. In this case, only finitely many steps of the optimization would moveWt in this direction, meaning that we must define a smaller space for convergence. The correct definition of this subspace turns out to be the following
V‖ := ∞⋂ t=0 ker [ ∞∏ s=t ( Id−2ηs N E[XsXTs ] )] (B.5)
= ∞⋂ t=0
{ u ∈ Rn ∣∣∣∣ limT→∞ T∏ s=t ( Id−2ηs N E[XsXTs ] ) u = 0 } .
With this re-definition of V|| and with Q‖ still denoting the orthogonal projection to V‖, Theorem B.1 holds verbatim and with the same proof. Note that if ηt → 0, V||colspan(E[XtXTt ]) is fixed in t, and (B.2) holds, this definition of V‖ reduces to that defined in (5.3).
Remark B.3. The condition (B.4) can be written in a more conceptual way as ∞∑ t=0 [ ‖XtXTt − E[XtXTt ]‖2F + η2t Tr [ Id ◦Var ( (E[Wt]Xt − Yt)XTt )]] <∞,
where we recognize that (E[Wt]Xt − Yt)XTt is precisely the stochastic gradient estimate at time t for the proxy loss Lt, evaluated at E [Wt], which is the location at time t for vanilla GD on Lt since taking expectations in the GD update equation (5.1) coincides with GD for Lt. Moreover, condition (B.4) actually implies condition (B.3) (see (B.12) below). The reason we state Theorem B.1 with both conditions, however, is that (B.4) makes explicit reference to the average E [Wt] of the augmented trajectory. Thus, when applying Theorem B.1 with this weaker condition, one must separately estimate the behavior of this quantity.
Theorem B.1 gave conditions on joint learning rate and data augmentation schedules under which augmented optimization is guaranteed to converge. Our next result proves rates for this convergence.
Theorem B.4. Suppose that ηt → 0 and that for some 0 < α < 1 < β1, β2 and C1, C2 > 0, we have
logE ∥∥∥∥∥( t∏
r=s
( Id−2ηr
N XrX
T r )) Q‖ ∥∥∥∥∥ 2
2
 < C1 − C2 ∫ t+1 s r−αdr (B.6)
as well as ‖Ξ∗t ‖F = O(t−β1) (B.7)
and
η2t Tr [ Id ◦Var(E[Wt]XtXTt − YtXTt )] = O(t−β2). (B.8)
Then, for any initialization W0, we have for any > 0 that
tmin{β1−1, β2−α 2 }− ‖WtQ‖ −W ∗∞‖F p→ 0.
Remark B.5. To reduce Theorem 5.2 to Theorem B.4, we notice that (5.9) and (5.10) mean that Theorem A.6 applies to Yt = Id−2ηt XtX T t N with at = 1 − Ω(t −α) and and b2t = O(t
−γ), thus implying (B.6).
The first step in proving both Theorem B.1 and Theorem B.4 is to obtain recursions for the mean and variance of the difference Wt − W ∗t between the time t proxy optimum and the augmented optimization trajectory at time t. We will then complete the proof of Theorem B.1 in §B.3 and the proof of Theorem B.4 in §B.4.

heading: B.2 RECURSION RELATIONS FOR PARAMETER MOMENTS
text: The following proposition shows that difference between the mean augmented dynamics E[Wt] and the time−t optimum W ∗t satisfies, in the sense of Definition A.1, one-sided decay of type ({At}, {Bt}) with
At = 2ηt N XtX T t , Bt = −Ξ∗t .
It also shows that the variance of this difference, which is non-negative definite, satisfies two-sided decay of type ({At}, {Ct}) with At as before and
Ct = 4η2t N2
[ Id ◦Var ( E[Wt]XtXTt − YtXTt )] .
In terms of the notations of Appendix A.1, we have the following recursions.
Proposition B.6. The quantity E[Wt]−W ∗t satisfies E[Wt+1]−W ∗t+1 = (E[Wt]−W ∗t ) (
Id−2ηt N
E[XtXTt ] ) − Ξ∗t (B.9)
and Zt := E[(Wt − E[Wt])T(Wt − E[Wt])] satisfies Zt+1 = E [ (Id−2ηt
N XtX
T t )Zt(Id− 2ηt N XtX T t )
] +
4η2t N2
[ Id ◦Var ( E[Wt]XtXTt − YtXTt )] .
(B.10)
Proof. Notice that E[XtXTt ]u = 0 if and only if XTt u = 0 almost surely, which implies that
W ∗t E[XtXTt ] = E[YtXTt ]E[XtXTt ]+E[XtXTt ] = E[YtXTt ]. Thus, the learning dynamics (5.1) yield
E[Wt+1] = E[Wt]− 2ηt N
( E[Wt]E[XtXTt ]− E[YtXTt ] ) = E[Wt]−
2ηt N (E[Wt]−W ∗t )E[XtXTt ].
Subtracting W ∗t+1 from both sides yields (B.9). We now analyze the fluctuations. Writing Sym(A) := A+AT, we have
E[Wt+1]TE[Wt+1] = E[Wt]TE[Wt] + 2ηt N
Sym ( E[Wt]TE[YtXTt ]− E[Wt]TE[Wt]E[XtXTt ] ) +
4η2t N2
( E[XtXTt ]E[Wt]TE[Wt]E[XtXTt ]+E[XtY Tt ]E[YtXTt ]−Sym(E[XtXTt ]E[Wt]TE[YtXTt ]) ) .
Similarly, we have that
E[WTt+1Wt+1] = E[WTt Wt] + 2ηt N Sym(E[WTt YtXTt −WTt WtXtXTt ])
+ 4η2t N2 E[XtXTt WTt WtXtXTt − Sym(XtXTt WTt YtXTt ) +XtY Tt YtXTt ].
Noting that Xt and Yt are independent of Wt and subtracting yields the desired.

heading: B.3 PROOF OF THEOREM B.1
text: First, by Proposition B.6, we see that E[Wt]−W ∗t has one-sided decay with
At = 2ηt XtX
T t
N and Bt = −Ξ∗t .
Thus, by Lemma A.2 and (B.1), we find that
lim t→∞
(E[Wt]Q‖ −W ∗t ) = 0, (B.11)
which gives convergence in expectation.
For the second moment, by Proposition B.6, we see that Zt has two-sided decay with
At = 2ηt XtX
T t
N and Ct = 4η2t N2
[ Id ◦Var ( E[Wt]XtXTt − YtXTt )] .
We now verify (A.7) and (A.8) in order to apply Lemma A.3.
For (A.7), for any > 0, notice that
E[‖As − E[As]‖2F ] = η2sE[‖XsXTs − E[XsXTs ]‖2F ] so by either (B.3) or (B.4) we may choose T1 > t so that ∑∞ s=T1
E[‖As − E[As]‖2F ] < 2 . Now choose T2 > T1 so that for T > T2, we have∥∥∥∥∥( T∏ r=T1 E[Id−Ar] ) Q‖ ∥∥∥∥∥ 2
2
<
2
1 ‖ ∏T1−1 s=t E[Id−As]‖2F + ∑T1−1 s=t E[‖As − E[As]‖2F ] .
For T > T2, we then have
E ∥∥∥∥∥( T∏ s=t (Id−As) ) Q‖ ∥∥∥∥∥ 2
2  ≤
∥∥∥∥∥( T∏ s=t E[Id−As] ) Q‖ ∥∥∥∥∥ 2 + T∑ s=t E ∥∥∥∥∥ s∏ r=t (Id−Ar) T∏ r=s+1 (Id−E[Ar])Q‖ ∥∥∥∥∥ 2
F
− ∥∥∥∥∥ s−1∏ r=t (Id−Ar) T∏ r=s (Id−E[Ar])Q‖ ∥∥∥∥∥ 2
F  =
∥∥∥∥∥( T∏ s=t E[Id−As] ) Q‖ ∥∥∥∥∥ 2
F
+ T∑ s=t E ∥∥∥∥∥ s−1∏ r=t (Id−Ar)(As − E[As]) T∏ r=s+1 (Id−E[Ar])Q‖ ∥∥∥∥∥ 2
F  ≤
∥∥∥∥∥ T1−1∏ s=t E[Id−As] ∥∥∥∥∥ 2 ∥∥∥∥∥F( T∏ r=T1 E[Id−Ar] ) Q‖ ∥∥∥∥∥ 2
2
+ T∑ s=t E[‖As − E[As]‖2F ] ∥∥∥∥∥( T∏ r=s+1 E[Id−Ar] ) Q‖ ∥∥∥∥∥ 2
2 ≤ (∥∥∥∥∥ T1−1∏ s=t E[Id−As] ∥∥∥∥∥ 2
F
+ T1−1∑ s=t E[‖As − E[As]‖2F ] )∥∥∥∥∥( T∏ r=T1 E[Id−Ar] ) Q‖ ∥∥∥∥∥ 2
2
+ T∑ s=T1 E[‖As − E[As]‖2F ]
< ,
which implies (A.7). Condition (A.8) follows from either (B.4) or (B.3) and the bounds
Tr(Ct) ≤ 8η2t N2
( ‖E[Wt](XtXTt − E[XtXTt ])‖2F + ‖YtXTt − E[YtXTt ]‖2F ) (B.12)
≤ 8η 2 t N2 ( ‖E[Wt]‖2‖XtXTt − E[XtXTt ]‖2F + ‖YtXTt − E[YtXTt ]‖2F ) ,
where in the first inequality we use the fact that ‖M1−M2‖2F ≤ 2(‖M1‖2F +‖M2‖2F ). Furthermore, iterating (B.9) yields ‖E[Wt]−W ∗t ‖F ≤ ‖W0−W ∗0 ‖F+ ∑∞ t=0 ‖Ξ∗t ‖F , which combined with (B.12) and either (B.3) or (B.4) therefore implies (A.8). We conclude by Lemma A.3 that
lim t→∞ QT‖ZtQ‖ = limt→∞ E[QT‖ (Wt − E[Wt]) T(Wt − E[Wt])Q‖] = 0. (B.13)
Together, (B.11) and (B.13) imply that WtQ‖ −W ∗t p→ 0. The conclusion then follows from the fact that limt→0W ∗t = W ∗ ∞. This complete the proof of Theorem B.1.

heading: B.4 PROOF OF THEOREM B.4
text: By Proposition B.6, E[Wt]−W ∗t has one-sided decay with
At = 2ηt N XtX T t , Bt = −Ξ∗t .
By Lemma A.7 and (B.6), E[At] satisfies
log ∥∥∥∥∥ t∏
r=s
( Id−2ηr 1
N E[XrXTr ]
) Q‖ ∥∥∥∥∥ 2 ≤ 1 2 logE ∥∥∥∥∥( t∏ r=s ( Id−2ηr XrX T r N )) Q‖ ∥∥∥∥∥ 2 2  < C1 2 − C2 2 ∫ t+1 s r−αdr.
Applying Lemma A.4 using this bound and (B.7), we find that
‖E[Wt]Q‖ −W ∗t ‖F = O(tα−β1).
Moreover, because ‖Ξ∗t ‖F = O(t−β1), we also find that ‖W ∗t −W ∗∞‖F = O(t−β1+1), and hence
‖E[Wt]Q‖ −W ∗∞‖F = O(t−β1+1).
Further, by Proposition B.6, E[(Wt − E[Wt])T(Wt − E[Wt])] has two-sided decay with
At = 2ηt N XtX T t , Ct = 4η2t N2
[ Id ◦Var ( E[Wt]XtXTt − YtXTt )] .
Applying Lemma A.5 with (B.6) and (B.8), we find that E [ ‖(Wt − E[Wt])Q‖‖2F ] = O(tα−β2).
By Chebyshev’s inequality, for any x > 0 we have P ( ‖WtQ‖ −W ∗∞‖F ≥ O(t−β1+1) + x ·O(t α−β2 2 ) ) ≤ x−2.
For any > 0, choosing x = tδ for small 0 < δ < we find as desired that
tmin{β1−1, β2−α 2 }− ‖WtQ‖ −W ∗∞‖F p→ 0,
thus completing the proof of Theorem B.4.

heading: C ANALYSIS OF NOISING AUGMENTATIONS
text: In this section, we give a full analysis of the noising augmentations presented in Section 4. Let us briefly recall the notation. As before, we consider overparameterized linear regression with loss
L(W ;D) = 1 N ||WX − Y ||2F ,
where the dataset D of size N consists of data matrices X,Y that each have N columns xi ∈ Rn, yi ∈ Rp with n > N. We optimize L(W ;D) by augmented gradient descent with additive Gaussian noise, which means that at each time t we replace D = (X,Y ) by a random dataset Dt = (Xt, Y ), where the columns xi,t of Xt are
xi,t = xi + σtGi, Gi ∼ N (0, 1) i.i.d.
We then take a step Wt+1 = Wt − ηt∇WL(Wt;Dt)
of gradient descent on the resulting randomly augmented loss L(W ;Dt) with learning rate ηt. A direct computation shows that the proxy loss
Lt = E [L(W ;Dt)] = L(W ;D) + σ2tN ||W || 2 F ,
which is strictly convex. Thus, the space
V‖ := column span of E[XtXTt ]
is simply all of Rn. Moreover, the proxy loss has a unique minimum, which is
W ∗t = Y X T (σ2tN Idn×n +XX T )−1.

heading: C.1 PROOF OF THEOREM 4.1
text: We first show convergence. For this, we seek to show that if σ2t , ηt → 0 with σ2t non-increasing and
∞∑ t=0 ηtσ 2 t =∞ and
∞∑ t=0 η2t σ 2 t <∞, (C.1)
then, Wt p→ Wmin. We will do this by applying Theorem 5.1, so we check that our assumptions imply the hypotheses of these theorems. For Theorem 5.1, we directly compute
E[YtXTt ] = Y XT and E[XtXTt ] = XXT + σ2tN · Idn×n
and
E[XtXTt Xt] = XXTX + σ2t (N + n+ 1)X E[XtXTt XtXTt ] = XXTXXT + σ2t ( (2N + n+ 2)XXT + Tr(XXT) Idn×n ) + σ4tN(N + n+ 1) Idn×n .
We also find that ‖Ξ∗t ‖F = |σ2t − σ2t+1|N ∥∥∥∥Y XT(XXT + σ2tN · Idn×n )−1(XXT + σ2t+1N · Idn×n )−1∥∥∥∥
F
≤ |σ2t − σ2t+1|N‖Y XT[(XXT)+]2‖F .
Thus, because σ2t is decreasing, we see that the hypothesis (5.5) of Theorem 5.1 indeed holds. Further, we note that
∞∑ t=0 η2tE [ ‖XtXTt − E[XtXTt ]‖2F + ‖YtXTt − E[YtXTt ]‖2F ] =
∞∑ t=0 η2t σ 2 t ( 2(n+ 1)‖X‖2F +N‖Y ‖2F + σ2tNn(n+ 1) ) = O ( ∞∑ t=0 η2t σ 2 t ) ,
which by (C.1) implies (B.3). Theorem 5.1 and the fact that limt→∞W ∗t = Wmin therefore yield that Wt
p→Wmin. For the rate of convergence, we aim to show that if ηt = Θ(t−x) and σ2t = Θ(t
−y) with x, y > 0, x+ y < 1, and 2x+ y > 1, then for any > 0, we have that
tmin{β, 1 2α}− ‖Wt −Wmin‖F p→ 0.
We now check the hypotheses for and apply Theorem B.4. For (B.6), notice that Yr = Id−2ηr XrX T r
N
satisfies the hypotheses of Theorem A.6 with ar = 1 − 2ηrσ2r and b2r = η2rσ 2 r
a2r
( 2(n + 1)‖X‖2F +
σ2rNn(n + 1) ) . Thus, by Theorem A.6 and the fact that ηt = Θ(t−x) and σ2t = Θ(t −y), we find
for some C1, C2 > 0 that
logE ∥∥∥∥∥ t∏
r=s
(Id−2ηr XrX
T r
N ) ∥∥∥∥∥ 2
2
 ≤ t∑ r=s b2r + 2 t∑ r=s log(1− 2ηrσ2r)
≤ C1 − C2 ∫ t+1 s r−x−ydr.
For (B.7), we find that
‖Ξ∗t ‖F ≤ |σ2t − σ2t+1|N‖Y XT[(XXT)+]2‖F = O(t−y−1).
Finally, for (B.8), we find that η2t Tr [ Id ◦Var ( E[Wt]XtXTt − YtXTt )] = O(t−2x−y).
Noting finally that ‖W ∗t −Wmin‖F = O(σ2t ) = O(t−y), we apply Theorem B.4 with α = x + y, β1 = y + 1, and β2 = 2x+ y to obtain the desired estimates. This concludes the proof of Theorem 4.1.

heading: D ANALYSIS OF SGD
text: This section gives the full analysis of the results for stochastic gradient descent with and without additive synthetic noise presented in Sections 6.1 and 6.2. Let us briefly recall the notation. As before, we consider overparameterized linear regression with loss
L(W ;D) = 1 N ||WX − Y ||2F ,
where the dataset D of size N consists of data matrices X,Y that each have N columns xi ∈ Rn, yi ∈ Rp with n > N.We optimize L(W ;D) by augmented SGD either with or without additive Gaussian noise. In the former case, this means that at each time t we replace D = (X,Y ) by a random batch Bt = (Xt, Yt) given by a prescribed batch size Bt = |Bt| in which each datapoint
in Bt is chosen uniformly with replacement from D, and the resulting data matrices Xt and Yt are scaled so that Lt(W ) = L(W ;D). Concretely, this means that for the normalizing factor ct := √ N/Bt we have
Xt = ctXAt and Yt = ctY At, (D.1)
whereAt ∈ RN×Bt has i.i.d. columnsAt,i with a single non-zero entry equal to 1 chosen uniformly at random. In this setting the minimum norm optimum for each t are the same and given by
W ∗t = W ∗ ∞ = Y X T(XXT)+,
which coincides with the minimum norm optimum for the unaugmented loss.
In the setting of SGD with additive noise at level σt, we take instead
Xt = ct(XAt + σtGt) and Yt = ctY At,
where ct and At are as before and Gt ∈ Rn×Bt has i.i.d. Gaussian entries. In this setting, the proxy loss is
Lt(W ) := 1 N E [ ‖ctY At − ctWXAt − ctσtWGt‖2F ] = 1 N ‖Y −WX‖2F + σ2t ‖W‖2F ,
which has ridge minimizer W ∗t = Y X T(XXT + σ2tN · Idn×n)−1.
We begin in §D.1 by treating the case of noiseless SGD. We then do the analysis in the presence of noise in §D.2.

heading: D.1 PROOF OF THEOREM 6.1
text: In order to apply Theorems B.1 and B.4, we begin by computing the moments of At as follows. Recall the notation diag(M) from Appendix A.1.
Lemma D.1. For any Z ∈ RN×N , we have that
E[AtATt ] = Bt N IdN×N and E[AtATt ZAtATt ] = Bt N diag(Z) + Bt(Bt − 1) N2 Z.
Proof. We have that
E[AtATt ] = Bt∑ i=1 E[Ai,tATi,t] = Bt N IdN×N .
Similarly, we find that
E[AtATt ZAtATt ] = Bt∑ i,j=1 E[Ai,tATi,tZAj,tATj,t]
= Bt∑ i=1 E[Ai,tATi,tZAi,tATi,t] + 2 ∑
1≤i<j≤Bt
E[Ai,tATi,tZAj,tATj,t]
= Bt N diag(Z) + Bt(Bt − 1) N2 Z,
which completes the proof.
Let us first check convergence in mean:
E[Wt]Q‖ →W ∗∞.
To see this, note that Lemma D.1 implies
E[YtXTt ] = Y XT E[XtXTt ] = XXT,
which yields that W ∗t = Y X T[XXT]+ = W ∗∞ (D.2)
for all t. We now prove convergence. Since all W ∗t are equal to W ∗ ∞, we find that Ξ ∗ t = 0. By (B.9) and Lemma D.1 we have
E[Wt+1]−W ∗∞ = (E[Wt]−W ∗∞) (
Id−2ηt N XXT
) ,
which implies since 2ηtN < λmax(XX T)−1 for large t that for some C > 0 we have
‖E[Wt]Q‖ −W ∗∞‖F ≤ ‖W0Q‖ −W ∗∞‖F t−1∏ s=0 ∥∥∥∥Q‖ − 2ηsN XXT ∥∥∥∥ 2
≤ C‖W0Q‖ −W ∗∞‖F exp ( − t−1∑ s=0 2ηs N λmin,V‖(XX T) ) . (D.3)
From this we readily conclude using (6.1) the desired convergence in mean E[Wt]Q‖ →W ∗∞.
Let us now prove that the variance tends to zero. By Proposition B.6, we find that Zt = E[(Wt − E[Wt])T(Wt − E[Wt])] has two-sided decay of type ({At}, {Ct}) with
At = 2ηt N XtX T t , Ct = 4η2t N2
[ Id ◦Var((E[Wt]Xt − Yt)XTt ) ] .
To understand the resulting rating of convergence, let us first obtain a bound on Tr(Ct). To do this, note that for any matrix A, we have
Tr (Id ◦Var[A]) = Tr ( E [ ATA ] − E [A]T E [A] ) .
Moreover, using the definition (D.1) of the matrix At and writing
Mt := E [Wt]X − Y, we find (
(E [Wt]Xt − Yt)XTt )T (E [Wt]Xt − Yt)XTt = XAtATtMTt MtAtATt XT
as well as E [( (E[Wt]Xt − Yt)XTt )]T E [(E[Wt]Xt − Yt)XTt ] = XE [AtATt ]MTt MtE [AtATt ]XT.
Hence, using the expression from Lemma D.1 for the moments ofAt and recalling the scaling factor ct = (N/Bt) 1/2, we find
Tr(Ct) = 4η2t Bt Tr
( X { diag ( MTt Mt ) − 1 N MTt Mt } XT ) .
Next, writing ∆t := E[Wt]−W ∗∞
and recalling (D.2), we see that Mt = ∆tX.
Thus, applying the estimates (D.3) about exponential convergence of the mean, we obtain
Tr(Ct) ≤ 8η2t Bt ∣∣∣∣∆tQ||∣∣∣∣22 ∣∣∣∣XXT ∣∣∣∣22 ≤ C 8η 2 t
Bt
∣∣∣∣XXT ∣∣∣∣2 2 ‖∆0Q‖‖2F exp ( − t−1∑ s=0 4ηs N λmin,V‖(XX T) ) . (D.4)
Notice now that Yr = Q‖ − Ar satisfies the conditions of Theorem A.6 with ar = 1 − 2ηr 1 N λmin,V‖(XX T) and b2r = 4η2r Bra2rN Tr ( X diag(XTX)X − 1NXX TXXT )
. By Theorem A.6 we then obtain for any t > s > 0 that
E ∥∥∥∥∥ t∏
r=s+1
(Q‖ −Ar) ∥∥∥∥∥ 2
2
 ≤ e∑tr=s+1 b2r t∏ r=s+1 ( 1− 2ηr 1 N λmin,V‖(XX T) )2 . (D.5)
By two-sided decay of Zt, we find by (D.4), (D.5), and (A.9) that
E[‖WtQ‖ − E[Wt]Q‖‖2F ] = Tr(Q‖ZtQ‖)
≤ e− 4 N λmin,V‖ (XX
T) ∑t−1 s=0 ηs
‖XXT‖22 N2 ‖∆0Q‖‖2FC t−1∑ s=0 8η2s Bs/N e 4ηs N λmin,V‖ (XX T)+ ∑t r=s+1 b 2 r . (D.6)
Since ηs → 0, we find that ηs NBs e 4ηs N λmin,V‖ (XX T) is uniformly bounded and that b2r ≤ 4 N λmin,V‖(XX T)ηr for sufficiently large r. We therefore find that for some C ′ > 0,
E[‖WtQ‖ − E[Wt]Q‖‖2F ] ≤ C ′ t−1∑ s=0 ηse − 4N λmin,V‖ (XX T) ∑s r=0 ηr ,
hence limt→∞ E[‖WtQ‖ − E[Wt]Q‖‖2F ] = 0 by Lemma A.8. Combined with the fact that E[Wt]Q‖ →W ∗∞, this implies that WtQ‖ p→W ∗∞.
To obtain a rate of convergence, observe that by (D.3) and the fact that ηt = Θ(t−x), for some C1, C2 > 0 we have ‖E[Wt]Q‖ −W ∗∞‖F ≤ C1 exp ( − C2t1−x ) . (D.7)
Similarly, by (D.6) and the fact that ηsBs/N <∞ uniformly, for some C3, C4, C5 > 0 we have E[‖WtQ‖ − E[Wt]Q‖‖2F ] ≤ C3 exp ( − C4t1−x ) t1−x
We conclude by Chebyshev’s inequality that for any a > 0 we have P ( ‖WtQ‖ −W ∗∞‖F ≥ C1 exp ( − C2t1−x ) + a · √ C3t 1 2− x 2 e−C4t 1−x/2 ) ≤ a−2.
Taking a = t, we conclude as desired that for some C > 0, we have
eCt 1−x ‖WtQ‖ −W ∗∞‖F p→ 0. This completes the proof of Theorem 6.1.

heading: D.2 PROOF OF THEOREM 6.2
text: We now complete our analysis of SGD with Gaussian noise. We will directly check that the optimization trajectory Wt converges at large t to the minimal norm interpolant W ∗∞ with the rates claimed in Theorem 6.2. We will deduce this from Theorem B.4. To check the hypotheses of this theorem, we will need expressions for its moments, which we record in the following lemma. Lemma D.2. We have
E[YtXTt ] = Y XT and E[XtXTt ] = XXT + σ2tN Idn×n . (D.8)

heading: Moreover,
text: E[YtXTt XtY Tt ] = c4tE[Y AtATt XTXAtATt Y T + σ2t Y AtGTt GtATt Y T]
= N
Bt Y diag(XTX)Y T + Bt − 1 Bt Y XTXY T + σ2tNY Y T
E[YtXTt XtXTt ] = c4tE[Y AtATt XTXAtATt XT + σ2t Y AtGTt GtATt XT
+ σ2t Y AtG T t XAtG T t + σ 2 t Y AtA T t X TGtG T t ]
= N
Bt Y diag(XTX)XT + Bt − 1 Bt Y XTXXT + σ2t (N + n+ 1 Bt/N )Y XT
E[XtXTt XtXTt ] = c4tE[XAtATt XTXAtATt XT + σ2tGtGTt XAtATt XT + σ2tXAtGTt GtATt XT
+ σ2tXAtA T t X TGtG T t + σ 2 tGtA T t X TGtA T t X T + σ2tXAtG T t XAtG T t
+ σ2tGtA T t X TXAtG T t + σ 4 tGtG T t GtG T t ]
= N
Bt X diag(XTX)XT + Bt − 1 Bt XXTXXT + σ2t (2N + n+ 2 Bt/N )XXT
+ σ2t N
Bt Tr(XXT) Idn×n +σ 4 tN(N +
n+ 1 Bt/N ) Idn×n .
Proof. All these formulas are obtained by direct, if slightly tedious, computation.
With these expressions in hand, we can readily check the of conditions Theorem B.4. First, we find using the Sherman-Morrison-Woodbury matrix inversion formula that
‖Ξ∗t ‖F = |σ2tN − σ2t+1N | ∥∥Y XT(XXT + σ2tN · Idn×n)−1(XXT + σ2t+1N · Idn×n)−1∥∥F
(D.9) ≤ N |σ2t − σ2t+1| ∥∥Y XT[(XXT)+]2∥∥ F .
Hence, assuming that σ2t = Θ(t −y), we see that condition (B.7) of Theorem B.4 holds with
β1 = −y − 1.
Next, let us verify that the condition (B.6) holds for an appropriate α. For this, we need to bound
logE ∣∣∣∣∣ ∣∣∣∣∣ t∏
r=s
( Id−2ηr
N XrX
T r )∣∣∣∣∣ ∣∣∣∣∣ 2
2
,
which we will do using Theorem A.6. In order to apply this result, we find by direct inspection of the formula E[XrXTr ] = XXT + σ2rN Idn×n that ∣∣∣∣∣∣∣∣E [Id−2ηrN XrXTr ]∣∣∣∣∣∣∣∣ 2 = 1− 2ηrσ2r := ar. Moreover, we have
E [∣∣∣∣∣∣∣∣Id−2ηrN XrXTr − E [ Id−2ηr N XrX T r ]∣∣∣∣∣∣∣∣2 2 ] = 4η2r N2 E [∣∣∣∣XrXTr − E [XrXTr ]∣∣∣∣22] .
Using the exact expressions for the resulting moments from Lemma D.2, we find
4η2r N2 E [∣∣∣∣XrXTr − E [XrXTr ]∣∣∣∣22] =
4η2r N2
[ 1 Bt Tr ( X(N diag(XTX)−XTX)XT ) + 2σ2t n+ 1 Bt/N Tr(XXT) + σ4t Nn(n+ 1) Bt/N ] ≤ Cη2r .
Thus, applying Theorem A.6, we find that
logE ∣∣∣∣∣ ∣∣∣∣∣ t∏
r=s
( Id−2ηr
N XrX
T r )∣∣∣∣∣ ∣∣∣∣∣ 2
2
≤ t∑
r=s
Cη2r log
( t∏
r=s
( 1− 2ηrσ2r )) ≤ t∑ r=s Cη2r − 2ηrσ2r .
Recall that, in the notation of Theorem 6.2, we have
ηr = Θ(r −x), σ2r = Θ(r −y).
Hence, since under out hypotheses we have x < 2y, we conclude that condition (B.6) holds with α = x+ y. Moreover, exactly as in Proposition B.6, we have
∆′t+1 = ∆ ′ t
( Id−2ηt N E [ XtX T t ]) + 2 N Ξ∗t , ∆ ′ t := E [Wt −W ∗t ] .
Since ||Ξ∗t ||F = O(t −y−1) and we already saw that ∣∣∣∣∣∣∣∣Id−2ηtN E [XtXTt ] ∣∣∣∣∣∣∣∣ 2 = 1− 2ηtσ2t , we may use the single sided decay estimates Lemma A.4 to conclude that
||∆′t||F = O(t x−1).
Finally, it remains to bound η2t Tr [ Id ◦Var(E[Wt]XtXTt − YtXTt )] .
A direct computation using Lemma D.2 shows
E [ ‖YtXTt − E[YtXTt ]‖2F ] = 1 Bt Tr ( Y (N diag(XTX)−XTX)Y T ) + σ2tN Tr(Y Y T).
Hence, again using D.2, we find η2t Tr [ Id ◦Var(E[Wt]XtXTt − YtXTt ) ] = η2t Tr ( 1 Bt E[Wt]X(N diag(XTX)−XTX)XTE[Wt]T
+ 2σ2t n+ 1
Bt/N E[Wt]XXTE[Wt]T + (σ2t
N Bt Tr(XXT) + σ4tN n+ 1 Bt/N )E[Wt]E[Wt]T ) − 2η2t Tr ( 1 Bt Y (N diag(XTX)−XTX)XTE[Wt]T + σ2t n+ 1 Bt/N Y XTE[Wt]T
) + η2t Tr ( 1 Bt Y (N diag(XTX)−XTX)Y T + σ2tNY Y T ) .
To make sense of this term, note that W ∗∞X = Y.
Hence, we find after some rearrangement that η2t Tr [ Id ◦Var(E[Wt]XtXTt − YtXTt ) ] ≤ Cη2t (σ2t + ||∆t|| 2 F ),
where we set ∆t := E [Wt −W ∗∞] .
Finally, we have
∆t ≤ ∆′t + ||W ∗t −W ∗∞||F = O(t x−1) + Θ(t−y) = Θ(t−y)
since we assumed that x+ y < 1. Therefore, we obtain η2t Tr [ Id ◦Var(E[Wt]XtXTt − YtXTt ) ] ≤ Cη2t σ2t = Θ(t−2x−y),
showing that condition (B.8) holds with β2 = 2x+ y. Applying Theorem B.4 completes the proof.

Review summaries:
.** Authors present a novel theoretical framework for assessing the effect of data augmentation (e.g. mini batch SGD), noise addition and the learning rate setup in gradient-based optimization with overparametrized models. Despite the analysis is only performed for linear regression, results extend the well-known Monro-Robbins theorem on rates of convergence. The manuscript is a first step for future analysis of the aforementioned techniques with other type of models and/or loss functions.

The paper considers stochastic gradient descent with noisy gradients. In contrast to the standard setting (e.g., gradient Langevin dynamics) where additive Gaussian noise is added to the model gradient, this work focuses on additive perturbations of data instances. As a result of this, the optimization objective changes throughout the training process because the data is no longer static/fixed but assumed to be sampled from some distribution governing the perturbation process (see Eq. 3.3).