# Revisiting and Advancing Fast Adversarial Training Through The Lens of Bi-
Level Optimization

Yihua Zhang1,∗ Guanhua Zhang2,∗ Prashant Khanduri3 Mingyi Hong3
Shiyu Chang2,4 Sijia Liu1,4
1Computer Science & Engineering, Michigan State University, USA
2Computer Science, University of California, Santa Barbara
3Electrical & Computer Engineering, University of Minnesota
4MIT-IBM Watson AI Lab
∗ Equal contribution

###### Abstract

Adversarial training (AT) is a widely recognized defense mechanism to gain the
robustness of deep neural networks against adversarial attacks. It is built on
min-max optimization (MMO), where the minimizer (_i.e._ , defender) seeks a
robust model to minimize the worst-case training loss in the presence of
adversarial examples crafted by the maximizer (_i.e._ , attacker). However,
the conventional MMO method makes AT hard to scale. Thus, Fast-AT [40] and
other recent algorithms attempt to simplify MMO by replacing its maximization
step with the single gradient sign-based attack generation step. Although easy
to implement, Fast-AT lacks theoretical guarantees, and its empirical
performance is unsatisfactory due to the issue of robust catastrophic
overfitting when training with strong adversaries. In this paper, we advance
Fast-AT from the fresh perspective of bi-level optimization (BLO). We first
show that the commonly-used Fast-AT is equivalent to using a stochastic
gradient algorithm to solve a linearized BLO problem involving a sign
operation. However, the discrete nature of the sign operation makes it
difficult to understand the algorithm performance. Inspired by BLO, we design
and analyze a new set of robust training algorithms termed Fast Bi-level AT
(Fast-BAT), which effectively defends sign-based projected gradient descent
(PGD) attacks without using any gradient sign method or explicit robust
regularization. In practice, we show that our method yields substantial
robustness improvements over multiple baselines across multiple models and
datasets. All code for reproducing the experiments in this paper is at
https://github.com/NormalUhr/Fast_BAT.

##  1 Introduction

Given the fact that machine learning (ML) models can be easily fooled by tiny
adversarial perturbations (also known as adversarial attacks) on the input
[15, 5, 28], training robust deep neural networks is now a major focus in
research. Nearly all existing effective defense mechanisms [26, 46, 35, 40,
45, 3] are built on the adversarial training (AT) recipe, introduced by [36]
and later formalized in [26] using _min-max optimization_ (MMO), where a
minimizer (_i.e._ defender) seeks to update model parameters against a
maximizer (_i.e._ attacker) that aims to increase the training loss by
perturbing the training data.

Despite the effectiveness of the AT-type defenses in various application
domains, the min-max nature makes them difficult to scale, because of the
multiple maximization steps required by the iterative attack generator at each
training step. The resulting prohibitive computation cost makes AT not
suitable in practical settings. For example, [41] used $128$ GPUs to run AT on
ImageNet. Thereby, how to speed up AT without losing accuracy and robustness
is now a grand challenge for adversarial defense [40].

Table 1: Performance overview of proposed Fast-BAT vs. the baselines Fast-AT
[40] and Fast-AT-GA [2] on CIFAR-10, CIFAR-100 and Tiny-ImageNet with
PreActResNet-18. All methods are robustly trained under two perturbation
budgets $\epsilon=8/255$ and $16/255$ over 20 epochs. We use the early-
stopping policy [33] to report the model with best robustness for each method.
The evaluation metrics include robust accuracy (RA) against PGD-50-10 attacks
(50-step PGD attack with 10 restarts) [26] at $\epsilon=8/255$ and $16/255$
(the test-time $\epsilon$ is the same as the train-time), RA against
AutoAttack (AA) [9] at $\epsilon=8/255$ and $16/255$, and computation time
(per epoch). The result $a$$\pm b$ represents mean $a$ and standard deviation
$b$ over 10 random trials. All experiments are run on a single GeForce RTX
3090 GPU.

CIFAR-10, PreActResNet-18
---
Method |  | SA (%)
---
($\epsilon=8/255$)
| RA-PGD (%)
---
($\epsilon=8/255$)
| RA-AA (%)
---
($\epsilon=8/255$)
| SA (%)
---
($\epsilon=16/255$)
| RA-PGD (%)
---
($\epsilon=16/255$)
| RA-AA (%)
---
($\epsilon=16/255$)
| Time
---
(s/epoch)
Fast-AT | 82.39$\pm 0.14$ | 45.49$\pm 0.21$ | 41.87$\pm 0.15$ | 44.15$\pm 7.27$ | 21.83$\pm 1.32$ | 12.49$\pm 0.33$ | 23.1
Fast-AT-GA | 79.71$\pm 0.24$ | 47.27$\pm 0.22$ | 43.24$\pm 0.27$ | 58.29$\pm 1.32$ | 26.01$\pm 0.16$ | 17.97$\pm 0.33$ | 75.3
Fast-BAT | 79.97$\pm 0.12$ | 48.83$\pm 0.17$ | 45.19$\pm 0.12$ | 68.16$\pm 0.25$ | 27.69$\pm 0.16$ | 18.79$\pm 0.24$ | 61.4
CIFAR-100, PreActResNet-18
Fast-AT |  52.62 $\pm 0.18$ | 24.66$\pm 0.21$ | 21.72$\pm 0.17$ | 21.32$\pm 3.27$ | 8.62$\pm 1.03$ | 6.22$\pm 0.61$ | 23.8
Fast-AT-GA | 50.06$\pm 0.27$ | 24.97$\pm 0.23$ | 21.82$\pm 0.21$ | 32.51$\pm 1.27$ | 12.27$\pm 0.36$ | 9.43$\pm 0.19$ | 77.1
Fast-BAT | 50.19$\pm 0.21$ | 26.49$\pm 0.20$ | 23.97$\pm 0.15$ | 39.29$\pm 0.53$ | 13.97$\pm 0.17$ | 11.32$\pm 0.22$ | 61.6
Tiny-ImageNet, PreActResNet-18
Fast-AT | 41.37$\pm 3.08$ | 17.05$\pm 3.25$ | 12.31$\pm 2.73$ | 31.38$\pm 0.19$ | 5.42$\pm 2.17$ | 3.13$\pm 0.24$ | 284.6
Fast-AT-GA | 45.52$\pm 0.24$ | 20.39$\pm 0.19$ | 16.25$\pm 0.17$ | 29.17$\pm 0.32$ | 6.79$\pm 0.27$ | 4.27$\pm 0.15$ | 592.7
Fast-BAT | 45.80$\pm 0.22$ | 21.97$\pm 0.21$ | 17.64$\pm 0.15$ | 33.78$\pm 0.23$ | 8.83$\pm 0.22$ | 5.52$\pm 0.14$ | 572.4

Recently, [40, 35, 45, 2] attempted to develop computationally-efficient
alternatives of AT, that is, the ‘fast’ versions of AT. To the best of our
knowledge, Fast-AT [40] and Fast-AT with gradient alignment (GA)
regularization, termed Fast-AT-GA [2], are the two state-of-the-art (SOTA)
‘fast’ versions of AT, since they achieve a significant reduction in
computational complexity and preserve accuracy and robustness to some extent.
To be specific, Fast-AT [40] replaces an iterative attack generator in AT with
a heuristic single-step attack generation method. Yet, Fast-AT suffers from
two main issues: (i) lack of stability, i.e., it has a large variance in
performance [25], and (ii) catastrophic overfitting, i.e., when training with
strong adversaries, the robustness of the resulting model can drop
significantly. To alleviate these problems, Andriushchenko and Flammarion [2]
proposed Fast-AT-GA by penalizing Fast-AT using an explicit GA regularization.
However, we will show that Fast-AT-GA encounters another problem: (iii) It
improves robust accuracy (RA) at the cost of a sharp drop in standard accuracy
(SA), leading to a poor accuracy-robustness tradeoff for large attack budget
(e.g., $\epsilon=16/255$ in Table 1). Moreover, (iv) there has been no
theoretical guarantee for the optimization algorithms used in Fast-AT and
Fast-AT-GA. Given the limitations (i)\- (iv), we ask:

How to design a ‘fast’ AT with improved stability, mitigated catastrophic
overfitting, enhanced RA-SA tradeoff, and some theoretical guarantees?

To address the above question, we formulate the AT problem as a unified bi-
level optimization (BLO) problem [10]. In the new formulation, the attack
generation is cast as a constrained lower-level optimization problem, while
the defense serves as the upper-level optimization problem. To the best of our
knowledge, this is the first work that makes a rigorous connection between
adversarial defense and BLO. Technically, we show that Fast-AT can be
interpreted as BLO with linearized lower-level problems. Delving into the
linearization of BLO, we propose a novel, theoretically-grounded ‘fast’ AT
framework, termed fast bi-level AT (Fast-BAT). Practically, Table 1 highlights
some achieved improvements over Fast-AT and Fast-AT-GA: When a stronger train-
time attack (i.e., $\epsilon=16/255$) is adopted, Fast-AT suffers from a large
degradation of RA and SA, together with higher variances than proposed Fast-
BAT. Although Fast-AT-GA outperforms Fast-AT, it still incurs a significant SA
loss at $\epsilon=16/255$. In contrast, Fast-BAT produces a more graceful SA-
RA tradeoff: As Fast-BAT achieves the best RA in all experimental settings,
its SA still remains at a high level, _e.g._ a significant improvement on SA
by $9.9\%$, $6.7\%$, and $4.6\%$ at $\epsilon=16/255$ for CIFAR-10, CIFAR-100,
and Tiny-ImageNet.

#### Contributions.

We summarize our contributions below.

① (Formulation-wise) We propose a unified BLO-based formulation for the robust
training problem. Within this formulation, we show the conventional Fast-AT
method is solving a low-level linearized BLO problem, rather than the vanilla
min-max problem. This key observation not only provides a new interpretation
of Fast-AT, but most importantly, also explains why Fast-AT is difficult to
possess strong theoretical guarantees.

② (Methodology-wise) We propose the new Fast-BAT algorithm based on our new
understanding of Fast-AT. The key enabling technique is to introduce a new
_smooth_ lower-level objective of BLO for robust training. In contrast to MMO,
BLO adopts a different optimization routine that requires implicit gradient
(IG) computation. We derive the closed-form of IG for Fast-BAT.

③ (Experiment-wise) We made a comprehensive experimental study to demonstrate
the effectiveness of Fast-BAT over SOTA baselines across datasets and model
types. Besides its merit in robustness enhancement, we also show its improved
stability, lifted accuracy-robustness trade-off, and mitigated catastrophic
overfitting.

##  2 Related work

#### Adversarial attack.

Adversarial attacks are techniques to generate malicious perturbations that
are imperceptible to humans but can mislead the ML models [15, 5, 9, 42, 4].
The adversarial attack has become a major approach to evaluate the robustness
of deep neural networks and thus, help build safe artificial intelligence in
many high-stakes applications such as autonomous driving [12, 24] and
surveillance [37, 43].

#### Adversarial defense and robust training at scale.

Our work falls into the category of robust training, which was mostly built on
MMO. For example, [26] established the framework of AT for the first time,
always recognized as one of the most powerful defenses [3]. Extended from AT,
TRADES [46] sought the optimal balance between robustness and generalization
ability. Further, AT-type defense has been generalized to the semi-/self-
supervised settings [6, 7] and integrated Eq. 2 with certified defense
techniques such as randomized smoothing [34].

Despite the effectiveness of AT and its variants, how to speed up AT without
losing performance remains an open question. Some recent works attempted to
impose algorithmic simplifications to AT, leading to fast but approximate AT
algorithms, such as ‘free’ AT [35], you only propagate once (YOPO) [45], Fast-
AT [40], and Fast-AT regularized by gradient alignment (termed Fast-AT-GA)
[2]. In particular, Fast-AT and Fast-AT-GA are the baselines most relevant to
ours due to their low computational complexity. However, their defense
performance is still unsatisfactory. For example, Fast-AT has poor training
stability [25] and suffers catastrophic overfitting when facing strong attacks
[2]. In contrast, Fast-AT-GA yields improved robustness but has a poor
accuracy-robustness tradeoff (see Table 1).

#### Bi-level optimization (BLO).

BLO is a unified hierarchical learning framework, where the objective and
variables of the upper-level problem depend on the lower-level problems. The
BLO problem in its most generic form is very challenging, and thus, the design
of algorithms and theory for BLO focuses on special cases [38, 39, 16, 14, 22,
21, 19]. In practice, BLO has been successfully applied to meta-learning [30],
data poisoning attack design [20], and reinforcement learning [8]. However, as
will be evident later, the existing BLO approach is not directly applied to
adversarial defense due to the presence of the constrained non-convex lower-
level problem (for attack generation). To the best of our knowledge, our work
makes a rigorous connection between adversarial defense and BLO for the first
time.

