# Machine-learning Kohn-Sham potential from dynamics in time-dependent Kohn-
Sham systems

Jun Yang  James Whitfield  Department of Physics and Astronomy, Dartmouth
College. Hanover, NH 03755

###### Abstract

A machine learning method to develop the energy functional and the Kohn-Sham
potential of a time-dependent Kohn-Sham(TDKS) system is proposed. The method
is based on the dynamics of the Kohn-Sham system, so no data of the exact
Kohn-Sham potential is required for training the model. We demonstrate the
results of our method with a 1d harmonic oscillator example and a 1d two-
electron example. We show that the machine-learned Kohn-Sham potential matches
the exact Kohn-Sham potential in the absence of memory effect. In the case
where there is a memory effect, our method can still capture the dynamics of
the Kohn-Sham system. The machine learning method developed in this article
provides an insight into making a better approximation to the energy
functional and the Kohn-Sham potential in the TDKS system.

##  I Introduction

One of the biggest challenges in time-dependent density functional
theory(TDDFT) is to find a good approximation to the Kohn-Sham potential
$v_{KS}(\bm{r},t)$ in the TDKS system. The Runge-Gross theoremRunge and Gross
(1984) guarantees that the Kohn-Sham potential is uniquely determined(up to a
purely time-dependent function) by the time-dependent electronic density
$n(\bm{r},t)$ for a given initial many-body state $\Psi(\bm{r},t_{0})$ and
initial Kohn-Sham orbitals $\Phi(\bm{r},t_{0})$. The associated Kohn-Sham
potential is often written as a sum of three terms
$v_{KS}[n](\bm{r},t)=v_{ext}(\bm{r},t)+v_{H}[n](\bm{r},t)+v_{xc}(\bm{r},t)$,
where the first term is the external potential, the second term is the Hartree
potential describing the classical electron-electron interaction. The first
two terms are numerically easy, which can be calculated explicitly given the
electronic density. While the last term, exchange-correlation(xc) potential,
involving all the complex and non-trivial effects, is hard to be calculated
numerically. So the real blocker of finding a good approximation to the Kohn-
Sham potential is finding a good approximation to the xc-potential.

Many approximation methods to the xc-potential have been developed with the
development of TDDFT, commonly used methods include local density
approximation(LDA), generalized gradient approximation(GGA), and meta-
generalized gradient approximation(meta-GGA)Kohn and Sham (1965); Perdew and
Burke (1996); Perdew  _et al._ (1999), etc. Besides these traditional explicit
constructions, recent works have provided approximation methods via machine
learning techniquesNagai  _et al._ (2018); Suzuki  _et al._ (2020), both works
require the exact xc-potential as the training set. So a natural question is,
can we develop a machine-learning-based method to approximate the energy
functional, thus the Kohn-Sham potential and the xc-potential, with access to
the data of Kohn-Sham orbitals only?

A similar question in classical mechanics has been extensively investigated by
researchers from the artificial intelligence community. Several different
types of neural networks have been developed to learn the potential from the
classical trajectories. The well-trained model can be further utilized to
predict the trajectories that are not contained in the training set. Most of
the neural networks in learning classical dynamics can be classified into two
categories, the Hamiltonian neural network(HNN)Greydanus  _et al._ (2019);
Tong  _et al._ (2021); Bertalan  _et al._ (2019); Han  _et al._ (2021); Chen
_et al._ and the Lagrangian neural network(LNN)Cranmer  _et al._ . HNN
incorporates Hamilton’s equations into the neural network. LNN leverages the
idea from the Lagrangian equation. Though promising results have been achieved
in learning classical mechanics, this method has not been generalized to the
quantum regime.

Our work generalizes the idea of HNN to the quantum regime. We propose a
machine-learning method to investigate the Kohn-Sham system under adiabatic
approximation, and give a positive answer to the question we put forward
above. In section II, we establish the connection between the Kohn-Sham
equations and Hamilton’s equations. The connection lies as the cornerstone for
the construction of our neural network. In section III, we present our results
with a harmonic oscillator example and a two-electron example. Conclusions and
summaries are offered in section IV.

##  II Kohn-Sham equations and Hamilton’s equations

Before discussing the Kohn-Sham equations, we will first go through the
relation between the time-dependent Schrodinger equation and Hamilton’s
equations. In fact, the Kohn-Sham equations in TDDFT are time-dependent
Schrödinger equations describing non-interacting particles. The requirement
for the Kohn-Sham equations in TDDFT is that the density reproduced from the
non-interacting system should be the same as the density of the original
interacting systemRunge and Gross (1984); van Leeuwen (1998).

###  II.1 Schrödinger equation and Hamilton’s equations

The Schrödinger equation reads (in atomic units throughout the article)

$\displaystyle i\frac{\partial}{\partial t}\ket{\psi(t)}=\hat{H}\ket{\psi(t)}$
(1)

The linear combination of basis states $\ket{\psi(t)}=\sum_{k}c_{k}(t)\ket{k}$
of the given state $\ket{\psi(t)}$ allows us to rewrite the Schrödinger
equation into the form of a vector differential equation

$\displaystyle i\dot{\bm{c}}(t)=\bm{Hc}(t)$ (2)

where $\bm{c}(t)=[c_{1}(t),\ldots,c_{n}(t),\ldots]^{T}$ is the coefficients
vector, $\bm{H}$ is the matrix representation of the Hamiltonian $\hat{H}$
with the matrix elements being $H_{ij}=\braket{i}{\hat{H}}{j}$. To uncover the
relation between the Schrödinger equation and Hamilton’s equations, we define
an energy functional $E[\bm{c}]=\bm{\bar{c}^{T}Hc}$ and separate the
coefficients vector into a real part and a imaginary part
$\bm{c}(t)=\frac{1}{\sqrt{2}}(\bm{q}(t)+i\bm{p}(t))$, where
$\bm{q}(t)=\sqrt{2}\mathrm{Re}\bm{c}(t),\bm{p}(t)=\sqrt{2}\mathrm{Im}\bm{c}(t)$.
With the definition of the energy functional, we obtain the following
equations with respect to the real and imaginary part of the coefficients (see
Appendix A for the details of the proof)

$\displaystyle\frac{\partial E[\bm{c}(t)]}{\partial
p_{k}(t)}=\dot{q}_{k}(t),\frac{\partial E[\bm{c}(t)]}{\partial
q_{k}(t)}=-\dot{p}_{k}(t),k=1,2,\ldots$ (3)

The form of these equations are identical to that of the Hamilton’s equation
in classical mechanics. For simplicity, we can also write Eq. 3 in a compact
form

$\displaystyle i\dot{\bm{c}}(t)=\frac{\partial
E[\bm{c}(t)]}{\partial\bar{\bm{c}}(t)}$ (4)

This equation connects the Schrödinger equation and Hamilton’s equation. In
the next section, we will see how Kohn-Sham equations are connected to
Hamilton’s equations with a similar treatment.

###  II.2 Kohn-Sham equations and Hamilton’s equation

To derive similar equations in the time-dependent Kohn-Sham system, we need to
go a bit beyond the Schrödinger equation in the above section and design a
different energy functional expression. To simplify things, we assume the
adiabatic approximation $v_{xc}[n](\bm{r},t)=v_{xc}[n(t)](\bm{r})$ in the
whole articleMarques and Gross (2004); Casida and Huix-Rotllant (2012); Li
_et al._ (2020), which is commonly used in TDDFT calculation. The adiabatic
approximation means the exchange-correlation potential depends on the
instantaneous electron density $n(t)$ only. There are two reasons for using
the adiabatic approximation. One is that the energy functional under adiabatic
approximation is a functional of the instantaneous electron density, as a
consequence the xc-potential can be calculated from the xc-energy functional
by $v_{xc}[n(t)](\bm{r})=\frac{\delta E_{xc}[n(t)]}{\delta n(t)}$. The other
consideration is that it is more convenient to construct a neural network that
maps the instantaneous electron density to the energy functional.

Although the adiabatic approximation is widely used in TDDFT calculation, it
often underestimates the complexity of the xc-potential. In fact, the exact
xc-potential at time $t$ is often dependent on the electron density at all
previous times, leading to the so-called memory effectLiao  _et al._ (2017);
Maitra and Burke (2001); Maitra  _et al._ (2002). It is worth noting that such
a complicated dependence on the history of the electron density is neglected
by the adiabatic approximation. Thus for a system where the memory effect is
not negligible, our machine learning method may not provide the exact Kohn-
Sham potential. In the results part of this article, we will confirm it with
examples.

With the adiabatic approximation, we can now discuss the relation between the
Kohn-Sham equations and the Hamilton’s equations. Notice that the Kohn-Sham
equations in TDDFT are time-dependent Schrödinger equations describing the
non-interacting particles,

$\displaystyle i\frac{\partial}{\partial t}\ket{\phi_{m}(t)}=$
$\displaystyle(-\frac{\nabla^{2}}{2}+v_{ext}(\bm{r},t)+v_{H}[n(t)](\bm{r})$
(5) $\displaystyle+v_{xc}[n(t)](\bm{r}))\ket{\phi_{m}(t)},m=1,2,\ldots$

where $m$ labels the $m$-th Kohn-Sham orbital. The corresponding energy
functional is given by

$\displaystyle E[n(t)]=$
$\displaystyle-\sum_{i}\braket{\phi_{i}(t)}{\frac{\nabla^{2}}{2}}{\phi_{i}(t)}$
(6) $\displaystyle+E_{H}[n(t)]+E_{ext}[n(t)]+E_{xc}[n(t)]$

Similar to the previous discussion, it’s possible to rewrite the Kohn-Sham
equations into vector differential equations in terms of the coefficients of
the Kohn-Sham orbitals.

$\displaystyle i\dot{\bm{c}}_{m}(t)=\bm{H}[\bm{c}]\bm{c}_{m}(t)$ (7)

where $\bm{c}_{m}(t)$’s are the coefficients of the $m$-th Kohn-Sham orbital,
and the matrix element of $\bm{H}[\bm{c}(t)]$ is given by
$\braket{i}{\hat{H}[n(t)]}{j}$.

Similarly, we can show that the Hamilton’s equations hold in Kohn-Sham system
under the adiabatic approximation as well, the details of the proof is shown
in Appendix B.

$\displaystyle i\dot{\bm{c}}_{m}(t)=\frac{\partial
E[n(t)]}{\partial\bm{\bar{c}}_{m}(t)},m=1,2,\ldots$ (8)

Eq. 8 along with Eq. 3 serves as the cornerstone for our machine learning
method, of which the details will be discussed in the next section.

###  II.3 Machine learning energy functional

From this part to the end, the discussion is limited to the 1-dimensional
system. In this section, we’ll demonstrate how we incorporate the equations
derived above into the neural networks.

Before discussing the technical details, we need to prepare the data set. We
pick $n$ basis functions, in our work, the sinc discrete variable
representation(DVR) basis $\braket{x}{k}=\sqrt{\Delta
x}\frac{\mathrm{sinc}\left[\frac{\pi(x-k\Delta x)}{\Delta
x}\right]}{\pi(x-k\Delta x)}$ is usedColbert and Miller (1992); Brown  _et
al._ (2020), where $\Delta x$ is the grid’s spacing. Then the coefficient
vectors calculated at $N$ different timestamp corresponding to the basis
functions are chosen as our data set.

The structure of our neural network is shown in Fig. 1. The input layer of the
neural network has 2$n$ entries, the first $n$ entries are for the scaled real
part of the coefficients at time $t$, the last $n$ entries are for the scaled
imaginary part of the coefficients. There are two hidden layers. There is one
scalar as the output layer representing the energy functional of the system in
our neural network. Let $\bm{y}^{(i)}$ be the variable vector of the $i$-th
layer, the $i$-th and the $i+1$-th layer are connected by

