You Only Propagate Once: Accelerating Adversarial Training via Maximal Principle

Introduction

Deep neural networks achieve state-of-the-art performance on many tasks 16, 8. However, recent works show that deep networks are often sensitive to adversarial perturbations 33, 25, 46, i.e., changing the input in a way imperceptible to humans while causing the neural network to output an incorrect prediction. This poses significant concerns when applying deep neural networks to safety-critical problems such as autonomous driving and medical domains. To effectively defend the adversarial attacks, 24 proposed adversarial training, which can be formulated as a robust optimization 36:

A major issue of the current adversarial training methods is their significantly high computational cost. In adversarial training, we need to solve the inner loop, which is to obtain the "optimal" adversarial attack to the input in every iteration. Such "optimal" adversary is usually obtained using multi-step gradient decent, and thus the total time for learning a model using standard adversarial training method is much more than that using the standard training. Considering applying 40 inner iterations of projected gradient descent (PGD 15) to obtain the adversarial examples, the computation cost of solving the problem (1) is about 40 times that of a regular training.

The main objective of this paper is to reduce the computational burden of adversarial training by limiting the number of forward and backward propagation without hurting the performance of the trained network. In this paper, we exploit the structures that the min-max objectives is encountered with deep neural networks. To achieve this, we formulate the adversarial training problem(1) as a differential game. Afterwards we can derive the Pontryagin’s Maximum Principle (PMP) of the problem.

From the PMP, we discover a key fact that the adversarial perturbation is only coupled with the weights of the first layer. This motivates us to propose a novel adversarial training strategy by decoupling the adversary update from the training of the network parameters. This effectively reduces the total number of full forward and backward propagation to only one for each group of adversary updates, significantly lowering the overall computation cost without hampering performance of the trained network. We name this new adversarial training algorithm as YOPO (You Only Propagate Once). Our numerical experiments show that YOPO achieves approximately 4 $\sim$5 times speedup over the original PGD adversarial training with comparable accuracy on MNIST/CIFAR10. Furthermore, we apply our algorithm to a recent proposed min max optimization objective "TRADES"43 and achieve better clean and robust accuracy within half of the time TRADES need.

To improve the robustness of neural networks to adversarial examples, many defense strategies and models have been proposed, such as adversarial training 24, orthogonal regularization 6, 21, Bayesian method 42, TRADES 43, rejecting adversarial examples 41, Jacobian regularization 14, 27, generative model based defense 12, 31, pixel defense 29, 23, ordinary differential equation (ODE) viewpoint 44, ensemble via an intriguing stochastic differential equation perspective 37, and feature denoising 40, 32, etc. Among all these approaches, adversarial training and its variants tend to be most effective since it largely avoids the the obfuscated gradient problem 2. Therefore, in this paper, we choose adversarial training to achieve model robustness.

Neural ODEs.

Recent works have built up the relationship between ordinary differential equations and neural networks 38, 22, 10, 5, 45, 35, 30. Observing that each residual block of ResNet can be written as $u_{n+1}=u_{n}+\Delta tf(u_{n})$, one step of forward Euler method approximating the ODE $u_{t}=f(u)$. Thus 19, 39 proposed an optimal control framework for deep learning and 5, 19, 20 utilize the adjoint equation and the maximal principle to train a neural network.

Decouple Training.

Training neural networks requires forward and backward propagation in a sequential manner. Different ways have been proposed to decouple the sequential process by parallelization. This includes ADMM 34, synthetic gradients 13, delayed gradient 11, lifted machines 1, 18, 9. Our work can also be understood as a decoupling method based on a splitting technique. However, we do not attempt to decouple the gradient w.r.t. network parameters but the adversary update instead.

2 Contribution

To the best of our knowledge, it is the first attempt to design NN–specific algorithm for adversarial defense. To achieve this, we recast the adversarial training problem as a discrete time differential game. From optimal control theory, we derive the an optimality condition, i.e. the Pontryagin’s Maximum Principle, for the differential game.

Through PMP, we observe that the adversarial perturbation is only coupled with the first layer of neural networks. The PMP motivates a new adversarial training algorithm, YOPO. We split the adversary computation and weight updating and the adversary computation is focused on the first layer. Relations between YOPO and original PGD are discussed.

We finally achieve about 4$\sim$ 5 times speed up than the original PGD training with comparable results on MNIST/CIFAR10. Combining YOPO with TRADES43, we achieve both higher clean and robust accuracy within less than half of the time TRADES need.

3 Organization

This paper is organized as follows. In Section 2, we formulate the robust optimization for neural network adversarial training as a differential game and propose the gradient based YOPO. In Section 3, we derive the PMP of the differential game, study the relationship between the PMP and the back-propagation based gradient descent methods, and propose a general version of YOPO. Finally, all the experimental details and results are given in Section 4.

