Title: AMATA: AN ANNEALING MECHANISM FOR ADVER-SARIAL TRAINING ACCELERATION

Abstract: Despite of the empirical success in various domains, it has been revealed that deep neural networks are vulnerable to maliciously perturbed input data that much degrade their performance. This is known as adversarial attacks. To counter adversarial attacks, adversarial training formulated as a form of robust optimization has been demonstrated to be effective. However, conducting adversarial training brings much computational overhead compared with standard training. In order to reduce the computational cost, we propose a simple yet effective modification to the commonly used projected gradient descent (PGD) adversarial training by increasing the number of adversarial training steps and decreasing the adversarial training step size gradually as training proceeds. We analyze the optimality of this annealing mechanism through the lens of optimal control theory, and we also prove the convergence of our proposed algorithm. Numerical experiments on standard datasets, such as MNIST and CIFAR10, show that our method can achieve similar or even better robustness with around 1/3 to 1/2 computation time compared with PGD.

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; Tramè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 = {(x 0 i , y i )} n i=1 with x 0 i ∈ R d as a normal or clean example in the d-dimensional input space and y i ∈ R C as its associated one-hot label, the objective of adversarial training is to solve the following min-max optimization problem:

EQUATION

where h θ : R d → R C 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 x i , within the -ball around a given normal example x 0 i 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 {x i } n i=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 x 0 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 x 0 to prevent the updated value fall outside the -ball of x 0 (Madry et al., 2017) :

