# Generalized description of the spatio-temporal biphoton state in spontaneous
parametric down-conversion

Baghdasar Baghdasaryan<EMAIL_ADDRESS>Theoretisch-
Physikalisches Institut, Friedrich Schiller University Jena, 07743 Jena,
Germany  Helmholtz-Institut Jena, 07743 Jena, Germany  Carlos Sevilla-
Gutiérrez  Fraunhofer Institute for Applied Optics and Precision Engineering
IOF, 07745 Jena, Germany  Fabian Steinlechner<EMAIL_ADDRESS>Fraunhofer Institute for Applied Optics and Precision Engineering IOF, 07745
Jena, Germany  Abbe Center of Photonics, Friedrich Schiller University Jena,
07745 Jena, Germany  Stephan Fritzsche  Theoretisch-Physikalisches Institut,
Friedrich Schiller University Jena, 07743 Jena, Germany  Helmholtz-Institut
Jena, 07743 Jena, Germany  Abbe Center of Photonics, Friedrich Schiller
University Jena, 07745 Jena, Germany

###### Abstract

Spontaneous parametric down-conversion (SPDC) is a widely used source for
photonic entanglement. Years of focused research have led to a solid
understanding of the process, but a cohesive analytical description of the
paraxial biphoton state has yet to be achieved. We derive a general expression
for the spatio-temporal biphoton state that applies universally across common
experimental settings and correctly describes the nonseparability of spatial
and spectral modes. We formulate a criterion on how to decrease the coupling
of the spatial from the spectral degree of freedom by taking into account the
Gouy phase of interacting beams. This work provides new insights into the role
of the Gouy phase in SPDC, and also into the preparation of engineered
entangled states for multidimensional quantum information processing.

##  I Introduction

Photon pairs generated via spontaneous parametric down-conversion (SPDC) have
provided an experimental platform for fundamental quantum science [1] and
figure prominently in applications in quantum information processing,
including recent milestone experiments in photonic quantum computing [2].

Several works in recent years have addressed the challenge of tailoring the
spectral and spatial properties of $signal$ and $idler$ photons generated via
SPDC in theory and experiment. In the spatial domain, that is, the transverse
momentum space, much of this work was motivated by the objective of improving
fiber coupling efficiency [3, 4] or the dimensionality of spatial entanglement
[5, 6, 7]. In the spectral domain, the motivation was usually to engineer pure
spectral states, which are crucial for protocols based on multiphoton
interference [8]. This has been performed either by tailoring the nonlinearity
of the crystal [9] or by using counterpropagating photon pair generation in
periodically poled waveguides [10]. The frequency degree of freedom (DOF) has
also been used to generate entangled states via spatial shaping of the pump
beam [11] or by transferring polarization into color entanglement [12]. The
spatial shaping of the pump beam has been also used in Hong-Ou-Mandel
interference experiments, in order to control the two-photon interference
behavior [13].

Closed expressions for the state emitted by SPDC in bulk crystals have been
derived using very special techniques and approximations, such as the
narrowband [14], thin-crystal [15, 16] or plane wave approximations [17],
where either the spectral or spatial biphoton state is considered. However,
from the $X$-shaped spatio-temporal correlations [18, 19], the spatial and
spectral properties of SPDC have been known to be coupled. The $X$-shaped
spatio-temporal correlation implies that if the twin photons are collected
from different positions, they are detected with a certain time delay. In
contrast, if the photons are detected at the same position, the time delay is
very short (a few nanoseconds) [20].

To date, models that address both spectrum and space together have been
limited to approximate phase matching functions [21] or numerical calculations
[22]. The work [18] investigated the quite general phase matching function,
but the pump beam was limited to monochromatic plane wave.

Here, we present a simple-to-use closed expression for the biphoton state. The
approach describes the full spectral and spatial properties of all interacting
beams and applies to a wide range of experimental settings. The analytical
treatment of the biphoton state decomposed into discrete Laguerre Gaussian
(LG) modes also provides a deeper insight into the role of the Guoy phase in
PDC. Especially, the spectral response of spatial modes in SPDC is determined
by the Gouy phase of the pump, signal and idler beams. We will also show that
the Gouy phase can be used to control the coupling strength of spatial and
spectral DOF in parametric down-conversion (PDC). Next to providing an
intuitive understanding, we also demonstrate the utility of the expression for
quantum state engineering in spatial DOF for multidimensional quantum
information processing.

##  II Theoretical Methods

Let us start with the basic expressions of SPDC. We can make use of the
paraxial approximation, since typical optical apparatuses support only
paraxial rays about a central axis. In the paraxial regime, the longitudinal
and transverse components of the wave vector can be treated separately
$\bm{k}=\bm{q}+k_{z}(\omega)\bm{z}$. Consequently, the biphoton state in the
momentum space can be represented by the following expression [23, 24, 25]

$\displaystyle\ket{\Psi}=\iint$ $\displaystyle
d\bm{q}_{s}\>d\bm{q}_{i}\>d\omega_{s}\>d\omega_{i}\>\Phi(\bm{q}_{s},\bm{q}_{i},\omega_{s},\omega_{i})$
$\displaystyle\hat{a}^{\dagger}_{s}(\bm{q}_{s},\omega_{s})\>\hat{a}^{\dagger}_{i}(\bm{q}_{i},\omega_{i})\ket{vac}.$
(1)

Equation (II) refers to the generation of photon pairs with energies
$\omega_{s,i}$ and transverse momenta $\bm{q}_{s,i}$ from the vacuum state
$\ket{vac}$. The biphoton mode function
$\Phi(\bm{q}_{s},\bm{q}_{i},\omega_{s},\omega_{i})$ contains the rich high-
dimensional spatio-temporal structure of SPDC that arises from the coupling
between the wave vectors of the pump, signal, and idler beams.

###  II.1 Biphoton state decomposed in Laguerre Gaussian basis

The transverse spatial [24, 26] and frequency DOF [27] have been successfully
used in continuous variable information processing. However, in practical
experimental settings, the continuous variable space is more often discretized
using a set of modes. The proper choice of a set reduces the number of
dimensions needed to describe the state. Moreover, discrete modes are easy to
manipulate and detect using efficient experimental techniques [28, 29]. Since
the projection of the orbital angular momentum (OAM) is conserved in SPDC
[30], it is convenient to decompose the biphoton state into LG modes
$\ket{p,\ell,\omega}=\int
d\bm{q}\,\mathrm{LG}_{p}^{\ell}(\bm{q})\,\hat{a}^{\dagger}(\bm{q},\omega)\ket{vac}$,
which are eigenstates of OAM [31]:

$\displaystyle\ket{\Psi}=$ $\displaystyle\iint\>d\omega_{s}\>d\omega_{i}\>$
$\displaystyle\sum_{p_{s},p_{i}=0}^{\infty}\>\sum^{\infty}_{\ell_{s},\ell_{i}=-\infty}C_{p_{s},p_{i}}^{\ell_{s},\ell_{i}}\ket{p_{s},\ell_{s},\omega_{s}}\ket{p_{i},\ell_{i},\omega_{i}},$
(2)

where the coincidence amplitudes are calculated from the overlap integral
$C^{\ell_{s},\ell_{i}}_{p_{s},p_{i}}=\braket{p_{s},\ell_{s},\omega_{s};p_{i},\ell_{i},\omega_{i}}{\Psi}$,

$\displaystyle C^{\ell_{s},\ell_{i}}_{p_{s},p_{i}}=\iint
d\bm{q}_{s}\>d\bm{q}_{i}\>\Phi(\bm{q}_{s},\bm{q}_{i},\omega_{s},\omega_{i})\>[\mathrm{LG}_{p_{s}}^{\ell_{s}}(\bm{q}_{s})]^{*}$
$\displaystyle\times[\mathrm{LG}_{p_{i}}^{\ell_{i}}(\bm{q}_{i})]^{*}.$ (3)

The angular distribution of an LG mode in the momentum space is given by

$\displaystyle\mathrm{LG}_{p}^{\ell}(\rho,\varphi)=$ $\displaystyle
e^{\frac{-\rho^{2}\,w^{2}}{4}}\,e^{i\ell\,\varphi}\,\sum_{u=0}^{p}\,T_{u}^{p,\ell}\,\rho^{2k+\mathinner{\\!\left\lvert\ell\right\rvert}}$
(4)

with $T_{u}^{p,\ell}$ being

$\displaystyle T_{u}^{p,\ell}=$
$\displaystyle\sqrt{\frac{p!\,(p+|\ell|)!}{\pi}}\,\biggr{(}\frac{w}{\sqrt{2}}\biggl{)}^{2u+|\ell|+1}\,\frac{(-1)^{p+u}(i)^{\ell}}{(p-u)!\,(\mathinner{\\!\left\lvert\ell\right\rvert}+u)!\,u!},$

and where $\rho$ and $\varphi$ stand for the cylindrical coordinates
$\bm{q}=(\rho,\varphi)$. The summations in Eq. (II.1) run over the LG mode
numbers $p$ and $\ell$ associated with the radial momentum and the OAM
projection, respectively. Except for the fact, that we now deal with
summations instead of integrations, this discretization will also help us to
understand the coupling of spatial and spectral DOF in the frame of the Gouy
phase of LG modes. Note that we discretize only the transverse spatial DOF,
but in principle, it is also possible to discretize the frequency DOF [32].

The construction of the biphoton state reduces to the calculation of the
coincidence amplitudes $C^{\ell_{s},\ell_{i}}_{p_{s},p_{i}}$, which in turn
depend on the mode function
$\Phi(\bm{q}_{s},\bm{q}_{i},\omega_{s},\omega_{i})$. A compact expression for
the mode function can be derived if the experimental setup and geometry is
fixed.

###  II.2 Geometry and mode function

Here, we consider the scenario when a coherent laser beam propagates along the
$z$ axis and is focused in the middle of a nonlinear crystal placed at $z=0$.
Signal and idler fields propagate close to the pump direction, known as the
quasicollinear regime. The crystal and the pump beam have typical transverse
cross sections in the order of millimeters and micrometers, respectively.
Hence, we assume that the crystal compared to the pump beam is infinitely
extended in the transverse direction, which enforces the conservation of the
transverse momentum, $\bm{q}_{p}=\bm{q}_{s}+\bm{q}_{i}$ [23]. Taking into
account also the energy conservation $\omega_{p}=\omega_{s}+\omega_{i}$, the
mode function can be written as [25]

$\displaystyle\Phi(\bm{q}_{s},\bm{q}_{i},\omega_{s},\omega_{i})=$
$\displaystyle
N_{0}\,\mathrm{V}_{p}(\bm{q}_{s}+\bm{q}_{i})\,\mathrm{S}_{p}(\omega_{s}+\omega_{i})$
$\displaystyle\times\int_{-L/2}^{L/2}dz\>\exp{\biggr{[}iz(k_{z,p}-k_{z,s}-k_{z,i})\biggl{]}},$
(5)

where $N_{0}$ is the normalization constant, $\mathrm{V}_{p}(\bm{q}_{p})$ is
the spatial and $\mathrm{S}_{p}(\omega_{p})$ the spectral distribution of the
pump beam, and $L$ is the length of the nonlinear crystal along the $z$ axis.

The important component of the mode function (II.2) is the phase mismatch in
the $z$ direction $\Delta k_{z}=k_{z,p}-k_{z,s}-k_{z,i}$, which characterizes
the differences in the energies and momenta of the signal and idler photons.
Therefore, careful calculation of $\Delta k_{z}$ is essential for the
quantitative description of SPDC, which we will do next.

Experimentally generated lights are usually not monochromatic and contain a
frequency distribution. Therefore, except for the central frequencies that
meet energy conservation condition $\omega_{0,p}=\omega_{0,s}+\omega_{0,i}$,
we expect a deviation from them, $\omega=\omega_{0}+\Omega$ with the
assumption $\Omega\ll\omega_{0}$. Furthermore, in the paraxial approximation,
the transverse component of the momentum is much smaller than the longitudinal
component $|\bm{q}|\ll k$. Hence, we can apply the Taylor series on $k_{z}$
(Fresnel approximation) to $|\bm{q}|/k$ and also to small $\Omega$:

$k_{z}=k(\Omega)\sqrt{1-\frac{|\bm{q}|^{2}}{k(\Omega)^{2}}}\approx
k+\frac{\Omega}{u_{g}}+\frac{G\Omega^{2}}{2}-\frac{|\bm{q}|^{2}}{2k},$

