Reweighting Augmented Samples by Minimizing the Maximal Expected Loss

1 INTRODUCTION . Deep neural networks have achieved state-of-the-art results in various tasks in natural language processing ( NLP ) tasks ( Sutskever et al. , 2014 ; Vaswani et al. , 2017 ; Devlin et al. , 2019 ) and computer vision ( CV ) tasks ( He et al. , 2016 ; Goodfellow et al. , 2016 ) . One approach to improve the generalization performance of deep neural networks is data augmentation ( Xie et al. , 2019 ; Jiao et al. , 2019 ; Cheng et al. , 2019 ; 2020 ) . However , there are some problems if we directly incorporate these augmented samples into the training set . Minimizing the average loss on all these samples means treating them equally , without considering their different implicit impacts on the loss . To address this , we propose to minimize a reweighted loss on these augmented samples to make the model utilize them in a cleverer way . Example reweighting has previously been explored extensively in curriculum learning ( Bengio et al. , 2009 ; Jiang et al. , 2014 ) , boosting algorithms ( Freund & Schapire , 1999 ) , focal loss ( Lin et al. , 2017 ) and importance sampling ( Csiba & Richtárik , 2018 ) . However , none of them focus on the reweighting of augmented samples instead of the original training samples . A recent work ( Jiang et al. , 2020a ) also assigns different weights on augmented samples . But weights in their model are predicted by a mentor network while we obtain the weights from the closed-form solution by minimizing the maximal expected loss ( MMEL ) . In addition , they focus on image samples with noisy labels , while our method can generally be applied to also textual data as well as image data . Tran et al . ( 2017 ) propose to minimize the loss on the augmented samples under the framework of Expectation-Maximization algorithm . But they mainly focus on the generation of augmented samples . ∗This work is done when Mingyang Yi is an intern at Huawei Noah ’ s Ark Lab . Unfortunately , in practise there is no way to directly access the optimal reweighting strategy . Thus , inspired by adversarial training ( Madry et al. , 2018 ) , we propose to minimize the maximal expected loss ( MMEL ) on augmented samples from the same training example . Since the maximal expected loss is the supremum over any possible reweighting strategy on augmented samples ’ losses , minimizing this supremum makes the model perform well under any reweighting strategy . More importantly , we derive a closed-form solution of the weights , where augmented samples with larger training losses have larger weights . Intuitively , MMEL allows the model to keep focusing on augmented samples that are harder to train . The procedure of our method is summarized as follows . We first generate the augmented samples with commonly-used data augmentation technique , e.g. , lexical substitution for textual input ( Jiao et al. , 2019 ) , random crop and horizontal flip for image data ( Krizhevsky et al. , 2012 ) . Then we explicitly derive the closed-form solution of the weights on each of the augmented samples . After that , we update the model parameters with respect to the reweighted loss . The proposed method can generally be applied above any data augmentation methods in various domains like natural language processing and computer vision . Empirical results on both natural language understanding tasks and image classification tasks show that the proposed reweighting strategy consistently outperforms the counterpart of without using it , as well as other reweighting strategies like uniform reweighting . 2 RELATED WORK . Data augmentation . Data augmentation is proven to be an effective technique to improve the generalization ability of various tasks , e.g. , natural language processing ( Xie et al. , 2019 ; Zhu et al. , 2020 ; Jiao et al. , 2019 ) , computer vision ( Krizhevsky et al. , 2014 ) , and speech recognition ( Park et al. , 2019 ) . For image data , baseline augmentation methods like random crop , flip , scaling , and color augmentation ( Krizhevsky et al. , 2012 ) have been widely used . Other heuristic data augmentation techniques like Cutout ( DeVries & Taylor , 2017 ) which masks image patches and Mixup ( Zhang et al. , 2018 ) which combines pairs of examples and their labels , are later proposed . Automatically searching for augmentation policies ( Cubuk et al. , 2018 ; Lim et al. , 2019 ) have recently proposed to improve the performance further . For textual data , Zhang et al . ( 2015 ) ; Wei & Zou ( 2019 ) and Wang ( 2015 ) respectively use lexical substitution based on the embedding space . Jiao et al . ( 2019 ) ; Cheng et al . ( 2019 ) ; Kumar et al . ( 2020 ) generate augmented samples with a pre-trained language model . Some other techniques like back translation ( Xie et al. , 2019 ) , random noise injection ( Xie et al. , 2017 ) and data mixup ( Guo et al. , 2019 ; Cheng et al. , 2020 ) are also proven to be useful . Adversarial training . Adversarial learning is used to enhance the robustness of model ( Madry et al. , 2018 ) , which dynamically constructs the augmented adversarial samples by projected gradient descent across training . Although adversarial training hurts the generalization of model on the task of image classification ( Raghunathan et al. , 2019 ) , it is shown that adversarial training can be used as data augmentation to help generalization in neural machine translation ( Cheng et al. , 2019 ; 2020 ) and natural language understanding ( Zhu et al. , 2020 ; Jiang et al. , 2020b ) . Our proposed method differs from adversarial training in that we adversarially decide the weight on each augmented sample , while traditional adversarial training adversarially generates augmented input samples . In ( Behpour et al. , 2019 ) , adversarial learning is used as data augmentation in object detection . The adversarial samples ( i.e. , bounding boxes that are maximally different from the ground truth ) are reweighted to form the underlying annotation distribution . However , besides the difference in the model and task , their training objective and the resultant solution are also different from ours . Sample reweighting . Minimizing a reweighted loss on training samples has been widely explored in literature . Curriculum learning ( Bengio et al. , 2009 ; Jiang et al. , 2014 ) feeds first easier and then harder data into the model to accelerate training . Zhao & Zhang ( 2014 ) ; Needell et al . ( 2014 ) ; Csiba & Richtárik ( 2018 ) ; Katharopoulos & Fleuret ( 2018 ) use importance sampling to reduce the variance of stochastic gradients to achieve faster convergence rate . Boosting algorithms ( Freund & Schapire , 1999 ) choose harder examples to train subsequent classifiers . Similarly , hard example mining ( Malisiewicz et al. , 2011 ) downsamples the majority class and exploits the most difficult examples . Focal loss ( Lin et al. , 2017 ; Goyal & He , 2018 ) focuses on harder examples by reshaping the standard cross-entropy loss in object detection . Ren et al . ( 2018 ) ; Jiang et al . ( 2018 ) ; Shu et al . ( 2019 ) use meta-learning method to reweight examples to handle the noisy label problem . Unlike all these existing methods , in this work , we reweight the augmented samples ’ losses instead of training samples . 3 MINIMIZE THE MAXIMAL EXPECTED LOSS . In this section , we derive our reweighting strategy on augmented samples from the perspective of maximal expected loss . We first give a derivation of the closed-form solution of the weights on augmented samples . Then we describe two kinds of loss under this formulation . Finally , we give the implementation details using the natural language understanding task as an example . 3.1 WHY MAXIMAL EXPECTED LOSS . Consider a classification task with N training samples . For the i-th training sample xi , its label is denoted as yxi . Let fθ ( · ) be the model with parameter θ which outputs the classification probabilities . ` ( · , · ) denotes the loss function , e.g . the cross-entropy loss between outputs fθ ( xi ) and the groundtruth label yxi . Given an original training sample xi , the set of augmented samples generated by some method isB ( xi ) . Without loss of generality , we assume xi ∈ B ( xi ) . The conventional training objective is to minimize the loss on every augmented sample z in B ( xi ) as min θ 1 N N∑ i=1 1 |B ( xi ) | ∑ ( z , yz ) ∈B ( xi ) ` ( fθ ( z ) , yz ) , ( 1 ) where yz is the label of z ∈ B ( xi ) , and can be different with yxi . |B ( xi ) | is the number of augmented samples in B ( xi ) , which is assumed to be finite . In equation ( 1 ) , for each given xi , the weights on its augmented samples are the same ( i.e. , 1/|B ( xi ) | ) . However , different samples have different implicit impacts on the loss , and we can assign different weights on them to facilitate training . Note that computing the weighted sum of losses of each augmented sample in B ( xi ) can be viewed as taking expectation of loss on augmented samples z ∈ B ( xi ) under a certain distribution . When the augmented samples generated from the same training sample are drawn from a uniform distribution , the loss in equation ( 1 ) can be rewritten as min θ Rθ ( PU ) = min θ 1 N N∑ i=1 [ Ez∼PU ( ·|xi ) [ ` ( fθ ( z ) , yz ) ] − λPKL ( PU ( · | xi ) ‖ PU ( · | xi ) ) ] , ( 2 ) where the Kullback–Leibler ( KL ) divergence KL ( PU ( · | xi ) ‖ PU ( · | xi ) ) equals zero . Here PU ( · | xi ) denotes the uniform distribution on B ( xi ) . When the augmented samples are drawn from a more general distribution PB ( · | · ) 1 instead of the uniform distribution , we can generalize PU ( · | · ) here to some other conditional distribution PB . min θ Rθ ( PB ) = min θ 1 N N∑ i=1 [ Ez∼PB ( ·|xi ) [ ` ( fθ ( z ) , yz ) ] − λPKL ( PB ( · | xi ) ‖ PU ( · | xi ) ) ] . ( 3 ) Remark 1 . When PB ( · | xi ) reduces to the uniform distribution PU ( · | xi ) for any xi , since KL ( PU ( · | xi ) ‖ PU ( · | xi ) ) = 0 , the objective in equation ( 3 ) reduces to the one in equation ( 1 ) . The KL divergence term in equation ( 3 ) is used as a regularizer to encourage PB close to PU ( see Remark 2 ) . From equation ( 3 ) , the conditional distribution PB determines the weights of each augmented sample in B ( xi ) . There may exist an optimal formulation of PB in some regime , e.g . corresponding to the optimal generalization ability of model . Unfortunately , we can not explicitly characterize such an unknown optimal PB . To address this , we borrow the idea from adversarial training ( Madry et al. , 2018 ) and minimize the maximal reweighted loss on augmented samples . Then , the model is guaranteed to perform well under any reweighting strategy , including the underlying optimal one . Specifically , let the conditional distribution PB be P∗θ = arg supPB Rθ ( PB ) . Our objective is to minimize the following reweighted loss min θ Rθ ( P∗θ ) = min θ sup PB Rθ ( PB ) . ( 4 ) The following Remark 2 discusses about the KL divergence term in equation ( 3 ) . 1In the following , we simplify PB ( · | · ) as PB if there is no obfuscation . Remark 2 . Since we take a supremum over PB in equation ( 4 ) , the regularizer KL ( PB ‖ PU ) encourages PB to be close to PU because it reaches the minimal value zero when PB = PU . Thus the regularizer controls the diversity among the augmented samples by constraining the discrepancy between PB and uniform distribution PU , e.g. , a larger λP promotes a larger diversity among the augmented samples . The following Theorem 1 gives the explicit formulation of Rθ ( P∗θ ) . Theorem 1 . Let Rθ ( PB ) and Rθ ( P∗θ ) be defined in equation ( 1 ) and ( 4 ) , then we have Rθ ( P∗θ ) = 1 N N∑ i=1 ∑ z∈B ( xi ) P∗θ ( z | xi ) ` ( fθ ( z ) , yz ) − λPP∗θ ( z | xi ) log ( |B ( xi ) |P∗θ ( z | xi ) ) , ( 5 ) where P∗θ ( z | xi ) = exp ( 1 λP ` ( fθ ( z ) , yz ) ) ∑ z∈B ( xi ) exp ( 1 λP ` ( fθ ( z ) , yz ) ) = Softmaxz ( 1 λP ` ( fθ ( B ( xi ) ) , yB ( xi ) ) ) , ( 6 ) where Softmaxz ( 1λP ` ( fθ ( B ( xi ) ) , yB ( xi ) ) ) represents the output probability of z for vector ( 1λP ` ( fθ ( z1 ) , yz1 ) , · · · , 1 λP ` ( fθ ( z|B ( xi ) | ) , y|B ( xi ) | ) ) . Remark 3 . If we ignore the KL divergence term in equation ( 3 ) , due to the equivalence of minimizing cross-entropy loss and MLE loss ( Martens , 2019 ) , the proposed MMEL also falls into the generalized Expectation-Maximization ( GEM ) framework ( Dempster et al. , 1977 ) . Specifically , given a training example , the augmented samples of it can be viewed as latent variable , and any reweighting on these augmented samples corresponds to a specific conditional distribution of these augmented samples given the training sample . In the expectation step ( E-step ) , we explicitly derive the closed-form solution of the weights on each of these augmented samples according to ( 6 ) . In the maximization step , since there is no analytical solution for deep neural networks , following ( Tran et al. , 2017 ) , we update the model parameters with respect to the reweighted loss by one step of gradient descent . The proof of this theorem can be found in Appendix A . From Theorem 1 , the loss of it decides the weight on each augmented sample z ∈ Bxi , and the weight is normalized by Softmax over all augmented samples in Bxi . The reweighting strategy allows more attention paid to augmented samples with higher loss values . The strategy is similar to those in ( Lin et al. , 2017 ; Zhao & Zhang , 2014 ) but they apply it on training samples .

Summary: This paper proposes a simple scheme for training with multiple augmentations of training data in one iteration and reweighting the instances by their relative loss. As authors note in their related works, the idea of reweighting examples based on their relative loss has been widely studied in a variety of machine learning problems. In contrast, this work proposes the reweighting only within augmentations of a single sample. They derive their particular reweighting scheme by proposing an alternative risk (Eq. 3). The new objective is a function of both model parameters and the distribution of augmentations. They propose to find the model parameters that minimize the alternative risk for the hardest distribution of augmentations that maximizes their alternative risk (Eq. 4). Then they consider the distribution of augmentations that are a function of model parameters and input and show that for fixed model parameters, the optimal distribution on a fixed finite set of augmentations is determined by the softmax on the loss of the model for each augmented input. In section 3.2, they propose two variations of their loss using the ground-truth label to evaluate the loss (hard loss) versus using the prediction of the model for the original raw input (soft loss). In section 3.3, they propose specific considerations for augmenting text data. They provide experiments on image and text data with ablations studies.