##  3 A Bi-Level Optimization View on Fast-AT

#### Preliminaries on AT and Fast-AT.

Let us consider the following standard min-max formulation for the robust
adversarial training problem [26]

$\displaystyle\begin{array}[]{l}\displaystyle\operatorname*{\text{minimize}}_{\bm{\theta}}~{}\mathbb{E}_{(\mathbf{x},y)\in\mathcal{D}}\left[\operatorname*{\text{maximize}}_{\bm{\delta}\in\mathcal{C}}\ell_{\mathrm{tr}}(\bm{\theta},\mathbf{x}+\bm{\delta},y)\right],\end{array}$
(2)

where $\bm{\theta}\in\mathbb{R}^{n}$ denotes model parameters; $\mathcal{D}$
is the training set consisting (a finite number) of labeled data pairs with
feature $\mathbf{x}$ and label $y$; $\bm{\delta}\in\mathbb{R}^{d}$ represents
adversarial perturbations subject to the perturbation constraint
$\mathcal{C}$, e.g.,
$\mathcal{C}=\\{\bm{\delta}\,|\,\|\bm{\delta}\|_{\infty}\leq\epsilon,\bm{\delta}\in[\mathbf{0},\mathbf{1}]\\}$
for $\epsilon$-toleration $\ell_{\infty}$-norm constrained attack (normalized
to $[\mathbf{0},\mathbf{1}]$); $\ell_{\mathrm{tr}}(\cdot)$ is the training
loss; $(\mathbf{x}+\bm{\delta})$ represents an adversarial example.

The standard solver for problem (2) is known as AT [26]. However, it has to
call an iterative optimization method (e.g., $K$-step PGD attack) to solve the
inner maximization problem of (2), which is computationally expensive. To
improve its scalability, the Fast-AT algorithm [40] was proposed to take a
single-step PGD attack for inner maximization. The algorithm backbone of Fast-
AT is summarized below.

####  Fast-AT algorithm

Let $\bm{\theta}_{t}$ be parameters at iteration $t$. The $(t+1)$th iteration
becomes [40]:

$\bullet$ Inner maximization by $1$-step PGD:

$\displaystyle\bm{\delta}\xleftarrow{}\mathcal{P}_{\mathcal{C}}\left(\bm{\delta}_{0}+\alpha\cdot\mathrm{sign}\left(\nabla_{\bm{\delta}}\ell_{\mathrm{tr}}(\bm{\theta}_{t},\mathbf{x}+\bm{\delta}_{0},y)\right)\right),$

where $\mathcal{P}_{\mathcal{C}}(\mathbf{a})$ denotes the projection operation
that projects the point $\mathbf{a}$ onto $\mathcal{C}$, i.e.,
$\mathcal{P}_{\mathcal{C}}(\mathbf{z})=\operatorname*{arg\,min}_{\bm{\delta}\in\mathcal{C}}\|\bm{\delta}-\mathbf{z}\|_{2}^{2}$,
$\bm{\delta}_{0}$ is a random uniform initialization within
$[\mathbf{0},\mathbf{1}]$, $\alpha>0$ is a proper learning rate (e.g.,
$1.25\epsilon$), and $\mathrm{sign}(\cdot)$ is the element-wise sign
operation.

$\bullet$ Outer minimization for model training: This can be conducted by any
standard optimizer, e.g.,

$\displaystyle\displaystyle\bm{\theta}_{t+1}\xleftarrow{}\bm{\theta}_{t}-\beta\nabla_{\bm{\theta}}\ell_{\mathrm{tr}}(\bm{\theta}_{t},\mathbf{x}+\bm{\delta},y)$

where $\beta>0$ is a proper learning rate (e.g., cyclic learning rate), and
$\bm{\delta}$ is provided from the inner maximization step.

A few remarks about Fast-AT are given below.

① Roughly speaking, Fast-AT is a simplification of AT using $1$-step PGD for
inner maximization. However, it is not entirely clear which problem Fast-AT is
actually solving. If we take a closer look at the algorithm, we will see that
the inner update only changes to the initial $\bm{\delta}_{0}$, not the most
recent $\bm{\delta}$. Clearly, this scheme is fundamentally different from the
classical gradient descent-ascent algorithm for min-max optimization [32],
which alternatively updates the inner and outer variables. _Therefore, it
remains elusive if Fast-AT is solving the original problem (2)_.

② Andriushchenko and Flammarion [2] demonstrated that Fast-AT could lead to
catastrophic overfitting when using strong adversaries for training. However,
there was no grounded theory to justify the pros and cons of Fast-AT. We will
show that BLO provides a promising solution.

#### BLO: Towards a unified formulation of robust training.

BLO (bi-level optimization) is an optimization problem that involves two
levels of optimization tasks, where the lower-level task is nested inside the
upper-level task. More precisely, it has the following generic form:

$\begin{array}[]{cl}\displaystyle\min_{\mathbf{u}\in\mathcal{U}}&\quad
f({\mathbf{u}},{\mathbf{v}}^{*}({\mathbf{u}}))\\\
\text{s.t.}&\quad{\mathbf{v}}^{*}({\mathbf{u}})=\arg\min_{\mathbf{v}\in\mathcal{V}}\;g(\mathbf{v},\mathbf{u})\end{array}$
(3)

where $\mathcal{U}$ and $\mathcal{V}$ are the feasible sets for the variables
$\mathbf{u}$ and $\mathbf{v}$, respectively; $f(\cdot)$ and $g(\cdot)$ are the
upper- and the lower-level objective functions, respectively. Intuitively, the
BLO (3) can be used to formulate the adversarial training problem as the
latter also involves two problems, one nested in the other. Importantly, it is
more powerful than the min-max formulation (2) as it allows the two problems
to have different objective functions. This flexibility provided by BLO is the
key to the generality of our proposed framework.

To make the above intuition precise, we use the upper-level problem to model
the training loss minimization, while the lower-level problem to model the
attack generation process, and consider the following BLO problem:

$\displaystyle\begin{array}[]{cl}\displaystyle\min_{\bm{\theta}}&\mathbb{E}_{(\mathbf{x},y)\in\mathcal{D}}[\ell_{\mathrm{tr}}(\bm{\theta},\mathbf{x}+\bm{\delta}^{*}(\bm{\theta};\mathbf{x},y);y)]\\\
\text{s.t.}&\bm{\delta}^{*}(\bm{\theta};\mathbf{x},y)=\displaystyle\operatorname*{arg\,min}_{\bm{\delta}\in\mathcal{C}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta};\mathbf{x},y),\end{array}$
(6)

where the training loss function $\ell_{\rm tr}(\cdot)$ has been defined in
(2), and $\ell_{\mathrm{atk}}(\cdot)$ denotes an attack objective. For the
notation simplicity, in the subsequent discussion, we will not indicate the
dependency of the functions $\ell_{\mathrm{tr}}$, $\ell_{\mathrm{atk}}$, and
$\bm{\delta}^{*}$ with respect to (w.r.t.) the data samples $(\mathbf{x},y)$.
The formulation (6) has two key differences from (2):

① When we choose $\ell_{\rm atk}=-\ell_{\rm tr}$, problem (6) becomes
equivalent to the min-max formulation (2). It follows that the BLO is suitable
to formulate the adversarial training problem. Moreover, we will see shortly
that the flexibility provided by choosing the lower-level objective
independently of the upper-level one enables us to interpret Fast-AT as
solving a certain special form of the BLO problem. Note that prior to our
work, it was not entirely clear what is the problem that Fast-AT is trying to
solve.

② BLO calls a different optimization routine from those to solve the original
min-max problem (2). As will be evident later, even if we set
$\ell_{\mathrm{atk}}=-\ell_{\mathrm{tr}}$ in (6), the BLO-enabled solver does
not simplify to Fast-AT (see more details in Appendix B). This is because for
a given data sample $(\mathbf{x},y)$, the gradient for the upper-level problem
of (6) yields:

$\displaystyle\displaystyle\frac{d\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))}{d\bm{\theta}}=\nabla_{\bm{\theta}}\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))+\underbrace{\frac{{d}\bm{\delta}^{*}(\bm{\theta})^{\top}}{{d}\bm{\theta}}}_{\text{{{\text{IG}}}}}\nabla_{\bm{\delta}}\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta})),$
(7)

where the superscript $\top$ denotes the transpose operation, and
$\nabla_{\bm{\theta}}\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))$
denotes the partial derivative w.r.t. the first input argument $\bm{\theta}$;
and
$\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}\in\mathbb{R}^{n\times
d}$, if exists, is referred to as implicit gradient (IG) because it is defined
through an implicitly constrained optimization problem
$\min_{\bm{\delta}\in\mathcal{C}}\ell_{\mathrm{atk}}$. The dependence on IG is
a ‘fingerprint’ of BLO (2) in contrast to AT or Fast-AT.

#### BLO-enabled interpretation of Fast-AT.

Next, we demonstrate how Fast-AT relates to BLO. Fast-AT can be interpreted as
an approximated stochastic gradient algorithm for solving the following lower-
level linearized BLO. That is to say, Fast-AT is not solving the original min-
max problem (2), but its linearized version below:

$\displaystyle\begin{array}[]{ll}\displaystyle\operatorname*{\text{minimize}}_{\bm{\theta}}&\mathbb{E}_{(\mathbf{x},y)\in\mathcal{D}}[\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))]\\\
\vspace*{3mm}\operatorname*{\text{subject
to}}&\bm{\delta}^{*}(\bm{\theta})=\operatorname*{arg\,min}_{\bm{\delta}\in\mathcal{C}}~{}[(\bm{\delta}-\mathbf{z})^{\top}\mathrm{sign}(\nabla_{\bm{\delta}=\mathbf{z}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta}))+(\lambda/2)\|\bm{\delta}-\mathbf{z}\|_{2}^{2}],\end{array}$
(10)

where $\mathbf{z}=\bm{\delta}_{0}$ and $\lambda=1/\alpha$. Our justification
for the above interpretation is elaborated on below.

① The simplified lower-level problem of (10) leads to the closed-form
solution:

$\displaystyle\bm{\delta}^{*}(\bm{\theta})=$
$\displaystyle\operatorname*{arg\,min}_{\bm{\delta}\in\mathcal{C}}\,(\lambda/2)\|\bm{\delta}-\mathbf{z}+(1/\lambda)\mathrm{sign}(\nabla_{\bm{\delta}=\mathbf{z}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta}))\|_{2}^{2}$
(11) $\displaystyle=$
$\displaystyle\mathcal{P}_{\mathcal{C}}\left(\mathbf{z}-(1/\lambda)\mathrm{sign}(\nabla_{\bm{\delta}=\mathbf{z}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta}))\right),$

which is exactly given by the $1$-step PGD attack with initialization
$\mathbf{z}$ and learning rate $(1/\lambda)$. In the linearization used in
(10), a quadratic regularization term (with regularization parameter
$\lambda$) is introduced to ensure the strong convexity of the lower-level
objective within the constraint set $\bm{\delta}\in\mathcal{C}$ so as to
achieve the unique minimizer (11). Note that imposing such a strongly convex
regularizer is also commonly used to stabilize the convergence of MMO (min-max
optimization) and BLO [29, 19]. If we set $\mathbf{z}=\bm{\delta}_{0}$ and
$\lambda=1/\alpha$, (11) precisely depicts the inner maximization step used in
Fast-AT.