where $u_{g}=1/(\partial k/\partial\Omega)$ is the group velocity and
$G=\partial/\partial\Omega\,(1/u_{g})$ is the group velocity dispersion,
evaluated at the respective central frequency. Here, we also assume that the
propagation is along a principal axis of the crystal, so we can ignore the
Poynting vector walk-off of extraordinary beams in the crystal. Next, we
insert the corresponding $k_{z}$ of the pump, signal, and idler into the phase
mismatch $\Delta k_{z}$ and arrive at

$\Delta
k_{z}=\Delta_{\Omega}+\rho_{s}^{2}\frac{k_{p}-k_{s}}{2k_{p}k_{s}}+\rho_{i}^{2}\frac{k_{p}-k_{i}}{2k_{p}k_{i}}-\frac{\rho_{s}\rho_{i}}{k_{p}}\cos{(\varphi_{i}-\varphi_{s})},$
(6)

where the frequency part $\Delta_{\Omega}$ is given by

$\displaystyle\Delta_{\Omega}=$
$\displaystyle\frac{\Omega_{s}+\Omega_{i}}{u_{g,p}}-\frac{\Omega_{s}}{u_{g,s}}-\frac{\Omega_{i}}{u_{g,i}}+\frac{G_{p}(\Omega_{s}+\Omega_{i})^{2}}{2}$
$\displaystyle-\frac{G_{s}\Omega_{s}^{2}}{2}-\frac{G_{i}\Omega_{i}^{2}}{2}.$
(7)

We used in Eq. (6) the relation
$\rho_{p}^{2}=\rho^{2}_{s}+\rho^{2}_{i}+2\rho_{s}\rho_{i}\cos{(\varphi_{i}-\varphi_{s})}$
and assumed momentum conservation for central frequencies $\Delta
k=k_{p}-k_{s}-k_{i}=0$. The condition $\Delta k=0$ ensures constructive
interference in the crystal between the pump, signal, and idler beams, which
is usually performed with birefringent crystals [25] or more recently by
periodic poling along the crystal axis, $k_{p}-k_{s}-k_{i}-2\pi/\Lambda=0$,
where $\Lambda$ is the poling period [33].

The remaining components of the mode function (II.2) that we should still fix
are the pump characteristics. We model the angular distribution of the pump
with an LG beam. The advantage of this choice is that an arbitrary paraxial
optical field can be expressed as a sum of LG beams
$\sum_{n}a_{n}\mathrm{LG}_{p_{n}}^{\ell_{n}}$ with $\sum_{n}|a_{n}|^{2}=1$ by
using their completeness relation. Thus, the theory developed for the LG pump
can be easily extended to SPDC with a particular pump. The amplitudes
$\eqref{expression}$ can then be upgraded to revised amplitudes
$\sum_{n}a_{n}C_{n}$, which follows from Eq. (II.1). Finally, the temporal
distribution is modeled with a Gaussian envelope of pulse duration $t_{0}$,
$\mathrm{S}_{p}(\omega_{p})=\exp{[-(\omega_{p}-\omega_{0,p})^{2}\,t_{0}^{2}/4]}\,t_{0}/\sqrt{\pi}$
[34], but which can be extended to any arbitrary pump spectrum.

###  II.3 Derivation of coincidence amplitudes

We can now substitute Eqs. (4)-(II.2) into Eq. (II.1) and calculate the
coincidence amplitudes:

$\displaystyle
C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}=N_{0}\sum_{u=0}^{p}\,\sum_{s=0}^{p_{s}}\,\sum_{i=0}^{p_{i}}\,T_{u}^{p,\ell}\>(T_{s}^{p_{s},\ell_{s}})^{*}\>(T_{i}^{p_{i},\ell_{i}})^{*}\>\int\,dz\,d\rho_{s}\,d\rho_{i}\,d\varphi_{s}\,d\varphi_{i}\,\Theta(z,\rho_{s},\rho_{i},\varphi_{i}-\varphi_{s})e^{i\ell\varphi_{s}}\,e^{i(-\ell_{s}\varphi_{s}-\ell_{i}\varphi_{i})},$
(8)

where we used a revised notation for the coincidence amplitudes
$C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}$ to indicate the mode numbers of
the pump. The function $\Theta(z,\rho_{s},\rho_{i},\varphi_{i}-\varphi_{s})$
is defined as

$\displaystyle\Theta(z,\rho_{s},\rho_{i},\varphi_{i}-\varphi_{s})=$
$\displaystyle[\rho^{2}_{s}+\rho^{2}_{i}+2\rho_{s}\rho_{i}\cos{(\varphi_{i}-\varphi_{s})}]^{\frac{2u+(\mathinner{\\!\left\lvert\ell\right\rvert}-\ell)}{2}}\,\rho_{s}^{\,\mathinner{\\!\left\lvert\ell_{s}\right\rvert}+2s+1}\,\rho_{i}^{\,\mathinner{\\!\left\lvert\ell_{i}\right\rvert}+2i+1}\,(\rho_{s}+\rho_{i}\,e^{i(\varphi_{i}-\varphi_{s})})^{\ell}$
$\displaystyle\times\exp{\biggl{[}-\frac{[\rho^{2}_{s}+\rho^{2}_{i}+2\rho_{s}\rho_{i}\cos{(\varphi_{i}-\varphi_{s})}]\,w^{2}}{4}-\frac{\rho_{s}^{2}\,w_{s}^{2}}{4}-\frac{\rho_{i}^{2}\,w_{i}^{2}}{4}\biggr{]}}\,\frac{t_{0}}{\sqrt{\pi}}e^{-\frac{t_{0}^{2}(\Omega_{s}+\Omega_{i})^{2}}{4}}$
$\displaystyle\times\exp{\biggr{[}iz\biggr{(}\Delta_{\Omega}+\rho_{s}^{2}\frac{k_{p}-k_{s}}{2k_{p}k_{s}}+\rho_{i}^{2}\frac{k_{p}-k_{i}}{2k_{p}k_{i}}-\cos{(\varphi_{i}-\varphi_{s})}\frac{\rho_{s}\rho_{i}}{k_{p}}\biggl{)}\biggl{]}}.$
(9)

In Eq. (9), the polar angle $\varphi$ of the pump beam has been expressed as a
function of signal and idler coordinates,

$e^{i\,\ell\,\varphi}=(\cos{\varphi}+i\sin{\varphi})^{\ell}=\frac{e^{i\ell\varphi_{s}}}{\rho_{p}^{\ell}}\,(\rho_{s}+\rho_{i}\,e^{i(\varphi_{i}-\varphi_{s})})^{\ell},$

by taking into account the conservation of transverse momentum,

$\bm{q}_{p}=\bm{q}_{s}+\bm{q}_{i}=\begin{pmatrix}\rho_{s}\cos{\varphi_{s}}+\rho_{i}\cos{\varphi_{i}}\\\
\rho_{s}\sin{\varphi_{s}}+\rho_{i}\sin{\varphi_{i}}\end{pmatrix}.$

The presentation of the coincidence amplitudes
$C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}$ in Eq. (8) with the function
$\Theta(z,\rho_{s},\rho_{i},\varphi_{i}-\varphi_{s})$ follows the goal to show
the OAM conservation in SDPC. To do so, we expand the function
$\Theta(z,\rho_{s},\rho_{i},\varphi_{i}-\varphi_{s})$ as superposition of
plane waves with the phases
$\exp{[i\ell^{{}^{\prime}}(\varphi_{i}-\varphi_{s})]}$(Fourier series with
complex coefficients),

$\Theta(z,\rho_{s},\rho_{i},\varphi_{i}-\varphi_{s})=\sum_{\ell^{{}^{\prime}}=-\infty}^{\infty}f_{\ell^{{}^{\prime}}}(z,\rho_{s},\rho_{i})e^{i\ell^{{}^{\prime}}(\varphi_{i}-\varphi_{s})}.$
(10)

We substitute expression (10) into Eq. (8) and perform the integration over
the polar angles $\varphi_{s}$ and $\varphi_{i}$:

$\displaystyle\sum^{\infty}_{\ell^{{}^{\prime}}=-\infty}f_{\ell^{{}^{\prime}}}(z,\rho_{s},\rho_{i})\int_{0}^{2\pi}\int_{0}^{2\pi}e^{i\ell\varphi_{s}}\,e^{i(-\ell_{s}\varphi_{s}-\ell_{i}\varphi_{i})}$
$\displaystyle\times
e^{i\ell^{{}^{\prime}}(\varphi_{i}-\varphi_{s})}d\varphi_{s}d\varphi_{i}\propto\delta_{\ell^{{}^{\prime}},\ell-\ell_{s}}\delta_{\ell^{{}^{\prime}},\ell_{i}}.$
(11)

As expected, the Kronecker delta functions appear in Eq. (11) which enforce
the conservation of OAM $\ell-\ell_{s}=\ell_{i}$. This conservation is not
valid out of the quasicollinear regime [35] because of the spin-orbital
angular momentum coupling in the nonparaxial regime [36]. In a non-collinear
regime, the total angular momentum should remain conserved, which can be a
future topic to study.

Going back to expression (8), we now calculate the integration over polar
coordinates $\varphi_{s,i}$ explicitly. For simplicity, we consider the
coincidence amplitudes $C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}$ for
positive OAM number of the pump beam $\ell\geq 0$. The coincidence amplitude
for $\ell<0$ is then given by
$C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}=(C_{p,p_{s},p_{i}}^{-\ell,-\ell_{s},-\ell_{i}})^{*}$,
which follows from Eq. (II.1). Furthermore, the two brackets on the first line
in Eq. (9) should be rewritten as finite sums by using the binomial formula.
For instance, the first bracket is written as

$\displaystyle[\rho^{2}_{s}+\rho^{2}_{i}+2\rho_{s}\rho_{i}\cos{(\varphi_{i}-\varphi_{s})}]^{u}=\sum_{m=0}^{u}\binom{u}{m}(\rho^{2}_{s}+\rho^{2}_{i})^{u-m}$
$\displaystyle\times[2\rho_{s}\rho_{i}\cos{(\varphi_{i}-\varphi_{s})}]^{m}.$

The cosine function can be expressed as the sum of two exponential functions
by using Euler’s formula, which should be again expressed as a Binomial sum.
After this step, the angular integration takes the form of the integral
representation of the Bessel function of the first kind [37]

$\frac{1}{2\pi}\int^{2\pi}_{0}e^{in\varphi\pm
iz\cos{(\varphi-\varphi^{\prime})}}d\varphi=(\pm
i)^{n}e^{in\varphi^{\prime}}J_{n}(z).$

Next, the sum representation of the Bessel function should be used

$J_{n}(z)=\sum_{k=0}^{\infty}\frac{(-1)^{k}}{k!\,\Gamma(k+n+1)}\biggl{(}\frac{z}{2}\biggr{)}^{2k+n},$
(12)

which transforms the integration over the radial coordinates into

$\int_{0}^{\infty}d\rho\,\rho^{n}e^{-a\,\rho^{2}}=\frac{\Gamma(\frac{n+1}{2})}{2a^{\frac{n+1}{2}}}.$

The final result is achieved via summing over $k$ from Eq. (12) by using the
definition of the Regularized hypergeometric function [38]. The coincidence
amplitudes read for $\ell\geq 0$ as

$\displaystyle C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}=$ $\displaystyle
N_{0}\,\pi^{3/2}\>t_{0}\>e^{-\frac{t_{0}^{2}(\Omega_{s}+\Omega_{i})^{2}}{4}}\>\delta_{\ell,\ell_{s}+\ell_{i}}$
$\displaystyle\sum_{u=0}^{p}\sum_{s=0}^{p_{s}}\sum_{i=0}^{p_{i}}T_{u}^{p,\ell}\>(T_{s}^{p_{s},\ell_{s}})^{*}\>(T_{i}^{p_{i},\ell_{i}})^{*}\>\sum_{n=0}^{\ell}\sum_{m=0}^{u}$
$\displaystyle\binom{\ell}{n}\binom{u}{m}\sum_{f=0}^{u-m}\sum_{v=0}^{m}\>\binom{u-m}{f}\binom{m}{v}\Gamma[h]\>\Gamma[b]$
$\displaystyle\int_{-L/2}^{L/2}dz\>e^{iz\,\Delta_{\Omega}}\>\frac{D^{d}}{H^{h}\>B^{b}}\>{{}_{2}}{\tilde{F}}_{1}\biggl{[}h,b,1+d,\frac{D^{2}}{H\,B}\biggl{]}$
(13)