$\displaystyle\bm{y}^{(i+1)}=\Phi(\bm{W}^{(i)}\bm{y}^{(i)}+\bm{b}^{(i)})$ (9)

where $\Phi(\cdot)$ is the non-linear activation function. In the numerical
experiments we carried out, $\Phi(x)=\tanh{x}$ and
$\Phi(x)=\mathrm{softplus}(x)$ were used.

[width=0.7]tikz/nn

Figure 1: The structure of our neural network.
$q_{i}=\sqrt{2}\mathrm{Re}c_{i}$’s are the scaled real part of the
coefficients, $p_{i}=\sqrt{2}\mathrm{Im}c_{i}$’s are the scaled imaginary part
of the coefficients. The scaled real and imaginary parts of the coefficients
at time $t$ are fed into the neural network as our inputs. There are two
hidden layers. Neurons are connected by non-linear activation functions (in
the numerical experiments we carried out, softplus and tanh were chosen). The
output layer is the energy functional of the Kohn-Sham system. The loss
function relies on the gradient (shown in the last layer in the figure) of the
energy functional, we compute it with auto differentiation functionality in
Pytorch.

The loss function is defined as

$\displaystyle\mathcal{L}=\sum_{i=1}^{n}\left(\left|\dot{q}_{i}-\frac{\partial
E}{\partial p_{i}}\right|^{2}+\left|\dot{p}_{i}+\frac{\partial E}{\partial
q_{i}}\right|^{2}\right)/n$ (10)

where the $\dot{q}_{i}$’s and the $\dot{p}_{i}$’s are obtained from the time
derivative of coefficients in the data set, $\frac{\partial H}{\partial
q_{i}}$’s and $\frac{\partial H}{\partial p_{i}}$’s are from the gradient of
the energy functional calculated using the auto-differentiation functionality
in modern machine learning frameworks. In our work, we chose PytorchPaszke
_et al._ (2019). The form of the loss function originates from the Hamilton’s
equations, which was originally proposed in Ref. De Wilde (1993) and Ref.
Greydanus  _et al._ (2019). The neural network is optimized using adaptive
moment estimation method(Adam)Kingma and Ba . By minimizing the loss function,
we can get the energy functional corresponding to the Kohn-Sham system.

We need to take a further step to calculate the Kohn-Sham potential from the
energy functional in the system. Since the sinc DVR is used in this article,
the corresponding matrix element of the potential term $V_{KS}^{ij}$ is given
by:

$\displaystyle
V_{KS}^{ij}=V_{KS}(x_{i})\delta^{ij}=\frac{1}{2}\left(\frac{\partial^{2}E[\bm{q},\bm{p}]}{\partial
q_{i}^{2}}+\frac{\partial^{2}E[\bm{q},\bm{p}]}{\partial
p_{i}^{2}}\right)-T^{ij}$ (11)

where $T^{ij}$ is the matrix element of the kinetic energy term whose value
isColbert and Miller (1992); Brown  _et al._ (2020)

$\displaystyle T^{ij}=\frac{(-1)^{(i-j)}}{2\Delta
x^{2}}\begin{cases}\frac{\pi^{2}}{3},&i=j\\\ \frac{2}{(i-j)^{2}},&i\neq
j\end{cases}$ (12)

The details of the proof of Eq. 11 are shown in Appendix. B.

##  III Results

###  III.1 One electron test

The first example we have chosen is the 1-dimensional harmonic oscillator. The
potential is time-independent $v(x)=\frac{1}{2}\omega^{2}x^{2}$, where
$\omega=1$. The eigenstates and their time derivatives of the Harmonic
oscillator have analytical expressions

$\displaystyle\phi_{n}(x,t)={\frac{1}{\sqrt{2^{n}\,n!}}}\left({\frac{1}{\pi}}\right)^{1/4}e^{-{\frac{x^{2}}{2}}-i(n+\frac{1}{2})t}H_{n}\left(x\right),$
(13)
$\displaystyle\dot{\phi}_{n}(x,t)={\frac{-i(n+\frac{1}{2})}{\sqrt{2^{n}\,n!}}}\left({\frac{1}{\pi}}\right)^{1/4}e^{-{\frac{x^{2}}{2}}-i(n+\frac{1}{2})t}H_{n}\left(x\right),$
(14) $\displaystyle n=0,1,2,\ldots,M$

where $H_{n}(x)$ is the $n$-th Hermite polynomial. We obtain the data set from
the lowest $M$ eigenstates. For each eigenstate in the data set, 150 evenly
spaced grid points in range $x\in[-6,6]$ are chosen with 2000 timestamps in
$t\in[0,4\pi]$. So the input dimension of the neural network is 300 (twice the
number of grid points), and the dimension of each hidden layer is 400. $M=15$
in this numerical experiment.

To verify the trained neural network produces the correct dynamics of the
system, we randomly pick a initial state that is not in the training set and
let it propagate under the given Hamiltonian. In the example given below, the
initial state is a linear combination of the lowest four eigenstates
$\ket{\psi(x,t=0)}=\frac{1}{\sqrt{2}}\ket{0}+\frac{1}{2}\ket{1}+\frac{1}{\sqrt{6}}\ket{2}+\frac{1}{\sqrt{12}}\ket{3}$.
Then we compare the machine-learned time-dependent density to the exact
density. To be precise, at time $t$, if the coefficients of the Kohn-Sham
orbitals are known, we can leverage them to compute the gradient of the energy
functional $\frac{\partial E}{\partial q_{i}}$ and $\frac{\partial E}{\partial
p_{i}}$, which are used to propagate the coefficients to the next timestamp
$t+dt$ with the Runge-Kutta (R-K) methodRunge (1895). In this sense, our
neural network can be viewed as a differential equation solver. The entire
pipeline of propagating the initial state with our neural network is shown in

[width=]tikz/flowchart

Figure 2: Pipeline of propagating the initial state with our neural network

The comparison is shown in Fig. 3. The machine learned density(black dashed
line) exhibits similar trend to the exact result at any point in the entire
position space at early timestamps, e.g. $t=1.0(a.u.)$. The error accumulates
with time. As shown in Fig. 3, the error grows larger at time $t=4.0(a.u.)$,
but still captures the general feature of the exact density. In the third plot
of Fig. 3, the results are completely different at longer time $t=15.0(a.u.)$.

Figure 3: The machine learned density and the exact density at different
timestamps. The black dashed line is the machine learned density, the red
solid line is the exact density. The three panels correspond to the result
result at timestamp $t=1.0(a.u.)$ (left), $t=4.0(a.u.)$ (middle) and
$t=15.0(a.u.)$ (right).

We use mean square error(MSE) to quantify the error in the density at each
timestamp. The result is shown in Fig. 4. It is clear from the figure that the
MSE grows in time. We notice the error in this example shows a periodic
pattern, this is due to the periodicity of the system. To avoid the error from
going too big, one can set a threshold and make corrections to the Kohn-Sham
orbitals when the error is larger than the threshold.

Figure 4: MSE in density accumulates with time. Periodic pattern in local
maximum and local minimum due to the periodicity of the system.

The machine-learned energy levels and the machine-learned Kohn-Sham potential
are also investigated in the study. We present the machine-learned energy
level in Fig. 5(a). Each red line corresponds to one single eigenstate at 2000
different timestamps. As shown in the figure, the energy is not changing with
time for a fixed initial state. Another interesting feature of our neural
network is that it can capture the correct level spacing $E_{n+1}-E_{n}=1$.

(a)

(b)

Figure 5: (a) is the figure of energy level. We manually shift the ground
state energy to 0 for comparison. The neural network generates the correct
level spacing for the harmonic oscillator test. (b) the machine learned Kohn-
Sham potential(black dots) versus the actual potential(red line).

Fig. 5(b) shows the machine-learned Kohn-Sham potential(black dots) versus the
exact potential(red line). The machine-learned Kohn-Sham potential captures
the structure of the exact quadratic potential around the center of the
position space, but it starts to deviate from the actual potential near the
boundary of the pre-selected finite region. One possible reason for this
phenomenon is that our training set only covers a fraction of the entire
Hilbert space, which means (i) the training set covers only a finite range of
the infinite position space, (ii) the training set includes only the lowest
$M$ eigenstates of the infinite many eigenstates. We show in Appendix C that
having more eigenstates in the training set results in more accurate machine-
learned Kohn-Sham potential.

###  III.2 Two electron test

We revisit the two-electron system previously studied inLacombe  _et al._
(2018); Suzuki  _et al._ (2017, 2020), and use the machine learning method to
study the Kohn-Sham potential in the system. The many-body Hamiltonian
governing the dynamics of the two-electron system is given by

$\displaystyle\hat{H}\left(x_{1},x_{2}\right)=\sum_{i=1,2}\left[-\frac{1}{2}\frac{\partial^{2}}{\partial
x_{i}^{2}}+v_{\mathrm{ext}}\left(x_{i}\right)\right]+W_{ee}\left(x_{1},x_{2}\right)$
(15)

where
$W_{ee}\left(x_{1},x_{2}\right)=\frac{1}{\sqrt{\left(x_{1}-x_{2}\right)^{2}+1}},v_{\mathrm{ext}}=-\frac{1}{\sqrt{\left(x_{1}+10\right)^{2}+1}}$.
The spatial part of the initial two-electron wavefunction is
$\Psi(x_{1},x_{2})=\frac{1}{\sqrt{2}}\left[\phi_{\mathrm{H}}\left(x_{1}\right)\phi_{\mathrm{WP}}\left(x_{2}\right)+\phi_{\mathrm{WP}}\left(x_{1}\right)\phi_{\mathrm{H}}\left(x_{2}\right)\right]$,
where $\phi_{H}(x)$ is the ground state hydrogen wavefunction,
$\phi_{WP}(x)=(2\alpha/\pi)^{\frac{1}{4}}e^{\left[-\alpha\left(x-x_{0}\right)^{2}+ip\left(x-x_{0}\right)\right]}$
with $\alpha=0.1,x_{0}=10.0,p=-1.5$. So we have one Gaussian wave packet
centered at $x_{0}=10.0$ moving leftward. As the system evolves, the Gaussian
wave packet collides with the static Hydrogen wave function, and then passes
through the Hydrogen wave function at longer times.

The two electron wavefunction $\Psi(x_{1},x_{2},t)$ can be computed
numerically. Based on this, we can calculate the electronic density
$n(x,t)=2\int dx_{2}\Psi^{\star}(x,x_{2},t)\Psi(x,x_{2},t)$ and the density
current $j(x,t)=2\mathrm{Im}\int
dx_{2}\Psi^{\star}(x,x_{2},t)\partial_{x}\Psi(x,x_{2},t)$. The two Kohn-Sham
orbitals are computed as
$\phi_{1}(x,t)=\phi_{2}(x,t)=\sqrt{\frac{n(x,t)}{2}}e^{i\int^{x}\frac{j\left(x^{\prime},t\right)}{n\left(x^{\prime}t\right)}dx^{\prime}}$.
Coefficients of the Kohn-Sham orbital are then obtained as our data set.

Similar to the discussion in the one electron test part, the system is
discretized evenly in the range $x\in[-24.78,21.97]$ with 200 grid points. The
propagation starts at $t=0$ to $t=15.0(a.u.)$ with time step length $\Delta
t=5\times 10^{-4}(a.u.)$. The input dimension of the neural network for the
two electron model is thus 400, and the dimension of each hidden layer is 400.

We show the time evolution of the density generated from the neural network in
Fig. 6. The time evolution of the machine-learned density(black dashed line)
captures the feature of the exact evolution(red solid line). As mentioned in
the previous section, the error accumulates over steps, so we calibrate the
initial condition every 500 time steps. We notice that the machine-learned
density overlaps well with the exact density at $t=5.0(a.u.)$ and
$t=13.0(a.u.)$, but shows a deviation at $t=8.5(a.u.)$ when the collision
between the two wave packets starts to happen.