② By substituting (11) into the upper-level problem of (10), we can then use
(7) to compute the stochastic gradients of the upper-level problem. If the
stochastic gradient can be precisely computed, then we can update the model
parameters $\bm{\theta}$ using SGD based on (7). That is,
$\scriptstyle{\bm{\theta}\xleftarrow{}\bm{\theta}-\beta}\textstyle{\frac{d\ell_{\mathrm{tr}}(\bm{\theta}_{t},\bm{\delta}^{*}(\bm{\theta}))}{d\bm{\theta}}}$
(with learning rate $\beta$). However, generally speaking, the IG function
$\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}$ involved in (7)
may not be differentiable, and even it is, the computation may not be easy.
For our case, $\bm{\delta}^{*}(\bm{\theta})$ expressed in (11) involves both a
projection and a sign operation, which can be particularly difficult to
compute. To proceed, let us make the following approximations. We assume that
the chain rule of the derivative of $\bm{\delta}^{*}(\bm{\theta})$ holds
w.r.t. $\bm{\theta}$, implying the differentiability of the projection
operation and the sign operation. Then, based on the closed-form of
$\bm{\delta}^{*}(\bm{\theta})$ in (11), IG is approximately equal to
$\displaystyle\textstyle\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}\scriptstyle=\mathbf{0}$,
where we use two facts: (1) The linearization point $\mathbf{z}$ is
independent of $\bm{\theta}$, i.e. $\mathbf{z}=\bm{\delta}_{0}$; And (2)
$\frac{d\mathrm{sign}(\cdot)}{d\bm{\theta}}=\mathbf{0}$ holds almost
everywhere. Clearly, the use of the gradient sign method simplifies the IG
computation. Further, the upper-level gradient is approximated by
$\frac{d\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))}{d\bm{\theta}}\scriptstyle\approx\nabla_{\bm{\theta}}\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))\mathrel{\mathop{:}}=\tilde{\mathbf{h}}(\bm{\theta})$,
and the upper-level optimization of problem (10) becomes
$\scriptstyle\bm{\theta}\xleftarrow{}\bm{\theta}-\beta\tilde{\mathbf{h}}(\bm{\theta})$,
which is the same as the outer minimization step in Fast-AT. In a nutshell,
the BLO solver of problem (10), which calls the IG computation based on (11),
eventually reduces to the Fast-AT algorithm.

The aforementioned analysis shows that Fast-AT can be viewed as using an
approximated stochastic gradient algorithm to solve the linearized BLO (10),
with the linearization point $\mathbf{z}$ and the regularization parameter
$\lambda$ set as $\mathbf{z}=\bm{\delta}_{0}$ and $\lambda=1/\alpha$. However,
since a series of approximations have been used when arriving at the
approximated gradient $\widetilde{\mathbf{h}}(\bm{\theta})$ used by Fast-AT,
it is no longer clear if the resulting algorithm can still sufficiently reduce
the objective function of the upper-level problem. Additionally, based on the
fact that the lower-level problem of (10) involves the discrete sign operator,
it is unlikely (if not impossible) that any approximated stochastic gradient-
based algorithms developed for it can possess any strong theoretical
guarantees.

##  4 Fast-BAT: Advancing Fast-AT by BLO

####  Fast-BAT and rationale.

The key take-away from (10) is that the conventional Fast-AT adopts the sign
of input gradient to linearize the lower-level attack objective. However, a
more natural choice is to use the first-order Taylor expansion for
linearization. By doing so, problem (10) can be modified to the following
form:

$\displaystyle\begin{array}[]{cl}\displaystyle\operatorname*{\text{minimize}}_{\bm{\theta}}&\mathbb{E}_{(\mathbf{x},y)\in\mathcal{D}}[\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))]\\\
\operatorname*{\text{subject
to}}&\bm{\delta}^{*}(\bm{\theta})=\operatorname*{arg\,min}_{\bm{\delta}\in\mathcal{C}}~{}[(\bm{\delta}-\mathbf{z})^{\top}(\nabla_{\bm{\delta}=\mathbf{z}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta}))+(\lambda/2)\|\bm{\delta}-\mathbf{z}\|_{2}^{2}],\end{array}$
(14)

Similar to (11), problem (14) can be solved as:

$\displaystyle\displaystyle\bm{\delta}^{*}(\bm{\theta})=\mathcal{P}_{\mathcal{C}}\left(\mathbf{z}-(1/\lambda)\nabla_{\bm{\delta}=\mathbf{z}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta})\right).$
(15)

In contrast to (11), the IG associated with (14) cannot be approximated by
zeros since the gradient sign operation is not present in (15). To compute IG,
the auto-differentiation (the chain rule) can be applied to the closed-form of
$\bm{\delta}^{*}(\bm{\theta})$. However, this will not give us an accurate and
generalizable IG solution since the projection operation
$\mathcal{P}_{\mathcal{C}}$ is not smooth everywhere and thus, the use of
chain rule does not yield a rigorous derivation. In the following subsection,
we address the IG challenge in a theoretically-grounded manner.

####  IG theory for Fast-BAT.

The problem of Fast-BAT (14) falls into a class of very challenging BLO
problems, which requires constrained lower-level optimization. The
unconstrained case is easier to handle since one can apply the implicit
function theory to the stationary condition of the lower-level problem to
obtain IG [19]. Yet, in our case with constrained problems, a stationary point
could violate the constraints, and thus, the stationary condition becomes non-
applicable. Fortunately, in problem (14), we are dealing with a special class
of lower-level constraints – linear constraints. Let us rewrite the
constraints below:

$\begin{split}\displaystyle\|\bm{\delta}\|_{\infty}\leq\epsilon,\bm{\delta}\in[\mathbf{-x},\mathbf{1-x}]\Longleftrightarrow{}\mathbf{B}\bm{\delta}\leq\mathbf{b},\end{split}$
(16)

where
$\scriptstyle{\mathbf{B}\mathrel{\mathop{:}}=}\begin{bmatrix}\scriptstyle{\mathbf{I}}\\\
\scriptstyle{-\mathbf{I}}\end{bmatrix},~{}\scriptstyle{\mathbf{b}\mathrel{\mathop{:}}=}\begin{bmatrix}\scriptstyle{\min\\{\epsilon\mathbf{1},\mathbf{1}-\mathbf{x}\\}}\\\
\scriptstyle{-\max\\{-\epsilon\mathbf{1},-\mathbf{x}\\}}\end{bmatrix}$. By
exploiting the above linearly constrained problem structure, we show that the
IG challenge associated with (14) can be addressed via Karush–Kuhn–Tucker
(KKT) conditions. We present our main theoretical results (Proposition 1 and
Corollary 1) below and refer readers to Appendix A for detailed derivation.

###### Proposition 1

For the BLO problem (14), let $g(\bm{\theta},\bm{\delta}^{*})$ denote its
lower-level objective function evaluated at $\bm{\delta}^{*}$ given
$\bm{\theta}$, then the analytical form of IG (implicit gradient) is given by

$\begin{array}[]{cc}\displaystyle\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}=-\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast})\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast})^{-1}+\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast})\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast})^{-1}\mathbf{B}_{0}^{\top}[\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast})^{-1}\mathbf{B}_{0}^{\top}]^{-1}\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast})^{-1},\end{array}$
(18)

where $\textstyle\bm{\delta}^{*}$ is given by (15) (the dependence on
$\textstyle\bm{\theta}$ is omitted for ease of notation),
$\textstyle\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{*})\in\mathbb{R}^{n\times
d}$ denotes the second-order partial derivatives evaluated at
$\textstyle(\bm{\theta},\bm{\delta}^{*})$, and $\textstyle\mathbf{B}_{0}$
denotes the sub-matrix of $\textstyle\mathbf{B}$ that only corresponds to the
active constraints at $\textstyle\bm{\delta}^{\ast}$, i.e., those in
$\textstyle\mathbf{B}\bm{\delta}^{\ast}\leq\mathbf{b}$ satisfied with
equality.

It is clear from (18) that the computation of IG requires the second-order
derivatives as well as matrix inversion. This is computationally intensive.
Recall from (14) that the Hessian matrix $\nabla_{\bm{\delta}\bm{\delta}}g$ of
the lower-level objective function is given by
$\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast})=\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}+\lambda\mathbf{I}$.
This inspires us to impose the Hessian-free approximation, i.e.,
$\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}=\mathbf{0}$. The rationale
behind the Hessian-free assumption is that ReLU-based neural networks commonly
lead to a piece-wise linear decision boundary w.r.t. the inputs [27, 1], and
thus, its second-order derivative (Hessian)
$\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}$ is close to zero. In Sec.
5.2 and Appendix D, we will empirically show that the Hessian-free assumption
is reasonable for both ReLU and non-ReLU neural networks. Thus, the Hessian
matrix is approximated as:

$\displaystyle\displaystyle\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))=\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}+\lambda\mathbf{I}=\mathbf{0}+\lambda\mathbf{I}.$
(19)

With (19), we can simplify closed-form of IG as below:

###### Corollary 1

With the Hessian-free assumption, namely
$\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}=\mathbf{0}$, the IG
(implicit gradient) of (14) is

$\displaystyle{\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}}=-(1/\lambda)\nabla_{\bm{\theta}\bm{\delta}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta}^{\ast})\mathbf{H}_{\mathcal{C}},$
(20)

where
$\scriptstyle\mathbf{H}_{\mathcal{C}}\mathrel{\mathop{:}}=\begin{bmatrix}\scriptstyle
1_{p_{1}<\delta_{1}^{*}<q_{1}}\mathbf{e}_{1}\cdots
1_{p_{1}<\delta_{d}^{*}<q_{d}}\mathbf{e}_{d}\end{bmatrix}\in\mathbb{R}^{d\times
d}$ and the function $1_{p_{i}<\delta_{i}^{*}<q_{i}}\in\\{0,1\\}$ denotes the
indicator over the constraint of
$\\{\delta_{i}\,|\,p_{i}<\delta_{i}^{*}<q_{i}\\}$, which returns $1$ if the
constraint is satisfied, $\delta_{i}^{*}$ denotes the $i$th entry of
$\bm{\delta}^{*}(\bm{\theta})$, $p_{i}=\max\\{-\epsilon,-x_{i}\\}$ and
$q_{i}=\min\\{\epsilon,1-x_{i}\\}$ characterize the boundary of the linear
constraint (16) for the variable $\delta_{i}$, and
$\mathbf{e}_{i}\in\mathbb{R}^{d}$ denotes the basis vector with the $i$th
entry being $1$ and others being $0$.

####  Fast-BAT algorithm and implementation.

Similar to Fast-AT or AT, the Fast-BAT algorithm also follows the principle of
alternating optimization. Specifically, it consists of the IG-based upper-
level gradient descent (7), interlaced with the lower-level optimal attack
(15). We summarize the Fast-BAT algorithm below.

$\bullet$ Lower-level solution: Obtain $\bm{\delta}^{*}(\bm{\theta}_{t})$ from
(15);

$\displaystyle\displaystyle\bm{\delta}^{*}(\bm{\theta})=\mathcal{P}_{\mathcal{C}}\left(\mathbf{z}-(1/\lambda)\nabla_{\bm{\delta}=\mathbf{z}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta})\right).$

$\bullet$ Upper-level model training: Integrating the IG (20) into (7), call
SGD to update the model parameters as:

$\displaystyle\begin{array}[]{ll}\displaystyle\bm{\theta}_{t+1}=\bm{\theta}_{t}-\alpha_{1}\nabla_{\bm{\theta}}\ell_{\mathrm{tr}}(\bm{\theta}_{t},\bm{\delta}^{\ast})-\alpha_{2}(-\frac{1}{\lambda})\nabla_{\bm{\theta}\bm{\delta}}\ell_{\mathrm{atk}}(\bm{\theta}_{t},\bm{\delta}^{\ast})\mathbf{H}_{\mathcal{C}}\nabla_{\bm{\delta}}\ell_{\mathrm{tr}}(\bm{\theta}_{t},\bm{\delta}^{\ast}),\end{array}$
(22)

where $\alpha_{1},\alpha_{2}>0$ are learning rates associated with the model
gradient and the IG-augmented descent direction.

It is clear from (22) that to train a robust model, Fast-BAT can be decomposed
into the regular Fast-AT update (i.e., the term multiplied by $\alpha_{1}$)
and the additional update that involves IG, (i.e., the term multiplied by
$\alpha_{2}$). To implement Fast-BAT, we highlight some key hyper-parameter
setups that are different from Fast-AT [40] and Fast-AT-GA [2].

#### Remark on implementation details

In what follows, we discuss the practical setups of hyperparameters used in
Fast-BAT; see more details and empirical justifications in Appendix C. In Sec.
5.1, we will show that a proper $\alpha_{2}$ helps mitigate catastrophic
overfitting. In practice, we set $\alpha_{2}/\lambda=0.1\alpha_{1}$ (see Table
A2 for empirical justifications). To specify the linearization point
$\mathbf{z}$ in (14), we investigate two types of linearization schemes: (1)
the random constant linearization (uniformly random linearization and random
corner linearization) and (2) the $1$-step perturbation warm-up-based
linearization ($1$-step sign-based PGD and $1$-step PGD w/o sign). We consider
these linearization schemes since they have computational complexities up to
the one-step attack generation. Empirically, we find that Fast-BAT using
“$1$-step PGD w/o sign" leads to the best defensive performance; see
justification in Table A4. We follow this experiment setup in the sequel.