and
$C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}=(C_{p,p_{s},p_{i}}^{-\ell,-\ell_{s},-\ell_{i}})^{*}$
for $\ell<0$. The function ${{}_{2}}{\tilde{F}}_{1}$ is known as the
regularized hypergeometric function [38]. The missing coefficients of Eq.
(II.3) are given by

$\displaystyle H$ $\displaystyle=$
$\displaystyle\frac{w_{p}^{2}}{4}+\frac{w_{s}^{2}}{4}-iz\frac{k_{p}-k_{s}}{2k_{p}k_{s}},\qquad
D=-\frac{w_{p}^{2}}{4}-iz\frac{1}{2k_{p}},$ $\displaystyle B$ $\displaystyle=$
$\displaystyle\frac{w_{p}^{2}}{4}+\frac{w_{i}^{2}}{4}-iz\frac{k_{p}-k_{i}}{2k_{p}k_{i}},\qquad
d=\ell_{i}+m-n-2v,$ $\displaystyle h$ $\displaystyle=$
$\displaystyle\frac{1}{2}(2+2s+\ell+\ell_{i}+2(-f+u)-2n-2v+\mathinner{\\!\left\lvert\ell_{s}\right\rvert}),$
$\displaystyle b$ $\displaystyle=$
$\displaystyle\frac{1}{2}(2+2f+2i+\ell_{i}+2m-2v+\mathinner{\\!\left\lvert\ell_{i}\right\rvert}),$

where $w_{p}$, $w_{s}$ and $w_{i}$ are the beam waists of the pump signal and
the idler beams, respectively.

Expression (II.3) for the coincidence amplitudes as a function of the pump
mode constitutes the main result of this work. It allows the spatial and
spectral emission profiles to be reconstructed mode by mode and is applicable
in any experimental setting that exhibits cylindrical symmetry. It can be
readily used to calculate many characteristics of SPDC: joint spectral
density, photon bandwidths, pair-collection probability, heralding ratio,
spectral and spatial correlation, etc. Previously, these could only be
achieved through numerical calculations or for special cases with a limited
scope of applicability. The experimental demonstration of Eq. (II.3) has
already been presented in Ref. [39], where we also showed how the coupling of
spatial and spectral DOF deteriorates the spatial entanglement but can be
compensated directly by a proper choice of the collection mode.

###  II.4 Gouy phase and spatio-temporal coupling

The spatio-temporal coupling encoded in Eq. (II.3) is a fundamental feature of
SPDC. However, the usual applications in quantum optics utilize either the
spatial or spectral DOF, neglecting the correlation between them.
Nevertheless, this coupling remains a fundamental issue in many protocols
based on entangled photon sources, where any distinguishability arising from
not-considered DOF reduces the coherence of the state. Next, we will
illustrate the utility of expression (II.3) in the frame of possible
decoupling of spatial and spectral DOF
$\Phi(\bm{q}_{s},\bm{q}_{i},\omega_{s},\omega_{i})=\Phi_{\bm{q}}(\bm{q}_{s},\bm{q}_{i})\Phi_{\omega}(\omega_{s},\omega_{i})$.
We will show that this decoupling is closely related to the Gouy phase of
interacting beams.

The role of the Gouy phase in nonlinear processes has been investigated
before. For instance, in SPDC, the change of the Gouy phase
$\psi_{G}(z)=(N+1)\arctan(z/z_{R})$ within the propagation distance has been
used to control the relative phase of two different LG modes of measurement
basis [40, 41]. Here, $N$ is the combined LG mode number
$N=2p+\mathinner{\\!\left\lvert\ell\right\rvert}$ and $z_{R}$ is the Rayleigh
length. In four-wave mixing (FWM), the conversion behavior between LG modes is
strongly affected by the Gouy phase [42]. The authors observed that the
existence of a relative Gouy phase between modes with different mode numbers
$N$ leads to a reduced FWM efficiency.

Here, we have a similar situation: pump, signal, and idler fields acquire
different Gouy phases along with propagation in the crystal due to different
mode numbers $N$, causing a reduced efficiency of mode down-conversion. We
expect intuitively that the shape of the spectrum of spatial modes is affected
by the relative Gouy phase of interacting beams. This is still a guess and
requires proof.

We consider for simplicity the scenario in which the Rayleigh lengths of the
three beams are equal $z_{R,p}=z_{R,i}=z_{R,s}$ and fixed. This condition
matches the Gouy angle $\arctan(z/z_{R})$ for all beams. Hence, the relative
Gouy phase can be written as

$\psi_{G,p}-\psi_{G,s}-\psi_{G,i}=(N_{p}-N_{s}-N_{i}-1)\arctan(z/z_{R}).$

This implies that the Gouy phase is fully defined by the relative mode number
$N_{R}=N_{p}-N_{s}-N_{i}$. If the Gouy phase is responsible for different
spectral dependencies of the coincidence amplitudes
$C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}(\Omega_{s}$,$\Omega_{i})$, the
shape of the spectrum should remain the same for fixed $N_{R}$. Assuming
$k_{p}=2k_{s}$, Eq. (II.3) transforms into

$C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}(\Omega_{s},\Omega_{i})\propto\int_{-L/2}^{L/2}dz\>e^{iz\,\Delta_{\Omega}}\>\frac{(i2z+k_{p}w_{p}^{2})^{N_{R}}}{(-i2z+k_{p}w_{p}^{2})^{N_{R}+1}}.$
(14)

We see from Eq. (14) that the spectral response of
$C_{p,p_{s},p_{i}}^{\ell,\ell_{s},\ell_{i}}(\Omega_{s},\Omega_{i})$ encoded
only in the term $e^{iz\,\Delta_{\Omega}}$ remains unaffected up to a constant
if $N_{R}$ is fixed. On the other hand, $N_{R}$ can be rewritten as

$N_{R}=\frac{\psi_{G,R}}{\arctan(z/z_{R})}+1,$ (15)

where $\psi_{G,R}$ is the relative Gouy phase
$\psi_{G,p}-\psi_{G,s}-\psi_{G,i}$. Therefore, it follows from Eqs. (14) and
(15) that the spectral response of coincidence amplitudes is determined by the
relative Gouy phase $\psi_{G,R}$ if the pump characteristics $z_{R}$, $w_{p}$,
and $k_{p}$ are fixed. This is what we wanted to prove. Note that the simple
form of Eq. (14) is due to the assumptions $k_{p}=2k_{s}$ and
$z_{R,p}=z_{R,i}=z_{R,s}$. The analytical proof for the general case requires
more effort, which we omit here.

This proof brings us a step forward in the decoupling problem of spatial and
spectral DOF: the decoupling can be achieved for a selected subspace of modes
that possess the same relative Gouy phase. So, if a state is engineered that
consists of modes with $N_{R}=const.$ assuming $z_{R,p}=z_{R,i}=z_{R,s}$, then
the modes contributing to the state have the same spectrum, i.e., the state is
separable. The question of decoupling of spatial and spectral DOF can be now
reformulated: How do we engineer a state only consisting of modes with the
same relative mode number $N_{R}$.

Figure 1: Normalized high-dimensional entangled states (a)
$\ket{\Psi^{\prime}_{4}}=\frac{1}{2}(\ket{0,0}+\ket{1,1}+\ket{2,2}+\ket{3,3})$
and (b) $\ket{\Psi_{4}}=\frac{1}{2}(\ket{0,1}+\ket{1,0}+\ket{2,3}+\ket{3,2})$.
The state $\ket{\Psi_{4}}$ is maximally entangled in the subspace
$\ell_{s},\ell_{i}=0,1,2,3$, but not $\ket{\Psi^{\prime}_{4}}$ because of
diagonal elements in the subspace. The state $\ket{\Psi^{\prime}_{4}}$ is
maximally entangled in a smaller subspace, namely in the subspace consisting
of only four modes $\\{\ket{0,0},\ket{1,1},\ket{2,2},\ket{3,3}\\}$. (c) The
modes of signal (idler) involved in the state $\ket{\Psi_{4}}$ have the same
spectrum (blue solid curve) compared to the modes out of the subspace shown in
the same color as corresponding bars (dotted curve $\ket{2,-1}$ and dashed
curve $\ket{3,-2}$). These curves correspond only to the spatial modes from
(b).

##  III Engineering high-dimensional entangled states in OAM basis

The state engineering in spatial DOF has been investigated theoretically in
the thin crystal regime [43, 44] and also implemented experimentally [45, 46].
In particular, three-, four-, and five-dimensional entangled states in OAM
basis have been generated in Ref. [45] using a superposition of LG beams for
the pump. The correct superposition for the pump has been determined with a
simultaneous perturbation stochastic approximation algorithm [47].

We show in this section how to calculate the correct superposition of LG beams
with Eq. (II.3), in order to generate entangled states in OAM basis including
the states from Ref. [45]. Our method is very straightforward and requires no
optimization algorithm. In comparison to Refs. [43, 44], our results can be
directly implemented in a real experiment, since we do not consider the thin
crystal approximation. State engineering in the thin crystal regime is
inefficient due to an infinite amount of spatial modes generated in the down-
conversion.

###  III.1 Determination of pump beam

We consider the four-dimensional subspace $\ell_{s},\ell_{i}=0,1,2,3$ and
$p_{s}=p_{i}=0$, which we refer to as $S_{4}$, with associated notation
$\ket{p_{s}=0,\ell_{s},\omega_{s}}\ket{p_{i}=0,\ell_{i},\omega_{i}}\mathrel{\mathop{\mathchar
58\relax}}=\ket{\ell_{s}(\omega_{s}),\ell_{i}(\omega_{i})}$. The goal is to
engineer a four-dimensional maximally entangled state in this subspace. We
model the pump beam as a superposition of LG beams,

$\displaystyle\mathrm{V}_{p}$ $\displaystyle=$
$\displaystyle\sum_{\ell}a_{\ell}\;\mathrm{LG}_{0}^{\ell},$

where the range of summation is determined with the possible minimal and
maximal OAM values in the subspace,
$\ell=[\mathrm{min}(\ell_{s}+\ell_{i}),\mathrm{max}(\ell_{s}+\ell_{i})]$. The
correct choice of the expansion amplitudes $a_{\ell}$ is now our task. Since
the pump function appears in Eq. (II.1) linearly, the corresponding state in
$S_{4}$ is given by

$\displaystyle\ket{\Psi_{4}}=\sum^{6}_{\ell=0}a_{\ell}\>\sum^{3}_{\ell_{s},\ell_{i}=0}C_{0,0,0}^{\ell,\ell_{s},\ell_{i}}\ket{\ell_{s},\ell_{i}}.$

The matrix representation of the state $\ket{\Psi_{4}}$ can clarify the right
choice of the coefficients $a_{\ell}$. The matrix consists of $16$ elements
and is given by the left-hand side of the following expression:

$\displaystyle\begin{pmatrix}{\color[rgb]{0,0,1}a_{0}}\,C_{0,0}&{\color[rgb]{1,0,0}a_{1}}\,C_{1,0}&{\color[rgb]{0,0,1}a_{2}}\,C_{2,0}&a_{3}\,C_{3,0}\\\
{\color[rgb]{1,0,0}a_{1}\,}C_{0,1}&{\color[rgb]{0,0,1}a_{2}}\,C_{1,1}&a_{3}\,C_{2,1}&{\color[rgb]{0,0,1}a_{4}}\,C_{3,1}\\\
{\color[rgb]{0,0,1}a_{2}}\,C_{0,2}&a_{3}\,C_{1,2}&{\color[rgb]{0,0,1}a_{4}}\,C_{2,2}&{\color[rgb]{1,0,0}a_{5}}\,C_{3,2}\\\
a_{3}\,C_{0,3}&{\color[rgb]{0,0,1}a_{4}}\,C_{1,3}&{\color[rgb]{1,0,0}a_{5}}\,C_{2,3}&{\color[rgb]{0,0,1}a_{6}}\,C_{3,3}\end{pmatrix}\rightarrow\begin{pmatrix}0&{\color[rgb]{1,0,0}1}&0&0\\\
{\color[rgb]{1,0,0}1}&0&0&0\\\ 0&0&0&{\color[rgb]{1,0,0}1}\\\
0&0&{\color[rgb]{1,0,0}1}&0\end{pmatrix},$ (16)