Figure 6: The machine learned density and the exact density at different
timestamps in the two electron test. The black dashed line is the machine
learned density, the red solid line is the exact density. The three panels
correspond to the result at timestamp $t=5.0(a.u.)$ (left), $t=8.5(a.u.)$
(middle) and $t=13.0(a.u.)$ (right).

We investigated how the error changes in time. The error is shown in Fig. 7.
The quantified error shows a spike at $t\approx 8.5(a.u.)$, which is
consistent with the our observations mentioned above, the error is larger at
around $t=8.5(a.u.)$ when the collision between the two wave packets starts to
happen.

Figure 7: MSE in density accumulates with time. A spike appears at $t\approx
8.5(a.u.).$

Although the neural network reproduces the correct dynamics of the system, we
need to mention that the machine-learned Kohn-Sham potential is not the same
as the actual one. We show the comparison in Fig. 8. One note here is that we
don’t compare the Kohn-Sham potential directly, instead, we compare the
exchange-correlation part of the Kohn-Sham potential. The machine-learned xc
potential(black dots) does not capture the features of the actual xc(red solid
line) potential. Similar phenomena were observed in Ref. Bhat  _et al._ ,
where the dynamics of the same system were reproduced correctly by adjoint
learning, but they mentioned that the machine learned xc potential did not
correspond to the actual one. The adiabatic local density approximation(ALDA)
xc potential(blue dash line) is also shown for comparison. Due to the lack of
memory effect, the ALDA xc potential does not capture the exact dynamics of
the two-electron system.

In our machine learning method, we attempted to learn the correct dynamics
under the adiabatic approximation. So the machine-learned xc potential is
expected to be different from both the exact xc potential and the ALDA xc
potential, but to remain the features of both potentials. The results shown in
Fig. 8 can show some evidences. In Fig. 8(a) and Fig. 8(b), the machine-
learned xc potential shows a peak at $x=-10$, which is in align with the ALDA
xc potential. The trend of machine-learned xc potential is similar to the
exact xc potential in the region of $x>0$.

(a)

(b)

(c)

Figure 8: The machine learned Kohn-Sham potential, the exact Kohn-Sham
potential and the ALDA potential at different timestamps in the two electron
test. The black dashed line is the machine learned density, the red solid line
is the exact density, and the blue dashed line is the ALDA potential. The
three panels correspond to the result at timestamp $t=5.0(a.u.)$ (left),
$t=8.5(a.u.)$ (middle) and $t=13.0(a.u.)$ (right).

We analyzed the entire procedure, the most important reason for the
discrepancy between the machine-learned Kohn-Sham potential and the actual one
is the memory effect. It has been proven that the Kohn-Sham potential at time
$t$ in this two-electron system depends on all densities at $t^{\prime}<t$ and
the choice of the initial Kohn-Sham orbitalsSuzuki  _et al._ (2017, 2020);
Lacombe  _et al._ (2018); Maitra  _et al._ (2002). In our work, we assume
adiabatic approximation, so the dynamics and the Kohn-Sham potential of the
system cannot be reproduced correctly at the same time.

##  IV Conclusions

In this article, we have described a new machine learning method to learn the
dynamics and the Kohn-Sham potential of the Kohn-Sham system. We have
demonstrated the method with two one-dimensional examples: a one-electron
harmonic oscillator model and a two-electron soft Coulomb model. In both
examples, the exact dynamics of the systems could be well reproduced from the
machine learning method. The machine-learned Kohn-Sham potential in the one-
electron test captures the general feature of the actual Kohn-Sham potential,
but it shows a discrepancy from the actual one in the two-electron test. We
have analyzed the possible reasons, the memory effect is the major source of
error. The memory effect requires considering the densities at previous
timestamps. To overcome this difficulty, we believe the neural networks
capable of handling time series(e.g. RNN, LSTM) are promisingSherstinsky ;
Hochreiter and Schmidhuber (1997).

##  V Acknowledgements

JY and JDW were supported by the U.S. Department of Energy, Office of Science,
Office of Advanced Scientific Computing Research, under the Quantum Computing
Application Teams program (Award 1979657). JDW was also supported by the NSF
(Grant 1820747) and additional funding from the DOE (Award A053685).

## Appendix A

In this part, we show the relation between the Schrödinger equation and the
Hamilton’s equations. For convenience, we want to write the Hamilton’s
equation into a compact form in terms of complex variables. The normal
expressions of the Hamilton’s equations are

$\displaystyle\frac{\partial H}{\partial p}=\dot{q}$ (16)
$\displaystyle\frac{\partial H}{\partial q}=-\dot{p}$ (17)

By rewriting $z=\frac{1}{\sqrt{2}}(q+ip)$ and
$\bar{z}=\frac{1}{\sqrt{2}}(q-ip)$, we can transform the derivatives into the
following form:

$\displaystyle\frac{\partial}{\partial z}=\frac{\partial}{\partial
q}\frac{\partial q}{\partial z}+\frac{\partial}{\partial p}\frac{\partial
p}{\partial z}=\frac{1}{\sqrt{2}}(\frac{\partial}{\partial
q}-i\frac{\partial}{\partial p})$ (18)
$\displaystyle\frac{\partial}{\partial\bar{z}}=\frac{\partial}{\partial
q}\frac{\partial q}{\partial\bar{z}}+\frac{\partial}{\partial p}\frac{\partial
p}{\partial\bar{z}}=\frac{1}{\sqrt{2}}(\frac{\partial}{\partial
q}+i\frac{\partial}{\partial p})$ (19)

Thus we can write the Hamilton’s equations into a tighter form,

$\displaystyle i\dot{z}=\frac{\partial H}{\partial\bar{z}}$ (20)

Now let’s look at the Schrödinger’s equation:

$\displaystyle i\frac{d}{dt}\ket{\Psi(t)}=\hat{H}\ket{\Psi(t)}$ (21)

Given the basis states $\ket{k},k=0,1,\ldots$, the wave function can be
expanded as

$\displaystyle\ket{\Psi(t)}=\sum_{k}c_{k}(t)\ket{k}$ (22)

Thus we obtain the differential equation of the coefficients vector
$\bm{c}(t)=[c_{1}(t),c_{2}(t),\cdots,c_{n}(t)]^{T}$

$\displaystyle i\dot{\bm{c}}=\bm{Hc}$ (23)

where $\bm{H}$ is the matrix representation of the Hamiltonian $\hat{H}$ with
its entries being $H_{ij}=\braket{i}{\hat{H}}{j}$.

By defining $H(\bm{c})=\bm{\bar{c}Hc}$, the Schrödinger’s equation can be
transformed as

$\displaystyle i\dot{\bm{c}}=\frac{\partial H(\bm{c})}{\partial\bar{\bm{c}}}$
(24)

which is in consistent with the classical Hamilton’s equation.

## Appendix B

We prove the Hamilton’s equations in Kohn-Sham system in this part. The Kohn-
Sham equations are given by

$\displaystyle i\dot{\phi}_{m}(t)=\hat{H}[n(t)]\phi_{m}(t),m=1,\cdots,N$ (25)

where $N$ is the number of electrons, and
$n(t)=\sum_{m=1}^{N}|\phi_{m}(x,t)|^{2}$.

The K-S Hamiltonian can be split into 4 parts,

$\displaystyle\hat{H}=\hat{T}_{s}+\hat{V}_{H}+\hat{V}_{ext}+\hat{V}_{xc}$ (26)

We can rewrite the Kohn-Sham equations as (suppose the basis state set
$\\{\ket{\varphi_{i}}\\}$ is given)

$\displaystyle i\dot{\bm{c}}_{m}=\bm{H}[\bm{c}]\bm{c}_{m}$ (27)

where the matrix element at $i$-th row $j$-th column is
$(\bm{H}[\bm{c}])_{ij}=\braket{\varphi_{i}}{\hat{T}_{s}+\hat{V}_{H}+\hat{V}_{ext}+\hat{V}_{xc}}{\varphi_{j}}$

Under the adiabatic approximation, the energy functional corresponding to the
system is $E[n]=T_{S}[n]+E_{ext}[n]+E_{H}[n]+E_{xc}[n]$.

Since the electronic density explicitly depends on the coefficients vector,

$\displaystyle
n(\bm{c})=\sum_{m,i,j}\bar{c}_{im}c_{jm}\bar{\varphi}_{i}(x)\varphi_{j}(x)$
(28)

we can express the energy functional in terms of coefficients vector $\bm{c}$

$\displaystyle
E(\bm{c})=T_{S}(\bm{c})+E_{ext}(\bm{c})+E_{H}(\bm{c})+E_{xc}(\bm{c})$ (29)

We anaylize the four terms one by one:

  1. 1.

Kinetic energy:

$\displaystyle
T_{s}[n]=\sum_{m}\sum_{i,j}\braket{\varphi_{i}}{\hat{T}_{s}}{\varphi_{j}}c_{jm}\bar{c}_{im}$
(30) $\displaystyle\frac{\partial
T_{s}[n]}{\partial\bar{c}_{im}}=\sum_{j}\braket{\varphi_{i}}{\hat{T}_{s}}{\varphi_{j}}c_{jm}$
(31)

  2. 2.

External energy:

$\displaystyle
E_{ext}[n]=\sum_{m}\sum_{i,j}\braket{\varphi_{i}}{\hat{V}_{ext}}{\varphi_{j}}c_{jm}\bar{c}_{im}$
(32) $\displaystyle\frac{\partial
E_{ext}[n]}{\partial\bar{c}_{im}}=\sum_{j}\braket{\varphi_{i}}{\hat{V}_{ext}}{\varphi_{j}}c_{jm}$
(33)

  3. 3.

Hartree energy:

$\displaystyle\frac{\partial E_{H}[n]}{\partial\bar{c}_{im}}$
$\displaystyle=\int\mathrm{d}x\int\mathrm{d}y\frac{\delta E_{H}[n]}{\delta
n(\bm{c})}(y)\delta(x-y)\frac{\partial n(\bm{c})}{\partial\bar{c}_{im}}(y)$
$\displaystyle=\sum_{j}\int\mathrm{d}x\
v_{H}(x)c_{jm}\bar{\varphi}_{i}(x)\varphi_{j}(x)=\sum_{j}\braket{\varphi_{i}}{\hat{V}_{H}}{\varphi_{j}}c_{jm}$
(34)

  4. 4.

Exchange-correlation energy:

$\displaystyle\frac{\partial E_{xc}[n]}{\partial\bar{c}_{im}}$
$\displaystyle=\int\mathrm{d}x\int\mathrm{d}y\frac{\delta E_{xc}[n]}{\delta
n(\bm{c})}(y)\delta(x-y)\frac{\partial n(\bm{c})}{\partial\bar{c}_{im}}(y)$
$\displaystyle=\sum_{j}\int\mathrm{d}x\
v_{xc}(x)c_{jm}\bar{\varphi}_{i}(x)\varphi_{j}(x)=\sum_{j}\braket{\varphi_{i}}{\hat{V}_{xc}}{\varphi_{j}}c_{jm}$
(35)

Putting together 3 terms, we have

$\displaystyle\frac{\partial
E(\bm{c})}{\partial\bm{\bar{c}}_{m}}=\bm{H}[\bm{c}]\bm{c}_{m}=i\dot{\bm{c}}_{m}$
(36)

This concludes the proof of Eq. 8.

We can take a further step to show Eq. 11
$\frac{\partial^{2}E(\bm{c})}{\partial\bar{c}_{mi}\partial\bar{c}_{mj}}-\braket{\varphi_{i}}{\hat{T}}{\varphi_{j}}=\braket{\varphi_{i}}{\hat{V}_{KS}}{\varphi_{j}}$.

## Appendix C