#### Remark on convergence analysis

We prove the convergence of Fast-BAT in Appendix E. It is worth mentioning
that the main analysis difficulty lies in the last term of the model updating
rule (22), which involves two coupled derivatives built upon the same mini-
batch.

##  5 Experiments

###  5.1 Experiment Setup

#### Datasets and model architectures.

We will evaluate the effectiveness of our proposal under CIFAR-10 [23],
CIFAR-100 [23], Tiny-ImageNet [11], and ImageNet [11]. Unless specified
otherwise, we will train DNN models PreActResNet (PARN)-18 [18] for all
datasets except ImageNet, and ResNet (RN)-50 [17] for ImageNet. As a part of
the ablation study, we also train larger models PARN-50 and WideResNet
(WRN)-16-8 [44] on CIFAR-10. Some preliminary ImageNet results are reported in
Appendix D.

#### Baselines.

We consider three baseline methods: Fast-AT [40], Fast-AT-GA [2], and PGD-2-AT
[26], i.e., the $2$-step PGD attack-based AT. The primal criterion of baseline
selection is computational complexity. The training time of all methods
including ours falls between the time of Fast-AT and Fast-AT-GA.

#### Training details.

We choose the training perturbation strength $\epsilon\ \in\\{8,16\\}/255$ for
CIFAR-10, CIFAR-100, and Tiny-ImageNet; and $\epsilon=2/255$ for ImageNet
following [40, 2]. Throughout the experiments, we utilize an SGD optimizer
with a momentum of $0.9$ and weight decay of $5\times 10^{-4}$. For CIFAR-10,
CIFAR-100 and Tiny-ImageNet, we train each model for 20 epochs in total, where
we use the cyclic scheduler to adjust the learning rate. The learning rate
linearly ascends from $0$ to $0.2$ within the first $10$ epochs and then
reduces to $0$ within the last $10$ epochs. Our batch size is set to 128 for
all settings. In the implementation of Fast-BAT, we follow the dataset-
agnostic hyperparameter scheme for $\lambda$, such that $\lambda=255/5000$ for
$\epsilon=8/255$ and $\lambda=255/2500$ for $\epsilon=16/255$ for CIFAR-10,
CIFAR-100 and Tiny-ImageNet. For ImageNet, we strictly follow the setup given
by [40] and we choose the train-time attack budget as $\epsilon=2/255$. For
each method, we use the early stopping method to pick the model with the best
robust accuracy, following [33]. All the experiments are conducted on a single
GeForce RTX 3090 GPU.All the baselines are trained with the recommended
configurations in their official GitHub repos. We refer readers to Appendix C
for more details on the training setup.

Table 2: SA, RA-PGD, and RA-AA of different robust training methods in the
setup (CIFAR-10, PARN-18 training with $\epsilon=8/255$) and (CIFAR-10,
PARN-18 training with $\epsilon=16/255$), respectively. All the results are
averaged over 10 independent trials with different random seeds.

CIFAR-10, PARN-18 trained with $\epsilon=8/255$
---
Method | SA (%) | RA-PGD (%) | RA-AA (%)
$\epsilon=8$ | $\epsilon=16$ | $\epsilon=8$ | $\epsilon=16$
Fast-AT | 82.39$\pm 0.44$ |  45.49 $\pm 0.41$ | 9.56 $\pm 0.26$ |  41.87 $\pm 0.15$ | 7.91 $\pm 0.06$
Fast-AT-GA | 79.71$\pm 0.44$ |  47.27 $\pm 0.42$ | 11.57$\pm 0.32$ |  43.24 $\pm 0.27$ | 9.48 $\pm 0.15$
PGD-2-AT | 81.97$\pm 0.41$ |  44.62 $\pm 0.39$ | 9.39 $\pm 0.32$ |  41.73 $\pm 0.20$ | 7.54$\pm 0.25$
Fast-BAT | 79.97 $\pm 0.12$ |  48.83 $\pm 0.17$ | 14.00 $\pm 0.21$ |  45.19 $\pm 0.12$ | 11.51 $\pm 0.20$
CIFAR-10, PARN-18 trained with $\epsilon=16/255$
Fast-AT | 44.15$\pm 7.27$ |  37.17 $\pm 0.74$ | 21.83 $\pm$1.32 |  31.66 $\pm 0.27$ | 12.49 $\pm 0.33$
Fast-AT-GA | 58.29 $\pm 1.32$ |  43.86 $\pm 0.67$ | 26.01 $\pm 0.16$ |  38.69 $\pm 0.56$ | 17.97 $\pm 0.33$
PGD-2-AT | 68.04 $\pm 0.30$ |  48.79 $\pm 0.31$ | 24.30 $\pm 0.46$ |  41.59 $\pm 0.22$ | 15.40 $\pm 0.29$
Fast-BAT | 68.16 $\pm 0.25$ |  49.05 $\pm 0.12$ | 27.69 $\pm 0.16$ |  43.64 $\pm 0.26$ | 18.79 $\pm 0.24$

Table 3: Performance of different robust training methods under different
model types. All the models are both trained and evaluated with the same
perturbation strength $\epsilon$.

Model | Method |  | SA(%)
---
($\epsilon=8/255$)
| RA-PGD(%)
---
($\epsilon=8/255$)
| SA(%)
---
($\epsilon=16/255$)
| RA-PGD(%)
---
($\epsilon=16/255$)
PARN-50 | Fast-AT | 73.15$\pm 6.10$ | 41.03$\pm 2.99$ | 43.86$\pm 4.31$ | 22.08$\pm 0.27$
| Fast-AT-GA | 77.40$\pm 0.81$ | 46.16$\pm 0.98$ | 42.28$\pm 6.69$ | 22.87$\pm 1.25$
| PGD-2-AT | 83.53$\pm 0.17$ | 46.17$\pm 0.59$ | 68.88$\pm 0.39$ | 22.37$\pm 0.41$
| Fast-BAT | 78.91$\pm 0.68$ | 49.18$\pm 0.35$ | 69.01$\pm 0.19$ | 24.55$\pm 0.06$
WRN-16-8 | Fast-AT | 84.39$\pm 0.46$ | 45.80$\pm 0.57$ | 49.39$\pm 2.17$ | 21.99$\pm 0.41$
| Fast-AT-GA | 81.51$\pm 0.38$ | 48.29$\pm 0.20$ | 45.95$\pm 13.65$ | 23.10$\pm 3.90$
| PGD-2-AT | 85.52$\pm 0.14$ | 45.47$\pm 0.14$ | 72.11$\pm 0.33$ | 23.61$\pm 0.16$
| Fast-BAT | 81.66$\pm 0.54$ | 49.93$\pm 0.36$ | 68.12$\pm 0.47$ | 25.63$\pm 0.44$

#### Evaluation details.

For adversarial evaluation, we report robust test accuracy (RA) of a learned
model against PGD attacks [26] (RA-PGD). Unless otherwise specified, we set
the test-time perturbation strength ($\epsilon$) the same as the train-time
value, and take $50$-step PGD with $10$ restarts all the datasets. Since
AutoAttack [9] is known as the strongest robust benchmark evaluation metric
(given as an ensemble attack), we also measure robust accuracy against
AutoAttack, termed RA-AA. Further, we measure the standard accuracy (SA)
against natural examples. Results are averaged over 10 independent trials. We
would like to highlight that all the methods share the same batch size, epoch
number, and training hardware. Thus, the time consumption per epoch reported
in Table 1 and Table 5 will serve as a fair indicator of the algorithm
complexity of different methods.

###  5.2 Results

#### Overall performance of Fast-BAT.

In Table 1, Table 2, and Table A1, we compare the overall performance of our
proposed Fast-BAT with baselines.

① We find that Fast-BAT consistently outperforms the other baselines across
the datasets and attack types. In Table 1, Fast-BAT improves the RA-PGD
performance consistently by over $1.5\%$ and RA-AA by around $1\%$ across all
the datasets on both attack strengths. For stronger attacks with a larger
perturbation budget the advantage of Fast-BAT is even clearer, e.g. a gain of
over $2\%$ in Tiny-ImageNet with $\epsilon=16/255$. On ImageNet, Fast-BAT
outperforms Fast-AT by 1.23% with $\epsilon=2/255$.

② Fast-BAT leads to a much better SA-RA trade-off compared with the baselines.
For example, in Table 1, the improvement in RA is not at cost of a huge drop
in SA. Instead, when models are trained with $\epsilon=16/255$, Fast-BAT even
enjoys a significant boost over Fast-AT-GA in SA by $9.9\%$, $6.7\%$, and
$4.6\%$ for CIFAR-10, CIFAR-100, and Tiny-ImageNet respectively.

#### Performance under different model architectures.

Besides PARN-18 reported above, Table 3 presents results on both deeper
(PARN-50) and wider (WRN-18-6) models. As we can see, Fast-BAT consistently
yields RA improvement over other methods. We also note that PGD-2-AT could be
a competitive baseline in terms of SA. In contrast to Fast-AT and Fast-AT-GA,
Fast-BAT is the only approach that yields an evident RA improvement over
PGD-2-AT.

(a) Without early stopping (b) With early stopping

Figure 1: RA-PGD of different robust training methods for CIFAR-10 with the
same training and evaluation attack strengths.

#### Mitigation of robust catastrophic overfitting.

As shown in [2], Fast-AT suffers robust catastrophic overfitting when the
train-time and test-time attack strength $\epsilon$ grows. Following [2],
Figure 1 presents two RA-PGD trajectories, i.e., training without early
stopping and training with early stopping, versus the train/test-time
$\epsilon$. As we can see, Fast-AT encounters a sharp RA drop when
$\epsilon>8$ when early stopping is not used, consistent with [2]. With early
stopping, the overfitting of RA can be alleviated to some extent for Fast-AT,
but its performance still remains the worst. Moreover, different from [2], we
find that PGD-2-AT yields resilient performance against catastrophic
overfitting. Our implementation gives a more positive baseline for PGD-2-AT,
since the implementation in [2] did not use random initialization to generate
train-time attacks. Furthermore, Figure 1 shows that Fast-BAT mitigates the
issue of catastrophic overfitting and yields improved RA over other methods.
We highlight that such an achievement is _free_ of any robustness stability
regularization, like gradient alignment used in Fast-AT-GA.

#### Gradient alignment for ‘free’.

As shown by [2], _catastrophic overfitting_ occurs with the local non-
linearity of deep networks, which can be measured by the gradient alignment
(GA) score:

$\displaystyle\displaystyle\mathbb{E}_{(\mathbf{x},y)\sim\mathcal{D},\bm{\eta}\sim\mathcal{U}({[-\epsilon,\epsilon]}^{d})}[\text{cos}(\nabla_{\mathbf{x}}\ell_{\mathrm{tr}}(\mathbf{x},y;\bm{\theta}),\nabla_{\mathbf{x}}\ell_{\mathrm{tr}}(\mathbf{x}+\bm{\eta},y;\bm{\theta}))],$

where $\mathcal{U}$ denotes the randomly uniform distribution.

Figure 2:  The training curve of GA score over different training methods on
CIFAR-10.

GA is a key performance indicator to measure the appearance of robustness
catastrophic overfitting, as catastrophic overfitting is always accompanied by
a sharp GA drop in the training trajectory of the robustly trained model. In
our paper, we calculate the GA for each method on the test set at the end of
each epoch throughout the training process. Figure 1 suggested that Fast-BAT
can mitigate overfitting without explicit GA regularization, and Figure 2
presents the GA score versus the training epoch number. As can be seen, Fast-
BAT automatically enforces GA and remains very close to Fast-AT-GA, which
maximizes GA with an explicit regularization. Therefore, a high GA score may
just be a necessary but not a sufficient condition for avoiding catastrophic
overfitting. As additional evidence, Fig. A1 shows that similar to Fast-AT-GA,
Fast-BAT has a flatter loss landscape than Fast-AT as well. Therefore, a
direct penalization on the input gradient norm may not achieve the state-of-
the-art model robustness.

Table 4: Robustness against adaptive attacks (RA-PGD) and transfer attacks
(RA-Transfer Attack). Naturally trained PARN-18, PARN-50, and WRN-16-8 serve
as source victim models for attack generation with PGD-20 ($\epsilon=8/255$)
and PARN-18 robustified by different methods as target model for transfer
attack evaluation.