where we used the notation
$C_{0,0,0}^{\ell_{s}+\ell_{i},\ell_{s},\ell_{i}}=C_{\ell_{i},\ell_{s}}$. The
state becomes maximally entangled in this subspace if the matrix has exactly
one entry of $1$ in each row and each column and $0$ elsewhere (permutation
matrix). The right-hand side of expression (16) is such a state that can be
engineered if we select $a_{1}=1/C_{0,1}\approx 1/C_{1,0}$,
$a_{5}=1/C_{2,3}\approx 1/C_{3,2}$ and $a_{0}=a_{2}=a_{3}=a_{4}=a_{6}=0$,
where we assumed degenerate SPDC $k_{p}\approx 2k_{s}$. This choice leads to
the state
$\ket{\Psi_{4}}=\frac{1}{2}(\ket{0,1}+\ket{1,0}+\ket{2,3}+\ket{3,2})$. Thus,
the state engineering is finished, where the coefficients of the pump
superposition $\\{a_{1},a_{5}\\}$ should be calculated with the expression
(II.3). In the same way, the state
$\ket{\Psi^{\prime}_{4}}=\frac{1}{2}(\ket{0,0}+\ket{1,1}+\ket{2,2}+\ket{3,3})$
from Ref. [45] can also be engineered, if we select
$\\{a_{0},a_{2},a_{4},a_{6}\\}$ to be equal to
$\\{1/C_{0,0},1/C_{1,1},1/C_{2,2},1/C_{3,3}\\}$ and $a_{1}=a_{3}=a_{5}=0$.

The states $\ket{\Psi^{\prime}_{4}}$ and $\ket{\Psi_{4}}$ are presented in
Figs. 1(a) and 1(b) with blue-colored bars on top. As we can see, the modes
contributing to the states $\ket{\Psi^{\prime}_{4}}$ and $\ket{\Psi_{4}}$
represent just a part of the full OAM emission (spiral bandwidth). Therefore,
the postselection should be the final step in the engineering process, where
undesirable modes are sorted out. Next, we should calculate the Schmidt number
and the purity of the presented states, in order to evaluate the efficiency of
the state preparation in the subspace $S_{4}$. We will use for all our
calculations the same experimental parameters as in Ref. [45]: $15$-$mm$-thick
periodically poled $\mathrm{KTiOPO}_{4}$ crystal designed for a collinear
frequency degenerate type-II phase matching, continuous-wave laser of
wavelength $405$ nm with beam waist $w_{p}=25$ $\mu m$ and detection modes of
radius $w_{s,i}=33$ $\mu m$.

###  III.2 Schmidt number and purity of subspace states

We compare first the azimuthal Schmidt numbers of the states $\ket{\Psi_{4}}$
and $\ket{\Psi^{\prime}_{4}}$ in the subspace $S_{4}$. Obviously, the diagonal
modes $\\{\ket{0,2},\ket{2,0},\ket{1,3},\ket{3,1}\\}$ in Fig. 1 (a) are non-
desirable and lead to a decrease of entanglement in $S_{4}$. Consequently, the
state $\ket{\Psi^{\prime}_{4}}$ has an azimuthal Schmidt number less than $4$,
$K=2.04$, while the Schmidt number of the state $\ket{\Psi_{4}}$ equals $4$.
Therefore, the preparation of the state $\ket{\Psi_{4}}$ is more efficient in
$S_{4}$ than for $\ket{\Psi^{\prime}_{4}}$. $K=4$ is necessary, but not a
sufficient condition for a four-dimensional state to be maximally entangled.
Additionally, the state should be pure. Hence, the state $\ket{\Psi_{4}}$ can
be called maximally entangled in $S_{4}$, if it is also spatially pure.

In order to calculate the spatial purity of $\ket{\Psi_{4}}$, we need the
reduced density matrix $\rho_{\bm{q}}$, which results from tracing over the
frequency $\rho_{\bm{q}}=\mathrm{Tr}_{\Omega}(\rho)$. The fact of a continuous
wave laser $\mathrm{S}_{p}(\omega_{p})\propto\delta(\omega_{p}-\omega_{0,p})$
sets the condition $\Omega_{s}=-\Omega_{i}\mathrel{\mathop{\mathchar
58\relax}}=\Omega$, which transforms Eq. (II.1) into

$\displaystyle\ket{\Psi}=$
$\displaystyle\iint\>d\Omega\>\sum^{\infty}_{\ell_{s},\ell_{i}=-\infty}C_{\ell_{s},\ell_{i}}(\Omega)\ket{\ell_{s},\Omega}\ket{\ell_{i},-\Omega}.$
(17)

Now, we calculate the density matrix $\rho=|\Psi\rangle\langle\Psi|$ and then
trace over the spectral domain, which yields:

$\rho_{\bm{q}}=\sum_{\ell_{s},\ell_{i}}\sum_{\tilde{\ell_{s}},\tilde{\ell_{i}}}A^{\tilde{\ell_{s}},\tilde{\ell_{i}}}_{\ell_{s},\ell_{i}}|\ell_{s},\ell_{i}\rangle\langle\tilde{\ell_{s}},\tilde{\ell_{i}}|,$
(18)

where $A^{\tilde{\ell_{s}},\tilde{\ell_{i}}}_{\ell_{s},\ell_{i}}=\int
d\Omega\>C_{\ell_{s},\ell_{i}}(\Omega)\>[C_{\tilde{\ell_{s}},\tilde{\ell_{i}}}(\Omega)]^{*}$
is the overlap integral of the spectra of the OAM modes. Equation (18) is very
useful to calculate the spatial purity in small subspaces.

We run summations in Eq. (18) over
$\ell_{s},\ell_{i},\tilde{\ell_{s}},\tilde{\ell_{i}}=0,1,2,3$, renormalize the
state, construct the density matrix of the subspace $\rho_{\bm{q},s}$ and
calculate the purity $\mathrm{Tr}(\rho^{2}_{\bm{q},s})$. Here, the subscript
$s$ indicates the consideration of the subspace $S_{4}$. In fact, the state
$\ket{\Psi_{4}}$ is spatially pure, $\mathrm{Tr}(\rho^{2}_{\bm{q},s})=1$. The
reason is very trivial: all modes that contribute to the state consist of only
positive OAM numbers, which leads to the same
$N_{R}=\mathinner{\\!\left\lvert\ell\right\rvert}-\mathinner{\\!\left\lvert\ell_{s}\right\rvert}-\mathinner{\\!\left\lvert\ell_{i}\right\rvert}=0$
for all modes due to $\ell=\ell_{s}+\ell_{i}$. Moreover, the experimental
parameters from Ref. [45] satisfy the condition $z_{R,p}\approx z_{R,i}\approx
z_{R,s}$. Hence, all modes have the same relative Gouy phase and consequently,
the same spectrum, which is presented in Fig. 1(c) with the blue curve. This
means that the spatial and spectral DOF are decoupled in $S_{4}$.
Interestingly, even though the authors did not consider the spectral DOF, the
prepared state from Ref. [45] is also separable in space and frequency in the
smaller subspace of only four modes
$\\{\ket{0,0},\ket{1,1},\ket{2,2},\ket{3,3}\\}$. We suppose that the
engineering of maximally entangled states in spatial DOF in a certain subspace
enforces automatic decoupling in spatial and spectral DOF in that subspace.

###  III.3 Purity and Schmidt number of the full biphoton state

Obviously, the subspace $S_{4}$ is a part of the full SPDC emission. The first
four OAM modes out of the subspace in Fig. 1(b), $\ket{2,-1}$ and
$\ket{3,-2}$, possess different spectra in contrast to the modes in $S_{4}$,
shown in Fig. 1(c) with dotted and dashed curves, respectively. The appearance
of modes with distinguishable spectra indicates the inseparability of spatial
and spectral DOF out of $S_{4}$. The more distinguishable modes contribute to
the state, the stronger the spatio-temporal coupling. This, in turn, leads to
reduced purity for the spatial biphoton state. Usually, narrowband filters are
used in front of detectors to increase the purity of the spatial state. On one
hand, the spectral filters improve the purity of the spatial state; on the
other hand, they reduce the rate of entangled photons.

We calculated the spatial purity
$\mathrm{Tr}(\rho^{2}_{\bm{q},\mathrm{full}})$ [21] of the full biphoton state
(II) depending on the filter bandwidth $\Delta\lambda$, in order to quantify
the influence of spectral filters on the biphoton state. We chose as a pump
the same beam, which leads to the state $\ket{\Psi_{4}}$. Very narrow filters
are required to end up with a more or less pure state, as we can see from Fig.
2. For instance, a typical spectral filter with a bandwidth of $1$ nm would
leave the state in a mixed state of purity $0.33$.

The Schmidt number of the full spatio-temporal biphoton state is also
different in comparison to the subspace state. The total Schmidt number can be
calculated from the reduced density matrix in space and frequency for the
signal by tracing over the idler
$\rho_{\mathrm{signal}}=\mathrm{Tr}_{\mathrm{idler}}(\rho)$[21]. The Schmidt
number is then given by $K=1/\mathrm{Tr}(\rho^{2}_{\mathrm{signal}})$ [48].
The number of both spatial and spectral Schmidt modes in the range of
frequencies $810\pm 10$ nm equals $140$, where $810$ nm is the central
frequency for signal and idler photons. In comparison, the number of Schmidt
modes generated only at central frequency $810$ $nm$ equals $5.8$.

Finally, a small remark about the thin crystal regime: The spatio-temporal
coupling is absent in the thin crystal regime $L\ll z_{R,p}$, since the
biphoton state is independent of the crystal features. The problem with this
regime is that it gives rise to a huge amount of spatial modes. Assume we keep
all parameters the same as in Ref. [45], but change the crystal length to
$L=1$ $\mu m$. The thin crystal regime is then well achieved according to Ref.
[16]. The state becomes spatially pure, but possesses a large amount of
Schmidt modes, $10^{7}$.

Figure 2: Purity of the spatial biphoton state depending on the bandwidth of
spectral filter.

##  IV CONCLUSION

In summary, we derived a closed analytical expression for the biphoton spatio-
temporal state in terms of the LG mode amplitudes. The expression readily
reveals the dependence of the modal decomposition on frequency and thus
correctly describes spectral-spatial coupling, a quintessential feature of
SPDC. The expression provides a new understanding of how the Gouy phase is
related to the decoupling of spatial and spectral DOF: the relative Gouy phase
of the interacting beams fully defines the shape of the spectrum of down-
converted photons.

Engineering the modal decomposition of the pump beam can be used to engineer a
high-dimensional OAM entanglement. State engineering can also be used to
decrease the coupling between the spatial and spectral DOF, leading to an
increase of the correlation stored in the spatial DOF. We thus hope that it
will aid experimenters in the design and quantitative modeling of challenging
experiments based on PDC.

The authors thank Egor Kovlakov and Darvin Wanisch for very helpful
discussions.