We show that having more eigenstates included in the data set results in more
accurate machine-learned Kohn-Sham potential. We examined this by training the
neural network on three different data sets. Each of the data sets includes
15, 10, and 5 eigenstates. The results are shown in Fig. 9. The three machine-
learned Kohn-Sham potentials in the three figures share a similar trend. All
of them capture the general feature of the actual Kohn-Sham potential and show
deviations on the boundary of the system. The gap between the machine-learned
potential and the exact potential closes when more eigenstates are included.

(a) 15 eigenstates

(b) 10 eigenstates

(c) 5 eigenstates

Figure 9: Comparison of the machine learned Kohn-Sham potential among
different data sets. From left to right, each plot shows the results of 15,
10, and 5 eigenstates in the data set.

We also have an interesting observation that the number of local maximums in
each plot is equal to the number of eigenstates in the data set.

## References

  * Runge and Gross (1984) E. Runge and E. K. U. Gross, Phys. Rev. Lett. 52, 997 (1984).
  * Kohn and Sham (1965) W. Kohn and L. J. Sham, Phys. Rev. 140, A1133 (1965).
  * Perdew and Burke (1996) J. P. Perdew and K. Burke, Int. J. Quantum Chem. 57, 309 (1996).
  * Perdew  _et al._ (1999) J. P. Perdew, S. Kurth, A. Zupan, and P. Blaha, Phys. Rev. Lett. 82, 2544 (1999).
  * Nagai  _et al._ (2018) R. Nagai, R. Akashi, S. Sasaki, and S. Tsuneyuki, J. Chem. Phys. 148, 241737 (2018).
  * Suzuki  _et al._ (2020) Y. Suzuki, R. Nagai, and J. Haruyama, Phys. Rev. A 101, 050501 (2020).
  * Greydanus  _et al._ (2019) S. Greydanus, M. Dzamba, and J. Yosinski, Advances in Neural Information Processing Systems 32 (2019).
  * Tong  _et al._ (2021) Y. Tong, S. Xiong, X. He, G. Pan, and B. Zhu, J. Comput. Phys. 437, 110325 (2021).
  * Bertalan  _et al._ (2019) T. Bertalan, F. Dietrich, I. Mezić, and I. G. Kevrekidis, Chaos: An Interdisciplinary Journal of Nonlinear Science 29, 121107 (2019).
  * Han  _et al._ (2021) C.-D. Han, B. Glaz, M. Haile, and Y.-C. Lai, Phys. Rev. Res. 3, 023156 (2021).
  * (11) Z. Chen, J. Zhang, M. Arjovsky, and L. Bottou, “Symplectic Recurrent Neural Networks,” arXiv:1909.13334 .
  * (12) M. Cranmer, S. Greydanus, S. Hoyer, P. Battaglia, D. Spergel, and S. Ho, “Lagrangian Neural Networks,” arXiv:2003.04630 .
  * van Leeuwen (1998) R. van Leeuwen, Phys. Rev. Lett. 80, 1280 (1998).
  * Marques and Gross (2004) M. A. L. Marques and E. K. U. Gross, Annu. Rev. Phys. Chem. 55, ; (2004), 15117259 .
  * Casida and Huix-Rotllant (2012) M. E. Casida and M. Huix-Rotllant, Annu. Rev. Phys. Chem. 63, 287 (2012).
  * Li  _et al._ (2020) X. Li, N. Govind, C. Isborn, A. E. DePrince, and K. Lopata, Chem. Rev. 120, 9951 (2020).
  * Liao  _et al._ (2017) S.-L. Liao, T.-S. Ho, H. Rabitz, and S.-I. Chu, Phys. Rev. Lett. 118, 243001 (2017).
  * Maitra and Burke (2001) N. T. Maitra and K. Burke, Phys. Rev. A 63, 042501 (2001).
  * Maitra  _et al._ (2002) N. T. Maitra, K. Burke, and C. Woodward, Phys. Rev. Lett. 89, 023002 (2002).
  * Colbert and Miller (1992) D. T. Colbert and W. H. Miller, J. Chem. Phys. 96, 1982 (1992).
  * Brown  _et al._ (2020) J. Brown, J. Yang, and J. D. Whitfield, J. Chem. Theory Comput. 16, 6014 (2020).
  * Paszke  _et al._ (2019) A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, in  _Advances in Neural Information Processing Systems 32_ (Curran Associates, Inc., 2019) pp. 8024–8035.
  * De Wilde (1993) P. De Wilde, Phys. Rev. E 47, 1392 (1993).
  * (24) D. P. Kingma and J. Ba, “Adam: A Method for Stochastic Optimization,” arXiv:1412.6980 .
  * Runge (1895) C. Runge, Mathematische Annalen 46, 167 (1895).
  * Lacombe  _et al._ (2018) L. Lacombe, Y. Suzuki, K. Watanabe, and N. T. Maitra, Eur. Phys. J. B 91, 1 (2018).
  * Suzuki  _et al._ (2017) Y. Suzuki, L. Lacombe, K. Watanabe, and N. T. Maitra, Phys. Rev. Lett. 119, 263401 (2017).
  * (28) H. S. Bhat, K. Collins, P. Gupta, and C. M. Isborn, “Dynamic Learning of Correlation Potentials for a Time-Dependent Kohn-Sham System,” arXiv:2112.07067 .
  * (29) A. Sherstinsky, “Fundamentals of Recurrent Neural Network (RNN) and Long Short-Term Memory (LSTM) Network,” arXiv:1808.03314 .
  * Hochreiter and Schmidhuber (1997) S. Hochreiter and J. Schmidhuber, Neural Comput. 9, 1735 (1997).


# Adversarial Robustness Across Representation Spaces

Pranjal Awasthi Equal contribution and corresponding authors. George Yu∗ Chun-
Sung Ferng  Andrew Tomkins  Da-Cheng Juan
Google Research
[pranjalawasthi, georgeyu, csferng, tomkins<EMAIL_ADDRESS>
###### Abstract

Adversarial robustness corresponds to the susceptibility of deep neural
networks to imperceptible perturbations made at test time. In the context of
image tasks, many algorithms have been proposed to make neural networks robust
to adversarial perturbations made to the input pixels. These perturbations are
typically measured in an $\ell_{p}$ norm. However, robustness often holds only
for the specific attack used for training. In this work we extend the above
setting to consider the problem of training of deep neural networks that can
be made simultaneously robust to perturbations applied in multiple natural
representations spaces. For the case of image data, examples include the
standard pixel representation as well as the representation in the discrete
cosine transform (DCT) basis. We design a theoretically sound algorithm with
formal guarantees for the above problem. Furthermore, our guarantees also hold
when the goal is to require robustness with respect to multiple $\ell_{p}$
norm based attacks. We then derive an efficient practical implementation and
demonstrate the effectiveness of our approach on standard datasets for image
classification.

##  1 Introduction

In recent years deep learning has enjoyed tremendous success in solving a
variety of machine learning tasks, even achieving or surpassing human level
performance in certain cases [14, 15]. At the same time important
vulnerabilities in these systems have also been discovered. One such example
is their susceptibility to imperceptible perturbations made to the input at
test time [23]. This has led to the new paradigm of adversarial machine
learning, i.e., making deep neural networks robust to test time perturbations.
There has been a flurry of recent works in this area with several proposed
defenses [18, 28, 6, 17] and methods to attack and evaluate these defenses [5,
3, 25]. When studying the design of networks robust to adversarial attacks
several aspects need to be considered such as a) what perturbations can the
adversary apply to the input, and b) what information does the adversary have
about the neural network? One widely-studied setting in the current literature
is white box attacks under $\ell_{p}$ norm perturbations [10, 18]. Here the
adversary has complete knowledge of the neural network and its parameters, and
given an input $x$ it can perturb it to $x^{\prime}$ such that
$\|x-x^{\prime}\|_{p}\leq\epsilon$ for some $p\geq 1$ specified apriori. In
the context of image data this corresponds to applying perturbations to the
input pixels. Current approaches for defending against such attacks are based
on studying variants of the following robust objective:

$\displaystyle\min_{\theta}\mathbb{E}_{(x,y)\sim
D}\big{[}\max_{x^{\prime}:\|x-x^{\prime}\|_{p}\leq\epsilon}L(f_{\theta}(x^{\prime}),y)\big{]}.$
(1)

Here $(x,y)$ is an example and label pair drawn from the data distribution,
$f$ is a neural network parameterized by weights $\theta$ and $L$ is a
standard loss function such as the cross entropy loss. As an example the
popular projected gradient descent (PGD) method [18] proposes to optimize the
above objective by alternately maximizing the inner objective via gradient
ascent and then performing the outer minimization via gradient descent. The
recent work of [21] combines the above objective with Gaussian smoothing to
achieve certified robustness guarantees, and another popular method namely the
TRADES algorithm [28] adds a regularization term requiring the predictions of
the network at $x$ and $x^{\prime}$ to be close to each other.

In this work we aim to address two main limitations of current approaches to
adversarial machine learning. The first concerns the choice of the
representation in which the adversary applies the perturbations. Using images
as an example, current approaches model the adversary as making small
magnitude changes in the pixel representation of the image. However, given
that the adversary has full access to the input $x$, apriori there is no
reason to restrict the perturbations to only the pixel representations. Real
data such as images have many other natural representations, such as the
discrete cosine transform (DCT) basis for images. One could envision an
adversary making changes to the input image in the DCT basis that are still
imperceptible but don’t satisfy the small $\ell_{p}$ norm property in the
pixel basis. Empirical attacks based on this have been shown to be successful
in recent works [4]. Hence it is important to consider adversarial robustness
in other representations for a model to be truly robust. Secondly, current
approaches fix a representation and the perturbation model, and design an
algorithm to achieve robustness for that specific setting. In general such
networks do not turn out to be robust to other types of attacks. For example a
network trained to be robust to $\ell_{\infty}$ norm perturbations in the
pixel representation may not be robust to $\ell_{1}$ norm perturbations.

Ideally, one would like to train networks that can be simultaneously robust to
multiple attack models in multiple representation spaces. At the same time it
is desirable to have a scalable solution with training cost not that much more
than standard adversarial training in a fixed attack model. This is precisely
the problem that we solve in this work. Our main contributions are listed
below.

  * •

We propose and motivate the problem of studying robustness to adversarial
perturbations in multiple representation spaces and under multiple attack
models.

  * •

We propose a min-max formulation of the above scenario and use ideas from the
theory of online learning, in particular the multiplicative weights update
method [13] to design an algorithm for our formulation and provide theoretical
guarantees to justify our approach.

  * •

We extend our theoretically principled algorithm to design a practical
implementation that can scale to multiple representation spaces and multiple
attack models with training cost not significantly more than that of standard
adversarial training for a fixed attack model and representation space. We
demonstrate the effectiveness of our algorithm for image classification tasks
on the MNIST [16] and the CIFAR-10 [14] datasets.

##  2 Related Work

There is a vast amount of literature on defenses and attacks for adversarial
robustness. See [25] for a survey. Here we discuss the works most relevant to
the results of the paper. As mentioned in the introduction most existing
defenses for adversarial robustness design customized solution for a fixed
attack model ($\ell_{p}$ norm) and representation space (pixel basis). These
methods are aimed at approximately solving the robust optimization objective
in (1). The FGSM method [10] solves the inner maximization problem via one
step of a gradient ascent whereas the PGD method [18] performs multiple
iterations of gradient ascent to better optimize the inner objective.
Typically this scales the cost of training linearly with the number of
iterations used in the inner maximization. There have been recent works aimed
at achieving the same performance as the PGD method but with faster training
time [22, 26].

The above approaches provide robustness to first order attacks that are of the
same type that are used in training. There has also been a lot of recent work
on provably certifying the robustness of neural networks via approaches such
as interval bound propagation [11], semi-definite programming [20], and
randomized smoothing [6, 17].