x k = x k-1 + α • sign(∇ x (h θ (x k-1 ), y) (2)

where α is the step size, (•) is the orthogonal projection function onto {x : x 0 -x ≤ }, and x k 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; K min : the minimum number of adversarial perturbations; K max : 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:

K t = K min + (K max -K min ) • t T

Compute adversarial perturbation step size:

α t = τ K t for each mini-batch x 0 B do for k = 1 to K t do

Compute adversarial perturbations:

x k B = x k-1 B + α t • sign(∇ x (h θ (x k B ), y) x k B = clip(x k B , x 0 B -, x 0 B + ) end for θ t+1 = θ t -η∇ θ (h θt (x K 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 (x 0 , y), since the nsample case is similar. The original min-max adversarial training problem is

min θ max {z: z-x 0 ≤ } (h θ (z), y).

(3)

The first relaxation is to replace the outer minimization with gradient descent so that we obtain the iteration

EQUATION

Then, the remaining maximization in each outer iteration step is replaced by an abstract algorithm

A u,θ : R d → R d (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 

EQUATION

Observe that Algorithm 1 is a particular realization of ( 6) where A θt,ut represents the inner PGD loop and u t = {α t , K t } 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η:

EQUATION

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 u s are the controls.

Our goal is two-fold: on a training interval [T 1 , T 2 ] 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:

EQUATION

where we have defined the shorthand u T1:

T2 = {u s : s ∈ [T 1 , T 2 ]} and L ∞ T1:T2 := L ∞ ([T 1 , T 2 ], G).

In this paper, we take Φ to be the DNN's prediction loss given the adversarial example (adversarial robustness), and R is set as γK t where γ is the coefficient for adversarial robustness and training time trade-off. This is to account for the fact that when K t 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 {u s }. The objective function taken as a sum serves to balance the adversarial robustness and computational cost, with γ a balancing coefficient.

2.3. PONTRYAGIN'S MAXIMUM PRINCIPLE:

In the last section, the problem of choosing hyper-parameters in the inner loops of adversarial training has been formulated as an optimal control problem in (8). Now, we show how this connection can help us design and validate algorithms.

A classical result in calculus of variations gives the following necessary conditions for optimality. For more background on the theory of calculus of variations and optimal control, we refer the reader to Boltyanskii et al. (1960) ; Bertsekas et al. (1995) .

Theorem 1 (Pontryagin's Maximum Principle (PMP) (Boltyanskii et al., 1960) ). Let u * T1:T2 ∈ L ∞ T1,T2 be a solution to (8). Suppose F (θ, u) is Lipschitz in θ and measurable in u. Define the Hamiltonian function

H(θ, p, u) = p T F (θ, u) -R(u) (9)

Then, there exists an absolutely continuous co-state process p * T1:T2 such that θ *

EQUATION

In short, the maximum principle says that a set of optimal parameter choices {u * s } (in our specific application, these are the optimal choices of inner-loop hyper-parameters as the training proceeds) must globally maximize the Hamiltonian defined above for each outer iteration. This statement is especially appealing for our application because unlike first-order gradient conditions, the PMP holds even when our hyper-parameters can only take a discrete set of values, or when there are nontrivial constraints amongst them. The maximization criterion holds generally under these conditions. Moreover, we now show that it gives us a quantitative measure of deviation from optimality, from which we can analyze and design algorithms.

Quantitative Measure of Sub-optimality. Given any hyper-parameter choice u T1:T2 over the training interval, let us define its "distance" from optimality as

EQUATION

where {θ u s , p u s : s ∈ [T 1 , T 2 ]} represents the solution of the equations ( 10) and ( 11) with u s in place of u * s . Observe that C(u T1:T2 ) ≥ 0 for any u T1:T2 with equality if and only if u T1:T2 satisfies the PMP for almost every s ∈ [T 1 , T 2 ]. Hence, C can be used as a measure of sub-optimality. When C is small, our annealing strategy {u s } is close to at least an locally optimal strategy, where as when it is large, our annealing strategy is far from an optimal one. One-step Approximation and Adversarial Training Criterion. Equation ( 13) requires information on the entire training interval and may be expensive to compute. In this paper, we use a one-step approximation where we take T 1 = t (current iteration) and T 2 = t + η with η 1. In this small interval, we can also take u s to be constant and thus equal to some u. This is in some sense a greedy approximation, where we evaluate in the immediate short term the optimality of a piece-wise constant choice of hyper-parameters. From (13) we then obtain via a Taylor expansion and our particular choices of Φ and R (See Appendix D)

EQUATION

with A θt,α,K denoting the inner PGD loop starting form x with K steps and step size α. Criterion ( 14) is a greedy version of the general criterion derived from the maximum principle. It can be used to either evaluate the near-term sub-optimality of some choice of hyper-parameters u, or to find an approximately optimal hyperparameter greedily by solving C(u, t) = 0 for u, which amounts to maximizing the first term. In this paper, we use Bayesian optimization2 to perform the maximization in ( 14) to evaluate and select strategies from the controllable space G.

Remark. By applying the method of successive approximations (MSA) (Chernousko & Lyubushin, 1982; Li et al., 2017; Li & Hao, 2018; E et al., 2019) , we can also obtain a solution of the PMP by iterating

EQUATION

Subsequent work will explore adaptive adversarial training acceleration based on this approach. We remark that although we appeal to control theory tools here, the application scenario is very different from prior work in Li et al. (2017) ; Li & Hao (2018); E et al. (2019) . For the latter, the controlled dynamics are idealized forward propagation through a deep neural network and the controls are the trainable weights of the network. In contrast, in the current setting the controlled dynamics are the evolution of weights under gradient descent in the outer loop and controls are the hyper-parameters for the adversarial perturbation inner loop.

Comparison with FOSC criterion: Wang et al. (2019) proposed an empirical criterion to measure the convergence of inner maximization:

EQUATION

There are some similarities between our criterion and FOSC when we do not consider the computational cost term R. For example, when the stationary saddle point is achieved, both our criterion and FOSC reach the minimum. However, our proposed criterion is quite different from FOSC in the following aspects:

1. Our criterion is derived from the optimal control theory, whereas FOSC is concluded from empirical observations. 2. Our criterion takes computation costs into consideration, whereas FOSC only considers the convergence of adversarial training. 3. Our criterion is based on the gradient of DNN parameters, whereas FOSC is based on the gradient of the input. Measuring the gradient of DNN parameters is arguably more suitable for considering robustness-training time trade-off as the variance of the DNN parameters is much larger than the input during training.

2.4. NUMERICAL JUSTIFICATION OF AMATA:

For simple linear-quadratic control problems, we can typically derive some analytical representations of an optimal control. However, for highly-nonconvex DNNs, we can only use a numerical form of the optimal control criterion (Equation 14) to analyze Amata for robustness and computational efficiency trade-off. We use the LeNet neural network3 for MNIST classification as an example.

As a sanity check, we first consider only adversarial robustness and neglect computational cost by setting γ = 0 and use the criterion to evaluate different adversarial training settings for PGD.

Recalling Algorithm 1, we will vary the other hyper-parameters in the inner loop: α (step-size of PGD) and K (number of PGD steps). In each case, we compute the optimal control criterion values on the training dataset and evaluate the adversarially trained networks' performance against PGD-40 attack on the test dataset.

Sanity check 1 In this case, we fix K = 5 and set α in [0.01, 0.02, 0.03, 0.04, 0.05].

Sanity check 2 In this case, we fix α = 0.01 and set K in [5, 10, 20, 30, 40, 50] .

The results are shown in Figure 1 (a)(b), where robustness is the prediction accuracy of the trained networks under PGD-40 attack on the test dataset.

From these results, we can observe that the robustness and the criterion value C are negatively correlated under varying α and K, as expected. Now, we will use the same setting as sanity checks but with different values of γ to numerically analyze the trade-off between robustness and training time for Amata and PGD.

Setting 1: fixed α, varying K and γ: in this case, we set α as 0.01, K in [5, 20, 40] and γ in [0.02, 0.04, 0.06, 0.08]. The result for Setting 1 is shown in Figure 1 (c). From this result, we can observe that with the increase of γ, reducing the computation time becomes more important. For example, when γ is 0.02, using PGD-40 is the best strategy (the lowest C value). When γ increases to 0.06 or 0.08, using PGD-20 achieves the best trade-off under the criterion. Next, we set γ as 0.04 and 0.08 and use the criterion to evaluate the Amata and compare it with other settings of PGD. It can be seen that Amata is better than PGD with fixed numbers of perturbation steps. Note that Amata also achieves better robustness than the highest robustness of PGD settings (PGD-40) with 91.47% adversarial accuracy compared to 90.56% adversarial accuracy.

Remark: Although computing the exact optimal control strategy for DNN adversarial training is inapplicable for real-time applications, with the criterion derived from the PMP, we are able to numerically compare the optimality of different adversarial training strategies. From this numerical evaluation, we have demonstrated that the proposed Amata algorithm is close to an optimal adversarial training strategy, or at least one that satisfies the maximum principle. We will show that our algorithm can achieve similar or even better adversarial accuracy much faster with empirical experiments on popular DNN models later in Section 3.

3. EXPERIMENTS:

To demonstrate the effectiveness of Amata, we conduct experiments on MNIST and CIFAR10. We find that the models trained with Amata have comparable or even better performance with that of the PGD adversarial training, but with much less computational cost. For our experiment, we use PyTorch 1.0.0 and a GTX1080 Ti GPU. We evaluate the adversarial trained networks against PGD and Carlini-Wagner (CW) attack (Carlini & Wagner, 2017) .

3.1. MNIST CLASSIFICATION:

We consider the standard MNIST classification using the "smallCNN" architecture that was used in (Zhang et al., 2019b) . We set τ to be 0.4 for Amata. In the experiment, We achieve 96% adversarial accuracy within 275 seconds while it takes PGD-40 1019 seconds to reach the same level. We also experimented with the FOSC adversarial training method in (Wang et al., 2019) and found that it took FOSC 3663 seconds to achieve the same level, which confirms our previous analysis that our optimal control criterion is more suitable for accelerating adversarial training. The error-time curve is shown in Figure 2 (a). Then, we run all methods for 57 epochs for full convergence. The robustness of convergence results are shown in Table 2 . From Table 2 , we can see that naively reducing the number of adversarial perturbations will hurt the robustness of DNNs. 

3.2. CIFAR10 CLASSIFICATION:

For the more complex CIFAR10 classification task, we use the PreAct-Res-18 network (Madry et al., 2017) . We set τ to be 20/255 for Amata. The proposed Amata method took 3045 seconds to achieve less than 55% robust error while for PGD-10, it took 6944 seconds. We also tested FOSC adversarial training and found that it took 8385 seconds to reach the same level, which is a bit longer than PGD-10. The error-time curve is shown in Figure 2 (b). We further run all methods for 100 epochs for convergence. The results are shown in Table 3 . 

4. CONCLUSION:

We have proposed a modification amounting to two lines of code to PGD adversarial training that achieves comparable or even better robustness with only 1/3 to 1/2 the computational cost over a variety of benchmarks. Moreover, a general optimal control formulation of adversarial training with hyper-parameters is developed to analyze this procedure and justify its superior performance through a numerical criterion based on the Pontryagin's maximum principle. A proof of convergence of our algorithm is also provided in Appendix C. As a point of future work, we will explore adaptive methods for adversarial training based on the optimal control formulation we introduced here. This approach can also lead to algorithms that can be combined with YOPO or parallel adversarial training methods for maximal efficiency. An example is shown in Appendix B.3.

B APPENDIX:ADDITIONAL EXPERIMENT RESULTS:



B.1 WIDERESNET-34 RESULTS:

We experiments using WideResNet-34 networks for CIFAR10 classification. We use the same experiment setting as the CIFAR10 classification task in the paper. We run all methods for 100 epochs for full convergence. The result is shown in Table 4 . 

B.3 AMATA+YOPO:

To demonstrate that our method can be incorporated in other acceleration method to further improve the performance, we implement the Amata in YOPO MNIST classification experiment4 . For Amata incorporation, we gradually increase the K and decrease σ in the codes that is similar to the case of modifying the PGD algorithm. We run Amata+YOPO and YOPO for 40 epochs as in the original setting. The results are shown in Figure 4 . From Figure 4 , we can see that Amata+YOPO takes 294 seconds to reach the adversarial accuracy of 94%, which is only around half of the time for YOPO to achieve the same adversarial accuracy. It is also worth noting that Amata+YOPO achieves even better adversarial accuracy when converged. From this example, we can see that Amata can be seamlessly incorporated in other adversarial training acceleration algorithm to provide further acceleration.

K 0 = K min and K T = K max to make the exponential decay scheme comparable with the linear scheme used in our paper. However, we tried different parameters of eta ranging from 0.1 to 5, but did not performance improvements. The results are shown in Figure 6 . The "Amata (Exp eta)" in the legend denotes the exponential decay scheme and the "Amata" in the legend denotes the original linear decay scheme. From Figure 6 , we can see that the the new exponential decay scheme cannot outperform the linear decay scheme. 

C APPENDIX:PROOF OF CONVERGENCE:

We provide a convergence analysis of our proposed Amata algorithm for solving the min-max problem. The proof of convergence in this paper largely follows (Wang et al., 2019; Sinha et al., 2018) . First, we will introduce some notations first for clarity. We denote x * i (θ) = arg max xi∈X i f (θ, x i ) where f (θ, x i ) is a short hand notation for the classification loss function, X i = {x| xx 0 i ≤ }, and fi (θ) = max xi∈X i f (θ, x i ), then xi (θ) is a δ-approximate solution by our algorithm to x * i (θ), if it satisfies that:

EQUATION

In addition, denote the objective function in Equation 1 by L S (θ) , and its gradient by ∇L S (θ) =

1 n n i=1 ∇ fi (θ). Let g(θ) = 1 |B| i∈B ∇ θ f (θ) be the stochastic gradient of L S (θ)

, where B is the mini-batch. Then, we have E[g(θ)] = ∇L S (θ). Let ∇ θ f (θ, x(θ)) be the gradient of f (θ, x(θ)) with respect to θ, and g(θ) = 1 |B| i∈B ∇ θ f (θ, xi (θ) be the approximate stochastic gradient of L S (θ). Before we prove the convergence of the algorithm, we have following assumptions. Assumption 1. The function f (θ, x) satisfies the gradient Lipschitz conditions:

sup x ∇ θ f (θ, x) -∇ θ f (θ * , x) 2 ≤ L θθ θ -θ * 2 sup θ ∇ θ f (θ, x) -∇ θ f (θ, x * ) 2 ≤ L θx x -x * 2 sup x ∇ x f (θ, x) -∇ x f (θ * , x * ) 2 ≤ L xθ θ -θ * 2

where L θθ , L θx , and L xθ are positive constants. Assumption 1 was used in (Sinha et al., 2018; Wang et al., 2019) . Assumption 2. The function f (θ, x) is locally µ-strongly concave in X = {x : xx 0 i ∞ ≤ } for all i ∈ [n], i.e., for any x 1 , x 2 ∈ X i :

f (θ, x 1 ) ≤ f (θ, x 2 ) + ∇ x f (θ, x 2 ), x 1 -x 2 - µ 2 x 1 -x 2 2 2

where µ is a positive constant which measures the curvature of the loss function. This assumption was used for analyzing distributional robust optimization problems (Sinha et al., 2018) .

Assumption 3. The variance of the stochastic gradient g(θ) is bounded by a constant σ 2 > 0:

E[ g(θ) -∇L S (θ) 2 2 ] ≤ σ 2

where ∇L S (θ) is the full gradient.

The Assumption 3 is commonly used for analyzing stochastic gradient optimization algorithms.

Theorem 2. Suppose Assumptions 1,2, and 3 holds. Denote ∆ = L S (θ 0 ) -min θ L S (θ). If the step size of outer minimization is η t = min(1/L, ∆ T Lσ 2 ). Then, we have:

1 T T -1 t=0 E[ ∇L S (θ t ) 2 2 ] ≤ 4σ L∆ T + 5L 2 θx δ 2 where L = L θx Lxθ/µ + L θθ .

The proof of this theorem can be found in the Appendix. Theorem 2 indicates that if the Amata finds solutions of the inner maximization problem closer enough to the maxima, Amata can converge at a sublinear rate.

Lemma 1. Under Assumption 1 and Assumption 2, we have L S (θ) is L-smooth where L = L θx L xθ /µ + L θθ , i.e., for any θ 1 and θ 2 it holds:

L S (θ 1 ) ≤ L S (θ 2 ) + ∇L S (θ 2 ), θ 1 -θ 2 + L 2 θ 1 -θ 2 2 2

The proof of Lemma 1 can be found in Wang et al. (2019) .

Lemma 2. Under Assumption 1 and Assumption 2, the norm of difference between the approximate stochastic gradient g(θ) and the stochastic gradient g(θ) is bounded, i.e., it holds that: g(θ) -g(θ) 2 ≤ L θx δ

Proof. We have

g(θ) -g(θ) 2 = 1 |B| i∈B ∇ θ f (θ, xi (θ)) -∇ fi (θ) 2 ≤ 1 |B| i∈B ∇ θ f (θ, xi (θ)) -∇ θ f (θ, x * i (θ)) 2 ≤ 1 |B| i∈B L θx xi (θ) -x * i (θ) 2 (19)

where the first inequality is from the triangle inequality, and the second inequality is from Assumption 1. Next, we insert Equation 18 into the above inequality then:

EQUATION

which completes the proof. Now we can prove the Theorem 2:

Proof. From Lemma 1, we have: Taking the telescope sum of the above equation from t = 0 to t = T -1, we have

L S (

T -1 t=0 η t 2 (1 -L η t 2 )E[ ∇L S (θ t ) 2 2 ] ≤ E[L S (θ 0 -θ T )] + T -1 t=0 η t 2 (1 + 3Lη t 2 )L 2 θx δ 2 + Lη 2 t σ 2

We set η t = min(1/L, ∆ T Lσ 2 ), we have

1 T T -1 t=0 E[ ∇L S (θ t ) 2 2 ] ≤ 4σ L∆ T + 5L 2 θx δ 2

Thus, we complete the proof.

D APPENDIX:APPROXIMATION OF THE OPTIMAL CONTROL CRITERION:

In this section we show how to derive the approximation ( 14) from ( 13). The primary assumption we make is that T 2 -T 1 = η 1, which allows one to use a local approximation for various functions to derive a simple-to-compute criterion. For this reason, the derivation here will be largely heuristic.

First, we assume that we consider a constant control u s ≡ u on the interval s ∈ [T 1 , T 2 ] ≡ [t, t + η]. Next, recall that the adversarial robustness which serves as our terminal loss function for the control problem is Φ(θ) = max z (h θ [A θ,z (x)], y). In practice, any control that sufficiently conducts the

:

Here we assume that the gradient with respect to θ is the partial derivative with respect to the parameters of the network h θ and θt in A θ t ,u t is held constant. This is the case for the PGD algorithm. Alternatively, we can also take the total derivative, but this leads to different algorithms.

Implementations can be found in https://github.com/hyperopt/hyperopt

Implementations can be found in https://github.com/pytorch/examples/blob/master/mnist/main.py

https://github.com/a1600012888/YOPO-You-Only-Propagate-Once/tree/82c5b902508224c642c8d0173e61435795c0ac42/experiments/M

[5][6][7][8][9][10]

Figures and Tables:

Figure 1: Sanity checks of the optimal control criterion (γ = 0).

Figure 2: Left: MNIST result. Training time against PGD-40 attack. We use Amata with the setting Kmin = 10 and Kmax = 40. Right: CIFAR10 result. Training time against PGD-20 attack. We use Amata with the setting Kmin = 2 and Kmax = 10.

Figure 3: MNIST result. Training epoch against PGD-40 attack. We use Amata with the setting Kmin = 10 and Kmax = 40.

Figure 6: MNIST result. Training epoch against PGD-40 attack. We use Amata with the setting Kmin = 2 and Kmax = 5.

Comparison of adversarial training strategies.

Table content: <table><tr><td>Strategy</td><td colspan="4">C(γ = 0.04) C(γ = 0.08) Robustness Training time (seconds)</td></tr><tr><td>Amata(K min = 5, K max = 40)</td><td>0.54</td><td>1.38</td><td>91.47%</td><td>697.73</td></tr><tr><td>Amata(K min = 10, K max = 40)</td><td>0.68</td><td>1.53</td><td>91.46%</td><td>760.16</td></tr><tr><td>PGD-10</td><td>7.82</td><td>7.70</td><td>68.07%</td><td>307.57</td></tr><tr><td>PGD-20</td><td>1.52</td><td>2.09</td><td>85.23%</td><td>567.11</td></tr><tr><td>PGD-40</td><td>1.20</td><td>2.19</td><td>90.56%</td><td>1086.31</td></tr></table>

MNIST adversarial training results.

Table content: <table><tr><td>Training methods</td><td colspan="4">Clean accuracy PGD-40 Attack CW Attack Time (Seconds)</td></tr><tr><td>ERM</td><td>99.52%</td><td>5.58%</td><td>0.10%</td><td>268.10</td></tr><tr><td>PGD-5</td><td>99.58%</td><td>80.99%</td><td>0.31%</td><td>709.21</td></tr><tr><td>PGD-10</td><td>99.57%</td><td>93.8%</td><td>1.06%</td><td>1234.30</td></tr><tr><td>PGD-40</td><td>99.51%</td><td>97.02%</td><td>93.49%</td><td>4331.81</td></tr><tr><td>FOSC(Wang et al., 2019)</td><td>99.51%</td><td>97.02%</td><td>81.79%</td><td>14928.53</td></tr><tr><td>Amata(Kmin = 5, Kmax = 40)</td><td>99.53%</td><td>97.11%</td><td>86.92%</td><td>2533.28</td></tr><tr><td>Amata(Kmin = 10, Kmax = 40)</td><td>99.45%</td><td>96.97%</td><td>94.37%</td><td>2675.64</td></tr></table>

CIFAR10 PreAct-Res-18 adversarial training results.

Table content: <table><tr><td>Training methods</td><td colspan="4">Clean accuracy PGD-20 Attack CW Attack Time (Seconds)</td></tr><tr><td>ERM</td><td>94.75%</td><td>0.0%</td><td>0.23%</td><td>2099.58</td></tr><tr><td>PGD-2</td><td>90.16%</td><td>31.70%</td><td>13.36%</td><td>6913.36</td></tr><tr><td>PGD-10</td><td>85.27%</td><td>47.31%</td><td>51.73%</td><td>23108.10</td></tr><tr><td>FOSC(Wang et al., 2019)</td><td>85.47%</td><td>48.04%</td><td>53.65%</td><td>26126.98</td></tr><tr><td>Amata(Kmin = 2, Kmax = 10)</td><td>85.52%</td><td>47.62%</td><td>52.94%</td><td>14308.96</td></tr></table>

CIFAR10 WideResNet-34 adversarial training results.We further plot the MNIST results in Figure2but with the x-axis denoting the number of training epochs. The results are shown in Figure3. From this figure, we can see that the time-saving of our method is mainly due to the savings in the adversarial examples generation step. This is expected as if we achieve faster convergence rates with regard to the number of epochs, it may indicate that the adversarial examples generated by the proposed mechanism is significantly weaker and the outer minimization step converges faster.

Table content: <table><tr><td>Training methods</td><td colspan="4">Clean accuracy PGD-20 Attack CW Attack Time (Seconds)</td></tr><tr><td>PGD-10</td><td>85.80%</td><td>45.38%</td><td>30.82%</td><td>150653.40</td></tr><tr><td>Amata(K min = 2, K max = 10)</td><td>85.77%</td><td>30.42%</td><td>52.94%</td><td>90945.15</td></tr><tr><td colspan="2">B.2 ADDITIONAL MNIST RESULTS</td><td/><td/><td/></tr></table>

θ t+1 ) ≤ L S (θ t ) + ∇L S (θ t , θ t+1 -θ t + L 2 θ t+1 -θ t 2 2 = L S (θ t ) + η t ∇L S (θ t ), ∇L S (θ t -g(θ t ) -η t ∇L S (θ t ) ∇L S (θ t ), ∇L S (θ t ) -g(θ t ) + ∇L S (θ t ), ∇L S (θ t ) -g(θ t ) + η t (1 -Lη t 2 ) ∇L S (θ t ),∇L S (θ t ) -g(θ t ) + ∇L S (θ), taking expectation on both sides of the inequality conditioned on θ t . Then we use Assumption 3 and Lemma 2 and simplify the above inequality: E[L S (θ t+1 ) -L S (θ t )|θ t ] ≤ -

Table content: <table><tr><td/><td/><td/><td/><td/><td/><td/><td/><td>2</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td>2</td></tr><tr><td>+</td><td colspan="2">Lη 2 t 2</td><td colspan="3">g(θ 2 )</td><td>2 2</td><td/></tr><tr><td colspan="8">= L S (θ t ) -η t (1 -</td><td>Lη t 2</td><td>) ∇L S (θ t )</td><td>2 2 + η t (1 -</td><td>Lη t 2</td><td>)•</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td>Lη 2 t 2</td><td>g(θ t ) -∇L S (θ t )</td><td>2 2</td></tr><tr><td colspan="8">= L S (θ t ) -η t (1 -</td><td>Lη t 2</td><td>) ∇L S (θ t )</td><td>2 2 + η t (1 -</td><td>Lη t 2</td><td>)•</td></tr><tr><td/><td/><td/><td/><td/><td/><td/><td/><td>Lη 2 t 2</td><td>g(θ 2 2</td></tr><tr><td colspan="4">≤ L S (θ t ) -</td><td colspan="2">η t 2</td><td colspan="2">(1 -</td><td>Lη t 2</td><td>) ∇L S (θ t )</td><td>2 2 + η t (1 -</td><td>Lη t 2</td><td>)•</td></tr><tr><td colspan="5">g(θ) -g(θ t )</td><td/><td colspan="3">2 2 + Lη 2 t ( g(θ η t 2</td><td>(1 -</td><td>Lη t 2</td><td>) ∇L S (θ t )</td><td>2 2</td></tr><tr><td/><td>+</td><td>η t 2</td><td colspan="2">(1 +</td><td colspan="2">3Lη t 2</td><td colspan="2">)L 2 θx δ 2 + Lη 2 t σ 2</td></tr></table>

Peer review (iclr2020):
The submission follows a recent line of work that formulates adversarial training (https://arxiv.org/abs/1412.6572) as a differentiable game and attempts to reduce the computational complexity of solving the inner maximization, corresponding to finding (an) adversarial example(s) for a given parameter setting. The authors propose a way to anneal the truncation of the inner iteration, and investigate a quantity to represent the suboptimality of the truncation. Strengths: - (From Table 3), the method is competitive with adversarial training (PGD; Madry et al. (2017)) as well as an improved method for adversarial robustness (FOSC; Wang et al. (2019)), while requiring half the computation time. - The writing is very clear. Weaknesses: - Significance: The technique AMATA itself, is a heuristic to gradually increase the length of the inner loop optimization in adversarial training. It is unrelated to the optimal control formulation. - Novelty: Much of the motivation from the perspective of an optimal control formulation of adversarial training is similar to https://arxiv.org/abs/1803.01299, including the PMP and successive approximation components; Eq. (13) of the submission is the continuous version of Theorem 2 in this paper. - Some experimental details are missing. - Minor clarity issues (see below). Questions for the authors: - Could the authors better contrast their work in reference to prior work (https://arxiv.org/abs/1803.01299)? - Could the authors clarify if they use random restarts for the PGD attacks? How was the number of iterations for the PGD attack (40) selected? Why was the iteration count of the PGD attack not varied? Clarity issues: - Please provide references for the optimal control formulation in 2.2. - Some references to variables in the text are not made precise by identifying the variable in question (e.g.: *a set of optimal parameter choices*; *requires information on the entire training interval*). - The function of the hyperparameters alpha, gamma, K are not well discussed in Section 2.4. - Reduce plot sizes and include more experimental details in the main text.