## References

  * Anwar  _et al._ [2021] A. Anwar, C. Perumangatt, F. Steinlechner, T. Jennewein, and A. Ling, Review of Scientific Instruments 92, 041101 (2021).
  * Zhong  _et al._ [2020] H.-S. Zhong, H. Wang, Y.-H. Deng, M.-C. Chen, L.-C. Peng, Y.-H. Luo, J. Qin, D. Wu, X. Ding, Y. Hu, P. Hu, X.-Y. Yang, W.-J. Zhang, H. Li, Y. Li, X. Jiang, L. Gan, G. Yang, L. You, Z. Wang, L. Li, N.-L. Liu, C.-Y. Lu, and J.-W. Pan, Science 370, 1460 (2020).
  * Grice  _et al._ [2011] W. P. Grice, R. S. Bennink, D. S. Goodman, and A. T. Ryan, Phys. Rev. A 83, 023810 (2011).
  * Srivastav  _et al._ [2021] V. Srivastav, N. H. Valencia, S. Leedumrongwatthanakun, W. McCutcheon, and M. Malik, “Characterising and tailoring spatial correlations in multi-mode parametric downconversion,” (2021), arXiv:2110.03462 [quant-ph] .
  * Krenn  _et al._ [2014] M. Krenn, M. Huber, R. Fickler, R. Lapkiewicz, S. Ramelow, and A. Zeilinger, Proceedings of the National Academy of Sciences 111, 6243 (2014).
  * Baghdasaryan and Fritzsche [2020] B. Baghdasaryan and S. Fritzsche, Phys. Rev. A 102, 052412 (2020).
  * Chen  _et al._ [2020] Y. Chen, W. Zhang, D. Zhang, X. Qiu, and L. Chen, Phys. Rev. Applied 14, 054069 (2020).
  * Caspani  _et al._ [2017] L. Caspani, C. Xiong, B. J. Eggleton, D. Bajoni, M. Liscidini, M. Galli, R. Morandotti, and D. J. Moss, Light: Science & Applications 6, e17100 (2017).
  * Graffitti  _et al._ [2018] F. Graffitti, P. Barrow, M. Proietti, D. Kundys, and A. Fedrizzi, Optica 5, 514 (2018).
  * Luo  _et al._ [2020] K.-H. Luo, V. Ansari, M. Massaro, M. Santandrea, C. Eigner, R. Ricken, H. Herrmann, and C. Silberhorn, Opt. Express 28, 3215 (2020).
  * Francesconi  _et al._ [2021] S. Francesconi, A. Raymond, N. Fabre, A. Lemaître, M. I. Amanti, P. Milman, F. Baboux, and S. Ducci, ACS Photonics 8, 2764 (2021).
  * Ramelow  _et al._ [2009] S. Ramelow, L. Ratschbacher, A. Fedrizzi, N. K. Langford, and A. Zeilinger, Phys. Rev. Lett. 103, 253601 (2009).
  * Walborn  _et al._ [2003] S. P. Walborn, A. N. de Oliveira, S. Pádua, and C. H. Monken, Phys. Rev. Lett. 90, 143601 (2003).
  * Miatto  _et al._ [2011] F. M. Miatto, A. M. Yao, and S. M. Barnett, Phys. Rev. A 83, 033816 (2011).
  * Yao [2011] A. M. Yao, New Journal of Physics 13, 053048 (2011).
  * Baghdasaryan  _et al._ [2021] B. Baghdasaryan, F. Steinlechner, and S. Fritzsche, Phys. Rev. A 103, 063508 (2021).
  * Valencia  _et al._ [2007] A. Valencia, A. Ceré, X. Shi, G. Molina-Terriza, and J. P. Torres, Phys. Rev. Lett. 99, 243601 (2007).
  * Gatti  _et al._ [2009] A. Gatti, E. Brambilla, L. Caspani, O. Jedrkiewicz, and L. A. Lugiato, Phys. Rev. Lett. 102, 223601 (2009).
  * Jedrkiewicz  _et al._ [2012] O. Jedrkiewicz, A. Gatti, E. Brambilla, and P. Di Trapani, Phys. Rev. Lett. 109, 243901 (2012).
  * Zhang and Zhang [2017] D. Zhang and Z. Zhang, Scientific Reports 7, 42373 (2017).
  * Osorio  _et al._ [2008] C. I. Osorio, A. Valencia, and J. P. Torres, New Journal of Physics 10, 113012 (2008).
  * Gatti  _et al._ [2012] A. Gatti, T. Corti, E. Brambilla, and D. B. Horoshko, Phys. Rev. A 86, 053803 (2012).
  * Saleh  _et al._ [2000] B. E. A. Saleh, A. F. Abouraddy, A. V. Sergienko, and M. C. Teich, Phys. Rev. A 62, 043816 (2000).
  * Walborn  _et al._ [2010] S. Walborn, C. Monken, S. Pádua, and P. Souto Ribeiro, Physics Reports 495, 87 (2010).
  * Karan  _et al._ [2020] S. Karan, S. Aarav, H. Bharadhwaj, L. Taneja, A. De, G. Kulkarni, N. Meher, and A. K. Jha, Journal of Optics 22, 083501 (2020).
  * Tasca  _et al._ [2011] D. S. Tasca, R. M. Gomes, F. Toscano, P. H. Souto Ribeiro, and S. P. Walborn, Phys. Rev. A 83, 052325 (2011).
  * Fabre  _et al._ [2022] N. Fabre, A. Keller, and P. Milman, Phys. Rev. A 105, 052429 (2022).
  * Bolduc  _et al._ [2013] E. Bolduc, N. Bent, E. Santamato, E. Karimi, and R. W. Boyd, Optics Letters 38, 3546 (2013).
  * Eckstein  _et al._ [2011] A. Eckstein, B. Brecht, and C. Silberhorn, Opt. Express 19, 13770 (2011).
  * Mair  _et al._ [2001] A. Mair, A. Vaziri, G. Weihs, and A. Zeilinger, Nature 412, 313 (2001).
  * Fickler  _et al._ [2012] R. Fickler, R. Lapkiewicz, W. N. Plick, M. Krenn, C. Schaeff, S. Ramelow, and A. Zeilinger, Science 338, 640 (2012), https://www.science.org/doi/pdf/10.1126/science.1227193 .
  * Gil-Lopez  _et al._ [2021] J. Gil-Lopez, Y. S. Teo, S. De, B. Brecht, H. Jeong, C. Silberhorn, and L. L. Sánchez-Soto, Optica 8, 1296 (2021).
  * Fradkin  _et al._ [1999] K. Fradkin, A. Arie, A. Skliar, and G. Rosenman, Applied Physics Letters 74, 914 (1999), https://doi.org/10.1063/1.123408 .
  * Grice and Walmsley [1997] W. P. Grice and I. A. Walmsley, Phys. Rev. A 56, 1627 (1997).
  * Molina-Terriza  _et al._ [2003] G. Molina-Terriza, J. P. Torres, and L. Torner, Optics Communications 228, 155 (2003).
  * Baghdasaryan  _et al._ [2019] B. Baghdasaryan, B. Böning, W. Paufler, and S. Fritzsche, Phys. Rev. A 99, 023403 (2019).
  * Yousif and Melka [1997] H. A. Yousif and R. Melka, Computer Physics Communications 106, 199 (1997).
  * [38] W. R. Inc., “Gauss hypergeometric function,” .
  * Sevilla-Gutiérrez  _et al._ [2022] C. Sevilla-Gutiérrez, V. R. Kaipalath, B. Baghdasaryan, M. Gräfe, S. Fritzsche, and F. Steinlechner, “Spectral properties of transverse laguerre-gauss modes in parametric down-conversion,” (2022).
  * Kawase  _et al._ [2008] D. Kawase, Y. Miyamoto, M. Takeda, K. Sasaki, and S. Takeuchi, Phys. Rev. Lett. 101, 050501 (2008).
  * de Brito  _et al._ [2021] F. de Brito, I. da Paz, B. Hiller, J. B. Araujo, and M. Sampaio, Physics Letters A 386, 126989 (2021).
  * Offer  _et al._ [2021] R. F. Offer, A. Daffurn, E. Riis, P. F. Griffin, A. S. Arnold, and S. Franke-Arnold, Phys. Rev. A 103, L021502 (2021).
  * Torres  _et al._ [2003] J. P. Torres, Y. Deyanova, L. Torner, and G. Molina-Terriza, Phys. Rev. A 67, 052313 (2003).
  * Bornman  _et al._ [2021] N. Bornman, W. Tavares Buono, M. Lovemore, and A. Forbes, Advanced Quantum Technologies 4, 2100066 (2021), https://onlinelibrary.wiley.com/doi/pdf/10.1002/qute.202100066 .
  * Kovlakov  _et al._ [2018] E. V. Kovlakov, S. S. Straupe, and S. P. Kulik, Phys. Rev. A 98, 060301(R) (2018).
  * Liu  _et al._ [2018] S. Liu, Z. Zhou, S. Liu, Y. Li, Y. Li, C. Yang, Z. Xu, Z. Liu, G. Guo, and B. Shi, Phys. Rev. A 98, 062316 (2018).
  * Spall [1998] J. Spall, IEEE Transactions on Aerospace and Electronic Systems 34, 817 (1998).
  * Nielsen and Chuang [2000] M. A. Nielsen and I. L. Chuang,  _Quantum Computation and Quantum Information_ (Cambridge University Press, 2000).


# Universal Pyramid Adversarial Training for Improved ViT Performance

Ping-yeh Chiang1, Yipin Zhou2, Omid Poursaeed2, Satya Narayan Shukla2, Ashish
Shah2, Tom Goldstein1, Ser-Nam Lim2

1Department of Computer Science
University of Maryland - College Park

2Meta AI

###### Abstract

Recently, Pyramid Adversarial training (Herrmann et al., 2022) has been shown
to be very effective for improving clean accuracy and distribution-shift
robustness of vision transformers. However, due to the iterative nature of
adversarial training, the technique is up to 7 times more expensive than
standard training. To make the method more efficient, we propose Universal
Pyramid Adversarial training, where we learn a single pyramid adversarial
pattern shared across the whole dataset instead of the sample-wise patterns.
With our proposed technique, we decrease the computational cost of Pyramid
Adversarial training by up to 70% while retaining the majority of its benefit
on clean performance and distribution-shift robustness. In addition, to the
best of our knowledge, we are also the first to find that universal
adversarial training can be leveraged to improve clean model performance.

##  1 Introduction

Human intelligence is exceptional at generalizing to previously unforeseen
circumstances. While deep learning models have made great strides with respect
to clean accuracy on a test set drawn from the same distribution as the
training data, a model’s performance often significantly degrades when
confronted with distribution shifts that are qualitatively insignificant to a
human. Most notably, deep learning models are still susceptible to adversarial
examples (perturbations that are deliberately crafted to harm accuracy) and
out-of-distribution samples (images that are corrupted or shifted to a
different domain).

Adversarial training has recently been shown to be a promising avenue for
improving both clean accuracy and robustness to distribution shifts. While
adversarial training was historically used for enhancing adversarial
robustness, recent works (Xie et al., 2020; Herrmann et al., 2022) found that
properly adapted adversarial training regimens could be used to achieve state-
of-the-art results (at the time of publication) on Imagenet (Xie et al., 2020)
and out-of-distribution robustness (Herrmann et al., 2022).

However, both proposed techniques (Herrmann et al., 2022; Xie et al., 2020)
use up to 7 times the standard training compute due to the sample-wise and
multi-step procedure for generating adversarial samples. The expensive cost
has prevented it from being incorporated into standard training pipelines and
more widespread adoptions. In this paper, we seek to improve the efficiency of
the adversarial training technique so that it can become more accessible for
practitioners and researchers.

Several prior works (Shafahi et al., 2019; Zhang et al., 2019; Mei et al.,
2022; Zheng et al., 2020; Wong et al., 2020) have proposed methods to increase
the efficiency of adversarial training in the context of adversarial
robustness, where they try to make models robust to deliberate malicious
attacks. Shafahi et al. (2019) proposed reusing the parameter gradient for
training during the sample-wise adversarial step for faster convergence.
Later, Wong et al. (2020) proposed making adversarial training more efficient
with a single-step adversary rather than the expensive multi-step adversary.
However, all prior works focus on the efficiency trade-off concerning
adversarial robustness rather than clean accuracy or out-of-distribution
robustness. In the setting of adversarial robustness, one often assumes a
deliberate and all-knowing adversary. Security is crucial, yet in reality,
deep learning systems already exhibit a significant number of errors without
adversaries, such as self-driving cars making mistakes in challenging
environments. Consequently, clean accuracy and robustness to out-of-
distribution data are typically prioritized in most industrial settings. Yet,
few works seek to improve the efficiency trade-off for the out-of-distribution
metric. Our work aims to fill this gap.

Figure 1: A graphical illustration of how our Universal Pyramid Adversarial
training is more efficient compared to Pyramid Adversarial training. In the
specific examples, the two-step Pyramid Adversarial training requires 4
forward and backward passes for a single example, where 2 passes are used for
generating the adversarial sample and 2 additional passes are used for
training. On the other hand, our proposed approach only requires 2 forward and
backward passes. These two passes are needed because the batch size has been
doubled since we optimize both the clean and the adversarial objectives.

By shifting the context from adversarial robustness to clean accuracy and out-
of-distribution robustness, we can free ourselves from certain constraints,
such as the need to train the model on _sample-wise_ adversaries, which is
very expensive to compute. Instead, we can leverage _universal_ perturbations,
which are shared across the whole dataset. By leveraging this simple idea, we
can generate adversarial samples for _free_ while getting more performance
improvement on clean accuracy compared to prior work (Herrmann et al., 2022).

In this paper, we focus our experiments on the Vision Transformer architecture
(Dosovitskiy et al., 2020). We focus on this architecture as it is the most
general and scalable architecture that applies to many domains, including
vision, language, and audio, while simultaneously achieving SOTA on many of
them. We believe focusing on this architecture will lead to more valuable
techniques for the community.

In summary, here are our three main contributions:

  * •

We propose Universal Pyramid Adversarial training that is 70% more efficient
than the multi-step approach while increasing ViT’s clean accuracy more than
Pyramid Adversarial training.

  * •

We evaluate our technique on 5 out-of-distribution datasets and find that
Universal Pyramid Adversarial training effectively increases the
distributional robustness and is competitive with Pyramid Adversarial Training
while being efficient.

  * •