Relatively little work exists on studying robustness to multiple types of
attacks simultaneously and in multiple representation spaces. The recent work
of [24] studies training classifiers that are simultaneously robust to
perturbations to the input pixels of different $\ell_{p}$ norms. However they
do not consider multiple representation spaces. Furthermore, their approach
does not come with theoretical guarantees and scales linearly with the number
of perturbations considered. In contrast our algorithm comes with theoretical
guarantees and has a training cost that is not much more than that of
adversarial training for a fixed attack. The recent work of [4] motivates the
problem of studying certified robustness in other representations such as the
DCT basis. However they do not consider training classifiers that are
simultaneously robust to multiple attack models.

##  3 Adversarial Robustness in Multiple Representations

In this section we motivate the need for studying adversarial robustness in
representations other than the one that is input to the network. Real world
data can be represented in many natural representations, each with their own
appealing properties. For instance, in the context of images, the DCT basis is
a popular choice and it is well known that signals when represented in this
basis are sparse. This has been exploited in recent works [4] to achieve
better robustness to $\ell_{\infty}$ perturbations in this representations.

Figure 1: The figure shows examples of images from the CIFAR-10 dataset with
their adversarially perturbed counterparts computed by launching a PGD based
attack in the DCT basis. The perturbed images, although imperceptible, are far
from the original images in the pixel basis in $\ell_{\infty}$ norm.

In the context of adversarial learning a fundamental question to ask is: what
constitutes an imperceptible perturbation? Is it enough for an adversarially
perturbed example to have a small $\ell_{p}$ norm in the pixel representation?
As one can imagine, this is not a sufficient condition for imperceptibility.
Many works have notice that images and their adversarial perturbations made in
the pixel basis have distinct spectral signatures when views in other bases
such as the discrete cosine transform (DCT). This has led to the proposal of
many learning systems for detecting pixel based adversarial attacks using
properties of images in other representations [2, 27, 9]. In particular, the
work of [19] shows that one can achieve high accuracy in detecting adversarial
perturbations made in the pixel representation by training a binary classifier
to separate real and perturbed images. Hence an adversary has to naturally
think about attacking the model in multiple representations simultaneously in
order to fool such systems.

Additionally, working in multiple representation spaces can help an adversary
craft stronger attacks. As an example, the recent work of [4] provides
examples where one can generate imperceptible examples by perturbing the image
in the DCT basis and at the same time the perturbed examples are far way from
the original image in the original pixel basis. Such an attack can fool
classifiers that are only trained for defending against small norm $\ell_{p}$
attacks in the pixel representation. We further illustrate this in Figure 1.
The figure shows examples of images from the CIFAR-10 dataset and
corresponding adversarial perturbations computed by launching a PGD based
adversarial attack in the DCT basis. For the case of CIFAR-10 it is generally
accepted that $\ell_{\infty}$ perturbations in the pixel basis upto a
magnitude of $\epsilon=0.03$ constitute imperceptible perturbations. However,
the adversarial images obtained in the Figure via working in the DCT basis,
while being imperceptible, have much higher $\ell_{\infty}$ distance from the
true images in the pixel representation.

From the above discussion we conclude that it is an important problem to
design classifiers that are simultaneously robust against adversarial attacks
in multiple representation spaces. Unfortunately, simply performing standard
adversarial training in a fixed space is not enough in order to achieve this
goal. As an example in Table 1 we show the performance of two neural networks,
one trained adversarially in the pixel representation and the other in the DCT
representation. As can be seen the trained networks have very poor robustness
against the attacks that were not considered during training.

|  Test w/pixel $\ell_{\infty}$ |  Test w/DCT $\ell_{\infty}$ |  Nat. Acc.
---|---|---|---
Train w/pixel $\ell_{\infty}$ |  $44.02$ $\pm 1.02$ |  $27.30$ $\pm 1.44$ |  $80.24$ $\pm 0.40$
Train w/DCT $\ell_{\infty}$ |  $11.80$ $\pm 0.68$ |  $51.92$ $\pm 0.43$ |  $74.92$ $\pm 0.61$
Table 1: The rows of the table correspond to two classifiers that have been
adversarially trained via the PGD method for $\ell_{\infty}$ robustness either
in the pixel basis or the DCT basis. The first two columns show the
adversarial accuracies achieved by the classifiers against $\ell_{\infty}$
attacks in the pixel and the DCT basis. The last column displays the natural
accuracy. As can be seen no classifier is simultaneously robust to both types
of attacks.

As a result of the above observations what is needed is a general algorithmic
approach for such scenarios. We next formulate and present such an approach.

##  4 Formulation and Algorithms

We next formulate the above scenario and design a near optimal algorithm for
simultaneously achieving robustness across multiple representation spaces. We
fix a canonical representation (say the pixel basis) and denote
$x\in\mathbb{R}^{d}$ as examples and $y$ being the label. We assume that the
example and label pairs $(x,y)$ are drawn from an unknown joint distribution
$D$. We then consider a given set of $k$ representation spaces with
corresponding maps being given by
$\mathcal{R}_{1},\mathcal{R}_{2},\dots,\mathcal{R}_{k}$. Hence, given an
example $x\in\mathbb{R}^{d}$ its representation in space $i$ is given by
$\mathcal{R}_{i}(x)\in\mathbb{R}^{d_{i}}$. It would be instructive to think of
$\mathcal{R}_{i}$ as the DCT basis although in general these maps could be
non-linear. The only assumption we require is that the maps be surjective,
i.e. $\mathcal{R}^{-1}_{i}$ exists. We will overload notation and denote by
$\mathcal{R}_{i}$ both the $i$th representation and the map corresponding to
it. For each $\mathcal{R}_{i}$ and any $x\in\mathbb{R}^{d}$, we denote by
$B_{i}(x)$ the set of allowed perturbations to $x$ in the representation space
$\mathcal{R}_{i}$. For example if we are modeling an $\ell_{\infty}$ attack of
radius $\epsilon$ in the representation $\mathcal{R}_{i}$ then we have
$B_{i}(x)=\\{z\in\mathbb{R}^{d_{i}}:\|\mathcal{R}_{i}(x)-z\|_{\infty}\leq\epsilon\\}$.
In this work we will be able to deal with very general perturbation sets.
Given a fixed representation space $\mathcal{R}_{i}$, the problem of learning
a robust classifier specific to $\mathcal{R}_{i}$ can be written as that of
minimizing:

$\displaystyle\min_{\theta}L_{i}(\theta)=\mathbb{E}_{(x,y)\sim
D}\big{[}\max_{z\in B_{i}(x)}L(f_{\theta}(\mathcal{R}^{-1}_{i}(z)),y)\big{]}.$
(2)

Then given $k$ representation spaces $\mathcal{R}_{1},\dots,\mathcal{R}_{k}$
our goal is to solve the following:

$\displaystyle\min_{\theta}\max_{i}L_{i}(\theta).$ (3)

The above min-max formulation lends itself naturally to techniques from online
learning. In particular, consider a two player game with the row player as the
one that chooses the network parameter $\theta$ and the column player as the
one that chooses the loss functions $L_{i}$ with the payoff for the column
player being $L_{i}(\theta)$. From the minimax theorem [13] for two player
games, we know that if for every distribution over the $k$ columns there
exists a good solution $\theta$, then there exists a distribution over
solutions that is simultaneously good for all the columns, i.e., the $k$ loss
functions. This immediately provides a way to solve the min-max formulation
via solving a simple cost sensitive adversarial optimization problem. Such
techniques have been widely used in the literature to solve a variety of
constrained problems in machine learning [1, 8]. Here we demonstrate their
applicability for adversarial robustness. There has also been recent work on
algorithms for solving optimization of the form
$\min_{\theta}\max_{\lambda\in\Lambda}\sum_{i}\lambda_{i}L_{i}(\theta)$ [7]
for $\Lambda$ being a convex set and functions $L_{i}$ being convex in
$\theta$. Our loss functions are non-convex in $\theta$ and hence we need
access to a cost sensitive optimization oracle to provide overall guarantees
for our formulation. We next define the adversarial cost sensitive
optimization problem.

######  Definition 4.1.

Given weights $w_{1},w_{2},\dots,w_{k}$ with $w_{i}\geq 0$ and non-negative
losses $L_{1},L_{2},\dots,L_{k}$ the adversarial cost sensitive optimization
corresponds to finding an approximately optimal solution $\hat{\theta}$ such
that

$\displaystyle\sum_{i=1}^{k}w_{i}L_{i}(\hat{\theta})\leq\min_{\theta}\sum_{i=1}^{k}w_{i}L_{i}({\theta})+\delta.$
(4)

Here $\delta$ quantifies the additive error in approximating the cost
sensitive objective.

We will show how to convert an algorithm for solving the adversarial cost
sensitive optimization problem above to provably optimize (3). The algorithm
is based on the popular multiplicative weights update method [13] and is
described in Figure 2.

Input: Training data $\\{(x_{1},y_{1}),\dots,(x_{m},y_{m})\\}$, Scaling factor
$\eta$. 1. Initialize $w_{i}=1$ for all $i\in[k]$. 2. For $t=1\dots,T$ do: •
Compute $\theta_{t}$ by approximately optimizing (4) with normalized weights
$\frac{w_{i}}{\sum_{j=1}^{k}w_{j}}$ as inputs. • For all $i$ set
$w_{i}=w_{i}\cdot e^{\eta{L}_{i}(\theta_{t})}$. 3. Output
$f_{\hat{\theta}}=\frac{1}{T}\sum_{t=1}^{T}f_{\theta_{t}}$.

Figure 2: An algorithm achieving robustness simultaneously across
representation spaces.

For the proposed algorithm we show the following guarantee

######  Theorem 4.2.

For a given set of non-negative losses bounded in $[0,R]$, if the adversarial
cost sensitive optimization in (4) can be solved to additive error $\delta$
for any setting of non-negative weights then the algorithm in Figure 2 when
run with $\eta=O(\epsilon/R)$ and $T=O(\frac{R^{2}\log k}{\epsilon^{2}})$
outputs a uniform distribution $P$ over solutions
$\theta_{1},\theta_{2},\dots,\theta_{T}$ such that

$\displaystyle\max_{i}\mathbb{E}_{\theta\sim
P}L_{i}(\theta)\leq\min_{\theta}\max_{i}L_{i}(\theta)+\epsilon+\delta.$ (5)

Furthermore if the loss function $L$ in (2) is convex in its first argument,
such as the cross-entropy loss, squared loss and hinge loss to name a few,
then the average hypothesis $f_{\hat{\theta}}$ satisfies

$\displaystyle\max_{i}L_{i}(f_{\hat{\theta}})\leq\min_{\theta}\max_{i}L_{i}(\theta)+\epsilon+\delta.$
(6)

Here $L_{i}(f_{\hat{\theta}})$ refers to the loss incurred by the ensembled
hypothesis as output by the algorithm in Figure 2.

The proof can be found in Appendix 8 in the supplementary material.

##  5 A Practical Implementation

While the algorithm in Figure 2 and the associated guarantee in Theorem 4.2
provide a principled way to approach the optimization in (3), we need to make
a number of modifications to the core algorithm in order to obtain a practical
and scalable implementation. In particular, we do not want the cost of
training the robust classifiers to scale linearly with $k$ the number of
representation spaces. We first discuss solving the adversarial cost sensitive
optimization in (4). In practice, each $L_{i}(\theta)$ itself represents a
hard optimization problem (of the form (2)). Luckily, there exists first order
algorithms such as the PGD method [18] to optimize each $L_{i}$ separately
that work well in practice. We now show how to combine them to solve (4). We
follow the methodology of stochastic optimization and proceed in epochs. In
each epoch, we sample a mini batch of $B$ data points, sample a loss $L_{i}$
with probability proportional to its current $w_{i}$ and then run the
corresponding PGD based algorithm for optimizing $L_{i}$ on the current mini
batch. After a few epochs of optimization we update the weights $w_{i}$ of the
losses as described in the algorithm in Figure 2. In order to evaluate the
losses for the weight update we use a separate validation set. This
significantly reduces the variance in our estimates.

