BulletTrain: Accelerating Robust Neural Network Training via Boundary Example Mining

1 Introduction . In the past decade , the performance and capabilities of deep neural networks ( DNNs ) have improved at an unprecedented pace . However , the reliability of DNNs is undermined by their lack of robustness against common and adversarial corruptions , raising concerns about their use in safety-critical applications such as autonomous driving ( Amodei et al. , 2016 ) . Towards improving the robustness of DNNs , a line of recent works proposed robust training by generating and learning from corrupted examples in addition to clean examples ( Cubuk et al. , 2019 ; Ding et al. , 2020 ; Hendrycks et al. , 2020 ; Madry et al. , 2018 ; Wang et al. , 2020 ; Zhang et al. , 2019b ) . While these methods have demonstrably enhanced model robustness , they are also very costly to deploy since the overall training time can be increased by up to ten times . Hard negative example mining techniques have shown success in improving the convergence rates of support vector machine ( SVM ) ( Joachims , 1999 ) , DNNs ( Shrivastava et al. , 2016 ) , metric learning ( Manmatha et al. , 2017 ) , and contrastive learning ( Kalantidis et al. , 2020 ) . Particularly , SVM with hard example mining maintains a working set of “ hard ” examples and alternates between training an SVM to convergence on the working set and updating the working set , where the “ hard ” examples are significant in that they violate the margins of the current model . Partly inspired by hard example mining , we hypothesize that there exists a small subset of examples that is critical to improving the robustness of the DNN model , so that reducing the computation of the remaining examples leads to a significant reduction in run time without compromising the robustness of the model . As depicted in Figure 1b , the essential idea is to predict and focus on the “ important ” subset of clean samples based on their geometric distance to the decision boundary ( i.e. , margin ) of the current model . For example , the correctly classified clean examples with a large margin are unlikely to help improving robustness as the corrupted examples generated from them are still likely to be appropriately categorized by the model . We refer to “ important ” examples as boundary examples since they are close to the decision boundary ( e.g. , in Figure 1b ) . After identifying the boundary 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . 1st iteration 2nd iteration 3rd iteration ( a ) Existing robust neural network training algorithms — Standard robust training approaches generate corrupted samples with equal effort for all samples ( ) in a mini-batch . ( b ) BulletTrain — The proposed scheme focuses on a small subset of “ important ” clean samples ( ) and allocates more computational cost on obtaining corrupted samples for those important ones . This figure illustrates the case that only generates the corrupted samples for the important samples . Figure 1 : Comparison between existing robust training approaches and BulletTrain on a synthetic binary classification task — The dash line is the current decision boundary of the linear model and the solid line represents the updated decision boundary after taking a SGD step . examples , the computational effort on generating corrupted samples can be reduced for the rest of the clean examples , thus accelerating the overall robust training process . In this work , we propose BulletTrain1 , a dynamic optimization that exploits boundary example mining to reduce the computational effort of robust DNN training . BulletTrain applies to robust training against either adversarial or common corruptions . It can substantially speed up state-of-the-art algorithms ( Hendrycks et al. , 2020 ; Wang et al. , 2020 ; Zhang et al. , 2019b ) with almost no loss in clean and/or robust accuracy . Figure 1 visualizes the situation where BulletTrain generates corrupted examples only for a small fraction of clean examples , while a standard robust training scheme requires obtaining corrupted examples for all examples . To the best of our knowledge , BulletTrain is the first work to exploit the idea of hard example mining in the context of accelerating robust DNN training . Unlike existing approaches which are limited to reducing the compute for all examples uniformly , BulletTrain distributes the compute for clean examples differently , focusing on a small set of boundary examples . BulletTrain is also generally applicable to robust training against common and adversarial corruptions . The end result is that BulletTrain can lower the running time of the adversarial training such as TRADES ( Zhang et al. , 2019b ) and MART ( Wang et al. , 2020 ) by more than 2.1× on CIFAR-10 ( Krizhevsky , 2009 ) and the robust training against common corruptions such as AugMix ( Hendrycks et al. , 2020 ) by 1.7× on CIFAR-10-C ( Hendrycks and Dietterich , 2019 ) without compromising the clean and robust accuracy . 2 Robust Training Overview . Preliminaries . We focus on K-class ( K ≥ 2 ) classification problems . Given a dataset { ( xi , yi ) } i=1 , ... , n with a clean example xi ∈ Rd and its associated ground truth label yi ∈ { 1 , ... , K } , a DNN classifier hθ parameterized by θ generates the logits output hθ ( xi ) for each class , where hθ ( xi ) = ( h 1 θ ( xi ) , ... , h K θ ( xi ) ) . The predicted label ŷi can be expressed as ŷi = arg maxk h k θ ( xi ) . To improve the robustness of a DNN classifier , state-of-the-art approaches ( Hendrycks and Dietterich , 2019 ; Madry et al. , 2018 ; Wang et al. , 2020 ; Zhang et al. , 2019b ) generate corrupted inputs x′i for the corresponding clean example xi at training time and minimize the loss of the classifier with respect to the corrupted input . For robustness against adversarial attacks , x′i can be obtained using 1The name of BulletTrain implies robust ( bullet-proof ) and high-speed training . Algorithm 1 : Standard robust DNN training algorithm . Input : batch size m , a N -step generation function GN for generating corrupted samples 1 while θ not converged do 2 Read mini-batch { x1 , ... , xm ; y1 , ... , ym } from training set 3 for i = 1 , ... , m do 4 x′i ← GN ( xi ) // GN ( · ) is an N -step procedure for generating perturbations 5 end 6 θ ← θ −∇θ ( L ( { hθ ( x1 ) , ... , hθ ( xm ) } , { hθ ( x′1 ) , ... , hθ ( x′m ) } ) ) // L is the surrogate loss 7 end adversarial attack methods such as projected gradient descent ( PGD ) ( Madry et al. , 2018 ) . For robustness to common corruptions and perturbations , x′i is generated by introducing different data augmentations ( Hendrycks and Dietterich , 2019 ) . Algorithm 1 shows pseudo-code for a standard robust training algorithm . Computational cost of robust training . Robust DNN training algorithms such as Algorithm 1 improve model robustness at the cost of paying extra computational overhead . There are two main sources of the overhead : 1 ) The procedure GN for generating corrupted input ; 2 ) Forward and backward propagate using the surrogate loss L with respect to both clean and corrupted samples . We denote the extra computational cost for robust neural network training as C [ GN , L ] . For adversarial robustness , PGD adversarial training generates the adversarially corrupted input by solving the inner maximization maxδ∈∆ L ( x + δ , y ; θ ) approximately using an N -step projected gradient descent method , where ∆ is the perceptibility threshold and is usually defined using an lp distance metric : ( ∆ = { δ : ‖δ‖p < } for some ≥ 0 ) . As a result , PGD adversarial training is roughly N times more expensive than ordinary DNN training , where N is usually between 7 and 20 . In addition , Hendrycks et al . ( 2020 ) propose to boost the robustness against common corruptions using Jensen-Shannon divergence ( JSD ) as the surrogate loss . JSD loss can be computed by first generating two corrupted examples ( x′ and x′′ ) , calculating a linear combination of the logits as M = ( hθ ( x ) + hθ ( x ′ ) + hθ ( x ′′ ) ) /3 , and then computing L ( x , x′ , x′′ ) = 13 ( KL [ hθ ( x ) ‖ M ] + KL [ hθ ( x′ ) ‖M ] + KL [ hθ ( x′′ ) ‖M ] ) , where KL stands for the Kullback–Leibler divergence . As a result , the surrogate loss for improving the robustness against common corruptions triples the cost of the ordinary training as it requires forward and backward propagation for both x , x′ , and x′′ . Existing efficient training algorithms . A number of recent efforts ( Andriushchenko and Flammarion , 2020 ; Shafahi et al. , 2019 ; Wong et al. , 2020 ) aimed at reducing the complexity of adversarial example generation , which constitutes the main source of computational overhead in adversarial training . For instance , Wong et al . ( 2020 ) proposed a variant of the single-step fast gradient sign method ( FGSM ) in place of multi-step methods , which drastically reduced the computational cost . However , it was later shown that the single-step approaches fail catastrophically against stronger adversaries ( Andriushchenko and Flammarion , 2020 ) . FGSM-based methods can also cause substantial degradation in robust accuracy ( Chen et al. , 2020 ) for recently proposed adversarial training methods such as TRADES ( Zhang et al. , 2019b ) . This result suggests that the adversarial examples generated using single-step FGSM are not challenging enough for learning a robust classifier . In addition , existing efficient training approaches can not be generalized to robust training against common corruptions . The state-of-the-art robust training algorithms against common corruptions ( Cubuk et al. , 2019 ; Hendrycks et al. , 2020 ) use data augmentations to create corrupted samples in a single pass , where the number of iterations can not be further reduced . 3 BulletTrain . 3.1 Motivation . As discussed in Section 2 , robust training can lead to a significant slowdown compared to ordinary DNN training . To reduce the overhead of robust training , our key insight is that different inputs contribute unequally to improving the model ’ s robustness ( which we later show ) , hence we can dynamically allocate different amounts of computation ( i.e. , different N ) for generating corrupted samples depending on their “ importance ” . In detail , we first categorize the set of all training samples into a disjoint union of three sets — outlier , robust , and boundary examples — as follows : • Outlier examples xO : the clean input x is misclassified by the DNN model hθ . xO = { x | arg max k=1 , ... , K hkθ ( x ) 6= y } ( 1 ) • Boundary examples xB : the clean input x is correctly classified while the corresponding corrupted input x′ is misclassified by the DNN model hθ . xB = { x | arg max k=1 , ... K hkθ ( x ) = y ∧ arg max k=1 , ... , K hkθ ( x ′ ) 6= y } ( 2 ) • Robust examples xR : both the clean input x and the corresponding corrupted input x′ = GN ( x ) are correctly classified by the DNN model hθ . xR = { x | arg max k=1 , ... K hkθ ( x ) = arg max k=1 , ... , K hkθ ( x ′ ) = y } ( 3 ) Note that our definition of the outlier , boundary , and robust examples is dependent on the current model hθ and can dynamically shift during training . That is , an outlier example may become a boundary example later on when the robustness of the model improves . Similarly , a robust example can also become a boundary or even outlier example when robustness improves in other regions of the training distribution . Robust training algorithms improve the robustness of a DNN model by increasing the classification margin for each training sample x so that its distance to the corresponding corrupted sample ( i.e. , ‖x− x′‖p ) is smaller than the margin . We hypothesize that this per-sample margin is an indicator of the amount of contribution that provided by this sample during robust training . • Outlier examples xO have a negative margin because they are misclassified by the current DNN model . Therefore , it is unnecessary to generate even “ harder ” corrupted examples and optimize the model using these “ harder ” examples . • Boundary examples xB have a positive margin to the current decision boundary but can be perturbed to cause misclassification . Training on a corrupted version of these samples is most helpful to increase the margin , thus making the model more robust . • Robust examples xR already have a margin large enough to tolerate common or adversarial corruption . As a result , training on corrupted versions of these examples would provide less benefit compared to training on corrupted boundary examples . To test our hypothesis , we conduct an empirical study of PGD adversarial training against l∞ attack ( = 0.3 ) on MNIST ( LeCun et al. , 2010 ) in a leave-one-out manner . The baseline PGD training achieves 90.6 % robust accuracy by setting the number of steps for all samples to be seven . In our experiment , we fix the number of PGD iterations ( N ) for one of xB , xO , and xR to be zero while varying N between zero and seven for the other two classes of data . The separation between xB , xO , and xR for each mini-batch of data is oracle , which is determined using ten PGD steps . As depicted in Figure 2a , PGD training fails to achieve a robust accuracy over 40 % when the number of steps for xB ( NB ) is zero ; thus demonstrating the importance of xB . When NO equals zero ( Fig . 2b ) , the highest robust accuracy is 1 % higher than that of the baseline , indicating that setting NO to zero does not impair the effectiveness of PGD training . Lastly , Figure 2c shows that PGD training can achieve a reasonable robust accuracy without paying the extra computational cost on xR . However , the highest robust accuracy achieved with this setup is 89.2 % , which is 1.4 % lower than the baseline . Therefore , setting a non-zero NR might help improving the robustness of the model in practice . This empirical observation reveals an opportunity to accelerate robust DNN training by focusing on the boundary examples . In doing so , the overhead of robust DNN training is then determined by the proportion of boundary examples among all training samples . As shown in Figure 3 , the fraction of xR increases while the fraction of xO decreases over epochs as the model becomes more robust . Moreover , the fraction of xB is relatively small compared to the other two classes of samples since boundary examples only present within the vicinity of the decision boundary . Therefore , we believe that the computational cost of robust training ( C [ GN , L ] ) can be reduced substantially by attending more to the boundary examples .

Summary: This paper introduces BulletTrain - an algorithm for speeding up robust training algorithms for both adversarial and corruption based robustness. This is an important problem to study because robust training (adversarial training in particular) is generally expensive which makes it hard to scale these algorithms on large datasets. The algorithm proposed in this paper, while being lightweight and easy to implement, shows strong empirical gains. Overall, I like this paper.