Differential Game Formulation and Gradient Based YOPO

Inspired by the link between deep learning and optimal control 20, we formulate the robust optimization (1) as a differential game 7. A two-player, zero-sum differential game is a game where each player controls a dynamics, and one tries to maximize, the other to minimize, a payoff functional. In the context of adversarial training, one player is the neural network, which controls the weights of the network to fit the label, while the other is the adversary that is dedicated to producing a false prediction by modifying the input.

The robust optimization problem (1) can be written as a differential game as follows,

2 Gradient Based YOPO

The Pontryagin’s Maximum Principle (PMP) is a fundamental tool in optimal control that characterizes optimal solutions of the corresponding control problem 7. PMP is a rather general framework that inspires a variety of optimization algorithms. In this paper, we will derive the PMP of the differential game (2), which motivates the proposed YOPO in its most general form. However, to better illustrate the essential idea of YOPO and to better address its relations with existing methods such as PGD, we present a special case of YOPO in this section based on gradient descent/ascent. We postpone the introduction of PMP and the general version of YOPO to Section 3.

Let us first rewrite the original robust optimization problem (1) (in a mini-batch form) as

The simplest way to solve the problem is to perform gradient ascent on the input data and gradient descent on the weights of the neural network as shown below. Such alternating optimization algorithm is essentially the popular PGD adversarial training 24. We summarize the PGD-$r$ (for each update on $\theta$) as follows, i.e. performing $r$ iterations of gradient ascent for inner maximization.

To reduce the total number of forward and backward propagation, we introduce a slack variable

and freeze it as a constant within the inner loop of the adversary update. The modified algorithm is given below and we shall refer to it as YOPO-$m$-$n$.

Another benefit of YOPO is that we take full advantage of every forward and backward propagation to update the weights, i.e. the intermediate perturbation $\eta^{j}_{i},j=1,\cdots,m-1$ are not wasted like PGD-$r$. This allows us to perform multiple updates per iteration, which potentially drives YOPO to converge faster in terms of the number of epochs. Combining the two factors together, YOPO significantly could accelerate the standard PGD adversarial training.

We would like to point out a concurrent paper 28 that is related to YOPO. Their proposed method, called "Free-$m$", also can significantly speed up adversarial training. In fact, Free-$m$ is essentially YOPO-$m$-1, except that YOPO-$m$-$1$ delays the weight update after the whole mini-batch is processed in order for a proper usage of momentum Momentum should be accumulated between mini-batches other than different adversarial examples from one mini-batch, otherwise overfitting will become a serious problem..

The Pontryagin’s Maximum Principle for Adversarial Training

In this section, we present the PMP of the discrete time differential game (2). From the PMP, we can observe that the adversary update and its associated back-propagation process can be decoupled. Furthermore, back-propagation based gradient descent can be understood as an iterative algorithm solving the PMP and with that the version of YOPO presented in the previous section can be viewed as an algorithm solving the PMP. However, the PMP facilitates a much wider class of algorithms than gradient descent algorithms 19. Therefore, we will present a general version of YOPO based on the PMP for the discrete differential game.

The PMP for continuous time differential game has been well studied in the literature 7. Here, we present the PMP for our discrete time differential game (2).

At the same time, the parameters of the first layer $\theta_{0}^{*}\in\Theta_{0}$ and the optimal adversarial perturbation $\eta^{*}_{i}$ satisfy

and the parameters of the other layers $\theta_{t}^{*}\in\Theta_{t},t\in[T]$ maximize the Hamiltonian functions

Proof is in the supplementary materials. ∎

From the theorem, we can observe that the adversary $\eta$ is only coupled with the parameters of the first layer $\theta_{0}$. This key observation inspires the design of YOPO.

2 PMP and Back-Propagation Based Gradient Descent

The classical back-propagation based gradient descent algorithm 17 can be viewed as an algorithm attempting to solve the PMP. Without loss of generality, we can let the regularization term $R=0$, since we can simply add an extra dynamic $w_{t}$ to evaluate the regularization term $R$, i.e.

We append $w$ to $x$ to study the dynamics of a new $(d_{t}+1)$-dimension vector and change $f_{t}(x,\theta_{t})$ to $(f_{t}(x,\theta_{t}),w+R_{t}(x,\theta_{t}))$. The relationship between the PMP and the back-propagation based gradient descent method was first observed by Li et al. 19. They showed that the forward dynamical system Eq.(3) is the same as the neural network forward propagation. The backward dynamical system Eq.(4) is the back-propagation, which is formally described by the following lemma.

To solve the maximization of the Hamiltonian, a simple way is the gradient ascent:

The update (8) is equivalent to gradient descent method for training networks19, 20.

3 YOPO from PMP’s View Point

Based on the relationship between back-propagation and the Pontryagin’s Maximum Principle, in this section, we provide a new understanding of YOPO, i.e. solving the PMP for the differential game. Observing that, in the PMP, the adversary $\eta$ is only coupled with the weight of the first layer $\theta_{0}$. Thus we can update the adversary via minimizing the Hamiltonian function instead of directly attacking the loss function, described in Algorithm 1.

For YOPO-$m$-$n$, to approximate the exactly minimization of the Hamiltonian, we perform $n$ times gradient descent to update the adversary. Furthermore, in order to make the calculation of the adversary more accurate, we iteratively pass one data point $m$ times. Besides, the network weights are optimized via performing the gradient ascent to Hamiltonian, resulting in the gradient based YOPO proposed in Section 2.2.

Experiments

To demonstrate the effectiveness of YOPO, we conduct experiments on MNIST and CIFAR10. We find that the models trained with YOPO have comparable performance with that of the PGD adversarial training, but with a much fewer computational cost. We also compare our method with a concurrent method "For Free"28, and the result shows that our algorithm can achieve comparable performance with around 2/3 GPU time of their official implementation.

We achieve comparable results with the best in within 250 seconds, while it takes PGD-40 more than 1250s to reach the same level. The accuracy-time curve is shown in Figuire 3(a). Naively reducing the backprop times of PGD-40 to PGD-10 will harm the robustness, as can be seen in Table 1. Experiment details can be seen in supplementary materials.

CIFAR10.

24 performs a 7-step PGD to generate adversary while training. As a comparison, we test YOPO-$3$-$5$ and YOPO-$5$-$3$ with a step size of 2/255. We experiment with two different network architectures.

Under PreAct-Res18, for YOPO-$5$-$3$, it achieves comparable robust accuracy with 24 with around half computation for every epoch. The accuracy-time curve is shown in Figuire 3(b).The quantitative results can be seen in Tbale 2. Experiment details can be seen in supplementary materials.

As for Wide ResNet34, YOPO-5-3 still achieves similar acceleration against PGD-10, as shown in Table 3. We also test PGD-3/5 to show that naively reducing backward times for this minmax problem 24 cannot produce comparable results within the same computation time as YOPO. Meanwhile, YOPO-3-5 can achieve more aggressive speed-up with only a slight drop in robustness.

2 YOPO for TRADES

TRADES43 formulated a new min-max objective function of adversarial defense and achieves the state-of-the-art adversarial defense results. The details of algorithm and experiment setup are in supplementary material, and quantitative results are demonstrated in Table 4.

Conclusion

In this work, we have developed an efficient strategy for accelerating adversarial training. We recast the adversarial training of deep neural networks as a discrete time differential game and derive a Pontryagin’s Maximum Principle (PMP) for it. Based on this maximum principle, we discover that the adversary is only coupled with the weights of the first layer. This motivates us to split the adversary updates from the back-propagation gradient calculation. The proposed algorithm, called YOPO, avoids computing full forward and backward propagation for too many times, thus effectively reducing the computational time as supported by our experiments.

Acknowledgement

We thank Di He and Long Chen for beneficial discussion. Zhanxing Zhu is supported in part by National Natural Science Foundation of China (No.61806009), Beijing Natural Science Foundation (No. 4184090) and Beijing Academy of Artificial Intelligence (BAAI). Bin Dong is supported in part by Beijing Natural Science Foundation (No. Z180001) and Beijing Academy of Artificial Intelligence (BAAI). Dinghuai Zhang is supported by the Elite Undergraduate Training Program of Applied Math of the School of Mathematical Sciences at Peking University.

References

Appendix A Proof Of The Theorems

In this section we give the full statement of the maximum principle for the adversarial training and present a proof. Let’s start from the case of the natural training of neural networks.

Then there exists co-state processes $p_{i}^{*}:={p^{*}_{i,t}:t=0,\cdots,T}$ such that the following holds for all $t\in[T]$ and $i\in[N]$:

At the same time, the parameter of the first layer $\theta_{0}^{*}\in\Theta_{0}$ and the best perturbation $\eta^{*}$ satisfy

while parameter of the other layers $\theta_{t}^{*}\in\Theta_{t},t=1,2,\cdots,T-1$ will maximize the Hamiltonian functions

We first propose PMP for discrete time dynamic system and utilize it directly gives out the proof of PMP for adversarial training.

There exists co-state processes $p_{i}^{*}:={p^{*}_{i,t}:t=0,\cdots,T}$ such that the following holds for all $t\in[T]$ and $i\in[N]$:

The parameters of the layers $\theta_{t}^{*}\in\Theta_{t},t=0,1,\cdots,T-1$ will maximize the Hamiltonian functions

Without loss of generality, we let $L=0$. The reason is that we can simply add an extra dynamic $w_{t}$ to calculate the regularization term $R$, i.e.

We append $w$ to $x$ to study the dynamic of a new $d_{t}+1$ dimension vector and modify $f_{t}(x,\theta)$ to $(f_{t}(x,\theta),w+R_{t}(x,\theta))$. Thus we only need to prove the case when $L=0$.

Now we begin the proof. Following the linearization lemma in 20, consider the linearized problem

The reachable states by the linearized dynamic system is denoted as

here $x_{t}^{\theta}$ denotes the the evolution of the dynamical system for $x_{t}$ under $\theta$. We also define

Setting $\theta_{s}=\theta_{s}^{*}$ for $s<t$ we have $\phi^{\theta}_{t+1}=f_{t}(x_{t}^{*},\theta_{t})$, which leads to $p_{t+1}^{*}\cdot(f_{t}(x_{t}^{*},\theta_{t})-x_{t+1}^{*})\leq 0$. This finishes the proof of the maximal principle on weight space $\Theta$.

We return to the proof of the theorem. The proof of the maximal principle on the weight space, i.e.

can be reached with the help of Lemma 2: replacing the dynamic start point $x_{i,0}$ in Eq.18 with $x_{i,0}+\eta_{i}^{*}$ makes this maximal principle a direct corollary of Lemma 2.

Next, we prove the Hamiltonian conidition for the adversary, i.e.

However in this time, all the layer parameters $\theta_{t}$ are fixed and $\eta_{i}$ is the control. From the above Lemma 2 we get

This finishes the proof for the adversarial control.

Appendix B Experiment Setup

Training against PGD-40 is a common practice to get sota results on MNIST. We adopt network architectures from 43 with four convolutional layers followed by three fully connected layers. Following 43 and 24, we set the size of perturbation as $\epsilon=0.3$ in an infinite norm sense. Experiments are taken on idle NVIDIA Tesla P100 GPUs. We train models for 55 epochs with a batch size of 256, longer than what convergence needs for both training methods. The learning rate is set to 0.1 initially, and is lowered by 10 times at epoch 45. We use a weight decay of $5e-4$ and a momentum of $0.9$. To measure the robustness of trained models, we performed a PGD-40 and CW4 attack with CW coefficient $c=5e2$ and $lr=1e-2$.

B.2 CIFAR-10

Following 24, we take Preact-ResNet18 and Wide ResNet-34 as the models for testing. We set the the size of perturbation as $\epsilon=8/255$ in an infinite norm sense. We perform a 20 steps of PGD with step size $2/255$ when testing. For PGD adversarial training, we train models for 105 epochs as a common practice. The learning rate is set to $5e-2$ initially, and is lowered by 10 times at epoch 79, 90 and 100. For YOPO-$m$-$n$, we train models for 40 epochs which is much longer than what convergence needs. The learning rate is set to $0.2/m$ initially, and is lowered by 10 times at epoch 30 and 36. We use a batch size of 256, a weight decay of $5e-4$ and a momentum of $0.9$ for both algorithm. We also test our model’s robustness under CW attack 4 with $c=5e2$ and $lr=1e-2$. The experiments are taken on idle NVIDIA GeForce GTX 1080 Ti GPUs.

B.3 TRADES

TRADES43 achieves the state-of-the-art results in adversarial defensing. The methodology achieves the 1st place out of the 1,995 submissions in the robust model track of NeurIPS 2018 Adversarial Vision Challenge. TRADES proposed a surrogate loss which quantify the trade-off in terms of the gap between the risk for adversarial examples and the risk for non-adversarial examples and the objective function can be formulated as

where $\Pi$ is projection operator. In the implementation of 43, after 10 such update iterations for each input data $x_{i}$, the update for weights is performed as

where $B$ is the batch size. We name this algorithm as TRADES-10, for it uses 10 iterations to update the adversary.

We name this algorithm as TRADES-YOPO-$m$-$n$. With less than half time of TRADES-$10$, TRADES-YOPO-$3$-$4$ achieves even better result than its baseline. Quantitative results is demonstrated in Table 4. The mini-batch size is $256$. All the experiments run for $105$ epochs and the learning rate set to $2e-1$ initially, and is lowered by $10$ times at epoch $70$, $90$ and $100$. The weight decay coefficient is $5e-4$ and momentum coefficient is $0.9$. We also test our model’s robustness under CW attack 4 with $c=5e2$ and $lr=5e-4$. Experiments are taken on idle NVIDIA Tesla P100 GPUs.