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: The paper studies adversarial training as a bi-level optimization problem. The authors show that Fast Adversarial Training can be viewed as a bi-level optimization problem where the lower-level problem is a linearization of the loss in the direction of the sign of the gradient. Motivated by this observation, the authors propose Fast Bi-level Adversarial Training, where the lower-level optimization linearizes the loss in the direction of the gradient (as opposed to the sign of the gradient). The authors then show analytically how to obtain the gradient of the upper-level problem under the assumption that the loss Hessian is equal to zero. This gradient is then used in an iterative manner similar to standard adversarial training. The proposed method is evaluated on CIFAR10 and ImageNet against multiple baselines including Fast-AT, Fast-AT-GA, and PGD-2. The empirical results show that compared to the baselines, within the same order of computational cost, the proposed method enjoys improved stability and mitigates the catastrophic overfitting present in other baselines.