To the best of our knowledge, we are the first to identify universal
adversarial training as a viable technique for improving clean performance and
out-of-distribution robustness on Imagenet 1K. In our ablations, , we found
that the pyramid structure is critical for the performance gain and plain
universal adversarial training is detrimental to performance, unlike Herrmann
et al. (2022), which found both instance wise adversarial training and pyramid
adversarial training to be beneficial.

##  2 Related Work

Improving the efficiency of adversarial training has been widely studied
(Shafahi et al., 2019; Zhang et al., 2019; Zheng et al., 2020; Wong et al.,
2020), but they have mainly been in the context of adversarial robustness.
Shafahi et al. (2019) proposed reusing the parameter gradients from the
adversarial step. By reusing the free parameter gradients for training, they
were able to achieve much faster convergence. Even though the proposed
approach was much more efficient, Shafahi et al. (2019) could not reach the
same robustness level as the original multi-step training on Imagenet-1K.
Zhang et al. (2019) proposed making the iterative attack cheaper by updating
the noise based on the Hamiltonian functions of the first few layers. Wong et
al. (2020) proposed using a single-step adversary for training instead of a
multi-step adversary. They found that random initialization and early stopping
could prevent adversarial over-fitting, where label leakage happens from using
adversaries with fewer steps. Zheng et al. (2020) proposed reusing the
adversarial perturbations between epochs with the observation that adversarial
noises are often transferable. The downside of the method is that the memory
requirement grows with the data size, which can be quite large given the size
of modern datasets. We differ from all the prior works in that we aim to
investigate the efficiency gain of adversarial training in the context of
clean accuracy. By focusing on clean accuracy and out-of-distribution
robustness, we gained more flexibility concerning the formulation of the min-
max problem.

Xie et al. (2020) was the first paper that showed adversarial training could
improve clean performance of convolutional networks. To achieve this, Xie et
al. (2020) employed split batchnorms (AdvProp) for adversarial and clean
samples. They argued that clean and adversarial samples have very different
distributions and that split batchnorms are needed to make optimization
easier. Before Xie et al. (2020), the community commonly believed that
adversarial training leads to a decrease in clean accuracy (Madry et al.,
2018).

In a similar line of work, Mei et al. (2022) proposed a faster variant of
AdvProp (Xie et al., 2020) that makes the training speed comparable to
standard training while retaining some of Advprop’s benefits. Even though the
proposed method is more efficient, it substantially trades off the performance
of the original multi-step method. Our work is similar to Mei et al. (2022)
because we also focus on improving the efficiency of the adversarial training
process, but we differ in that we focus on ViT architecture with Pyramid
Adversarial training where their proposed approach is not applicable. Also, we
can achieve a performance gain that is _comparable or better_ than the multi-
step approach, whereas the previous method trades off performance for
efficiency.

Several recent approaches have shown that adversarial training can be used to
improve the performance of vision transformers. Bai et al. (2022) showed that
ViT relies more on low-frequency signals than high-frequency signals. By
adversarially training the model on high-frequency signals, Bai et al. (2022)
further boosted ViT’s performance. Mao et al. (2022) showed that by converting
images to discrete tokens, adversarial training could further increase the
performance of ViTs. Later, Herrmann et al. (2022) showed that by
incorporating the pyramid structure into standard adversarial training, they
could boost the performance of ViT, where the split batchnorms idea introduced
in Xie et al. (2020) were not directly applicable to ViT models. In our work,
we focus on the Pyramid Adversarial training technique proposed by Herrmann et
al. (2022) since it is the best-performing method that achieves SOTA on
multiple fronts while being applicable to ViT, a more modern architecture. The
main drawback of Herrmann et al. (2022) is the significantly higher training
time which can go up to 7x of the standard training time. In this work, we
propose Universal Pyramid Adversarial training to improve the efficiency of
Pyramid Adversarial training while retaining its effectiveness.

While universal adversarial samples have been used in prior work Shafahi et
al. (2020) for training to defend against universal adversarial attacks, our
proposed approach differs from them in that it leverages these samples to
improve clean model performance. Benz et al. (2021) finds that universal
perturbations tend to slide images into some classes more than others. They
find that by updating universal perturbations in a class-wise manner, they can
achieve better robustness compared to Shafahi et al. (2019). Both prior works
Shafahi et al. (2020); Benz et al. (2021) show that universal adversarial
training consistently decreases the performance of the model, similar to
standard adversarial training. Our ablation study shows that incorporating
both clean loss and pyramid structures are crucial for the performance gain
observed in Universal Pyramid Adversarial training. Without our proposed
modifications, universal adversarial training consistently decreases the clean
performance of the model. To the best of our knowledge, we are the first to
show that universal adversarial training can be leveraged for improved model
performance.

##  3 Method

In this section, we will go over the formulation of the proposed adversarial
training objective, the pyramid structure that we leveraged from Herrmann et
al. (2022), and our more efficient Universal Pyramid Adversarial training.

###  3.1 Adversarial Training

Adversarial training remains one of the most effective methods for defending
against adversarial attacks Bai et al. (2021). Adversarial training is aimed
at solving the following min-max optimization problem:

$\min_{\theta}E_{(x,y)\sim D}\left[\max_{\delta\in
B}L(f(x+\delta;\theta),y)\right],$ (1)

where $\theta$ is the model parameter, $\delta$ is the adversarial
perturbation, $L$ is the loss function, $D$ is the data distribution, and $B$
is the constraint for the adversarial perturbation, which is often an
$\ell_{\infty}$ ball. The inner objective seeks to find an adversarial
perturbation within the constraint, and the outer objective aims to minimize
the worst-case loss by optimizing the model parameters. While the method
effectively improves robustness to adversarial attacks, it often reduces clean
performance. However, the loss of clean performance is often not acceptable
for most practical applications.

Since our goal is to improve performance as opposed to worst case robustness,
we train the model on the following formulation (similar to Xie et al. (2020);
Herrmann et al. (2022)) instead where the clean loss is optimized in addition
to the adversarial loss:

$\small\min_{\theta}\underset{(x,y)\sim
D}{E}\left[L(f(x;\theta),y)+\lambda\max_{\delta\in
B}L(f(x+\delta;\theta),y)\right].$ (2)

Here, the $\lambda$ controls the trade-off between adversarial and clean loss.
However, adding clean loss alone is often not sufficient for improving the
performance of the model Xie et al. (2020). Additional techniques such as
split batchnorms Xie et al. (2020) and pyramid structures Herrmann et al.
(2022) are necessary for performance gain.

The main problem with the adversarial training formulation is that the inner
maximization is often expensive to compute, requiring several steps to
approximate accurately Madry et al. (2018). Specifically, for each iteration
of the adversarial step, one needs a full forward and backward pass on all of
the examples in a batch (see Figure 1). For example, if five steps are used,
which is the setting in both Xie et al. (2020); Herrmann et al. (2022), then
five forward and backward passes are needed. The generation of adversarial
samples is already five times more expensive than regular training. In
addition, one needs to use both the clean and the generated adversarial
samples for training, so now we have doubled the batch size. The larger batch
size increases the cost by another factor of two. When training with a
$5$-step adversary, the total computational cost will be seven times more
expensive than standard training. In Section 3.3, we describe our proposed
Universal Pyramid Adversarial training, where we can substantially reduce the
computational cost.

###  3.2 Pyramid Structure

Adversarial training alone even when coupled with the clean loss does not
typically increase performance of the model Xie et al. (2020). In order to
increase clean accuracy, certain techniques have to be used. Here we leverage
the pyramid structure from Herrmann et al. (2022). Herrmann et al. (2022)
aimed to endow the adversarial perturbation with more structure so that the
adversary can make larger edits without changing an image’s class. The pyramid
adversarial noise is parameterized with different levels of scales as follows:

$\vspace{-2mm}\delta=\sum_{s\in S}m_{s}\cdot C_{B}(\delta_{s}),\vspace{-2mm}$
(3)

where $C_{B}$ clips the noise within the constraint set $B$, $S$ is all of the
scales used, $m_{s}$ is the multiplicative constant, and $\delta_{s}$ is
perturbation at scale $s$. For the $\delta_{s}$ at a given scale, $s\times s$
number of pixels within a square tile share a single parameter giving greater
structure to the noise. Since the larger scale can often tolerate more
changes, larger $m_{s}$ at the coarser scales allow us to update the coarser
noises more quickly relative to the granular noise.

###  3.3 Universal Pyramid Adversarial Training

While Pyramid Adversarial training Herrmann et al. (2022) is effective at
increasing clean model performance, it is seven times more expensive compared
to standard training. To address this, we propose Universal Pyramid
Adversarial training, an efficient adversarial training approach to improve
model performance on clean and out-of-distribution data. Our proposed approach
learns a universal adversarial perturbation with pyramid structure, thus
unifying both the effectiveness of Pyramid Adversarial training and the
efficiency of universal adversarial training Shafahi et al. (2020).
Specifically, we attempt to solve the following objective:

$\small\min_{\theta}\max_{\delta\in B}\underset{(x,y)\sim
D}{E}\left[L(f(x;\theta),y)+\lambda L(f(x+\delta;\theta),y)\right].$ (4)

With this objective, we only have to solve for a _single_ universal
adversarial pattern that can be shared across the whole dataset, and we do not
have to optimize a new adversary for each sample. Even though the objective
looks similar, they are not the same. Due to Jensen’s inequality, Equation 4
is always strictly upper-bounded by Equation 2. We have described the complete
method in Algorithm 1. This yields up to 70% saving compared to the 5-step
sample-wise approach (see Table 1). Further, we update the universal
adversarial pattern during the backward pass of training, where we can get the
gradients of $\delta$ for free (see Figure 1 for an illustration of how
universal adversarial training can help save compute). However, since we still
need to train the model on twice the number of samples, our proposal is still
twice as expensive as standard training, but it is already 33% cheaper than
the fastest (one-step) sample-wise adversarial training approach.

More concretely, the generation step for the one-step sample-wise adversarial
training costs a single forward-backward pass, and the training step is twice
as expensive as standard training. Overall, one-step sample-wise adversarial
training is 3x the cost of standard training making our method 33% faster.
This is because in case of the one-step adversarial training, the gradient
from the first generation step cannot be reused because the patterns are
randomly initialized and the induced gradient is different from the clean
training gradient.

Algorithm 1 Universal Pyramid Adversarial Training

$\tau$ is the step size of adversarial attack

$\lambda$ controls the regularization strength

Initialize $\delta_{s}$ as zeros

for epoch = 1 … $N_{ep}$ do

for $x,y$ $\in$ $D$ do

$\delta\leftarrow\sum_{s\in S}m_{s}C_{B}(\delta_{s})$

$\mathcal{L}(x,y,\delta;\theta)\leftarrow\ell(x,y;\theta)+\lambda\ell(x+\delta,y;\theta)$

$\theta\leftarrow\theta-lr\cdot\nabla_{\theta}\mathcal{L}(x,y,\delta;\theta)$

\\\Update model parameters with some optimizer

for $s\in S$ do

$\delta_{s}\leftarrow\delta_{s}+\tau\cdot
sign(\nabla_{\delta_{s}}\mathcal{L}(x,y,\delta;\theta))$

\\\Update noise with the free gradients

end for

end for

end for

Augmentation | Method | Radius | Radius Schedule | # of Steps | Training Time (hrs) | Top 1 Accuracy | Gain
---|---|---|---|---|---|---|---
Weak | Baseline | - | - | - | 12.7 | 72.90% | -
Pyramid Adversarial TrainingHerrmann et al. (2022) | 6/255 | - | 1 | 36.5 | 73.52% | 0.62%
6/255 | - | 2 | 49.3 | 73.95% | 1.05%
6/255 | - | 3 | 61.8 | 74.68% | 1.78%
6/255 | - | 4 | 75.4 | 74.80% | 1.90%
6/255 | - | 5 | 88.9 | 74.18% | 1.28%
Universal Pyramid Adversarial Training | 8/255 | Yes | - | 26.6 | 74.87% | 1.97%
8/255 | - | - | 26.6 | 74.57% | 1.67%
Strong | Baseline | - | - | - | 12.7 | 79.85% | -
Pyramid Adversarial TrainingHerrmann et al. (2022) | 6/255 | - | 1 | 36.5 | 79.46% | -0.38%
6/255 | - | 2 | 49.3 | 79.87% | 0.03%
6/255 | - | 3 | 61.8 | 80.19% | 0.35%
6/255 | - | 4 | 75.4 | 80.04% | 0.22%
6/255 | - | 5 | 88.9 | 80.10% | 0.26%
Universal Pyramid Adversarial Training | 8/255 | Yes | - | 26.6 | 80.15% | 0.31%
8/255 | - | - | 26.6 | 80.28% | 0.44%

