Data-Efficient Augmentation for Training Neural Networks

1 INTRODUCTION . Data augmentation expands the training data by applying transformations , such as rotations or crops for images , to the original training examples . Due to its effectiveness , data augmentation is a key component in achieving nearly all state-of-the-art results in deep learning applications ( Shorten & Khoshgoftaar , 2019 ) . The most effective data augmentation techniques often search over a ( possibly large ) space of transformations to find sequences of transformations that speeds up training the most ( Cubuk et al. , 2019 ; 2020 ; Luo et al. , 2020 ; Wu et al. , 2020 ) . In addition , multiple augmented examples are usually generated for a single data point to obtain better results , increasing the size of the training data by orders of magnitude . As a result , state-of-the-art data augmentation techniques become computationally prohibitive for large real-world problems ( c.f . Fig . 1 ) . To make data augmentation more efficient and scalable , an effective approach is to carefully select a small subset of the training data such that augmenting only the subset have similar training dynamics to that of full data augmentation . If such a subset can be quickly found , it would directly lead to a significant reduction in storage and training costs , and lower costs incurred from selecting and tuning the optimal set of transformations to apply . Despite the efficiency and scalability that it can provide , this direction has remained largely unexplored . Existing studies are limited to fully training a network and subsampling data points based on its loss or influence for augmentation in subsequent training runs ( Kuchnik & Smith , 2018 ) . However , this method is prohibitive for large datasets , provides a marginal improvement over augmenting random subsets , and does not provide any theoretical guarantee for the performance of the network trained on the augmented subsets . A major challenge in finding the most effective data points for augmentation is to theoretically understand how data augmentation affects the optimization and generalization of neural networks . Existing theoretical results are mainly limited to simple linear classifiers and analyze data augmentation as enlarging the span of the training data ( Wu et al. , 2020 ) , providing a regularization effect ( Bishop , 1995 ; Dao et al. , 2019 ; Wager et al. , 2013 ; Wu et al. , 2020 ) , enlarging the margin of a linear classifier ( Rajput et al. , 2019 ) , or having a variance reduction effect ( Chen et al. , 2019 ) . However , such tools do not provide insights on the effect of data augmentation on training deep neural networks . Here , we study the effect of label invariant data augmentation modeled by small additive perturbations ( Rajput et al. , 2019 ) on training dynamics of overparameterized neural networks . In particular , we rely on recent results that characterize training dynamics of neural networks based on the alignment of the labels/residuals with the singular subspace of the network ’ s Jacobian matrix containing all its first-order partial derivatives ( Arora et al. , 2019 ) . We show that label invariant data augmentation enlarges smaller singular values of the information space , and prove that this will speed up training . Next , we develop a rigorous method to iteratively find small weighted subsets ( coresets ) that when augmented , closely capture the alignment of the full augmented data with the label/residual , at every point during the training . Augmenting the coresets guarantees similar training dynamics to that of full data augmentation . Our key observation is that early in training , this alignment is best captured by data points that are highly representative of their classes . However towards the end of training when the network converges , data points with maximum loss best capture this alignment . Data augmentation has been empirically shown to mainly affect the initial phase of training ( Golatkar et al. , 2019 ) , which crucially determines the final basin of convergence ( Fort et al. , 2020 ) . Better selection of points by our method during this initial phase explains the superior accuracy improvement resulted by augmenting our coresets vs. the max-loss and the better final generalization performance . Importantly , we show that our coresets can be provably trained on even in absence of full data , and even when a high fraction of labels are noisy . We demonstrate the effectiveness of our approach applied to training ResNet20 , ResNet32 , WideResNet on CIFAR10 , CIFAR10-IB and SVHN , compared to random and max-loss baselines ( Kuchnik & Smith , 2018 ) . We show that augmenting coresets found by our approach outperforms the stateof-the-art even in absence of the full data . For instance , when small augmented subsets of size 30 % found by our approach are appended to CIFAR10 , we attain 75 % of the improvement in test accuracy compared to augmenting the full dataset while enjoying a 3.4x speedup in training time . Even in the absence of full data , training and augmenting tiny coresets of size 1 % can achieve 74.7 % accuracy on CIFAR10/ResNet20 while providing a 39x speedup compared to using the full dataset . On CIFAR10 with 50 % noisy labels , augmenting 50 % of the training data outperforms full data augmentation . 2 PROBLEM FORMULATION . We begin by formally describing the problem of learning from augmented data . Consider a dataset Dtrain = ( Xtrain , ytrain ) , where Xtrain = ( x1 , · · · , xn ) ∈ Rd×n is the set of n normalized data points xi ∈ [ 0 , 1 ] d , from the index set V , i.e. , i ∈ V = { 1 , · · · , n } , and ytrain = ( y1 , · · · , yn ) ∈ { y ∈ { ν1 , ν2 , · · · , νC } } with { νj } Cj=1 ∈ [ 0 , 1 ] . Assume at every step t during training , we have a set of augmented examples Dtaug generated by a set of label-invariant transformations . In particular , following ( Rajput et al. , 2019 ) we model data augmentation as an arbitrary bounded additive perturbation , with ‖ ‖≤ 0 . Formally , for a given upper bound 0 and the set of all possible transformations A , we study the transformations selected from S ⊆ A satisfying S = { Ti ∈ A | ‖Ti ( x ) − x‖≤ 0 ∀x ∈Xtrain } . ( 1 ) Under the smoothness constraint of images , where adjacent pixels have close intensities , such as small translations , crops , rotations , and for other pixel-wise augmentation methods such as sharpening , blurring , and color distortions ( Cubuk et al. , 2020 ) , 0 is small . This model is also especially suitable for modelling state-of-the-art augmentation techniques such as structured adversarial perturbation ( Luo et al. , 2020 ) , in which pixel intensities are changed minimally . While in our analysis we focus on small perturbations ε0 , our experiments use variety of strong and weak augmentations . Although other transformations such as data synthesis ( Baluja & Fischer , 2017 ; Mirza & Osindero , 2014 ) , semantic augmentation ( Wang et al. , 2019 ) , large translations , crops , rotations , and flips are still valid under the additive perturbation model , they can be more effectively modelled using matrices of linear transforms ( Wu et al. , 2020 ) , as we analyze in Appendix B.3 . In practice , multiple augmentations are generated for each example xi , and each augmented data point can be a combination of multiple transformations , e.g . random cropping and rotating followed by horizontal flipping . The set of augmentations at iteration t generating r augmented examples per data point can be specified , with abuse of notation , as Dtaug = { ⋃r i=1 ( T t i ( Xtrain ) , ytrain ) } , where |Dtaug|= rn and T ti ( Xtrain ) transforms all the training data points with the set of transformations T ti ⊂ S at iteration t. We denoteXtaug = { ⋃r i=1 T t i ( Xtrain ) } and ytaug = { ⋃r i=1 ytrain } . Let f ( W , x ) be an arbitrary neural network with m vectorized ( trainable ) parametersW ∈ Rm . We assume that the network is trained using ( stochastic ) gradient descent with learning rate η to minimize the squared loss L over the original and augmented training examples Dt = { Dtrain ∪ Dtaug } with associated index set V t , at every iteration t : L ( W t , X ) : = 1 2 ∑ i∈V t Li ( W t , xi ) : = 1 2 ∑ ( xi , yi ) ∈Dt ‖f ( W t , xi ) − yi‖22 . ( 2 ) The gradient update at iteration t is given by W t+1 = W t − η∇L ( W t , X ) , ∇L ( W t , X ) = J T ( W t , X ) ( f ( W t , X ) − y ) , ( 3 ) where Xt = { Xtrain ∪Xtaug } and yt = { ytrain ∪ ytaug } are the set of original and augmented examples and their labels , J ( W , X ) ∈ Rn×m is the Jacobian matrix associated with f , and rt = f ( W t , X ) − y is the residual . We further assume that J is smooth with Lipschitz constant L : ‖J ( W , xi ) − J ( W , xj ) ‖≤ L‖ xi − xj‖ ∀ xi , xj ∈X . ( 4 ) This trivially holds for linear models , and whenW is bounded , it holds for deep ReLU , and generally for networks with any activation φ with bounded derivatives φ′ and φ′′ ( Jordan & Dimakis , 2020 ) . Under this assumption , augmentation as defined in Eq . ( 1 ) results in bounded perturbations to the Jacobian matrix . I.e. , for any transformation Tj ∈ S , we have ‖J ( W , xi ) −J ( W , Tj ( xi ) ) ‖≤ L 0 . Using the shorthand notations J = J ( W , Xtrain ) and J̃ = J ( W , Tj ( Xtrain ) ) , we obtain J̃ = J +E , where E is the perturbation matrix with ‖E‖2≤ ‖E‖F≤ √ nL 0 . 3 DATA AUGMENTATION SPEEDS UP LEARNING . In this section , we analyze the effect of data augmentation on training dynamics of neural networks , and show that data augmentation can provably speed up learning . To do so , we leverage the recent results that characterize the training dynamics based on properties of neural network Jacobian matrix and the corresponding Neural Tangent Kernel ( NTK ) ( Jacot et al. , 2018 ) defined as Θ = J ( W , X ) J ( W , X ) T . Formally ( Arora et al. , 2019 ) : rt = n∑ i=1 ( 1− ηλi ) ( uiuTi ) rt−1 = n∑ i=1 ( 1− ηλi ) t ( uiuTi ) r0 , ( 5 ) where Θ = UΛUT = ∑ i=1 λiuiu T i is the eigendecomposition of the NTK . Although the constant NTK assumption holds only in the infinite width limit , Lee et al . ( 2019 ) found close empirical agreement between the NTK dynamics and the true dynamics for wide but practical networks , such as wide ResNet architectures ( Zagoruyko & Komodakis , 2016 ) . Eq . ( 5 ) shows that training dynamics depend on the alignment of the NTK with the residual vector at every iteration t. In particular , shrinkage of residuals along the directions associated with larger eigenvalues of the NTK is fast and happens early during the training , while learning along the space associated with the small eigenvalues is slow and happens later . In the following , we prove that for small perturbations 0 , data augmentation speeds up training by enlarging smaller eigenvalues of the NTK , while decreasing larger eigenvalues with a high probability . Intuitively , this can be characterized as decreasing the learning rate for dimensions with larger gradient and slightly increasing the learning in dimensions with smaller gradients , and having a regularization effect by slightly perturbing the eigenvectors .

Summary: Deep learning uses augmentations to improve generalization performance. Using all augmentations for a dataset may slow down training. A subset selection technique is proposed (e.g., "coreset") using insights from the Neural Tangent Kernel (NMT) framework such that an alignment between the NMT Jacobian and the residuals is preserved. This alignment score is used to select the coreset using submodular optimization, which allows the model to be trained on a subset of augmented data (e.g., 0.1% to 30%) while preserving most augmentation benefits. The speedup of the method is reported to be up to 6.3x.