Next we consider approximating the output $f_{\hat{\theta}}$. Notice that the
guarantee of Theorem 4.2 applies to an ensemble of $T$ neural networks
provided by parameters $\theta_{1},\theta_{2},\dots,\theta_{T}$. Maintaining
this ensemble requires a high storage cost and makes the final output model
impractically big. We first notice that if the losses $L_{i}$ were convex,
then the guarantee of Theorem 4.2 will also hold for the average parameter,
i.e., $\hat{\theta}=\frac{1}{T}\sum_{t}\theta_{t}$. To get a practical
implementation we make a near convexity assumption on the losses and simply
take the average of the model weights. Furthermore, in our experiments we
observe that taking the average of the last few model parameters performs
better than the uniform average of all the model weights. Fixing these choices
leads to a scalable variant as shown in Figure 3.

Input: Training data $\\{(x_{1},y_{1}),\dots,(x_{m},y_{m})\\}$, Validation
data $\\{(x_{m+1},y_{m+1}),\dots,(x_{m+s},y_{m+s})\\}$, mini batch size $B$,
time steps $T$, update frequency $r$, window size $h$, Scaling factor $\eta$.
1. Initialize $w_{i}=1$ for all $i\in[k]$. 2. For $t=1\dots,T$ do: • Repeat
for $r$ epochs: – Get the next mini batch of size $B$. Sample loss $L_{i}$
with probability $p_{i}=\frac{w_{i}}{\sum_{j=1}^{k}w_{j}}$. – Run the PGD
based algorithm to optimize $L_{i}$ on the mini batch. • For all $i$ set
$w_{i}=w_{i}\cdot e^{\eta{L_{i}^{\text{val}}}(\theta_{t})}$. Here
$L^{\text{val}}$ is the loss evaluated on the validation set. 3. Output
$\hat{\theta}=\frac{1}{h}\sum_{t=T-h+1}^{T}\theta_{t}$.

Figure 3: A scalable variant of the algorithm in Figure 2.

##  6 Experimental Evaluation

We next demonstrate the effectiveness of our approach on the task of learning
a classifier for image classification that is simultaneously robust to
adversarial attacks in multiple representations spaces.

### Datasets.

We perform the experimental evaluation on two public datasets namely the MNIST
dataset [16] and the CIFAR-10 dataset [14]. The MNIST dataset consists of
$60,000$ training images with each being a $28\times 28\times 1$ tensor. The
CIFAR-10 dataset consists of $50,000$ training images each of dimensionality
$32\times 32\times 3$. Both the datasets consist of $10,000$ test images and
correspond to a multi class classification problem with $10$ class labels. In
each case we reserve $10\%$ of the training data to be used as the validation
set in the Algorithm from Figure 3. This validation set will be used to
evaluate the loss $L^{\text{val}}$ in the algorithms.

### Representation Spaces and Attack Models.

To demonstrate the scalability of our approach we consider two different
representation spaces namely the pixel basis and the DCT basis. In each
representation space we consider three types of $\ell_{p}$ norm based attacks
for $p=1,2,\infty$. Hence, in total we have $6$ loss functions $L_{i}$ to
optimize as in (3). All our experiments are conducted on a ResNet-50 deep
neural network that is a popular architecture for training on image
classification tasks [12]. To compute adversarial examples in the pixel basis
for norm bounded $\ell_{\infty}$ and $\ell_{2}$ perturbations we use the
standard PGD based attack as proposed in [18]. For computing a norm bounded
$\ell_{1}$ perturbation we use the sparse ascent algorithm namely the SLIDE
method as proposed in [24]. To compute an adversarial attack in the DCT basis
we append the ResNet-50 architecture with a linear DCT transformation follows
by an inverse DCT transformation as shown in Figure 4. Notice that both the
DCT and the Inverse DCT are fixed linear layers and in the absence of any
perturbations the output of the network in Figure 4 is exactly the same as the
original ResNet-50 network.

Using the modified architecture we first compute the DCT representation of the
image and then launch an adversarial perturbation in the DCT basis using
either the PGD method (for $\ell_{2},\ell_{\infty}$ attacks) or the SLIDE
method (for $\ell_{1}$ attacks). In this way we get the perturbed image after
taking the inverse DCT transform of the perturbed example $x^{\prime}$ as
shown in Figure 1. After computing the adversarial perturbation and passing it
through the inverse DCT transform we clip the pixel values in $[0,1]$ to make
the example a valid input for the ResNet-50 network.

Figure 4: The modified network architecture for computing an adversarial
perturbation in the DCT basis.

### Baselines.

We next describe the baselines that we use when comparing our proposed
approach. The problem of being simultaneously robust to multiple adversarial
attacks has been largely ignored in the literature so far. The recent work of
[24] studies being robust in pixel basis to different $\ell_{p}$ norm based
attacks. However the proposed method is not scalable to a large number of
attacks.

We instead compare our proposed algorithm with following two baseline
heuristics. We choose these heuristics due to their simplicity and more
importantly due to their scalability.

Round Robin. The round robin heuristic sketched in Figure 5 follows the same
outline as the algorithm in Figure 3 except that instead of maintaining and
updating weights, it simply picks the loss function $L_{i}$ to apply to a mini
batch in a fixed order. It is easy to see that this method scales very well.

Input: Training data $\\{(x_{1},y_{1}),\dots,(x_{m},y_{m})\\}$. Input: Mini
batch size $B$, time steps $T$. 1. Initialize $index=1$. 2. For $t=1\dots,T$
do: • Get the next mini batch of size $B$. • Set $i=index$. • Use PGD to
optimize $L_{i}$ on the mini batch to get $\theta_{t}$. •
$index=(index+1)\text{ mod }k+1$. 3. Output $\hat{\theta}=\theta_{T}$.

Figure 5: The round robin heuristic.

Greedy. The greedy heuristic also follows the same outline as the algorithm in
Figure 3. However, for each time step it chooses the loss with the worst error
(on the validation set) to apply next. See Figure 6.

Input: Training data $\\{(x_{1},y_{1}),\dots,(x_{m},y_{m})\\}$. Input:
Validation data: $\\{(x_{m+1},y_{m+1}),\dots,(x_{m+s},y_{m+s})\\}$. Input:
Mini batch size $B$, time steps $T$, update frequency $r$. 1. For $t=1\dots,T$
do: • Set $i=\arg\max_{j}L^{\text{val}}_{j}$. • Repeat for $r$ epochs: – Get
the next mini batch of size $B$. – Use PGD to optimize $L_{i}$ on the mini
batch to get $\theta_{t}$. 2. Output $\hat{\theta}=\theta_{T}$.

Figure 6: The greedy heuristic. | Greedy | Round Robin |  Mult. Weights ($h=1$) |  Mult. Weights ($h=3$)
---|---|---|---|---
Pixel ($\ell_{\infty}$) | $36.03\pm 6.69$ | $29.99\pm 1.49$ | ${35.70\pm 6.36}$ | $\mathbf{39.13\pm 1.74}$
Pixel ($\ell_{2}$) | $69.03\pm 1.58$ | $\mathbf{71.92\pm 0.36}$ | $69.51\pm 2.21$ | $\mathbf{71.27\pm 0.22}$
Pixel ($\ell_{1}$) | $45.84\pm 2.85$ | $\mathbf{53.45\pm 0.70}$ | $44.22\pm 4.79$ | $46.90\pm 1.46$
DCT ($\ell_{\infty}$) | ${44.72\pm 8.60}$ | $\mathbf{45.63\pm 1.75}$ | $39.44\pm 6.76$ | $42.27\pm 2.77$
DCT ($\ell_{2}$) | $68.99\pm 1.58$ | $\mathbf{72.01\pm 0.28}$ | $69.66\pm 2.21$ | $71.24\pm 0.12$
DCT ($\ell_{1}$) | $37.51\pm 6.78$ | $\mathbf{41.44\pm 0.82}$ | $39.62\pm 5.71$ | $\mathbf{42.96\pm 1.12}$
Min. Accuracy | $34.62\pm 5.60$ | $29.99\pm 1.49$ | $\mathbf{35.70\pm 6.36}$ | $\mathbf{39.13\pm 1.74}$
Union Attack | $31.95\pm 4.55$ | $29.70\pm 1.42$ | $\mathbf{33.17\pm 5.95}$ | $\mathbf{36.18\pm 1.89}$
Nat. Acc. | $78.56\pm 1.46$ | $\mathbf{81.80\pm 0.38}$ | $79.64\pm 0.51$ | $\mathbf{80.66\pm 0.89}$
Table 2: Comparison of the adversarial accuracies achieved on the CIFAR-10
dataset by the greedy algorithm, the round robin algorithm and our proposed
algorithm in Figure 3.

### Hyperparameter Configurations.

Next we discuss the hyperparameters we use when computing the adversarial
perturbations for the different $\ell_{p}$ norm based attacks. When running
our proposed algorithm in Figure 3 and the greedy heuristic, we set $T=40$,
$r=5$ for CIFAR-10 ($200$ epochs total), and $T=20$, $r=3$ for MNIST ($60$
epochs total). We train the round robin heuristic for the same number of
epochs.

For the case of $\ell_{\infty}$ and $\ell_{2}$ attacks, during training we use
$10$ steps of gradient ascent to optimize the inner maximization in (2).
During evaluation we again run the PGD based attack on our model across all
the representation spaces and use $40$ steps of the PGD method to solve the
inner maximization in (1). For the case of $\ell_{1}$ attacks we use $20$
iterations of the SLIDE method during training to compute adversarial
perturbations and $100$ iterations of the method during evaluation. We
experiment with both running the PGD method with $20$ random restarts, and a
simpler attack with no restarts. The experiments we report here are for the
latter case. The qualitative conclusions of our experiments remain the same
when using $20$ random restarts. See the supplementary material for details.

Dataset |  Pixel ($\ell_{\infty}$) |  DCT ($\ell_{\infty}$) | Round Robin | Greedy | Mult. Weights
---|---|---|---|---|---
MNIST | $2.25$ | $2.09$ | $3.36$ | $2.67$ | $2.90$
CIFAR-10 | $7.66$ | $7.06$ | $10.61$ | $8.70$ | $9.51$
Table 3: Training time in hours (wall clock time) for the baselines and our
proposed method on the MNIST and the CIFAR-10 datasets. The first two columns
represent the training time for optimizing a single loss, i.e.,
$\ell_{\infty}$ attack in the pixel and the DCT basis respectively. The next
three columns represent the training time of the three methods when optimizing
over all the $6$ losses simultaneously. The reported numbers are averaged over
$5$ runs.

For the MNIST dataset we perturbation magnitudes of $0.4,1$ and $5$ for
$\ell_{\infty}$, $\ell_{2}$ and $\ell_{1}$ norm based attacks respectively.
The corresponding magnitudes for the CIFAR-10 dataset are $0.06,0.1$ and
$7.84$. We keep the perturbation magnitudes the same across both the pixel and
the DCT basis. In our experiments when performing gradient ascent for $\ell$
steps to compute a perturbation, we use a step size of
$2.5\frac{\epsilon}{\ell}$, where $\epsilon$ is the perturbation magnitude.

### Metrics.

For each of our trained classifiers we report the adversarial accuracy for
each of the $6$ individual attacks launched separately on the trained model.
In addition we also report the worst adversarial accuracy among the $6$
attacks on the same model. Notice that this is the metric that our proposed
algorithm in Figure 3 aims to optimize. Finally, we also report the accuracy
of our trained models on a union attack, i.e., for each example we produce all
$6$ adversarial perturbations and consider the attack successful if any one of
them succeeds in making the prediction of the model incorrect. Finally, notice
that our proposed algorithm in Figure 3 has a parameter $h$ namely the window
size. We report results for $h=1$ and $h=3$. These correspond to either using
the parameters of the last time step or using the model averaged over the last
three time steps.