Table 1: Comparison of the effectiveness of universal and sample-wise pyramid
adversarial training. Universal Pyramid Adversarial training improves
performance compared to sample-wise Pyramid Adversarial training while being
much more efficient.

###  3.4 Radius Schedule

In our experiments, we find that a radius schedule occasionally benefits
performance. The radius dictates the extent to which an adversary is permitted
to alter the image, as measured by the $\ell_{\infty}$ distance between the
original and perturbed images. A larger radius permits greater perturbations,
thus strengthening the adversary, while a smaller radius restricts the
perturbation, rendering the adversary weaker. We propose this schedule as we
observe that a more aggressive/larger radius tends to promote faster
convergence at the beginning, but these images are very far out of
distribution, resulting in poor performance. By using a linearly decreasing
radius schedule, we are sometimes able to get a considerable performance boost
while maintaining fast convergence. Precisely, we calculate the radius at a
given epoch as follows

$r(e)=r_{\text{start}}+(r_{\text{end}}-r_{start})\frac{\max(e-e_{start},0)}{e_{end}-e_{start}},$
(5)

where $r_{start},r_{end}$ are the starting and ending radius with
$r_{start}>r_{end}$, $e_{start},e_{end}$ are the starting and ending epochs
for the radius schedule, and $r(e)$ is the radius at a given epoch $e$.

##  4 Experiments

###  4.1 Experimental Set-up

In all of our experiments, we focus on the training setup in Beyer et al.
(2022) since it allows us to achieve a competitive 79.8% on Imagenet-1K with a
ViT-S/16. The setup allows us to study ViT in a computationally feasible
setting.

Following Beyer et al. (2022), we use the AdamW optimizer, with a batch size
of 1024, a learning rate of 0.001 with a linear warm-up for the first 8
epochs, and weight decay of 0.1. We train the model for a total of 300 epochs
across all settings. For augmentation, we apply a simple inception crop and
horizontal flip. For experiments with strong data augmentation, we apply
RandomAugment Cubuk et al. (2020) of level 10 and MixUp Zhang et al. (2017) of
probability 0.2. We used strong data augmentation in all of our experiments
except for the first part of experiments in Table 1.

For Pyramid Adversarial training, following Herrmann et al. (2022) we use
$S=[32,16,1]$, $M=[20,10,1]$, and radius of 6/255. For step size, we simply
divide the radius by the number of steps used. For our proposed Universal
Pyramid Adversarial training, we use the same $M$ and $S$ as Pyramid
Adversarial training, but with a radius of 8/255. When a radius schedule is
used, we linearly decrease the radius by 90% in increments starting from epoch
30.

In addition to Imagenet-1K, we also evaluated our models on five out-of-
distribution datasets: Imagenet-C (Hendrycks & Dietterich, 2019), Imagenet-A
(Hendrycks et al., 2021), Imagenet-Rendition, Imagenet Sketch (Wang et al.,
2019), and Stylized Imagenet (Geirhos et al., 2019). Using a diverse set of
out-of-distribution datasets, we can more thoroughly evaluate the model’s
robustness to unexpected distribution shifts.

Method | Radius | Training Time (hrs) | Clean $\uparrow$ | C $\downarrow$ | A $\uparrow$ | Rendition $\uparrow$ | Sketch $\uparrow$ | Stylized $\uparrow$
---|---|---|---|---|---|---|---|---
Universal Pyramid Adversarial Training | 8/255 | 26.6 | 80.28% | 41.02% | 29.08% | 35.45% | 16.89% | 17.74%
Gain/ Loss Relative to |  |  |  |  |  |  |  |
\- Regular Training | - | 12.7 | 0.44% | -1.16% | 1.56% | 1.80% | 1.81% | 0.68%
\- Pyramid Adversarial Training 1 steps | 6/255 | 36.5 | 0.82% | 0.62% | 2.60% | 0.11% | 0.33% | -1.26%
\- Pyramid Adversarial Training 2 steps | 6/255 | 49.3 | 0.41% | 1.25% | 0.59% | -0.71% | -0.06% | -2.21%
\- Pyramid Adversarial Training 3 steps | 6/255 | 61.8 | 0.09% | 2.10% | -1.01% | -2.29% | -0.77% | -3.22%
\- Pyramid Adversarial Training 4 steps | 6/255 | 75.4 | 0.24% | 1.80% | -1.25% | -2.95% | -1.13% | -3.28%
\- Pyramid Adversarial Training 5 steps | 6/255 | 88.9 | 0.18% | 2.17% | -0.96% | -3.17% | -0.91% | -3.47%
Universal Pyramid Adversarial training | 12/255 | 26.6 | 80.04% | 40.37% | 28.21% | 37.05% | 17.42% | 19.53%
Gain/ Loss Relative to |  |  |  |  |  |  |  |
\- Regular Training | - | 12.7 | 0.20% | -1.81% | 0.69% | 3.39% | 2.34% | 2.48%
\- Pyramid Adversarial Training 1 steps | 6/255 | 36.5 | 0.58% | -0.02% | 1.73% | 1.70% | 0.87% | 0.53%
\- Pyramid Adversarial Training 2 steps | 6/255 | 49.3 | 0.17% | 0.60% | -0.28% | 0.88% | 0.48% | -0.41%
\- Pyramid Adversarial Training 3 steps | 6/255 | 61.8 | -0.15% | 1.46% | -1.88% | -0.70% | -0.23% | -1.42%
\- Pyramid Adversarial Training 4 steps | 6/255 | 75.4 | 0.00% | 1.15% | -2.12% | -1.36% | -0.59% | -1.48%
\- Pyramid Adversarial Training 5 steps | 6/255 | 88.9 | -0.06% | 1.52% | -1.83% | -1.58% | -0.38% | -1.67%

Table 2: Evaluating models trained with Universal Pyramid Adversarial training
(across two radius) on 5 additional out-of-distribution datasets and comparing
them with regular and Pyramid Adversarial training. Universal Pyramid
Adversarial training consistently boosts out-of-distribution robustness and is
competitive with 1-step and 2-step Pyramid Adversarial training. We report the
gain/loss relative to our method, i.e., the values in the colored cells are
calculated by the following formula: (our method - the alternative method).
For all columns except the Imagenet-C column, positive numbers mean that our
method performs better and vice versa. We report mean Corruption Error (mCE)
for Imagenet-C where lower is better. For the rest of the datasets, we simply
report the top-1 accuracy. The complete table can be found in the Appendix.

###  4.2 Experimental Results

Overall, we find that Universal Pyramid Adversarial training effectively
increases the clean accuracy and out-of-distributional robustness similar to
the original Pyramid Adversarial training while being much more efficient.

#### Clean Accuracy

Here we first analyze the effectiveness of our proposed Universal Pyramid
Adversarial training when applied to ViT with weak data augmentation. Pyramid
Adversarial training, as expected, significantly increases the performance of
the ViT by up to 1.9% (Table 1) when a 4-step attack is used. As we increase
the step count to 5, the benefit of Pyramid Adversarial training starts
diminishing. On the other hand, our Universal Pyramid Adversarial training
increased the performance even further. With the radius schedule, we can
obtain a performance increase of 1.97%, exceeding the performance benefit of
the Pyramid Adversarial training for all step counts. Without the radius
schedule, we still obtain a competitive gain of 1.67%. In addition to better
performance, our method is much more efficient than the original Pyramid
Adversarial training. In Table 1, we also reported the training time of each
method on 8 Nvidia A100 GPUs in hours, and our approach is 70% faster compared
to the 5-step Pyramid Adversarial training.

To further verify our proposal’s effectiveness, we analyze our method’s
performance when coupled with strong data augmentations. The setting with
strong data augmentation is a more challenging setting since all components of
the training pipelines are heavily tuned. It is worth noting that our baseline
ViT-S performs comparably with the baseline ViT-B/16 in Herrmann et al.
(2022); Steiner et al. (2021). Again, we continue to see the benefit of
Universal Pyramid Adversarial training compared to standard training.

In the more challenging setting, we see less benefit from both approaches, as
expected. As we increase the number of steps for the Pyramid Adversarial
training, the accuracy first increases, reaching the maximum at 3 steps, and
starts decreasing with 4 or more steps. On the other hand, our Universal
Pyramid Adversarial training achieves more performance gain than all of the
step count tested while being more efficient.

#### Out-of-Distribution Robustness

We found that Universal Pyramid Adversarial training effectively increases
models’ out-of-distribution robustness and is comparable to 1-step and 2-step
Pyramid Adversarial training. In Table 2, we see that our Universal Pyramid
Adversarial training consistently increases models’ performance on all five
out-of-distribution datasets relative to the baseline. When compared with
Pyramid Adversarial training, we find that Universal Pyramid Adversarial
training with a radius of 12/255 consistently improves performance with
respect 1-step Pyramid Adversarial training and is comparable with the 2-step
Pyramid Adversarial training. Note that both 1-step and 2-step Pyramid
Adversarial training are already 50% and 100% more expensive than our proposed
Universal Pyramid Adversarial training. However, unlike the case of clean
accuracy, Universal Pyramid Adversarial training still underperforms relative
to the more costly Pyramid Adversarial training with 3 or more steps.

| Radius |  |  |  |  |
---|---|---|---|---|---|---
| 2/255 | 4/255 | 6/255 | 8/255 | 10/256 | 12/256
Baseline | 79.85% |  |  |  |  |
Universal Adversarial Training (w/o Pyramid Structure) | 79.73% | 79.81% | 79.64% | 79.59% | 79.65% | 79.75%
\- Gain Relative to Baseline | -0.12% | -0.04% | -0.21% | -0.26% | -0.20% | -0.10%
Universal Pyramid Adversarial Training | 79.89% | 80.23% | 80.15% | 80.28% | 80.13% | 80.04%
\- Gain Relative to Baseline | 0.04% | 0.38% | 0.30% | 0.43% | 0.28% | 0.19%

Table 3: Comparing universal adversarial training with and without the pyramid
structure. We find that the pyramid structure is indeed crucial for the
performance gain observed. Without the pyramid structure, universal
adversarial training consistently decreases the model’s performance relative
to the baseline.

Method | Radius | Top 1 Accuracy | Gain
---|---|---|---
Baseline | - | 79.84% | -
Uni. Pyramid Adv. | 2/255 | 79.89% | 0.05%
4/255 | 80.23% | 0.39%
6/255 | 80.15% | 0.31%
8/255 | 80.28% | 0.44%
10/255 | 80.13% | 0.29%
12/255 | 80.04% | 0.20%

Table 4: Universal Pyramid Adversarial training consistently increases the
performance of ViT across a range of hyperparameters but achieves the best
performance at radius 8/255.

###  4.3 Ablations

In this section, we ablated several components of Universal Pyramid
Adversarial training, including its sensitivity to the selected radius, the
importance of the pyramid structure, and the benefit of incorporating clean
loss.

#### Sensitivity to Radius

Hyperparameter sensitivity is crucial for a method’s practicality, and we find
that our Universal Pyramid Adversarial training is consistent and stable with
respect to the selected radius. In Table 4, we see that Universal Pyramid
Adversarial training consistently increases the model’s performance across a
wide range of radii from 2/255 to 12/255. This consistency is important
because it allows us to benefit from the method without finely tuning the
radius. We also find that the performance varies in a predictable upside-down
U-shape. As we increase the radius, we see that the performance steadily
increases until radius 8/255 after which the performance decreases. The way
that performance changes with respect to the radius gives practitioners a
clear signal on whether to increase or decrease the radius and makes
hyperparameter tuning easier.

| Top 1 Accuracy |
---|---|---
Baseline | 79.84% |
Uni. Pyramid Adv. |  |
w/o clean loss | 78.08% ( -1.76%) |
w/ clean loss | 80.28% ( +0.44%) |

Table 5: Comparing Universal Pyramid Adversarial training with and without
adding the clean loss. We found that the addition of clean loss is critical
for performance improvements.

#### Pyramid Structure

We also found that the pyramid structure is crucial for performance gain with
our proposed universal adversarial approach. We experimented with naively
combining clean loss with universal adversarial training as in Shafahi et al.
(2020) as an additional regularizer. However, in Table 3, we see that without
using the pyramid structure, the model consistently performs worse after
adding the adversarial samples.

(a) Baseline

(b) Pyramid Adversarial training

(c) Universal Pyramid Adversarial training