Method |  RA-PGD(%) | RA-Transfer Attack(%)
---|---|---
PARN-18 | PARN-50 | WRN-16-8
Fast-AT | 45.44$\pm 0.06$ | 76.35$\pm 0.12$ | 76.94$\pm 0.14$ | 77.23$\pm 0.21$
PGD-2-AT | 44.71$\pm 0.04$ | 77.56$\pm 0.14$ | 78.64$\pm 0.12$ | 78.84$\pm 0.17$
Fast-AT-GA | 47.31$\pm 0.05$ | 77.34$\pm 0.13$ | 78.34$\pm 0.13$ | 78.53$\pm 0.12$
Fast-BAT(Ours) | 48.67$\pm 0.05$ | 78.03$\pm 0.15$ | 79.93$\pm 0.12$ | 79.21$\pm 0.15$

#### Sanity check for obfuscated gradients

As pointed out by [3], model robustness could be overestimated due to
obfuscated gradients. The model with obfuscated gradients could have
‘obfuscated’ stronger resilience to white-box (adaptive) attacks than black-
box (transfer) attacks. To justify the validity of Fast-BAT, Table 4
summarizes the comparison between our proposal and the other baselines when
facing adaptive and transfer attacks. Firstly, for all the methods, RA
increases if the transfer attack is applied, implying that the transfer attack
is weaker than the adaptive attack. This is desired in the absence of
obfuscated gradients. Moreover, Fast-BAT consistently outperforms the other
baselines against both adaptive and transfer attacks. The absence of
obfuscated gradients can also be justified by the flatness of the adversarial
loss landscape in Figure. A1.

#### The validity of the Hessian-free assumption on ReLU-based neural
networks.

Table 5: Performance of Hessian-free and Hessian-aware Fast-BAT on CIFAR-10.
We train and evaluate with the same attack budgets $\epsilon=8/255$ and
$\epsilon=16/255$ to show the influence brought by Hessian matrix.

Method |  | SA(%)
---
($\epsilon=8/255$)
| RA-PGD(%)
---
($\epsilon=8/255$)
| SA(%)
---
($\epsilon=16/255$)
| RA-PGD(%)
---
$(\epsilon=16/255)$
| Time
---
(s/epoch)
Hessian-free | 79.97 $\pm 0.12$ | 48.83 $\pm 0.17$ | 68.16 $\pm 0.25$ | 27.69 $\pm 0.16$ | 61.4
Hessian-aware | 79.62 $\pm 0.17$ | 49.13 $\pm 0.14$ | 67.82 $\pm 0.23$ | 27.82 $\pm 0.19$ | 82.6

In Corollary 1, the Hessian-free assumption,
i.e.$\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}=0$, was made to
simplify the computation of implicit gradient. To justify this assumption we
conduct experiments to compare the Hessian-free Fast-BAT with the Hessian-
aware version. In Hessian-aware Fast-BAT, the implicit gradient is calculated
based on (34). In Table 5, the results do not indicate much difference when
Hessian is used. However, the extra computations required to evaluate the
Hessian heavily slows down Fast-BAT as around $30\%$ more time is needed.
Therefore, the Hessian-free assumption is reasonable and also necessary in
terms of the efficiency of the algorithm. We also justify this assumption on
some non-ReLU neural networks and for more results please refer to Appendix D.

##  6 Conclusion

In this paper, we introduce a novel bi-level optimization-based fast
adversarial training framework, termed Fast-BAT. The rationale behind
designing Fast-BAT lies in two aspects. First, from the perspective of
implicit gradients, we show that the existing Fast-AT framework is equivalent
to the lower-level linearized BLO along the sign direction of the input
gradient. Second, we show that Fast-BAT is able to achieve improved stability
of performance, mitigated catastrophic overfitting, and enhanced accuracy-
robustness trade-off. To the best of our knowledge, we for the first time
establish the theory and the algorithmic foundation of BLO for adversarial
training. Extensive experiments are provided to demonstrate the superiority of
our method to state-of-the-art accelerated AT baselines.

## References

  * Alfarra et al. [2020] Motasem Alfarra, Adel Bibi, Hasan Hammoud, Mohamed Gaafar, and Bernard Ghanem.  On the decision boundaries of deep neural networks: A tropical geometry perspective.  _arXiv preprint arXiv:2002.08838_ , 2020.
  * Andriushchenko and Flammarion [2020] Maksym Andriushchenko and Nicolas Flammarion.  Understanding and improving fast adversarial training.  In _NeurIPS_ , 2020.
  * Athalye et al. [2018a] Anish Athalye, Nicholas Carlini, and David Wagner.  Obfuscated gradients give a false sense of security: Circumventing defenses to adversarial examples.  _arXiv preprint arXiv:1802.00420_ , 2018a.
  * Athalye et al. [2018b] Anish Athalye, Logan Engstrom, Andrew Ilyas, and Kevin Kwok.  Synthesizing robust adversarial examples.  In _International Conference on Machine Learning_ , pages 284–293, 2018b.
  * Carlini and Wagner [2017] Nicholas Carlini and David Wagner.  Towards evaluating the robustness of neural networks.  In _IEEE Symposium on S &P_, 2017.
  * Carmon et al. [2019] Yair Carmon, Aditi Raghunathan, Ludwig Schmidt, Percy Liang, and John C Duchi.  Unlabeled data improves adversarial robustness.  _arXiv preprint arXiv:1905.13736_ , 2019.
  * Chen et al. [2020] T. Chen, S. Liu, S. Chang, Y. Cheng, L. Amini, and Z. Wang.  Adversarial robustness: From self-supervised pretraining to fine-tuning.  In _CVPR_ , 2020.
  * Chen et al. [2019] Zhangyu Chen, Dong Liu, Xiaofei Wu, and Xiaochun Xu.  Research on distributed renewable energy transaction decision-making based on multi-agent bilevel cooperative reinforcement learning.  2019\.
  * Croce and Hein [2020] Francesco Croce and Matthias Hein.  Reliable evaluation of adversarial robustness with an ensemble of diverse parameter-free attacks.  In _International Conference on Machine Learning_ , pages 2206–2216. PMLR, 2020.
  * Dempe [2002] Stephan Dempe.  _Foundations of bilevel programming_.  Springer Science & Business Media, 2002.
  * Deng et al. [2009] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _Computer Vision and Pattern Recognition, 2009. CVPR 2009. IEEE Conference on_ , pages 248–255. IEEE, 2009.
  * Deng et al. [2020] Yao Deng, Xi Zheng, Tianyi Zhang, Chen Chen, Guannan Lou, and Miryung Kim.  An analysis of adversarial attacks and defenses on autonomous driving models.  In _2020 IEEE International Conference on Pervasive Computing and Communications (PerCom)_ , pages 1–10. IEEE, 2020.
  * Engstrom et al. [2018] Logan Engstrom, Andrew Ilyas, and Anish Athalye.  Evaluating and understanding the robustness of adversarial logit pairing.  _arXiv preprint arXiv:1807.10272_ , 2018.
  * Ghadimi and Wang [2018] Saeed Ghadimi and Mengdi Wang.  Approximation methods for bilevel programming.  _arXiv preprint:1802.02246_ , 2018.
  * Goodfellow et al. [2014] Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy.  Explaining and harnessing adversarial examples.  _arXiv preprint arXiv:1412.6572_ , 2014.
  * Gould et al. [2016] Stephen Gould, Basura Fernando, Anoop Cherian, Peter Anderson, Rodrigo Santa Cruz, and Edison Guo.  On differentiating parameterized argmin and argmax problems with application to bi-level optimization.  _arXiv preprint arXiv:1607.05447_ , 2016.
  * He et al. [2016a] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Proceedings of the IEEE conference on computer vision and pattern recognition_ , pages 770–778, 2016a.
  * He et al. [2016b] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Identity mappings in deep residual networks.  In _European conference on computer vision_ , pages 630–645. Springer, 2016b.
  * Hong et al. [2020] Mingyi Hong, Hoi-To Wai, Zhaoran Wang, and Zhuoran Yang.  A two-timescale framework for bilevel optimization: Complexity analysis and application to actor-critic.  _arXiv preprint arXiv:2007.05170_ , 2020.
  * Huang et al. [2020] W Ronny Huang, Jonas Geiping, Liam Fowl, Gavin Taylor, and Tom Goldstein.  Metapoison: Practical general-purpose clean-label data poisoning.  _arXiv preprint arXiv:2004.00225_ , 2020.
  * Ji et al. [2020] Kaiyi Ji, Junjie Yang, and Yingbin Liang.  Bilevel optimization: Nonasymptotic analysis and faster algorithms.  _arXiv preprint arXiv:2010.07962_ , 2020.
  * Khanduri et al. [2021] Prashant Khanduri, Siliang Zeng, Mingyi Hong, Hoi To Wai, Zhaoran Wang, and Zhuoran Yang.  A near-optimal algorithm for stochastic bilevel optimization via double-momentum.  In _Thirty-Fifth Conference on Neural Information Processing Systems_ , 2021.
  * Krizhevsky and Hinton [2009] A. Krizhevsky and G. Hinton.  Learning multiple layers of features from tiny images.  _Master’s thesis, Department of Computer Science, University of Toronto_ , 2009.
  * Kumar et al. [2020] K Naveen Kumar, C Vishnu, Reshmi Mitra, and C Krishna Mohan.  Black-box adversarial attacks in autonomous vehicle technology.  In _2020 IEEE Applied Imagery Pattern Recognition Workshop (AIPR)_ , pages 1–7. IEEE, 2020.
  * Li et al. [2020] Bai Li, Shiqi Wang, Suman Jana, and Lawrence Carin.  Towards understanding fast adversarial training.  _arXiv preprint arXiv:2006.03089_ , 2020.
  * Madry et al. [2018] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu.  Towards deep learning models resistant to adversarial attacks.  In _International Conference on Learning Representations_ , 2018.
  * Moosavi-Dezfooli et al. [2019] Seyed-Mohsen Moosavi-Dezfooli, Alhussein Fawzi, Jonathan Uesato, and Pascal Frossard.  Robustness via curvature regularization, and vice versa.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , pages 9078–9086, 2019.
  * Papernot et al. [2016] Nicolas Papernot, Patrick McDaniel, Somesh Jha, Matt Fredrikson, Z Berkay Celik, and Ananthram Swami.  The limitations of deep learning in adversarial settings.  In _Security and Privacy (EuroS &P), 2016 IEEE European Symposium on_, pages 372–387. IEEE, 2016.
  * Qian et al. [2019] Qi Qian, Shenghuo Zhu, Jiasheng Tang, Rong Jin, Baigui Sun, and Hao Li.  Robust optimization over multiple domains.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 33, pages 4739–4746, 2019.
  * Rajeswaran et al. [2019] Aravind Rajeswaran, Chelsea Finn, Sham Kakade, and Sergey Levine.  Meta-learning with implicit gradients.  _arXiv preprint arXiv:1909.04630_ , 2019.
  * Ramachandran et al. [2017] Prajit Ramachandran, Barret Zoph, and Quoc V Le.  Searching for activation functions.  _arXiv preprint arXiv:1710.05941_ , 2017.
  * Razaviyayn et al. [2020] Meisam Razaviyayn, Tianjian Huang, Songtao Lu, Maher Nouiehed, Maziar Sanjabi, and Mingyi Hong.  Nonconvex min-max optimization: Applications, challenges, and recent theoretical advances.  _IEEE Signal Processing Magazine_ , 37(5):55–66, 2020.
  * Rice et al. [2020] Leslie Rice, Eric Wong, and Zico Kolter.  Overfitting in adversarially robust deep learning.  In _International Conference on Machine Learning_ , pages 8093–8104. PMLR, 2020.
  * Salman et al. [2019] Hadi Salman, Greg Yang, Jerry Li, Pengchuan Zhang, Huan Zhang, Ilya Razenshteyn, and Sebastien Bubeck.  Provably robust deep learning via adversarially trained smoothed classifiers.  _arXiv preprint arXiv:1906.04584_ , 2019.
  * Shafahi et al. [2019] Ali Shafahi, Mahyar Najibi, Mohammad Amin Ghiasi, Zheng Xu, John Dickerson, Christoph Studer, Larry S Davis, Gavin Taylor, and Tom Goldstein.  Adversarial training for free!  In _Advances in Neural Information Processing Systems_ , pages 3353–3364, 2019.
  * Szegedy et al. [2014] Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus.  Intriguing properties of neural networks.  _International Conference on Learning Representations_ , 2014.
  * Thys et al. [2019] Simen Thys, Wiebe Van Ranst, and Toon Goedemé.  Fooling automated surveillance cameras: adversarial patches to attack person detection.  In _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition Workshops_ , pages 0–0, 2019.
  * Vicente et al. [1994] Luis Vicente, Gilles Savard, and Joaquim Júdice.  Descent approaches for quadratic bilevel programming.  _Journal of Optimization Theory and Applications_ , 81(2):379–399, 1994.
  * White and Anandalingam [1993] Douglas J White and G Anandalingam.  A penalty function approach for solving bi-level linear programs.  _Journal of Global Optimization_ , 3(4):397–419, 1993.
  * Wong et al. [2020] Eric Wong, Leslie Rice, and J. Zico Kolter.  Fast is better than free: Revisiting adversarial training.  In _International Conference on Learning Representations_ , 2020.
  * Xie et al. [2019] Cihang Xie, Yuxin Wu, Laurens van der Maaten, Alan L Yuille, and Kaiming He.  Feature denoising for improving adversarial robustness.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 501–509, 2019.
  * Xu et al. [2019] Kaidi Xu, Sijia Liu, Pu Zhao, Pin-Yu Chen, Huan Zhang, Quanfu Fan, Deniz Erdogmus, Yanzhi Wang, and Xue Lin.  Structured adversarial attack: Towards general implementation and better interpretability.  In _ICLR_ , 2019.
  * Xu et al. [2020] Kaidi Xu, Gaoyuan Zhang, S. Liu, Quanfu Fan, Mengshu Sun, Hongge Chen, Pin-Yu Chen, Yanzhi Wang, and Xue Lin.  Adversarial T-Shirt! evading person detectors in a physical world.  In _ECCV_ , 2020.
  * Zagoruyko and Komodakis [2016] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  _arXiv preprint arXiv:1605.07146_ , 2016.
  * Zhang et al. [2019a] Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, and Bin Dong.  You only propagate once: Accelerating adversarial training via maximal principle.  _arXiv preprint arXiv:1905.00877_ , 2019a.
  * Zhang et al. [2019b] Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric P Xing, Laurent El Ghaoui, and Michael I Jordan.  Theoretically principled trade-off between robustness and accuracy.  _ICML_ , 2019b.

