Revisiting and Advancing Fast Adversarial Training Through the lens of Bi-Level Optimization

1 INTRODUCTION . Given the fact that machine learning ( ML ) models can be easily fooled by tiny adversarial perturbations ( also known as adversarial attacks ) on the input ( Goodfellow et al. , 2014 ; Carlini & Wagner , 2017 ; Papernot et al. , 2016 ) , learning robust deep neural networks ( DNNs ) is now a major focus in research . Nearly all existing effective defense mechanisms ( Madry et al. , 2018 ; Zhang et al. , 2019b ; Shafahi et al. , 2019 ; Wong et al. , 2020 ; Zhang et al. , 2019a ; Athalye et al. , 2018a ) are built on the adversarial training ( AT ) recipe , first developed in ( Szegedy et al. , 2014 ) and later formalized in ( Madry et al. , 2018 ) using min-max optimization . In contrast to standard model training using empirical risk minimization , AT ( Madry et al. , 2018 ) calls min-max optimization . That is , a minimizer ( i.e . defender ) seeks to update model parameters against a maximizer ( i.e . attacker ) that aims to worsen the training loss by perturbing each training example . The AT-type defenses have been widely adopted in various application domains including image classification ( Goodfellow et al. , 2014 ; Madry et al. , 2018 ; Kurakin et al. , 2017 ) , object detection ( Zhang & Wang , 2019 ) , natural language processing ( Miyato et al. , 2016 ; Zhu et al. , 2019 ) , and healthcare ( Finlayson et al. , 2019 ; Mahmood et al. , 2019 ) . Despite their effectiveness , the min-max optimization nature makes them difficult to scale . This is because multiple maximization steps ( required by an iterative attack generator ) are needed at every model training step in AT . The resulting prohibitive computation cost prevents AT from a feasible solution to enhance adversarial robustness when computing resource is limited . For example , Xie et al . ( 2019 ) used 128 GPUs to make AT practical on ImageNet . Thereby , how to speed up AT without losing accuracy and robustness is now a grand challenge for adversarial defense . Very recently , some work attempted to develop computationally-efficient alternatives of AT , which we call ‘ fast ’ versions of AT ( Shafahi et al. , 2019 ; Zhang et al. , 2019a ; Wong et al. , 2020 ; Andriushchenko & Flammarion , 2020 ) . To the best of our knowledge , FAST-AT ( Wong et al. , 2020 ) and FAST-AT with gradient alignment ( GA ) regularization , termed FAST-AT-GA ( Andriushchenko & Flammarion , 2020 ) , are the two state-of-the-art ( SOTA ) ‘ fast ’ versions of AT , since they achieve a significant reduction in computation complexity and preserve accuracy and robustness to some extent . To be specific , FAST-AT ( Wong et al. , 2020 ) replaces an iterative attack generator used in AT with a heuristics-based single-step attack generation method . Thus , it merely takes computation cost comparable to standard model training . However , FAST-AT suffers two main issues : ( i ) lack of stability , i.e. , large variance in performance ( Li et al. , 2020 ) , and ( ii ) robustness catastrophic overfitting , i.e. , a large drop of robustness when training with strong adversaries ( Andriushchenko & Flammarion , 2020 ) . To alleviate these problems , Andriushchenko & Flammarion ( 2020 ) proposed FAST-AT-GA by penalizing FAST-AT using an explicit robust regularization given by GA . However , we will show that FAST-AT-GA encounters a new problem ( iii ) : FAST-AT-GA hampers standard accuracy , making a poor accuracy-robustness tradeoff at large attack budget ( = 16/255 ) , i.e . the improvement on RA is at cost of a sharp drop on SA . Given the limitations ( i ) - ( iii ) , we ask : How to design a theoretically-grounded ‘ fast ’ version of AT with improved stability , mitigated catastrophic overfitting , and enhanced accuracy-robustness tradeoff ? To address above question , paper we revisit and advance AT through the lens of bi-level optimization ( BLO ) ( Dempe , 2002 ) , where we cast the attack generation problem as a lower-level optimization problem with constraints and the defense as an upper-level optimization problem in the objective . To the best of our knowledge , this is the first work to make a solid connection between adversarial defense and BLO . Technically , we show that FAST-AT can be interpreted as BLO with linearized lower-level problems . Delving into linearization of BLO , we propose a novel , theoretically-grounded ‘ fast ’ AT framework , fast bi-level AT ( FAST-BAT ) . Practically , Table 1 highlights some achieved improvements over FAST-AT and FAST-AT-GA : When a stronger train-time attack ( i.e. , = 16/255 vs. 8/255 ) is adopted , FAST-AT suffers a large degradation of robust accuracy ( RA ) and standard accuracy ( SA ) , together with higher variances than proposed FAST-BAT . Although FAST-AT-GA outperforms FAST-AT , it still incurs a significant SA loss ( over 21 % ) at = 16/255 . By contrast , FAST-BAT yields a more graceful SA-RA tradeoff : 9 % improvement of SA without loss of RA . Different from FAST-AT-GA , FAST-BAT achieves above improvements in stability , RA and SA without resorting to any extra robust regularization and thus takes less computation cost . Contributions . We summarize our contributions below . ¬ We propose a new formulation of adversarially robust training through the lens of BLO , yielding a novel and theoretically-grounded interpretation of FAST-AT . We propose a new systematic and effective fast BLO-oriented AT framework , termed FAST-BAT , with rigorously-established theory and algorithm . ® We conduct extensive experiments on FAST-BAT , showing its improved stability , mitigated catastrophic overfitting , and enhanced accuracy-robustness tradeoff ; see illustrations in Table 1 . 2 RELATED WORK . Adversarial attack . Adversarial attacks are techniques to generate malicious perturbations that are imperceptible to humans but can mislead the machine learning ( ML ) models ( Goodfellow et al. , 2014 ; Carlini & Wagner , 2017 ; Croce & Hein , 2020 ; Xu et al. , 2019 ; Athalye et al. , 2018b ) . A popular threat model that an adversary used is known as ` p-norm ball constrained attack ( p ∈ { 0 , 1 , 2 , ∞ } ) . This is also the focus of this paper . The adversarial attack has become a major approach to evaluate the robustness of deep neural networks ( DNNs ) and thus , help build safe artificial intelligence in many high stakes applications such as autonomous driving ( Deng et al. , 2020 ; Kumar et al. , 2020 ) , surveillance ( Thys et al. , 2019 ; Xu et al. , 2020 ) , and healthcare ( Finlayson et al. , 2019 ) . Adversarial defense and robust training at scale . Our work falls into the category of robust training , which was mostly built upon min-max optimization . For example , Madry et al . ( 2018 ) established the framework of AT for the first time , which has been recognized as one of the most powerful defenses ( Athalye et al. , 2018a ) . Extended from AT , TRADES ( Zhang et al. , 2019b ) sought the optimal balance between robustness and generalization ability . Further , AT-type defense has been generalized to the semi-/self-supervised settings ( Carmon et al. , 2019 ; Chen et al. , 2020 ) and integrated 1 with certified defense techniques such as randomized smoothing ( Salman et al. , 2019 ) . Despite the effectiveness of AT and its variants , they need to take high computation costs . How to speed up AT without losing performance remains an open question . Some recent works attempted to impose algorithmic simplifications to AT , leading to fast but approximate AT algorithms , such as ‘ free ’ AT ( Shafahi et al. , 2019 ) , you only propagate once ( YOPO ) ( Zhang et al. , 2019a ) , FASTAT ( Wong et al. , 2020 ) , and FAST-AT regularized by gradient alignment ( termed FAST-AT-GA ) ( Andriushchenko & Flammarion , 2020 ) . In particular , FAST-AT and FAST-AT-GA are the baselines most relevant to ours since they were designed with the least computation complexity . However , their defense performance is far from satisfactory . For example , FAST-AT has poor training stability ( Li et al. , 2020 ) and suffers catastrophic overfitting when facing strong attacks ( Andriushchenko & Flammarion , 2020 ) . In contrast to FAST-AT , FAST-AT-GA yields improved robustness but has a poor accuracy-robustness tradeoff ( e.g. , Table 1 ) . In this paper , we aim to advance the algorithm foundation of ‘ fast robust training ’ through the lens of BLO ( bi-level optimization ) . We will show that the proposed FAST-BAT can lead to stable robust learning without suffering catastrophic overfitting and graceful tradeoff between accuracy and robustness . Bi-level optimization ( BLO ) . BLO is a unified hierarchical learning framework , where the objective and variables of an upper-level problem depend on the optimizer of certain lower-level problems . The BLO problem in its most generic form is a class of very challenging problems , and thus , the design of algorithms and theory for BLO focuses on special cases ( Vicente et al. , 1994 ; White & Anandalingam , 1993 ; Gould et al. , 2016 ; Ghadimi & Wang , 2018 ; Ji et al. , 2020 ; Hong et al. , 2020 ) . In practice , some successful applications of BLO to ML have been witnessed in meta-learning ( Rajeswaran et al. , 2019 ) , data poisoning attack design ( Huang et al. , 2020 ) , and reinforcement learning ( Chen et al. , 2019 ) . However , as will be evident later , the existing BLO approach is not directly applied to adversarial defense due to the presence of the constrained nonconvex lower-level problem ( for attack generation ) . To the best of our knowledge , our work makes a rigorous connection between adversarial defense and BLO for the first time . 3 A BI-LEVEL OPTIMIZATION VIEW ON FAST-AT . Preliminaries on FAST-AT . FAST-AT is designed for solving the adversarial training problem ( Madry et al. , 2018 ) given below minimize θ E ( x , y ) ∈D [ maximize δ∈C ` tr ( θ , x+ δ , y ) ] , ( 1 ) where θ ∈ Rn denotes model parameters , D is the training set consisting of labeled data pairs with feature x and label y , δ ∈ Rd represents adversarial perturbations subject to the perturbation constraint C , e.g. , C = { δ | ‖δ‖∞ ≤ , δ ∈ [ 0,1 ] } for -toleration ` ∞-norm constrained attack ( normalized to [ 0,1 ] ) , ( x+ δ ) is then called adversarial example , and ` tr ( · ) represents a training loss . The standard solver to problem ( 1 ) is known as AT ( Madry et al. , 2018 ) . However , it has to call an iterative optimization method ( e.g. , K-step PGD attack ) to solve the inner maximization problem of ( 1 ) . As a result , AT is computationally intensive . To improve its scalability , FAST-AT that only takes the single-step PGD attack for inner maximization was proposed and successfully implemented in ( Wong et al. , 2020 ) . The algorithm backbone of FAST-AT is summarized below .

Summary: This work focuses on the problem of speeding up adversarial training in the $\ell_\inf$ threat model. The work first describes two previous works designed to speed up adversarial training: - Fast-AT: carefully perform single-step PGD (otherwise known as FGSM) to train $\ell_\inf$ robust models - from Wong, Rice and Kolter 2020 - Fast-AT-GA: single-step PGD + a gradient alignment loss function - from Andriushchenko and Flammarion 2020.  The work then lists a problem with Fast-AT: - Fast-AT experiences catastrophic overfitting (i.e. provides no robustness on held-out data) for $\epsilon = 16/255$ (Note: Fast-AT-GA does not have this catastrophic overfitting issue)  Finally, the authors present a new method based on bi-level optimization to perform fast $\ell_\inf$-robust learning. They evaluate it and compare with Fast-AT-GA using a 20 epoch schedule, finding that for $\epsilon = 16/255$, the technique can provide the same robustness except with standard accuracy ~68% instead of ~59%.