Figure 2: Analyzing the loss landscape of the final trained models. We
employed the filter normalization method from Li et al. (2017) for
visualization. Pyramid Adversarial training seems to induce minima that are
sharper compared to the baseline and Universal Pyramid Adversarial training.
However, despite landing the model in sharper minima, Pyramid Adversarial
training still produces a better performing model than the baseline. On the
other hand, Universal Pyramid Adversarial training does not seem to change the
sharpness of minima while attaining the best performance. This observation
shows that both adversarial pyramid approaches do not improve performance
through flattening the loss landscape.

#### Clean Loss

In addition to using pyramid structure, we found that incorporating clean loss
to the Universal Pyramid Adversarial training is vital to obtain the
performance gain we see. In Table 5, we see that removing the clean loss
results in a performance decrease of 1.76%, making the model much worse than
the baseline.

###  4.4 Analysis

In this section, we try to understand the similarities between universal and
sample-wise adversarial pyramid training, given their similar benefits and
formulation. We analyzed the attack strength, the noise pattern, and the loss
landscape and found that despite their similarities, they are quite different
in many aspects. The findings point to the need to further understand the
mechanism that both universal and sample-wise adversarial pyramid training use
to increase model performance.

(a) Attack Strength on the Training Set

(b) Attack Strength on the Validation Set

Figure 3: Comparing the strength of the universal adversary with the sample-
wise adversary. The x-axis is the number of epochs, and the y-axis is the
increased error rate after adversarial attack. We see that instance wise
adversary is much stronger compared to universal adversary even though
universal adversary achieves competitive performance on clean and ood tasks.

#### Attack Strength

Given that the Universal Pyramid Adversarial training achieves similar
performance gain compared to the Pyramid Adversarial training, one may expect
that their attack strength is similar. However, we found that the sample-wise
adversary is significantly stronger than the universal adversary even though
the universal adversary uses a larger radius (8/255 vs. 6/255). In Figure 3,
the multi-step adversary consistently achieves a much higher adversarial error
rate throughout the training process. The observation shows that we don’t
necessarily need to make the attack very strong to benefit from adversarial
training.

#### Qualitative Differences in Perturbation Pattern

In Figure 4, we visualize the universal and sample-wise adversarial patterns
used during training. We find the perturbations to have qualitatively
different patterns despite their similar effectiveness in improving clean
accuracy.

For the perturbation at the coarser scales, universal perturbation has a more
diverse color than sample-wise perturbation. The diversity may be because
universal perturbations need to transfer between images, and large brightness
changes may be effective in removing some information from the samples. On the
other hand, the sample-level perturbations may exploit the color cues to move
an image to the adversarial class by consistently changing the color of the
image.

For the perturbation at the pixel level, sample-wise perturbation is much more
salient compared to universal perturbation. We can see that the sample-wise
perturbations have some resemblance to objects. Even though the universal
perturbations have some salient patterns, they are less obvious than the
pattern from the sample-wise perturbations.

(a)

(b)

Figure 4: Adversarial patterns generated by sample-wise pyramid adversarial
training in 4(a) and universal pyramid adversarial training in 4(b). We
plotted each level of the pyramid separately, so that we can visually inspect
the differences between each level. We found that (1) universal pyramid seems
to rely on more diverse perturbation values at the coarser scale and (2)
sample-wise pyramid adversarial seems to rely more on the pixel level
perturbations according to the more salient pixel level perturbations.

#### Loss Landscape

To understand how Pyramid Adversarial training and Universal Pyramid
Adversarial training improve the performance of a model, we visualize the loss
landscape of models trained with both approaches to see whether it achieves
the performance by implicitly inducing a flatter minimum Li et al. (2017).
Surprisingly, we found this not to be the case. Sample-wise Pyramid
Adversarial training produces sharper minima compared to regular training and
yet has better performance (see Figure 2). On the other hand, Universal
Pyramid Adversarial training does not noticeably change the sharpness of the
minimum and yet produces the greatest performance improvement. This finding
suggests that both adversarial pyramid approaches rely on different underlying
mechanisms to improve the model’s performance compared to an optimizer, such
as SAM Foret et al. (2020) that explicitly searches for flatter minima.

##  5 Conclusion

In this paper, we propose Universal Pyramid Adversarial training to improve
the clean performance and out-of-distribution robustness of ViT. It obtains
similar accuracy gain as sample-wise Pyramid Adversarial training while being
up to 70% faster than the original approach. To the best of our knowledge, we
are also the first to identify universal adversarial training as a possible
technique for improving the model’s clean accuracy. We hope that the proposed
method will help make the adversarial technique more accessible to
practitioners and future researchers.

## References

  * Bai et al. (2022) Jiawang Bai, Li Yuan, Shu-Tao Xia, Shuicheng Yan, Zhifeng Li, and Wei Liu.  Improving vision transformers by revisiting high-frequency components.  _arXiv preprint arXiv:2204.00993_ , 2022.
  * Bai et al. (2021) Tao Bai, Jinqi Luo, Jun Zhao, Bihan Wen, and Qian Wang.  Recent advances in adversarial training for adversarial robustness.  _arXiv preprint arXiv:2102.01356_ , 2021.
  * Benz et al. (2021) Philipp Benz, Chaoning Zhang, Adil Karjauv, and In So Kweon.  Universal adversarial training with class-wise perturbations.  In _2021 IEEE International Conference on Multimedia and Expo (ICME)_ , pp. 1–6, 2021.  doi: 10.1109/ICME51207.2021.9428419.
  * Beyer et al. (2022) Lucas Beyer, Xiaohua Zhai, and Alexander Kolesnikov.  Better plain vit baselines for imagenet-1k.  _arXiv preprint arXiv:2205.01580_ , 2022.
  * Cubuk et al. (2020) Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le.  Randaugment: Practical automated data augmentation with a reduced search space.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops_ , pp. 702–703, 2020.
  * Dosovitskiy et al. (2020) Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  _arXiv preprint arXiv:2010.11929_ , 2020.
  * Foret et al. (2020) Pierre Foret, Ariel Kleiner, Hossein Mobahi, and Behnam Neyshabur.  Sharpness-aware minimization for efficiently improving generalization.  _arXiv preprint arXiv:2010.01412_ , 2020.
  * Geirhos et al. (2019) Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A Wichmann, and Wieland Brendel.  Imagenet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness.  In _International Conference on Learning Representations_ , 2019.  URL https://openreview.net/forum?id=Bygh9j09KX.
  * Hendrycks & Dietterich (2019) Dan Hendrycks and Thomas Dietterich.  Benchmarking neural network robustness to common corruptions and perturbations.  _arXiv preprint arXiv:1903.12261_ , 2019.
  * Hendrycks et al. (2021) Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song.  Natural adversarial examples.  _CVPR_ , 2021.
  * Herrmann et al. (2022) Charles Herrmann, Kyle Sargent, Lu Jiang, Ramin Zabih, Huiwen Chang, Ce Liu, Dilip Krishnan, and Deqing Sun.  Pyramid adversarial training improves vit performance.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 13419–13429, 2022.
  * Li et al. (2017) Hao Li, Zheng Xu, Gavin Taylor, and Tom Goldstein.  Visualizing the loss landscape of neural nets.  _CoRR_ , abs/1712.09913, 2017.  URL http://arxiv.org/abs/1712.09913.
  * 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.  URL https://openreview.net/forum?id=rJzIBfZAb.
  * Mao et al. (2022) Xiaofeng Mao, Yuefeng Chen, Ranjie Duan, Yao Zhu, Gege Qi, Shaokai Ye, Xiaodan Li, Rong Zhang, and Hui Xue.  Enhance the visual representation via discrete adversarial training.  _arXiv preprint arXiv:2209.07735_ , 2022.
  * Mei et al. (2022) Jieru Mei, Yucheng Han, Yutong Bai, Yixiao Zhang, Yingwei Li, Xianhang Li, Alan Yuille, and Cihang Xie.  Fast advprop.  _arXiv preprint arXiv:2204.09838_ , 2022.
  * 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!  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Shafahi et al. (2020) Ali Shafahi, Mahyar Najibi, Zheng Xu, John Dickerson, Larry S Davis, and Tom Goldstein.  Universal adversarial training.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 34, pp. 5636–5643, 2020.
  * Steiner et al. (2021) Andreas Steiner, Alexander Kolesnikov, Xiaohua Zhai, Ross Wightman, Jakob Uszkoreit, and Lucas Beyer.  How to train your vit? data, augmentation, and regularization in vision transformers.  _arXiv preprint arXiv:2106.10270_ , 2021.
  * Wang et al. (2019) Haohan Wang, Songwei Ge, Zachary Lipton, and Eric P Xing.  Learning robust global representations by penalizing local predictive power.  In _Advances in Neural Information Processing Systems_ , pp. 10506–10518, 2019.
  * Wong et al. (2020) Eric Wong, Leslie Rice, and J. Zico Kolter.  Fast is better than free: Revisiting adversarial training.  _CoRR_ , abs/2001.03994, 2020.  URL https://arxiv.org/abs/2001.03994.
  * Xie et al. (2020) Cihang Xie, Mingxing Tan, Boqing Gong, Jiang Wang, Alan L Yuille, and Quoc V Le.  Adversarial examples improve image recognition.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 819–828, 2020.
  * Zhang et al. (2019) Dinghuai Zhang, Tianyuan Zhang, Yiping Lu, Zhanxing Zhu, and Bin Dong.  You only propagate once: Accelerating adversarial training via maximal principle.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Zhang et al. (2017) Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  _arXiv preprint arXiv:1710.09412_ , 2017.
  * Zheng et al. (2020) Haizhong Zheng, Ziqi Zhang, Juncheng Gu, Honglak Lee, and Atul Prakash.  Efficient adversarial training with transferable adversarial examples.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 1181–1190, 2020.

##  Appendix A Further Details for Table 2

Method | Radius | Clean $\uparrow$ | C $\downarrow$ | A $\uparrow$ | Rendition $\uparrow$ | Sketch $\uparrow$ | Stylized $\uparrow$
---|---|---|---|---|---|---|---
Universal Pyramid Adversarial Training | 12/255 | 80.04% | 40.37% | 28.21% | 37.05% | 17.42% | 19.53%
Universal Pyramid Adversarial Training | 8/255 | 80.28% | 41.02% | 29.08% | 35.45% | 16.89% | 17.74%
Regular Training |  | 79.84% | 42.18% | 27.52% | 33.65% | 15.08% | 17.06%
Pyramid Adversarial Training 1 steps | 6/255 | 79.46% | 40.39% | 26.48% | 35.34% | 16.56% | 19.00%
Pyramid Adversarial Training 2 steps | 6/255 | 79.87% | 39.76% | 28.49% | 36.17% | 16.95% | 19.95%
Pyramid Adversarial Training 3 steps | 6/255 | 80.19% | 38.91% | 30.09% | 37.75% | 17.66% | 20.96%
Pyramid Adversarial Training 4 steps | 6/255 | 80.04% | 39.22% | 30.33% | 38.40% | 18.02% | 21.02%
Pyramid Adversarial Training 5 steps | 6/255 | 80.10% | 38.84% | 30.04% | 38.63% | 17.80% | 21.21%

Table 6: Here we provide further details for the performance of each of the
adversarial training methods on different corruption datasets.

##  Appendix B Visualization of Attention Operations

Figure 5: The first column displays the image, followed by attention
visualizations for the baseline model, UPAT model, and PAT model in the next
three columns. Despite exhibiting similar attention patterns to the baseline
model, the UPAT model consistently outperforms it. In contrast, the PAT model
demonstrates sparse attention, indicating that UPAT and PAT models improve
performance through different mechanisms.

##  Appendix C Radius Ablation with Respect to Imagenet V2

Method | Radius | Top 1 Accuracy (Imagenet v2) | Gain
---|---|---|---
Baseline | - | 79.81% | -
Uni. Pyramid Adv. | 2/255 | 79.93% | 0.12%
4/255 | 79.91% | 0.10%
6/255 | 79.94% | 0.13%
8/255 | 80.24% | 0.43%
10/255 | 79.91% | 0.10%
12/255 | 80.20% | 0.39%

Table 7: During training, we lacked a separate validation set, raising
concerns that our model’s hyperparameters may overfit to the test accuracy. To
mitigate this, we employed additional evaluation using ImageNet-V2, an
independent dataset not used for hyperparameter selection in our experiments.
Ultimately, we observed similar evaluation results on ImageNet-V2 and the test
set, alleviating concerns of hyperparameter overfitting.