##  Appendix A Proof of Proposition 1 and Corollary 1

Proof: Upon defining
$g(\bm{\theta},\bm{\delta})=(\bm{\delta}-\mathbf{z})^{\top}\nabla_{\bm{\delta}=\mathbf{z}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta})+(\lambda/2)\|\bm{\delta}-\mathbf{z}\|_{2}^{2}$,
we repeat (14) as

$\displaystyle\begin{array}[]{ll}\displaystyle\operatorname*{\text{minimize}}_{\bm{\theta}}&\mathbb{E}_{(\mathbf{x},y)\in\mathcal{D}}[\ell_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))]\\\
\operatorname*{\text{subject
to}}&\bm{\delta}^{*}(\bm{\theta})=\displaystyle\operatorname*{arg\,min}_{\mathbf{B}\bm{\delta}\leq\mathbf{b}}g(\bm{\theta},\bm{\delta}),\end{array}$
(25)

where we have used the expression of linear constraints in (16).

Our goal is to derive the IG
$\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}$ shown in (7). To
this end, we first build implicit functions by leveraging KKT conditions of
the lower-level problem of (25). We say $\bm{\delta}^{*}(\bm{\theta})$ and
$\bm{\lambda}^{*}(\bm{\theta})$ (Lagrangian multipliers) satisfy the KKT
conditions:

$\displaystyle\begin{split}\text{Stationarity:}\qquad&\nabla_{\bm{\delta}}g(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))+\mathbf{B}^{\top}\bm{\lambda}^{\ast}(\bm{\theta})=\mathbf{0},\\\
\text{Complementary slackness :
}\qquad&\bm{\lambda}^{\ast}(\bm{\theta})\cdot(\mathbf{B}\bm{\delta}^{\ast}(\bm{\theta})-\mathbf{b})=\mathbf{0}\\\
\text{Dual
feasibility:}\qquad&\bm{\lambda}^{\ast}(\bm{\theta})\geq\mathbf{0}\end{split}$
(26)

where $\cdot$ denotes the elementwise product.

Active constraints & definition of $\mathbf{B}_{0}$: Let $\mathbf{B}_{0}$
denote the sub-matrix of $\mathbf{B}$ and $\mathbf{b}_{0}$ the sub-vector of
$\mathbf{b}$, which consists of only the active constraints at
$\bm{\delta}^{\ast}(\bm{\theta})$, i.e., those satisfied with the equality
$\mathbf{B}_{0}\bm{\delta}^{\ast}(\bm{\theta})=\mathbf{b}_{0}$ (corresponding
to nonzero dual variables). The determination of active constraints is done
given $\bm{\theta}$ at each iteration.

With the aid of $(\mathbf{B}_{0},\mathbf{b}_{0})$, KKT (26) becomes

$\displaystyle\nabla_{\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))+\mathbf{B}_{0}^{\top}\bm{\lambda}^{\ast}(\bm{\theta})=\mathbf{0},~{}~{}\text{and}~{}~{}\mathbf{B}_{0}\bm{\delta}^{\ast}(\bm{\theta})-\mathbf{b}_{0}=\mathbf{0},$
(27)

where the nonzero $\bm{\lambda}^{\ast}(\bm{\theta})$ only corresponds to the
active constraints. We take derivatives w.r.t. $\bm{\theta}$ of (27), and thus
obtain the following

$\displaystyle\frac{d\nabla_{\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{\top}}{d\bm{\theta}}+\nabla_{\bm{\theta}}\bm{\lambda}^{\ast}(\bm{\theta})^{\top}\mathbf{B}_{0}=\mathbf{0}$
$\displaystyle\Longrightarrow$
$\displaystyle\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))+\underbrace{\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}}_{\text{IG}}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))+\nabla_{\bm{\theta}}\bm{\lambda}^{\ast}(\bm{\theta})^{\top}\mathbf{B}_{0}=\mathbf{0},$
(28) and
$\displaystyle~{}\underbrace{\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}}_{\text{IG}}\mathbf{B}_{0}^{\top}=\mathbf{0},$
(29)

where
$\nabla_{\bm{\theta}\bm{\delta}}\in\mathbb{R}^{|\bm{\theta}|\times|\bm{\delta}|}$
denotes second-order partial derivatives (recall that $|\bm{\theta}|=n$ and
$|\bm{\delta}|=d$). According to (28), we have

$\displaystyle{\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}}=-[\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))+\nabla_{\bm{\theta}}\bm{\lambda}^{\ast}(\bm{\theta})^{\top}\mathbf{B}_{0}]\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}.$
(30)

Substituting the above into (29), we obtain

$\displaystyle\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}+\nabla_{\bm{\theta}}\bm{\lambda}^{\ast}(\bm{\theta})^{\top}\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}=\mathbf{0},$
(31)

which yields:

$\displaystyle\nabla_{\bm{\theta}}\bm{\lambda}^{\ast}(\bm{\theta})^{\top}=-\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}[\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}]^{-1},$
(32)

and thus,

$\displaystyle\nabla_{\bm{\theta}}\bm{\lambda}^{\ast}(\bm{\theta})^{\top}\mathbf{B}_{0}=-\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}[\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}]^{-1}\mathbf{B}_{0}.$
(33)

Substituting (33) into (30), we obtain the IG

$\displaystyle{\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}}=-\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}-\nabla_{\bm{\theta}}\bm{\lambda}^{\ast}(\bm{\theta})^{\top}\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}$
$\displaystyle=-\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}$
$\displaystyle+\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}[\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}\mathbf{B}_{0}^{\top}]^{-1}\mathbf{B}_{0}\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))^{-1}.$
(34)

The above completes the proof of Proposition 1. $\square$

To further compute (34), the Hessian matrix
$\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}$ is needed. Recall from
the definition of the lower-level objective that the Hessian matrix is given
by

$\displaystyle\nabla_{\bm{\delta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))=\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}+\lambda\mathbf{I}=\mathbf{0}+\lambda\mathbf{I}.$
(35)

Here we used the assumption that
$\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}=\mathbf{0}$. The rationale
behind that is neural networks commonly leads to a piece-wise linear decision
boundary w.r.t. the inputs [27, 1], and thus, its second-order derivative
(Hessian) $\nabla_{\bm{\delta}\bm{\delta}}\ell_{\mathrm{atk}}$ is close to
zero.

Based on the simplification (35), we have

$\displaystyle{\frac{d\bm{\delta}^{*}(\bm{\theta})^{\top}}{d\bm{\theta}}}=$
$\displaystyle-(1/\lambda)\nabla_{\bm{\theta}\bm{\delta}}g(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\underbrace{\left(\mathbf{I}-\mathbf{B}_{0}^{\top}[\mathbf{B}_{0}\mathbf{B}_{0}^{\top}]^{-1}\mathbf{B}_{0}\right)}_{\mathrel{\mathop{:}}=\mathbf{H}_{\mathcal{C}}}$
$\displaystyle-(1/\lambda)\nabla_{\bm{\theta}\bm{\delta}}\ell_{\mathrm{atk}}(\bm{\theta},\bm{\delta}^{\ast}(\bm{\theta}))\mathbf{H}_{\mathcal{C}},$
(36)

where we have used the fact that
$\nabla_{\bm{\theta}\bm{\delta}}{{g}}=\nabla_{\bm{\theta}\bm{\delta}}\ell_{\mathrm{atk}}$.

What is $\mathbf{H}_{\mathcal{C}}$ in (36)? Since
$\mathbf{B}=\begin{bmatrix}\mathbf{I}\\\ -\mathbf{I}\end{bmatrix}$, we can
obtain that $\mathbf{B}_{0}\mathbf{B}_{0}^{\top}=\mathbf{I}$ and
$\mathbf{B}_{0}^{\top}\mathbf{B}_{0}$ is a sparse diagonal matrix with
diagonal entries being $0$ or $1$. Thus, $\mathbf{H}_{\mathcal{C}}$ can be
first simplified to

$\displaystyle\mathbf{H}_{\mathcal{C}}=\mathbf{I}-\mathbf{B}_{0}^{\top}\mathbf{B}_{0}.$
(37)

Clearly, $\mathbf{H}_{\mathcal{C}}$ is also a diagonal matrix with either $0$
or $1$ diagonal entries. The $1$-valued diagonal entry of
$\mathbf{H}_{\mathcal{C}}$ corresponds to the inactive constraints in
$\mathbf{B}\bm{\delta}^{*}(\bm{\theta})<\mathbf{b}$, i.e., those satisfied
with strict inequalities in
$\\{\|\bm{\delta}\|_{\infty}\leq\epsilon,\mathbf{0}\leq\bm{\delta}\leq\mathbf{1}\\}$.
This can be expressed as

$\displaystyle\mathbf{H}_{\mathcal{C}}=\begin{bmatrix}1_{p_{1}\leq\delta_{1}^{*}\leq
q_{1}}\mathbf{e}_{1},\ldots,1_{p_{1}\leq\delta_{d}^{*}\leq
q_{d}}\mathbf{e}_{d}\end{bmatrix}$ (38)

where $1_{p_{i}\leq\delta_{i}^{*}\leq q_{i}}\in\\{0,1\\}$ denotes the
indicator function over the constraint $\\{p_{i}\leq\delta_{i}^{*}\leq
q_{i}\\}$ and returns $1$ if the constraint is satisfied, $\delta_{i}^{*}$
denotes the $i$th entry of $\bm{\delta}^{*}(\bm{\theta})$,
$p_{i}=\max\\{-\epsilon,-x_{i}\\}$ and $q_{i}=\min\\{\epsilon,1-x_{i}\\}$, and
$\mathbf{e}_{i}\in\mathbb{R}^{d}$ denotes the basis vector with the $i$th
entry being $1$ and others being $0$s.

Based on the definition of $g$, (36) and (38), we can eventually achieve the
desired IG formula (20). The proof of Corollary 1 is now complete. $\square$

##  Appendix B Discussion on case $\ell_{\mathrm{atk}}=-\ell_{\mathrm{tr}}$

We provide an in-depth explanation on the fact that even if we set
$\ell_{\text{atk}}=-\ell_{\text{tr}}$, the optimization routine given by (7)
to solve problem (6) does not reduce to the ordinary IG-absent gradient
descent to solve problem (2) because of the presence of lower-level
constraints.

  * •

In the absence of the constraint $\bm{\delta}\in\mathcal{C}$, if we set
$\ell_{\text{atk}}=-\ell_{\text{tr}}$, then solving problem (6) via IG-
involved descent method (7) will reduce to the ordinary IG-absent method that
solves problem (2).

This is a known BLO result (_e.g._ [14]) and can be readily proven using the
stationary condition. To be specific, based on the stationary condition of
unconstrained lower-level optimization, we have
$\nabla_{\bm{\delta}}\ell_{\text{atk}}(\bm{\theta},\bm{\delta}^{*})=0$. Since
$\ell_{\text{atk}}=-\ell_{\text{tr}}$, we have
$\nabla_{\bm{\delta}}\ell_{\text{tr}}(\bm{\theta},\bm{\delta}^{*})=0$. As a
result, the second term in (7) becomes $\mathbf{0}$ and solving problem (6)
becomes identical to solving the min-max problem (2).

  * •

In the presence of the constraint $\bm{\delta}\in\mathcal{C}$, the stationary
condition cannot be applied since the stationary point may not be a feasible
point in the constraint. In other words,
$\nabla_{\bm{\delta}}\ell_{\text{atk}}(\bm{\theta},\bm{\delta}^{*})=0$ does
not hold in the case of $\ell_{\text{atk}}=-\ell_{\text{tr}}$. As a matter of
fact, one has to resort to KKT conditions instead of the stationary condition
for a constrained lower-level problem. Similar to our proof in Proposition 1,
the implicit gradient (and thus the second term of (7)) cannot be omitted in
general. This makes the optimization routine to solve problem (6) different
from solving problem (2).

##  Appendix C Detailed Experiment settings

###  C.1 Training Set-up

For CIFAR-10, we summarize the training setup for each method. 1) Fast-AT: We
use FGSM with an attack step size of $1.25\epsilon$ to generate perturbations;
2) PGD-2-AT: 2-step PGD attacks111We use random initialization to generate
perturbations for PGD, while in the paper of Fast-AT-GA [2], 2-step PGD is
initialized at zero point, which we believe will underestimate the effect of
PGD-2-AT with an attack step size of $0.5\epsilon$ is implemented; 3) Fast-AT-
GA: The gradient alignment regularization parameter is set to the recommended
value for each $\epsilon$; 4) Fast-BAT: We select $\lambda$ from $255/5000$ to
$255/2000$ for different $\epsilon$. At the same time, we adjust $\alpha_{2}$
accordingly, so that the coefficient of the second term in (22), namely
$\alpha_{2}/\lambda$ always equals to $0.1\alpha_{1}$.

