Data augmentation as stochastic optimization

1 INTRODUCTION . 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 ) . 2 RELATED WORK . In addition to the extensive empirical work on data augmentation cited elsewhere in this article , we briefly catalog other theoretical work on data augmentation and learning rate schedules . The latter were first considered in the seminal work 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 . 3 DATA AUGMENTATION AS STOCHASTIC OPTIMIZATION . 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 .

Summary: 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).