Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

You Only Propagate Once: Accelerating Adversarial Training via Maximal Principle

Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, Bin Dong

TL;DR

This work reframes adversarial training as a discrete-time differential game and derives a Pontryagin’s Maximum Principle to reveal that adversarial perturbations only couple to the network’s first layer. Leveraging this insight, the authors propose YOPO, which decouples adversary updates from full backpropagation and confines most computations to the first layer, drastically reducing training cost. Empirical results on MNIST and CIFAR-10 show 4–5× speedups with comparable robustness to PGD-based adversarial training, and the method also enhances TRADES performance while halving training time. The approach provides a principled, efficient route to robust training with practical impact for large-scale deployment.

Abstract

Deep learning achieves state-of-the-art results in many tasks in computer vision and natural language processing. However, recent works have shown that deep networks can be vulnerable to adversarial perturbations, which raised a serious robustness issue of deep networks. Adversarial training, typically formulated as a robust optimization problem, is an effective way of improving the robustness of deep networks. A major drawback of existing adversarial training algorithms is the computational overhead of the generation of adversarial examples, typically far greater than that of the network training. This leads to the unbearable overall computational cost of adversarial training. In this paper, we show that adversarial training can be cast as a discrete time differential game. Through analyzing the Pontryagin's Maximal Principle (PMP) of the problem, we observe that the adversary update is only coupled with the parameters of the first layer of the network. This inspires us to restrict most of the forward and back propagation within the first layer of the network during adversary updates. This effectively reduces the total number of full forward and backward propagation to only one for each group of adversary updates. Therefore, we refer to this algorithm YOPO (You Only Propagate Once). Numerical experiments demonstrate that YOPO can achieve comparable defense accuracy with approximately 1/5 ~ 1/4 GPU time of the projected gradient descent (PGD) algorithm. Our codes are available at https://https://github.com/a1600012888/YOPO-You-Only-Propagate-Once.

You Only Propagate Once: Accelerating Adversarial Training via Maximal Principle

TL;DR

This work reframes adversarial training as a discrete-time differential game and derives a Pontryagin’s Maximum Principle to reveal that adversarial perturbations only couple to the network’s first layer. Leveraging this insight, the authors propose YOPO, which decouples adversary updates from full backpropagation and confines most computations to the first layer, drastically reducing training cost. Empirical results on MNIST and CIFAR-10 show 4–5× speedups with comparable robustness to PGD-based adversarial training, and the method also enhances TRADES performance while halving training time. The approach provides a principled, efficient route to robust training with practical impact for large-scale deployment.

Abstract

Deep learning achieves state-of-the-art results in many tasks in computer vision and natural language processing. However, recent works have shown that deep networks can be vulnerable to adversarial perturbations, which raised a serious robustness issue of deep networks. Adversarial training, typically formulated as a robust optimization problem, is an effective way of improving the robustness of deep networks. A major drawback of existing adversarial training algorithms is the computational overhead of the generation of adversarial examples, typically far greater than that of the network training. This leads to the unbearable overall computational cost of adversarial training. In this paper, we show that adversarial training can be cast as a discrete time differential game. Through analyzing the Pontryagin's Maximal Principle (PMP) of the problem, we observe that the adversary update is only coupled with the parameters of the first layer of the network. This inspires us to restrict most of the forward and back propagation within the first layer of the network during adversary updates. This effectively reduces the total number of full forward and backward propagation to only one for each group of adversary updates. Therefore, we refer to this algorithm YOPO (You Only Propagate Once). Numerical experiments demonstrate that YOPO can achieve comparable defense accuracy with approximately 1/5 ~ 1/4 GPU time of the projected gradient descent (PGD) algorithm. Our codes are available at https://https://github.com/a1600012888/YOPO-You-Only-Propagate-Once.

Paper Structure

This paper contains 26 sections, 5 theorems, 41 equations, 3 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Works
  3. Adversarial Defense.
  4. Neural ODEs.
  5. Decouple Training.
  6. Contribution
  7. Organization
  8. Differential Game Formulation and Gradient Based YOPO
  9. The Optimal Control Perspective and Differential Game
  10. Gradient Based YOPO
  11. The Pontryagin’s Maximum Principle for Adversarial Training
  12. PMP
  13. PMP and Back-Propagation Based Gradient Descent
  14. YOPO from PMP's View Point
  15. Experiments
  16. ...and 11 more sections

Key Result

Theorem 1

(PMP for adversarial training) Assume $\ell_i$ is twice continuous differentiable, $f_t(\cdot,\theta),R_t(\cdot,\theta)$ are twice continuously differentiable with respect to $x$; $f_t(\cdot,\theta),R_t(\cdot,\theta)$ together with their $x$ partial derivatives are uniformly bounded in $t$ and $\the At the same time, the parameters of the first layer $\theta_0^*\in\Theta_0$ and the optimal adversa

Figures (3)

  • Figure 1: Our proposed YOPO expolits the structure of neural network. To alleviate the heavy computation cost, YOPO focuses the calculation of the adversary at the first layer.
  • Figure 2: Pipeline of YOPO-$m$-$n$ described in Algorithm \ref{['alg:A']}. The yellow and olive blocks represent feature maps while the orange blocks represent the gradients of the loss w.r.t. feature maps of each layer.
  • Figure 3: Performance w.r.t. training time

Theorems & Definitions (9)

  • Theorem 1
  • proof
  • Lemma 1
  • Theorem 2
  • Theorem
  • proof
  • Lemma 2
  • proof
  • Remark