For ImageNet, we set $\epsilon$ to $2/255$ , and we strictly follow the
training setting adopted by [40]. In Fast-BAT, we fix $\lambda$ at $255/3000$
and adopt the same $\alpha_{2}$ selection strategy as CIFAR-10.

#### Parameter for Fast-AT-GA

Regarding Fast-AT-GA with different model types, we adopt the same
regularization parameter recommended in its official repo222Fast-AT-GA:
https://github.com/tml-epfl/understanding-fast-adv-training/blob/master/sh
intended for PreActResNet-18 (namely 0.2 for $\epsilon=8/255$ and 2.0 for
$\epsilon=16/255$).

###  C.2 The choice of initialization point $\mathbf{z}$

To specify $\mathbf{z}$ in (14), we investigate two classes of linearization
schemes. The first class is random constant linearization, including:
“uniformly random linearization”, i.e., $\mathbf{z}=\bm{\delta}_{0}$ similar
to Fast-AT, and “random corner linearization" under the $\epsilon$-radius
$\ell_{\infty}$-ball, i.e., $\mathbf{z}\in\\{-\epsilon,\epsilon\\}^{d}$. The
second class is $1$-step perturbation warm-up-based linearization, including
the other two specifications: “$1$-step sign-based PGD" , and “$1$-step PGD
w/o sign" . We consider this linearization schemes as their computation
complexities are less than or close to the complexity of one-step attack
generation. As a result, Fast-BAT takes comparable computation cost to the
baselines Fast-AT, PGD-2-AT and Fast-AT-GA. Empirically, we find that Fast-BAT
using “$1$-step PGD w/o sign" leads to the best defensive performance; see
justification in Table A4. We follow this experiment setup in the sequel.

##  Appendix D Additional Experimental Results

#### Results on ImageNet

Fast-AT Fast-AT-GA Fast-BAT

Figure A1:  Visualization of adversarial loss landscapes of Fast-AT, Fast-AT-
GA and Fast-BAT trained using the ResNet-18 model on the CIFAR-10 dataset. The
losses at are calculated w.r.t. the same image example ID #001456, and the
landscape is obtained by tracking the loss changes w.r.t. input variations
following [13]. That is, the loss landscape is generated by
$z=\text{loss}(I+x\cdot\mathbf{r_{1}}+y\cdot\mathbf{r_{2}})$, where $I$
denotes an image, and the $x$-axis and the $y$-axis correspond to linear
coefficients associated with the sign-based attack direction
$\mathbf{r_{1}}=\text{sign}(\nabla_{I}\text{loss}(I))$ $\mathbf{\hat{x}}$ and
a random direction $\mathbf{r_{2}}\sim\text{Rademacher}(0.5)$, respectively.
Table A1: SA and RA on ImageNet.

Method | SA (%) | RA-PGD (%)
---|---|---
Fast-AT | 60.90 | 43.43
Fast-BAT | 60.18 | 44.64

We train DNN models ResNet (RN)-50 [17] for ImageNet and we choose the
training perturbation strength $\epsilon=2/255$ strictly following [40, 2]. We
remark that when evaluating on ImageNet, we only compare ours with Fast-AT
since as shown in Table 6 of [2], the other baseline methods did not show any
improvement over Fast-AT at the attack budget $\epsilon=2/255$. RA-PGD stands
for the robustness against PGD-50-10 (50-step PGD attack with 10 restarts)
with $\epsilon=2/255$. Table A1 shows the performance comparison between Fast-
AT and Fast-BAT. We can see Fast-BAT outperforms Fast-AT by 1.23% when facing
attacks with $\epsilon=2/255$.

#### Sensitivity to regularization parameter $\lambda$

Table A2: Performance of Fast-BAT with different parameter $\lambda$. We train
and evaluate with the same attack budget $\epsilon=16/255$ on CIFAR-10 to show
the influence of $\lambda$.

CIFAR-10, PreActResNet-18, $\epsilon=16/255$
---
1/$\lambda$ (/255) | 500 | 1000 | 1500 | 2500 | 3500
SA (%) | 83.20 | 75.06 | 69.31 | 68.16 | 64.37
RA-PGD (%) | 19.02 | 21.42 | 23.34 | 27.69 | 25.32

In Table A2, we show the sensitivity of Fast-BAT to the regularization
parameter $\lambda$. All the parameters remain the same as the default
setting, except that for different $\lambda$. We always adjust $\alpha_{2}$ so
that $\alpha_{2}/\lambda=0.1\alpha_{1}$ holds. Note $1/\lambda$ also serves as
the attack step in (15). As $\lambda$ decreases, the improvement in robust
accuracy is evidently strengthened, and there is an obvious trade-off between
robust accuracy (SA) and standard accuracy (RA). At a certain level of
$\lambda$, namely when $\lambda\leq 255/3500$, RA starts to converge and stop
surging.

#### Sensitivity to different $\alpha_{2}$ choices

Table A3: Performance of Fast-BATwith different $\alpha_{2}$ choices on
CIFAR-10. Models are trained and evaluated with the same attack budget
($\epsilon=16/255$). Here $\alpha_{1}$ is set as the cyclic learning rate and
is not a constant value. $\alpha_{2}$ is always set proportionate to
$\alpha_{1}$ for simplicity.

|  $\alpha_{2}$ (CIFAR-10,
---
PreActResNet18,
$\epsilon=16/255)$
0.025$\alpha_{1}$ | 0.0167$\alpha_{1}$ | 0.0125$\alpha_{1}$ | 0.008$\alpha_{1}$
SA (%) | 75.06 | 69.31 | 68.16 | 57.92
RA-PGD (%) | 21.42 | 23.34 | 27.69 | 20.53

We consider the case of robust training with the large $\epsilon$ choice
($16/255$). As we can see from Table A3, if $\alpha_{2}$ is set too small
($\alpha_{2}=0.008\alpha_{1}$), then both SA and RA will drop significantly.
Here $\alpha_{1}$ is set as the cyclic learning rate and thus not a constant
parameter. However, in the $\alpha_{2}$ interval
$[0.0125\alpha_{1},0.025\alpha_{1}]$, we observed a tradeoff between standard
accuracy (SA) and robust accuracy (RA): That is, the improvement in RA
corresponds to a loss in SA. In our experiments, we choose $\alpha_{2}$ when
the tradeoff yields the best RA without suffering a significant drop of SA
(which still outperforms the baseline approaches).

#### Sensitivity of linearization schemes

Fast-BAT needs a good linearization point $\mathbf{z}$ in (14). In
experiments, we adopt the perturbation generated by 1-step PGD without sign as
our default linearization scheme. In Table A4, we show the performance of the
other possible linearization options. We find that 1-step PGD without sign
achieves the best robust accuracy among all the choices. This is not spurring
since this linearization point choice is consistent with the first-Taylor
expansion that we used along the direction of the input gradient without the
sign operation involved. By contrast, Fast-BAT linearized with uniformly
random noise suffers from catastrophic overfitting and reaches a rather low
standard accuracy (SA). Fast-BAT with other linearizations also yields a worse
SA-RA trade-off than our proposal.

Table A4: Performance of Fast-BAT with different linearization schemes.
Besides 1-step PGD without sign (PGD w/o Sign), we further generate
linearization point with the following methods: uniformly random noise
$[-\epsilon,\epsilon]^{d}$ (Uniformly Random); uniformly random corner
$\\{-\epsilon,\epsilon\\}^{d}$ (Random Corner); and perturbation from 1-step
PGD attack with $0.5\epsilon$ as attack step (PGD).

CIFAR-10, PreActResNet-18, $\epsilon=16/255$
---
| Linearization
---
Method
| PGD
---
w/o Sign
| Uniformly
---
Random
| Random
---
Corner
PGD
SA (%) | 69.31 | 43.42 | 62.19 | 75.30
RA-PGD (%) | 23.34 | 21.25 | 16.5 | 19.42

#### The validity of the Hessian-free assumption on non-ReLU based neural
networks.

The Hessian-free assumption is based on the fact that the commonly used ReLU
activation function is piece-wise linear w.r.t. input. We further conduct
experiments to verify the feasibility of such an assumption on models with
non-ReLU activation functions. We choose two commonly used activation
functions, Swish [31] and Softplus, as alternatives for the non-smooth ReLU
function. We compare the results both calculating Hessian as well as the
Hessian-free version to see if the Hessian-free assumption still holds for the
non-ReLU neural network. The results are shown in Table A5. As we can see, the
use of Hessian does not affect performance much. A similar phenomenon can be
observed across different $\epsilon$ and different model activation functions
(ReLU, Softplus, and Swish). However, the introduction of Hessian leads to an
increase in time consumption by more than 30%. Therefore, we can draw the
conclusion that the Hessian-free assumption is reasonable across different
activation function choices.

Ablation studies. In Appendix D, we present additional empirical studies
including 1) the sensitivity analysis of the linearization hyperparameter
$\lambda$, 2) the choice of the linearization point,  and 3) the sensitivity
analysis of $\alpha_{2}$.

Table A5: Performance of Fast-AT and Fast-BAT with different activation
functions on CIFAR-10. ReLU, Swish and Softplus are taken into consideration.
For Fast-BAT, we compare the Hessian-free and Hessian-aware version to verify
the influence of Hessian matrix. The results are averaged over 3 independent
trials.

