Amata: An Annealing Mechanism for Adversarial Training Acceleration

1 INTRODUCTION . Recently , the revival of deep neural networks has led to breakthroughs in various fields , including computer vision , natural language processing , game playing , etc . Despite these advancements , deep neural networks were found to be vulnerable to malicious perturbations on the original input data . While the perturbations remain almost imperceptible to humans , they can lead to wrong predictions over the perturbed examples ( Szegedy et al. , 2013 ; Goodfellow et al. , 2014 ; Akhtar & Mian , 2018 ) . These maliciously crafted examples are known as adversarial examples , which have caused serious concerns over the reliability and security of deep learning systems , particularly when deployed in the life-critical scenarios , such as autonomous driving systems and health/medical domains . Several defense mechanisms have been proposed , such as input reconstruction ( Meng & Chen , 2017 ; Song et al. , 2018 ) , input encoding ( Buckman et al. , 2018 ) , and adversarial training ( Goodfellow et al. , 2014 ; Tramèr et al. , 2017 ; He et al. , 2017 ; Madry et al. , 2017 ) . Among these methods , adversarial training is the most effective defense method so far . Adversarial training can be posed as a robust optimization problem ( Ben-Tal & Nemirovski , 1998 ) , where a min-max optimization problem is solved ( Madry et al. , 2017 ; Kolter & Wong , 2017 ) . For example , given a C-class dataset S = { ( x0i , yi ) } ni=1 with x0i ∈ Rd as a normal or clean example in the d-dimensional input space and yi ∈ RC as its associated one-hot label , the objective of adversarial training is to solve the following min-max optimization problem : min θ 1 N N∑ i=1 max ‖zi−x0i‖≤ ` ( hθ ( zi ) , yi ) ( 1 ) where hθ : Rd → RC is the deep neural network ( DNN ) function , ` is the loss function and is the maximum perturbation constraint . The inner maximization problem is to find an adversarial example xi , within the -ball around a given normal example x0i that maximizes the classification loss ` . On the other hand , the outer minimization problem is to find model parameters that minimizes the loss ` on the adversarial examples { xi } ni=1 that are generated from the inner maximization . The inner maximization problem is commonly solved by projected gradient descent ( PGD ) . PGD perturbs a normal example x0 by iteratively updating it in the steepest ascent direction for a total of K times . Each ascent step is modulated by a small step size and a projection step back onto the -ball of x0 to prevent the updated value fall outside the -ball of x0 ( Madry et al. , 2017 ) : xk = ∏ ( xk−1 + α · sign ( ∇x ` ( hθ ( xk−1 ) , y ) ) ( 2 ) where α is the step size , ∏ ( · ) is the orthogonal projection function onto { x′ : ‖x0 − x′‖ ≤ } , and xk is the adversarial example at k-th step . However , a major problem prohibiting adversarial training to be practically applicable is the huge computational burden associated with the inner maximization steps : we need to iteratively solve the inner maximization problem to find good adversarial examples for DNN to be robust . Recently , to accelerate adversarial training , a few methods have been proposed . For example , YOPO estimated the gradient on the input by only propagating the first layer ( Zhang et al. , 2019a ) , and parallel adversarial training utilized multiple graphics computation units ( GPUs ) for acceleration ( Bhat & Tsipras , 2019 ; Shafahi et al. , 2019 ) . A common drawback of these methods are their implementation complexity or their need for multiple GPUs for acceleration . On the other hand , an empirical observation made by ( Wang et al. , 2019 ) , indicated that we might not need to find good solutions to the inner maximization at the initial stages of adversarial training to achieve even better robustness . In line of such observations , in this paper we propose a simple yet effective Annealing Mechanism for Adversarial Training Acceleration , which we call Amata , that gradually controls the degree of which the inner maximization is solved as the training proceeds . This annealing algorithm only takes 1/3 to 1/2 the time to achieve comparable or even better robustness with the addition of only two lines of code . To motivate the present approach , we develop a general formulation of adversarial training as an optimal control problem , from which an approximate optimality criterion can be derived based on the Pontryagin ’ s maximum principle . This criterion is related to but different from the empirical FOSC criterion ( Wang et al. , 2019 ) that improves the robustness of adversarial training by determining the number of adversarial training steps according to the empirical criterion . We will also demonstrate in our experiments that unlike FOSC , algorithms derived from our proposed criterion leads to acceleration . Our contributions are as follows : 1 . We propose an adversarial training algorithm , Amata , that gradually increases the number of iterations and decreases the step size for the inner maximization . We show that comparable or even better robustness can be achieved with much reduced computational overhead . We also prove the convergence of our algorithm . 2 . To obtain more theoretical insight , we develop a general formulation of adversarial training subject to hyper-parameters in the inner maximization loop as an optimal control problem . This includes the PGD-based methods considered in this work , but potentially encompasses a larger class of annealed or adaptive adversarial training algorithms . 3 . Using the Pontryagin ’ s maximum principle from optimal control theory , we provide a principled criterion to measure the near-term optimality of an annealing schedule for inner maximization , thereby guiding and validating our proposed algorithm in balancing the trade-off between robustness and computational cost . 2 ACCELERATING ADVERSARIAL TRAINING BY ANNEALING MECHANISM . In this section , we will first introduce the proposed adversarial training algorithm , which aims to balance the computational cost and the accuracy of solving inner maximization . A proof of the convergence of the algorithm can be found in the Appendix . We then develop a general optimal control formulation of adversarial training and derive a criterion based on the Pontryagin ’ s maximum principle ( Boltyanskii et al. , 1960 ) . The criterion is shown to be effective in quantifying adversarial training performance both on accuracy and computational efficiency , justifying the superiority of the proposed algorithm . 2.1 PROPOSED ANNEALING ADVERSARIAL TRAINING ALGORITHM . The proposed adversarial training algorithm Amata is found in Algorithm 1 . Compared with PGD , we only need to add two lines of code that are shown in blue . The intuition behind Amata is that , at the initial stage , the neural network focus on learning features , which might not require very accurate adversarial examples . Therefore , we only need coarse approximations of the inner maximization problem solutions . With this consideration , we set a small number of update steps but with a large step size for inner maximization , and then gradually increase K and decreases α to improve the quality of inner maximization solutions . This adaptive annealing mechanism would largely reduce the computational cost in the early iterations while still maintaining reasonable accuracy for the entire optimization . Algorithm 1 Amata : an annealing mechanism for adversarial training acceleration Input : T : training epochs ; Kmin : the minimum number of adversarial perturbations ; Kmax : the maximum number of adversarial perturbations ; θ : parameter of neural network to be adversarially trained ; B : mini-batch ; α : adversarial training time step ; η : learning rate of neural network parameters . τ : constant , maximum perturbation : . Initialization θ = θ0 for t = 0 to T − 1 do Compute the annealing number of adversarial perturbations : Kt = Kmin + ( Kmax −Kmin ) · t T Compute adversarial perturbation step size : αt = τKt for each mini-batch x0B do for k = 1 to Kt do Compute adversarial perturbations : xkB = x k−1 B + αt · sign ( ∇x ` ( hθ ( x k B ) , y ) xkB = clip ( x k B , x 0 B − , x0B + ) end for θt+1 = θt − η∇θ ` ( hθt ( xK t B ) , y ) end for end for Collect θT as the parameter of adversarially-trained neural network . In the following section , we develop a general formulation of adversarial training based on optimal control theory and derive a novel criterion to quantify the optimality of a training strategy , taking into account the trade-off between accuracy and efficiency . Importantly , we show that our proposal , Amata , performs favorably under this criterion , thus providing theoretical justification for our approach . 2.2 OPTIMAL CONTROL FORMULATION OF ADVERSARIAL TRAINING . In essence , the PGD-based adversarial training algorithm ( Madry et al. , 2017 ) represents a sequence of relaxations of the original min-max problem ( 1 ) . First , the outer minimization over θ is replaced by gradient descent . Then , a full solution of the inner maximization is replaced by a number of PGD steps in the steepest ascent direction . Consequently , a natural question is how to choose the number of steps , the step size or any other hyper-parameters associated with this relaxed version of the inner maximization , and how their choices affect the performance of the overall algorithm . This is the central point of analysis in this paper and the basis of the algorithm proposed . It turns out that this question can be systematically formulated in the framework of optimal control theory ( Bertsekas , 1995 ) . We now introduce the general setup of our problem and establish connections with optimal control , and in particular the classical maximum principle in the calculus of variations . For simplicity of presentation , let us consider just one fixed input-label pair ( x0 , y ) , since the nsample case is similar . The original min-max adversarial training problem is min θ max { z : ‖z−x0‖≤ } ` ( hθ ( z ) , y ) . ( 3 ) The first relaxation is to replace the outer minimization with gradient descent so that we obtain the iteration θt+1 = θt − η∇θ max { z : ‖z−x0‖≤ } ` ( hθt ( z ) , y ) . ( 4 ) Then , the remaining maximization in each outer iteration step is replaced by an abstract algorithm Au , θ : Rd → Rd ( 5 ) which solves the inner maximization approximately . Here , we assume that the algorithm depends on the parameters of our neural network , as well as hyper-parameters u which takes values in a closed subset G of a Euclidean space . No further assumptions is placed on G , which may be a continuum , a discrete set , or even a finite set . This relaxation leads to the following general iterations1 θt+1 = θt − η∇θ ` ( hθt ( Aθt , ut ) , y ) . ( 6 ) Observe that Algorithm 1 is a particular realization of ( 6 ) where Aθt , ut represents the inner PGD loop and ut = { αt , Kt } are the hyper-parameters we pick at each t. Nevertheless , Equation ( 6 ) represents a general formulation of an inner-outer loop adversarial training algorithm . For small η , we can replace ( 6 ) by an ordinary differential equation with the identification s ≈ tη : θ̇s = −∇θ ` ( hθs ( Aθs , us ) , y ) . ( 7 ) Formulated this way , it is immediately clear that we can approach this problem from a control perspective : the dynamics of the trainable parameters are given by ( 7 ) and the hyper-parameter choices at each inner algorithm us are the controls . Our goal is two-fold : on a training interval [ T1 , T2 ] in the outer loop , we want to minimize some loss under adversarial training measured by a real-valued function Φ ( θ ) while minimizing training cost associated with each inner algorithm loop under the hyper-parameter u , which is measured by another real-valued function R ( u ) . We thus arrive at the general optimal control formulation of our problem : min uT1 : T2∈L ∞ T1 : T2 Φ ( θT ) + ∫ T2 T1 R ( us ) ds subject to : θ̇t = F ( θs , us ) where F ( θs , us ) : = −∇θ ` ( hθs ( Aθs , us ) , y ) , ( 8 ) where we have defined the shorthand uT1 : T2 = { us : s ∈ [ T1 , T2 ] } and L∞T1 : T2 : = L ∞ ( [ T1 , T2 ] , G ) . In this paper , we take Φ to be the DNN ’ s prediction loss given the adversarial example ( adversarial robustness ) , and R is set as γKt where γ is the coefficient for adversarial robustness and training time trade-off . This is to account for the fact that when Kt increases , the cost of the inner loop training increases accordingly . The integral over t of R is taken so as to account for the total computational cost corresponding to a choice of hyper-parameters { us } . The objective function taken as a sum serves to balance the adversarial robustness and computational cost , with γ a balancing coefficient .

Summary: Deep learning algorithms are known to be prone to adversarial examples. This paper proposes a simple modification for adversarial training in order to improve the robustness of the algorithms. The adversarial training of the neural networks involve two steps: the outer loop where the loss is being minimized, the inner loop where the adversarial examples are being found by using PGD. The inner PGD optimization makes training more expensive, because for every gradient step, algorithm needs to perform several steps of PGD. This paper proposes a simple modification to the PGD that is being used in the inner loop. The proposed modification involves the increasing the number of adversarial training steps and decreasing the adversarial training step size gradually as the training proceeds. Then, they analyze the modifications from the theory of optimal control. Their experiments show that their method can achieve similar robustness to the other existing approaches with less computations.