You Only Propagate Once: Accelerating Adversarial Training via Maximal Principle

TL;DR

This paper introduces YOPO, a novel adversarial training method that significantly reduces computational costs by restricting most network updates to the first layer, while maintaining comparable robustness.

Contribution

The paper formulates adversarial training as a differential game and applies Pontryagin's Maximal Principle to develop YOPO, which minimizes forward and backward passes during adversary updates.

Findings

YOPO achieves similar robustness as PGD with 4-5 times less GPU time.

Restricting updates to the first layer greatly reduces computational overhead.

The approach maintains high accuracy in adversarial defense.

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.