Setting |  | SA (%)
---
($\epsilon=8/255$)
| RA-PGD (%)
---
($\epsilon=8/255$)
| SA (%)
---
($\epsilon=16/255$)
| RA-PGD (%)
---
($\epsilon=16/255$)
| Time
---
(s/epoch)
Fast-AT-ReLU | 82.39$\pm 0.14$ | 45.49$\pm 0.21$ | 44.15$\pm 7.27$ | 21.83$\pm 1.32$ | 23.1
| Fast-BAT-ReLU
---
Hessian-free
79.97 $\pm 0.12$ | 48.83 $\pm 0.17$ | 68.16 $\pm 0.25$ | 27.69 $\pm 0.16$ | 61.4
| Fast-BAT-ReLU
---
Hessian-aware
79.62 $\pm 0.17$ | 49.13 $\pm 0.14$ | 67.82 $\pm 0.23$ | 27.82 $\pm 0.19$ | 82.6
Fast-AT-Softplus | 81.29 $\pm 0.16$ | 47.26 $\pm 0.24$ | 45.39 $\pm 3.27$ | 22.40 $\pm 0.75$ | 23.3
| Fast-BAT-Softplus
---
Hessian-free
79.48 $\pm 0.18$ | 49.67 $\pm 0.21$ | 68.57 $\pm 0.27$ | 25.59 $\pm 0.15$ | 61.7
| Fast-BAT-Softplus
---
Hessian-aware
79.59 $\pm 0.21$ | 49.74 $\pm 0.12$ | 68.63 $\pm 0.23$ | 25.54 $\pm 0.19$ | 82.8
Fast-AT-Swish | 75.61 $\pm 0.15$ | 44.43 $\pm 0.18$ | 52.03 $\pm 4.29$ | 23.08 $\pm 2.23$ | 23.1
| Fast-BAT-Swish
---
Hessian-free
73.89 $\pm 0.14$ | 45.90 $\pm 0.23$ | 62.59 $\pm 0.29$ | 23.81 $\pm 0.17$ | 61.7
| Fast-BAT-Swish
---
Hessian-aware
73.93 $\pm 0.16$ | 45.97 $\pm 0.19$ | 62.49 $\pm 0.27$ | 23.99 $\pm 0.17$ | 82.6

##  Appendix E Convergence Analysis

Let us consider we have data samples
$\\{\mathbf{x}_{i},\mathbf{y}_{i}\\}_{i=1}^{N}$, where $N$ denotes the number
of the training data. Then the goal of Fast-BAT algorithm is to solve:

$\begin{split}&\min_{\bm{\theta}}L_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))=\frac{1}{N}\sum_{i=1}^{N}\ell_{\text{tr}}(\bm{\theta};\mathbf{x}_{i}+\bm{\delta}_{i}^{*}(\bm{\theta}),y_{i})\\\
\bm{\delta}^{*}(\bm{\theta})\in&\operatorname*{arg\,min}_{\bm{\delta}_{i}\in\mathcal{C}}L_{\text{atk}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))\frac{1}{N}\sum_{i=1}^{N}\ell_{\text{atk}}(\bm{\theta};\mathbf{x}_{i}+\bm{\delta}_{i}^{*}(\bm{\theta}),y_{i})\end{split}$
(39)

Let us denote the batch size as $b$. Then the problem can be reformulated as :

$\begin{split}&\min_{\bm{\theta}}L^{b_{t}}_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))=\frac{1}{b_{t}}\sum_{i=1}^{b_{i}}\ell_{\text{tr}}(\bm{\theta};\mathbf{x}_{i}+\bm{\delta}_{i}^{*}(\bm{\theta}),y_{i})\\\
\bm{\delta}^{*}(\bm{\theta})\in&\operatorname*{arg\,min}_{\bm{\delta}_{i}\in\mathcal{C}}L_{\text{atk}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))\underbrace{\frac{1}{b_{t}}\sum_{i=1}^{b_{t}}\ell_{\text{atk}}(\bm{\theta};\mathbf{x}_{i}+\bm{\delta}_{i}^{*}(\bm{\theta}),y_{i})}_{L_{\mathrm{atk}}^{b_{t}}(\bm{\theta},\bm{\delta}_{i}(\bm{\theta}))}\end{split}$
(40)

The way to computes the gradient (7) is equal to the case where $b=N$.

$\displaystyle\frac{dL^{N}_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))}{d\bm{\theta}}=\nabla_{\bm{\theta}}L^{N}_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))+\frac{{d}\bm{\delta}^{*}(\bm{\theta})^{\top}}{{d}\bm{\theta}}\nabla_{\bm{\delta}}L^{N}_{\mathrm{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))$
(41)

If we approximate the gradient
$\nabla_{\bm{\theta}}L_{\text{tr}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))}$
using a batch size of $b$, then we have the following assumptions:

  * •

Bias assumption:

$\displaystyle\mathbb{E}[\nabla_{\bm{\theta}}L^{b}_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))]=\nabla_{\bm{\theta}}L_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))+\beta(b),$
(42)

where $\beta(b)$ is the bias and the expectation is taken w.r.t. batches. Note
that $\beta(b)=0$ for $b=N$.

  * •

Variance assumption:

$\displaystyle\mathbb{E}\|\nabla_{\bm{\theta}}L^{b}_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))-\nabla_{\bm{\theta}}L_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))-\beta(b)\|^{2}=\sigma^{2}[1+\|\nabla_{\bm{\theta}}L_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))\|^{2}],$
(43)

Note that the variance equals 0 for $b=N$.

Now the Fast-BAT algorithm with batch size $b$ can be stated as follows:

  * •

Initialization: $\bm{\theta}_{0}\in\mathbb{R}^{n}$, step size
$\\{\alpha_{t}\\}_{t=0}^{T-1}$, batch size $\\{\alpha_{b}\\}_{t=0}^{T-1}$.

  * •

for $t=1$ to $T$:

    * –

Choose batsh size $b_{t}$ at each iteration $t\in[T]$.

    * –

$L_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))\approx
L^{b_{t}}_{\text{tr}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))$,
$L_{\text{atk}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))\approx
L^{b_{t}}_{\text{atk}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))$

    * –

Update:

$\displaystyle\bm{\theta}_{t}=\bm{\theta}_{t-1}-\alpha_{t-1}\nabla_{\bm{\theta}}L^{b_{t}}_{\text{tr}}(\bm{\theta}_{t-1},\bm{\delta}_{t-1}^{*}(\bm{\theta}_{t-1}))$
(44)

where

$\displaystyle\nabla_{\bm{\theta}}L^{b_{t}}_{\text{tr}}(\bm{\theta}_{t-1},\bm{\delta}_{t-1}^{*}(\bm{\theta}_{t-1}))=\nabla_{\bm{\theta}}L^{b_{t}}_{\mathrm{tr}}(\bm{\theta}_{t-1},\bm{\delta}^{*}(\bm{\theta}_{t-1}))+\nabla_{\bm{\theta}}\bm{\delta}^{*}(\bm{\theta}_{t-1})^{\top}\nabla_{\bm{\delta}}L^{b_{t}}_{\mathrm{tr}}(\bm{\theta}_{t-1},\bm{\delta}^{*}(\bm{\theta}_{t-1})).$
(45)

Note $\nabla_{\bm{\theta}}\bm{\delta}^{*}(\bm{\theta}_{t-1})$ is computed
using (34) by replacing:

      * *

$L_{\text{atk}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))$ with
$L^{b_{t}}_{\text{atk}}(\bm{\theta},\bm{\delta}^{*}(\bm{\theta}))$

      * *

$B_{0}$ with $B_{0}$ computed with (40).

Assuming smoothness of $L_{\mathrm{tr}(\bm{\theta})}$ w.r.t. $\bm{\theta}$, we
get:

$\displaystyle L_{\mathrm{tr}}(\bm{\theta}_{t+1})$ $\displaystyle\leq
L_{\mathrm{tr}}(\bm{\theta}_{t})+\left\langle\nabla_{\bm{\theta}}(\bm{\theta}_{t}),\bm{\theta}_{t+1}-\bm{\theta}_{t}\right\rangle+\frac{L}{2}\|\bm{\theta}_{t+1}-\bm{\theta}_{t}\|_{2}^{2}$
$\displaystyle=L_{\mathrm{tr}}(\bm{\theta}_{t})-\alpha_{t}\left\langle\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t}),\nabla_{\bm{\theta}}L^{b_{t}}_{\mathrm{tr}}(\bm{\theta}_{t})\right\rangle+\frac{\alpha_{t}^{2}L}{2}\|\nabla_{\bm{\theta}}L^{b_{t}}_{\mathrm{tr}}(\bm{\theta}_{t})\|^{2}_{2}$
(46)

Taking expectation w.r.t. random samples, we get:

$\displaystyle\mathbb{E}[L_{\mathrm{tr}}(\bm{\theta}_{t+1})]$
$\displaystyle\leq\mathbb{E}[L_{\mathrm{tr}}(\bm{\theta}_{t})]-\alpha_{t}\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|_{2}^{2}]-\alpha_{t}\mathbb{E}[\left\langle\nabla_{\bm{\theta}}(\bm{\theta}_{t}),\beta(b_{t})\right\rangle]+\alpha_{t}^{2}\sigma^{2}L\mathbb{E}[1+\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|^{2}_{2}]$
$\displaystyle\hskip
170.71652pt+2\alpha_{t}^{2}L\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|^{2}_{2}]+2\alpha_{t}^{2}L\mathbb{E}[\|\beta(b_{t})\|^{2}_{2}]$
$\displaystyle\leq\mathbb{E}\left[L_{\mathrm{tr}}(\bm{\theta}_{t})-\left[\frac{\alpha_{t}}{2}-\alpha_{t}^{2}\sigma^{2}L\right]\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|_{2}^{2}+\left[\frac{\alpha_{t}}{2}+2\alpha_{t}^{2}L\right]\|\beta(b_{t})\|^{2}_{2}+\alpha^{2}_{t}L\sigma^{2}\right]$

Using $\alpha_{t}\leq\frac{1}{4L(2+\sigma^{2})}$, we get:

$\displaystyle\mathbb{E}[L_{\mathrm{tr}}(\bm{\theta}_{t+1})]\leq\mathbb{E}\left[L_{\mathrm{tr}}(\bm{\theta}_{t})-\frac{\alpha_{t}}{4}\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|^{2}_{2}+\alpha_{t}\|\beta(b_{t})\|^{2}+\alpha_{t}^{2}L\sigma^{2}\right].$
(47)

Finally, after rearranging the terms, we get:

$\displaystyle\frac{\alpha_{t}}{4}\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|_{2}^{2}]\leq\mathbb{E}[(L_{\mathrm{tr}}(\bm{\theta}_{t})-L_{\mathrm{tr}}(\bm{\theta}_{t+1}))+\alpha_{t}\|\beta(b_{t})\|^{2}+\alpha_{t}^{2}L\sigma^{2}]$
(48)

Summing over $t=1$ to $T$ and multiplying by $1/T$, we get:

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\frac{\alpha_{t}}{4}\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|^{2}_{2}]\leq\frac{L_{\mathrm{tr}}(\bm{\theta}_{0})-L_{\mathrm{tr}}(\bm{\theta}^{*})}{T}+\frac{1}{T}\sum_{t=1}^{T}\alpha_{t}\|\beta(b_{t})\|^{2}+\frac{1}{T}\sum_{t=1}^{T}\alpha_{t}^{2}\sigma_{t}^{2}.$
(49)

Taking $\alpha_{t}=\alpha$ and $b_{t}=b$, we get:

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|_{2}^{2}]\leq\frac{4[L_{\mathrm{tr}}(\bm{\theta}_{0})-L_{\mathrm{tr}}(\bm{\theta}^{*})]}{\alpha
T}+4\|\beta(b)\|^{2}+4\alpha\sigma^{2}.$ (50)

Recall if $b=N$, we have $\sigma^{2}=0$ and $\beta(b)=0$, we can further get:

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|_{2}^{2}\leq\frac{4[L_{\mathrm{tr}}(\bm{\theta}_{0})-L_{\mathrm{tr}}(\bm{\theta}^{*})]}{\alpha
T}.$ (51)

Case I: We can choose $\alpha=\frac{1}{8L}$ and get:

$\displaystyle\frac{1}{T}\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t})\|_{2}^{2}=\mathcal{O}(\frac{1}{T}).$
(52)

Case II: We can choose $\alpha=\sqrt{\frac{1}{T}}$:

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t}\|^{2}_{2})]\leq\frac{4[L_{\mathrm{tr}}(\bm{\theta}_{0})-L_{\mathrm{tr}}(\bm{\theta}^{*})]}{\sqrt{T}}+\underbrace{\frac{4\sigma^{2}}{\sqrt{T}}}_{\mathcal{O}(\frac{1}{\sqrt{T}})}+\underbrace{4\|\beta(b)\|^{2}}_{\mathcal{O}(1)}$

Now if $\|\beta(b)\|^{2}=\mathcal{O}(\frac{1}{\sqrt{T}})$, then

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t}\|^{2}_{2}]\leq\mathcal{O}(\frac{1}{\sqrt{T}}),$
(53)

otherwise we get in the worst case:

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}[\|\nabla_{\bm{\theta}}L_{\mathrm{tr}}(\bm{\theta}_{t}\|^{2}_{2}]\leq\mathcal{O}(\frac{1}{\sqrt{T}})+\mathcal{O}(1).$
(54)