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: .** 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.