| Greedy | Round Robin |  Mult. Weights ($h=1$) |  Mult. Weights ($h=3$)
---|---|---|---|---
Pixel ($\ell_{\infty}$) | $63.12\pm 19.50$ | $64.40\pm 25.12$ | $\mathbf{67.73\pm 14}$ | $\mathbf{66.56\pm 13.45}$
Pixel ($\ell_{2}$) | $22.51\pm 17.66$ | $10.23\pm 8.40$ | $\mathbf{71.23\pm 2.68}$ | $\mathbf{71.77\pm 4.25}$
Pixel ($\ell_{1}$) | $64.33\pm 30.73$ | $43.14\pm 23.57$ | $\mathbf{73.20\pm 10.54}$ | $\mathbf{73.66\pm 9.47}$
DCT ($\ell_{\infty}$) | $61.53\pm 21.45$ | $59.47\pm 32.83$ | $\mathbf{60.65\pm 10.60}$ | $\mathbf{60.32\pm 11.34}$
DCT ($\ell_{2}$) | $25.80\pm 18.82$ | $12.28\pm 12.16$ | $\mathbf{84.58\pm 2.52}$ | $\mathbf{84.46\pm 4.18}$
DCT ($\ell_{1}$) | $66.00\pm 31.83$ | $38.23\pm 21.47$ | $\mathbf{72.39\pm 8.43}$ | $\mathbf{73.13\pm 7.33}$
Min. Accuracy | $22.13\pm 17.17$ | $9.76\pm 8.69$ | $\mathbf{57.64\pm 7.83}$ | $\mathbf{57.57\pm 8.46}$
Union Attack | $12.32\pm 9.80$ | $3.72\pm 4.96$ | $\mathbf{35.30\pm 4.48}$ | $\mathbf{35.87\pm 6.69}$
Nat. Acc. | $77.16\pm 37.64$ | $60.93\pm 32.31$ | $\mathbf{91.00\pm 10.50}$ | $\mathbf{91.43\pm 9.61}$
Table 4: Comparison of the adversarial accuracies achieved on the MNIST
dataset by the greedy algorithm, the round robin algorithm and our proposed
algorithm in Figure 3.

### Results.

The performance of our algorithm as compared to the baseline is shown in Table
2 for the CIFAR-10 dataset and in Table 4 for the MNIST dataset. We make a few
observations. In both the cases the performance of the multiplicative weights
update based algorithm is significantly better than the baseline on the
minimum accuracy metric and the union attack metric. This difference is
significantly higher for the MNIST dataset where both the greedy and the round
robin heuristics are unstable and have much higher variances. The round robin
heuristic switches among different losses much more often and pays unnecessary
attention to the adversaries which it already covers well. The greedy
heuristic, on the other hand, switches less frequent than our proposed
algorithm. But greedy fails to address the runner-up adversary which may be
almost as difficult as the chosen one, causing instability. Furthermore, we
also notice that using the average of the last three model parameters in the
multiplicative weights based method achieves slightly better performance than
simply using the parameters of the last time step.

### Comparison of Training Times.

We next demonstrate the scalability of our proposed algorithm. Table 3 shows
the training time of our method, measured in wall clock time, as compared to
the baselines when optimizing over all $6$ loss functions. Moreover, the first
two columns represent the training times for optimizing a single loss function
($\ell_{\infty}$ attack) in either the pixel or the DCT basis. As can be seen
our the training cost of our approach scales sublinearly with the number of
representation spaces.

### On the Convexity Assumption.

Recall that the guarantees of Theorem 4.2 apply to the algorithm in Figure 2
that requires one to produce a hypothesis that is an ensemble of the
intermediate trained models. In other words, ideally one should be averaging
the post-softmax outputs of the trained models. If the loss functions were
convex then one could replace the ensembling with simply averaging the model
weights and retain the theoretical guarantees. Even though we have non-convex
losses we still make the near convexity assumption and average the model
weights to produce a scalable implementation. In Table 6 and Table 6 we
compare the performance of our weight averaging strategy with that of the
ideal one that ensembles the models. As can be seen the loss in making the
near convexity assumption is negligible and justifies our implementation in
Figure 3.

##  7 Discussion

In this work we motivated the problem of designing neural networks that are
simultaneously robust to multiple types of adversarial attacks in multiple
representation spaces. We provided a theoretically sound algorithm with
training cost that grows sublinearly with the number of representation spaces.
Our implementation is scalable and significantly outperforms strong baselines,
with similar training cost to standard adversarial training.

Several future directions emerge from this work. Notice that in our proposed
algorithm we use the PGD based method of [18] to optimize the individual
losses. There has been very recent work proposing faster training methods that
achieve similar performance to that of the PGD method [22, 26]. It would be
interesting to incorporate them in our framework to drive down the training
cost even further. The benefits of this could be significant as the number of
representation spaces grows.

While our theory applies to an ensemble of neural networks, as our experiments
indicate, in practice simply averaging the model weights does as well as
ensembling. This justifies our near convexity assumption. It would be
interesting to study this behavior further and provide a more formal
theoretical justification. Finally, we hope that future work on adversarial
robustness will consider evaluating robustness of classifiers in multiple
representation spaces.

|  Mult. Weights ($h=3$) (ensemble) |  Mult. Weights ($h=3$) (wt. avg.)
---|---|---
Pixel ($\ell_{\infty}$) |  ${39.13\pm 1.74}$ |  ${40.52\pm 1.28}$
Pixel ($\ell_{2}$) |  $71.27\pm 0.22$ |  $71.41\pm 0.24$
Pixel ($\ell_{1}$) |  $46.90\pm 1.46$ |  $48.14\pm 0.99$
DCT ($\ell_{\infty}$) |  $42.27\pm 2.77$ |  $43.95\pm 2.05$
DCT ($\ell_{2}$) |  $71.24\pm 0.12$ |  $71.38\pm 0.20$
DCT ($\ell_{1}$) |  $42.96\pm 1.12$ |  $44.49\pm 0.58$
Min. Accuracy |  ${39.13\pm 1.74}$ |  ${40.52\pm 1.28}$
Union Attack |  ${36.18\pm 1.89}$ |  ${37.76\pm 1.30}$
Nat. Acc. |  $80.66\pm 0.89$ |  ${80.46\pm 0.81}$

Table 5: Comparison of the adversarial accuracies achieved on the CIFAR-10
dataset by our proposed algorithm in Figure 3 when using the average of the
last three model parameters (convexity assumption) vs. ensembling the outputs
of the last three models.

|  Mult. Weights ($h=3$) (ensemble) |  Mult. Weights ($h=3$) (wt. avg.)
---|---|---
Pixel ($\ell_{\infty}$) |  ${66.56\pm 13.45}$ |  ${66.61\pm 13.77}$
Pixel ($\ell_{2}$) |  ${71.77\pm 4.25}$ |  ${71.89\pm 3.79}$
Pixel ($\ell_{1}$) |  ${73.66\pm 9.47}$ |  ${73.38\pm 9.45}$
DCT ($\ell_{\infty}$) |  ${60.32\pm 11.34}$ |  ${59.87\pm 11.32}$
DCT ($\ell_{2}$) |  ${84.46\pm 4.18}$ |  ${84.34\pm 4.13}$
DCT ($\ell_{1}$) |  ${73.13\pm 7.33}$ |  ${73.08\pm 7.38}$
Min. Acc. |  ${57.57\pm 8.46}$ |  ${57.23\pm 8.51}$
Union Attack |  ${35.87\pm 6.69}$ |  ${35.54\pm 6.69}$
Nat. Acc. |  ${91.43\pm 9.61}$ |  ${91.44\pm 9.62}$

Table 6: Comparison of the adversarial accuracies achieved on the MNIST
dataset by our proposed algorithm in Figure 3 when using the average of the
last three model parameters (convexity assumption) vs. ensembling the outputs
of the last three models.

## References

  * [1] Alekh Agarwal, Alina Beygelzimer, Miroslav Dudík, John Langford, and Hanna Wallach.  A reductions approach to fair classification.  arXiv preprint arXiv:1803.02453, 2018.
  * [2] Naveed Akhtar, Jian Liu, and Ajmal Mian.  Defense against universal adversarial perturbations.  In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, pages 3389–3398, 2018.
  * [3] 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, 2018.
  * [4] Pranjal Awasthi, Himanshu Jain, Ankit Singh Rawat, and Aravindan Vijayaraghavan.  Adversarial robustness via robust low rank representations.  Advances in Neural Information Processing Systems, 33, 2020.
  * [5] Nicholas Carlini and David Wagner.  Towards evaluating the robustness of neural networks.  In 2017 ieee symposium on security and privacy (sp), pages 39–57. IEEE, 2017.
  * [6] Jeremy M Cohen, Elan Rosenfeld, and J Zico Kolter.  Certified adversarial robustness via randomized smoothing.  arXiv preprint arXiv:1902.02918, 2019.
  * [7] Corinna Cortes, Mehryar Mohri, Javier Gonzalvo, and Dmitry Storcheus.  Agnostic learning with multiple objectives.  Advances in Neural Information Processing Systems, 33, 2020.
  * [8] Andrew Cotter, Heinrich Jiang, and Karthik Sridharan.  Two-player games for efficient non-convex constrained optimization.  In Algorithmic Learning Theory, pages 300–332. PMLR, 2019.
  * [9] Gintare Karolina Dziugaite, Zoubin Ghahramani, and Daniel M Roy.  A study of the effect of jpg compression on adversarial images.  arXiv preprint arXiv:1608.00853, 2016.
  * [10] Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy.  Explaining and harnessing adversarial examples.  arXiv preprint arXiv:1412.6572, 2014.
  * [11] Sven Gowal, Krishnamurthy Dvijotham, Robert Stanforth, Rudy Bunel, Chongli Qin, Jonathan Uesato, Relja Arandjelovic, Timothy Mann, and Pushmeet Kohli.  On the effectiveness of interval bound propagation for training verifiably robust models.  arXiv preprint arXiv:1810.12715, 2018.
  * [12] 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, 2016.
  * [13] Satyen Kale.  Efficient algorithms using the multiplicative weights update method.  Princeton University, 2007.
  * [14] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * [15] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  Communications of the ACM, 60(6):84–90, 2017.
  * [16] Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner.  Gradient-based learning applied to document recognition.  Proceedings of the IEEE, 86(11):2278–2324, 1998.
  * [17] Mathias Lecuyer, Vaggelis Atlidakis, Roxana Geambasu, Daniel Hsu, and Suman Jana.  Certified robustness to adversarial examples with differential privacy.  In 2019 IEEE Symposium on Security and Privacy (SP), pages 656–672. IEEE, 2019.
  * [18] Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu.  Towards deep learning models resistant to adversarial attacks.  arXiv preprint arXiv:1706.06083, 2017.
  * [19] Jan Hendrik Metzen, Tim Genewein, Volker Fischer, and Bastian Bischoff.  On detecting adversarial perturbations.  arXiv preprint arXiv:1702.04267, 2017.
  * [20] Aditi Raghunathan, Jacob Steinhardt, and Percy S Liang.  Semidefinite relaxations for certifying robustness to adversarial examples.  In Advances in Neural Information Processing Systems, pages 10877–10887, 2018.
  * [21] 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.
  * [22] 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.
  * [23] Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus.  Intriguing properties of neural networks.  arXiv preprint arXiv:1312.6199, 2013.
  * [24] Florian Tramèr and Dan Boneh.  Adversarial training and robustness for multiple perturbations.  In Advances in Neural Information Processing Systems, pages 5866–5876, 2019.
  * [25] Florian Tramer, Nicholas Carlini, Wieland Brendel, and Aleksander Madry.  On adaptive attacks to adversarial example defenses.  arXiv preprint arXiv:2002.08347, 2020.
  * [26] Eric Wong, Leslie Rice, and J Zico Kolter.  Fast is better than free: Revisiting adversarial training.  arXiv preprint arXiv:2001.03994, 2020.
  * [27] Dong Yin, Raphael Gontijo Lopes, Jon Shlens, Ekin Dogus Cubuk, and Justin Gilmer.  A fourier perspective on model robustness in computer vision.  Advances in Neural Information Processing Systems, 32:13276–13286, 2019.
  * [28] Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric P Xing, Laurent El Ghaoui, and Michael I Jordan.  Theoretically principled trade-off between robustness and accuracy.  arXiv preprint arXiv:1901.08573, 2019.

##  8 Appendix: Proof of Theorem 4.2

###### Proof.

Given the weights $w_{1},w_{2},\dots w_{k}$ for the $k$ losses, we denote by
$\mathbf{p}$ the normalized probabilities, i.e.,
$p_{i}=w_{i}/\sum_{j=1}^{k}w_{j}$. Let
$\mathbf{p}_{1},\mathbf{p}_{2},\dots,\mathbf{p}_{T}$ be the sequence of
probability vectors produced by the column players namely the player that
chooses among the $k$ losses. Denote by $p_{t,j}$ the $j$th coordinate of the
vector $\mathbf{p}_{t}$. Since the player is performing multiplicative weights
updates, by the standard guarantee of the multiplicative weights update (see
Theorem 2 in [13]) we get that for any $i\in[k]$ the following holds

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\sum_{j=1}^{k}p_{t,j}L_{j}(\theta_{t})\geq(1-\eta)\frac{1}{T}\sum_{t=1}^{T}L_{i}(\theta_{t})-2R\frac{\log
k}{\eta T}.$ (7)

Setting $\eta=O(\epsilon/R)$ and $T=O(R^{2}\frac{\log k}{\epsilon^{2}})$ we
get that

$\displaystyle\frac{1}{T}\sum_{t=1}^{T}\sum_{j=1}^{k}p_{t,j}L_{j}(\theta_{t})\geq\frac{1}{T}\sum_{t=1}^{T}L_{i}(\theta_{t})-\epsilon.$
(8)

Hence, denoting by $P$ the uniform distribution over the $T$ parameters, we
get that

$\displaystyle\mathbb{E}_{\theta\sim
P}L_{i}(\theta)\leq\frac{1}{T}\sum_{t=1}^{T}\sum_{j=1}^{k}p_{t,j}L_{j}(\theta_{t})+\epsilon.$
(9)

Next we will write down the right hand side in terms of the optimal value of
the objective as in (1). Given the guarantee that we can solve the adversarial
cost sensitive optimization upto additive error $\delta$, we have that for any
$t\in[T]$,

$\displaystyle\sum_{j=1}^{k}p_{t,j}L_{j}(\theta_{t})\leq\min_{\theta}\sum_{j=1}^{k}p_{t,j}L_{j}(\theta)+\delta.$
(10)

Substituting into (9) we get

$\displaystyle\mathbb{E}_{\theta\sim P}L_{i}(\theta)$
$\displaystyle\leq\frac{1}{T}\sum_{t=1}^{T}\min_{\theta}\sum_{j=1}^{k}p_{t,j}L_{j}(\theta)+\epsilon+\delta$
(11)
$\displaystyle\leq\min_{\theta}\frac{1}{T}\sum_{t=1}^{T}\sum_{j=1}^{k}p_{t,j}L_{j}(\theta)+\epsilon+\delta.$
(12)

Next for any $j\in[k]$ define
$\tilde{p}_{j}=\frac{1}{T}\sum_{t=1}^{T}p_{t,j}$. Then we can rewrite the
above as

$\displaystyle\mathbb{E}_{\theta\sim P}L_{i}(\theta)$
$\displaystyle\leq\frac{1}{T}\sum_{t=1}^{T}\min_{\theta}\sum_{j=1}^{k}p_{t,j}L_{j}(\theta)+\epsilon+\delta$
$\displaystyle\leq\min_{\theta}\sum_{j=1}^{k}\tilde{p}_{j}L_{j}(\theta)+\epsilon+\delta.$
(13)

It is easy to check that $\tilde{p}_{j}\in[0,1]$ and
$\sum_{j=1}^{k}\tilde{p}_{j}=1$. Hence we get that
$\sum_{j=1}^{k}\tilde{p}_{j}L_{j}(\theta)\leq\min_{\theta}\max_{j}L_{j}(\theta)$.
Substituting into (8) we get the guarantee of the theorem that for all
$i\in[k]$,

$\displaystyle\mathbb{E}_{\theta\sim P}L_{i}(\theta)$
$\displaystyle\leq\min_{\theta}\max_{j}L_{j}(\theta)+\epsilon+\delta.$ (14)

Next we prove the consequence when the loss function $L$ in (2) is convex in
its first argument. This is true for commonly used loss functions in practice
such as the cross entropy loss, squared loss and the hinge loss. We denote by
$f_{\theta_{t}}$ the network corresponding to the parameter $\theta_{t}$ and
by $f_{\hat{\theta}}$ the ensembled network i.e.,

$f_{\hat{\theta}}=\frac{1}{T}\sum_{t=1}^{T}f_{\theta_{t}}.$

By expanding out $\mathbb{E}_{\theta\sim P}L_{i}(\theta)$ we get

$\displaystyle\mathbb{E}_{\theta\sim P}L_{i}(\theta)$
$\displaystyle=\frac{1}{T}\sum_{t=1}^{T}L_{i}(f_{\theta_{t}})$ (15)
$\displaystyle=\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}_{(x,y)\sim
D}\big{[}\max_{z\in B_{i}(x)}L(f_{\theta_{t}}(\mathcal{R}^{-1}(z)),y)\big{]}$
(16) $\displaystyle\geq\max_{z\in
B_{i}(x)}\frac{1}{T}\sum_{t=1}^{T}\mathbb{E}_{(x,y)\sim
D}\big{[}L(f_{\theta_{t}}(\mathcal{R}^{-1}(z)),y)\big{]}$ (17)
$\displaystyle=\max_{z\in B_{i}(x)}\mathbb{E}_{(x,y)\sim
D}\mathbb{E}_{\theta\sim P}L(f_{\theta}(\mathcal{R}^{-1}(z)),y)$ (18)
$\displaystyle\geq\max_{z\in B_{i}(x)}\mathbb{E}_{(x,y)\sim
D}L(\mathbb{E}_{\theta\sim P}f_{\theta}(\mathcal{R}^{-1}(z)),y)$ (19)
$\displaystyle=\max_{z\in B_{i}(x)}\mathbb{E}_{(x,y)\sim
D}L(f_{\hat{\theta}}(\mathcal{R}^{-1}(z)),y).$ (20)

Here in the last inequality we have used the fact that $L$ is convex in its
first argument. Substituting back into (14) we get that $f_{\hat{\theta}}$
satisfies that for all $i\in[k]$,

$\displaystyle L_{i}(f_{\hat{\theta}})$
$\displaystyle\leq\min_{\theta}\max_{j}L_{j}(\theta)+\epsilon+\delta.$ (21)

∎

##  9 Appendix: Further Experiments

| Greedy | Round Robin |  Mult. Weights ($h=1$) |  Mult. Weights ($h=3$)
---|---|---|---|---
Pixel ($\ell_{\infty}$) | $34.16\pm 6.57$ | $28.28\pm 1.53$ | ${34.01\pm 6.38}$ | $\mathbf{37.24\pm 1.89}$
Pixel ($\ell_{2}$) | $64.74\pm 1.73$ | $\mathbf{67.33\pm 0.27}$ | $64.97\pm 3.03$ | $\mathbf{66.94\pm 0.28}$
Pixel ($\ell_{1}$) | $45.26\pm 3.11$ | $\mathbf{53.08\pm 0.72}$ | $43.46\pm 4.99$ | $46.25\pm 1.44$
DCT ($\ell_{\infty}$) | ${42.64\pm 9.02}$ | $\mathbf{43.62\pm 1.74}$ | $37.17\pm 6.92$ | $39.93\pm 3.01$
DCT ($\ell_{2}$) | $64.66\pm 1.71$ | $\mathbf{67.26\pm 0.38}$ | $65.03\pm 2.88$ | $66.98\pm 0.30$
DCT ($\ell_{1}$) | $37.07\pm 6.84$ | $\mathbf{41.01\pm 0.81}$ | $39.16\pm 5.86$ | $\mathbf{42.52\pm 1.17}$
Min. Accuracy | $32.50\pm 5.31$ | $28.28\pm 1.53$ | $\mathbf{34.01\pm 6.38}$ | $\mathbf{37.24\pm 1.89}$
Union Attack | $30.29\pm 4.24$ | $28.07\pm 1.48$ | $\mathbf{31.61\pm 6.07}$ | $\mathbf{34.63\pm 2.02}$
Nat. Acc. | $78.56\pm 1.46$ | $\mathbf{81.80\pm 0.38}$ | $79.64\pm 0.51$ | $\mathbf{80.66\pm 0.89}$
Table 7: Comparison of the adversarial accuracies achieved on the CIFAR-10 dataset by the greedy algorithm, the round robin algorithm and our proposed algorithm in Figure 3. | Greedy | Round Robin |  Mult. Weights ($h=1$) |  Mult. Weights ($h=3$)
---|---|---|---|---
Pixel ($\ell_{\infty}$) | $53.50\pm 24.05$ | $59.60\pm 24.55$ | $\mathbf{60.88\pm 17}$ | $\mathbf{60.01\pm 16.61}$
Pixel ($\ell_{2}$) | $7.50\pm 4.48$ | $6.04\pm 4.27$ | $\mathbf{63.32\pm 8.25}$ | $\mathbf{63.47\pm 7.72}$
Pixel ($\ell_{1}$) | $63.60\pm 30.36$ | $41.98\pm 22.87$ | $\mathbf{72.23\pm 10.54}$ | $\mathbf{72.51\pm 9.68}$
DCT ($\ell_{\infty}$) | $53.80\pm 24.53$ | $53.24\pm 22.41$ | $\mathbf{55.29\pm 12.16}$ | $\mathbf{54.81\pm 13.09}$
DCT ($\ell_{2}$) | $9.18\pm 5.19$ | $5.90\pm 5.71$ | $\mathbf{79.97\pm 5.92}$ | $\mathbf{79.50\pm 6.50}$
DCT ($\ell_{1}$) | $65.45\pm 31.95$ | $36.81\pm 21.10$ | $\mathbf{71.74\pm 8.49}$ | $\mathbf{72.51\pm 7.33}$
Min. Accuracy | $7.31\pm 4.38$ | $5.19\pm 4.72$ | $\mathbf{52.53\pm 11.31}$ | $\mathbf{53.29\pm 12.73}$
Union Attack | $3.37\pm 3.10$ | $1.97\pm 3.03$ | $\mathbf{27.65\pm 7.06}$ | $\mathbf{28.09\pm 9.38}$
Nat. Acc. | $77.16\pm 37.64$ | $60.93\pm 32.31$ | $\mathbf{91.00\pm 10.50}$ | $\mathbf{91.43\pm 9.61}$
Table 8: Comparison of the adversarial accuracies achieved on the MNIST
dataset by the greedy algorithm, the round robin algorithm and our proposed
algorithm in Figure 3.

In this section we present the results of evaluating our trained
multiplicative weights method based algorithm from Figure 3, as well as the
greedy and the round robin heuristics, against the PGD based attack where the
PGD method is run with $20$ random restarts in order to find an adversarial
example. Tables 7 shows the results for the CIFAR-10 dataset and Table 8 shows
the results for the MNIST dataset. Similar to the results presented in Section
6, the multiplicative weights method significantly outperforms the baselines
on both the minimum accuracy metric and accuracy against a union attack.