# Principal Component Analysis for Spatial Phase Reconstruction in Atom
Interferometry

Stefan Seckmeyer<EMAIL_ADDRESS>Holger Ahlers  Present
Address: DLR Institute for Satellite Geodesy and Inertial Sensing, 30167
Hannover, Germany  Jan-Niclas Kirsten-Siemß  Matthias Gersemann  Ernst M.
Rasel  Sven Abend  Naceur Gaaloul<EMAIL_ADDRESS>Gottfried
Wilhelm Leibniz Universität Hannover, Institut für Quantenoptik, Welfengarten
1, 30167 Hannover, Germany

###### Abstract

Atom interferometers are sensitive to a wide range of forces by encoding their
signals in interference patterns of matter waves. To estimate the magnitude of
these forces, the underlying phase shifts they imprint on the atoms must be
extracted. Up until now, extraction algorithms typically rely on a fixed model
of the patterns’ spatial structure, which if inaccurate can lead to systematic
errors caused by, for example, wavefront aberrations of the used lasers. In
this paper we employ an algorithm based on Principal Component Analysis, which
is capable of characterizing the spatial phase structure and per image phase
offsets of an atom interferometer from a set of images. The algorithm does so
without any prior knowledge about the specific spatial pattern as long as this
pattern is the same for all images in the set. On simulated images with atom
projection noise we show the algorithm’s reconstruction performance follows
distinct scaling laws, i.e., it is inversely-proportional to the square-root
of the number atoms or the number of images respectively, which allows a
projection of its performance for experiments. We also successfully extract
the spatial phase patterns of two experimental data sets from an atom
gravimeter. This algorithm is a first step towards a better understanding and
complex spatial phase patterns, e.g., caused by inhomogeneous laser fields in
atom interferometry.

##  I Introduction

Atom interferometry allows for high precision experiments in several fields
such as gravimetry [1, 2, 3], gravity cartography [4] and inertial navigation
[5, 6]. Other use-cases include measurements of fundamental constants like the
fine structure constant [7, 8], the gravitational constant [9] and the
universality of free fall [10]. In atom interferometry, a cloud of atoms is
usually coherently split and recombined via atom-light interactions with laser
pulses. The signal to be measured is translated into a relative phase shift
between the interferometer arms as depicted in Fig. 1a. However, such a phase
difference is not a quantum mechanical observable and therefore cannot be
measured directly.

To measure the phase of an atom interferometer, it is projected into an
interferogram, which can be detected after the recombination of the two matter
waves via a final laser pulse (Fig. 1a). Figure 1b illustrates a result of
such a measurement of atomic populations, e.g., obtained via fluorescence or
absorption imaging [11, 12]. These two images also show the typical structure
of such an interferogram having two so-called ports. In a light interferometer
the ports’ brightness is given by the light intensity. For atom
interferometers it is determined by the number of atoms interacting with the
detection laser. Their relative brightness between the ports directly depends
on the phase shift introduced during the sequence (see Fig. 1a).

In order to reconstruct the phase difference from its projection into
intensity, one needs to recover their functional relationship. To do so,
interferograms with different phase offsets are recorded over multiple runs to
cover a whole oscillation period, which results in a phase scan as shown in
Fig. 1c. Alternatively, intentionally creating a phase shear imprint across
the atomic clouds and fitting a fixed model has also been used to perform a
phase measurement in a single shot [13, 14, 10].

It is typically assumed that the ports of an interferometer have a phase
difference of $\pi$ and that the phase in each port is spatially homogeneous.
However, there are multiple processes which can introduce an unwanted spatial
phase pattern in addition to the desired signal phase into each of the ports.
Early studies on spatial phase imprints in atom interferometry [15, 16]
focused on atom-atom interactions and waveguides. These optical and magnetic
waveguides often suffer from imperfections (e.g. speckles), which can lead to
unwanted spatially dependent phase patterns. The quality of atom
interferometric measurements depends directly on the quality of the light used
to manipulate the atoms. Imperfections of the optical potentials are one of
the leading systematic effects of current and probably future measurements.
They are imprinted into the spatial phase profile of the final recorded
interferometer images and usually referred to as wavefront aberrations [17,
18, 19].

Currently, the primary mitigation strategy is to minimize these aberrations as
much as possible [20]. That is because they are difficult to accurately
measure and post-correct for in terms of systematic shifts [19, 8, 7], which
is why standard analysis algorithms do not generally take them into account.
Thus, this strategy introduces significant systematic errors, if the wavefront
aberrations are strong enough to have a significant impact on the
measurements.

In the fields of light and acoustic interferometry a very similar problem
appears in phase shifting interferometry [21, 22, 23], which extracts the
phase from a set of phase-shifted interferograms. One way to account for the
spatial phase imprints of wavefront aberrations is the use of statistical
algorithms like Principal Component Analysis (PCA) and Independent Component
Analysis, which have been examined in [24, 25, 26, 27, 28].

In the context of atom interferometry, PCA was used to enhance the contrast
and shape of shear-interferometer fringes from atom interferometry output
images for analysis [14, 13].

In this paper, we present a PCA-based, spatial phase reconstruction (PSPR)
algorithm that extracts an arbitrary spatial phase map and the phase offsets
between different images of an image set. The algorithm shares similarities
with the one presented in [29] and we successfully apply it to both simulated
atom interferometry data as well as actual measurements. The paper is
organized as follows. In Section II, we introduce the PSPR algorithm and
discuss the underlying assumptions in contrast to commonly used phase
extraction methods in atom interferometry. Section III establishes an accuracy
baseline for PSPR by analyzing synthetic data of a two-port atom
interferometer in the absence of noise. We do so by studying synthetic
interferometry images affected by three exemplary phase patterns. Each of
these patterns could be reconstructed from the synthetic images at almost the
numerical machine precision. In Section IV, we introduce atomic shot noise
into our interferometer simulation and show that the performance of the phase
reconstruction scales inversely proportional with the square root of the
number of atoms per image as expected for uncorrelated particles. Finally, we
showcase the application of the PSPR to experimental data in Section V and
show its robustness under real-world conditions. We extract the spatial phase
patterns for two different interferometer configurations and verify that the
coefficients exhibit structures similar to the results from the simulations.
To conclude, in Section VI we discuss the potential of PSPR for precision
interferometry in general and for automation and data compression in
particular for transportable atom interferometry sensors deployed on ground or
in space.

##  II PCA-based Spatial Phase Reconstruction (PSPR) Algorithm

(a)

aaaaaa$\phi=0.4\pi$ |  aaaaaa$\phi=2\pi$
---|---
|
(b)

(c)

Figure 1:  a) Scheme of an atom interferometer. The phase shift $\phi$ can be
caused by multiple effects and is abstracted by a box on one of the arms. b)
Ideal images of the two ports of an atom interferometer for two different
phase shifts $\phi$. The ports are illustrated by the marked integration
regions for simple box integrals. The integration is the estimation of the
intensity in each port. c) A phase scan: The population of the two ports as a
function of a deliberately added phase shift $\phi$. To estimate the phase
shift a cosine can be fitted to this oscillation in the populations.

###  II.1 The Interferometer Models Underlying the Extraction Algorithms

Every phase extraction algorithm is based on an assumed model of how the
interferometer works. Any deviation of the actual experiment from that model
causes a systematic error in the results of the algorithm.

In almost all phase extraction algorithms, one needs to vary the phase
difference of an interferometer to reconstruct the phase that was projected
into the interferograms. Varying the phase difference between the two arms
causes the intensity to oscillate, which allows to distinguish between
background intensity and the intensity of the interferometer output ports.
This phase is usually varied between different experimental runs and therefore
images. It is also possible to vary the phase via an imprint across the atomic
cloud to allow the phase extraction from a single image, as done in phase
shear interferometry [14, 13].

As a baseline for the models typically used in current state-of-the-art atom
interferometers [30, 31], we focus on taking multiple images and changing a
phase offset in the last beam splitter. Assuming the phase in each of the
ports is constant over the whole area of the port, the phase offset determines
how many atoms end up in the left or right port as shown in Figure 1b. In
order to extract amplitudes from these images, the ratio of the number of
atoms in each port needs to be determined. One approach is to simply define an
area in the image and integrate the intensity inside it, to determine the
ratio of atom numbers between the areas (see Fig. 1b). Defining this
integration area is, however, a difficult trade-off because the clouds become
very dilute at the edges. On the one hand, choosing a too large integration
area reduces the signal-to-noise ratio because the outer areas of the clouds
are dominated by background noise. On the other hand, if the integrated region
is too small, information is discarded. Selecting a shape for the integration
area that more closely reflects the spatial distribution of the atoms can
provide a better alternative that allows using the whole atom cloud. This,
however, can introduce an additional systematic error, if the assumed cloud
shape deviates from the actual data.

Both the box integral as well as fitting a specific shape on each cloud rely
on estimating the number of atoms in each port, also called populations.
Plotting these populations as a function of the phase shift in the last beam
splitter results in oscillations like the ones depicted in Fig. 1c. The
populations of each port are fitted to a cosine, in order to determine the
contrast and the phase shift of the interferometer. This population
oscillation is shifted by the phase shift caused by the measured quantity.
Each population value can only be uniquely mapped into an interval from
$[0,\pi]$ which results in a dynamic range of $\pi$. Doing such a cosine fit
requires to reliably determine and control the phase shift of each image
relative to the other images or to fix the ambiguity introduced by phase jumps
by correlating the atom interferometer with a classical interferometer [5, 32,
20].

PSPR, on the other hand, does not rely on the existence of a port structure at
all. It only requires a minimum of two pixels which have a fixed phase offset
$\neq\pi$ between each other over the whole image sequence. This is formalized
as defining a set of images that are assumed to follow the form

$\displaystyle I(i,x,y)$ $\displaystyle=A(x,y)+B(x,y)\ \cos[\Phi(i,x,y))]$ (1)
$\displaystyle\Phi(i,x,y)$ $\displaystyle=\theta(i)+\gamma(x,y).$ (2)

The image set has $n_{\mathrm{images}}$ images which are indexed with $i$ and
each image has $xy$ many pixels with $x$ indexing the pixel row and $y$
indexing the pixel column. $A(x,y)$ is the background intensity, $B(x,y)$ is
an intensity envelope defining the signal strength of each pixel and
$\gamma(x,y)$ is the free-form spatial phase profile of the interferometer.
The intensities and the phase profile are assumed to be the same in all
analyzed images in the input set. $\theta(i)$ is the sum of the phase shift
from the measured quantity and the artificially introduced phase shift in
image $i$.

The form of Eq. 1 assumes that the intensity of each pixel is only ever
influenced by a single phase $\Phi(i,x,y)$, so it requires non-overlapping
ports. Both the intensity and the phase do not, however, have any
predetermined spatial form and the order in which the images are recorded does
not matter either. From such a set of images, PSPR allows to reconstruct
$\theta(i)$ and $\gamma(x,y)$ with a dynamic range of $2\pi$.

###  II.2 PCA Based Spatial Phase Reconstruction (PSPR)

The PSPR reconstruction steps are visually displayed in Fig. 2 and follow a
similar logic than the one of reference [29]. In the following, we introduce
subscripts to properly distinguish the input phases $\theta_{\mathrm{inp}}(i)$
and $\gamma_{\mathrm{inp}}(x,y)$ from the reconstructed phases
$\theta_{\mathrm{rec}}(i)$ and $\gamma_{\mathrm{rec}}(x,y)$. The input model
of Eq. 1 can be rewritten as a Fourier decomposition yielding linear
combinations of terms depending only on $i$ or $x,y$:

$\displaystyle I(i,x,y)=A(x,y)$
$\displaystyle+B(x,y)\cos[\theta_{\mathrm{inp}}(i)]\cos[\gamma_{\mathrm{inp}}(x,y)]-B(x,y)\sin[\theta_{\mathrm{inp}}(i)]\sin[\gamma_{\mathrm{inp}}(x,y)].$
(3)

As shown above, the image model has two linearly independent degrees of
freedom. These two degrees of freedom can be extracted using PCA. It
decomposes any data set into the average intensity
$\bar{I}(x,y)=\mathrm{mean}[I(i,x,y)]$ and $n_{\mathrm{comp}}$ linearly
independent orthonormal components $\mathrm{pc}_{j}(x,y)$ with associated
coefficients $w_{j}(i)$. As depicted in step 2 in Fig. 2, PCA generates only
two $\mathrm{pc}_{j}(x,y)$ and $w_{j}(i)$ respectively, since there are only
two linearly independent components in any set of images generated with Eq. 1:

$\displaystyle
I(i,x,y)=\bar{I}(x,y)+w_{1}(i)\mathrm{pc}_{1}(x,y)+w_{2}(i)\mathrm{pc}_{2}(x,y).$
(4)

The principal components $\mathrm{pc}_{j}(x,y)$ are chosen by PCA such that
they are orthonormal and maximize the variance of the coefficients $w_{j}(i)$.
Although $\cos[\gamma_{\mathrm{inp}}(x,y)]$ and
$\sin[\gamma_{\mathrm{inp}}(x,y)]$ are orthogonal to each other analytically,
they are not orthogonal for almost any choice of $A(x,y)$, $B(x,y)$,
$\theta_{\mathrm{inp}}(i)$ and $\gamma_{\mathrm{inp}}(x,y)$. Likewise, the
mean $\bar{I}(x,y)$ does not generally equal $A(x,y)$.

This causes plotting $w_{1}(i)$ against $w_{2}(i)$ or $\mathrm{pc}_{1}(x,y)$
against $\mathrm{pc}_{2}(x,y)$ (with $B(x,y)=1$) to result in ellipses instead
of the ideal circles. If $B(x,y)\neq 1$ the ellipse formed by
$\mathrm{pc}_{1}(x,y)$ and $\mathrm{pc}_{2}(x,y)$ becomes distorted because
each point has a potentially different distance from the origin. Therefore, an
ellipse fit [33, 34] is performed on the generated $w_{1}(i)$ and $w_{2}(i)$
instead of fitting $\mathrm{pc}_{1}(x,y)$ and $\mathrm{pc}_{2}(x,y)$ as done
in [29].

The ellipse fit (step 3 in Fig. 2) has the free fit parameters
${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}x_{c}}$,
${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}y_{c}}$,
${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}a}$,
${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}b}$
and
${\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta}$:

$\displaystyle
w_{1}(i)={\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}x_{c}}+{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}a}\cos{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta}\cos
t_{i}-{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}b}\sin{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta}\sin
t_{i}$ (5) $\displaystyle-
w_{2}(i)={\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}y_{c}}+{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}a}\sin{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta}\cos
t_{i}+{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}b}\cos{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta}\sin
t_{i}$ (6)

The negation in front of $w_{2}(i)$ stems from the subtraction in Eq. 3. This
negation needs to be assigned to either the coefficients or the components
during reconstruction in order to have a consistent sign between
$\theta_{\mathrm{rec}}(i)$ and $\gamma_{\mathrm{rec}}(x,y)$.

Correcting for the ellipse distortion in step 4 transforms $w_{1}(i)$ and
$w_{2}(i)$ onto an approximate unit circle.

$\displaystyle\begin{pmatrix}w^{\prime}_{1}(i)\\\
w^{\prime}_{2}(i)\end{pmatrix}$
$\displaystyle=\underbrace{\begin{pmatrix}1/{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}a}&0\\\
0&1/{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}b}\end{pmatrix}\times}_{\text{scale}}\underbrace{\begin{pmatrix}\cos(-{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})&-\sin(-{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})\\\
\sin(-{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})&\cos(-{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})\end{pmatrix}\times}_{\text{rotate}}\left[\begin{pmatrix}w_{1}(i)\\\
-w_{2}(i)\end{pmatrix}\smash[b]{\underbrace{-\begin{pmatrix}{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}x_{c}}\\\
{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}y_{c}}\end{pmatrix}}_{\text{translate}}}\right]$
(7) $\displaystyle\begin{pmatrix}\mathrm{pc}^{\prime}_{1}(x,y)\\\
\mathrm{pc}^{\prime}_{2}(x,y)\end{pmatrix}$
$\displaystyle=\underbrace{\begin{pmatrix}{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}a}&0\\\
0&{\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}b}\end{pmatrix}\times}_{\text{scale}}\underbrace{\begin{pmatrix}\cos({\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})&-\sin({\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})\\\
\sin({\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})&\cos({\color[rgb]{0,0,0}\definecolor[named]{pgfstrokecolor}{rgb}{0,0,0}\pgfsys@color@gray@stroke{0}\pgfsys@color@gray@fill{0}\vartheta})\end{pmatrix}\times}_{\text{rotate}}\begin{pmatrix}\mathrm{pc}_{1}(x,y)\\\
\mathrm{pc}_{2}(x,y)\end{pmatrix}$ (8)

Computing the angle of these corrected components
$\mathrm{pc}^{\prime}_{1}(x,y)$ and $\mathrm{pc}^{\prime}_{2}(x,y)$ and
coefficients $w^{\prime}_{1}(i)$ and $w^{\prime}_{2}(i)$ on their respective
unit circle in step 5 reconstructs the full phase $\Phi_{\mathrm{rec}}(i,x,y)$
for each image:

$\displaystyle\theta_{\mathrm{rec}}(i)$
$\displaystyle=\mathrm{arctan2}(w^{\prime}_{2}(i),w^{\prime}_{1}(i))$ (9)
$\displaystyle\gamma_{\mathrm{rec}}(x,y)$
$\displaystyle=\mathrm{arctan2}(\mathrm{pc}^{\prime}_{2}(x,y),\mathrm{pc}^{\prime}_{1}(x,y))$
(10) $\displaystyle\Phi_{\mathrm{rec}}(i,x,y)$
$\displaystyle=\theta_{\mathrm{rec}}(i)+\gamma_{\mathrm{rec}}(x,y).$ (11)

Figure 2: Steps of the PSPR algorithm. Analysis begins with a set of images
$I(i,x,y)$ (1), those are transformed with PCA (2), an ellipse is fit to the
coefficients $w_{1}(i)$ and $w_{2}(i)$ (3), the coefficients $w_{j}(i)$ and
components $\mathrm{pc}_{j}(x,y)$ are corrected for the ellipse distortion (4)
and the phase is extracted via arctan2 (5).

In addition to the already mentioned [34, 28] the implementation of the
algorithm, simulation and visualization in this paper use the python libraries
numpy [35] and jax [36, 37]. The bokeh library [38] was used for plotting,
colorcet [39] and palettes from [40] were used as a basis for the color
schemes.

###  II.3 Comparison between Inputs, Outputs and Multiple Data Sets

The first fundamental building block for phase comparisons is the difference
between two angles. When viewing the two angles on a circle there are two
distances between them, one going left-handed and the other going right-
handed. Given two phases $\varphi_{1}$ and $\varphi_{2}$, Eq. 12 returns
$\varphi_{\mathrm{diff}}\in[-\pi,\pi)$ which is defined as the shorter of the
two possible distances and the sign indicating which direction was taken
around the circle:

$\displaystyle\varphi_{\mathrm{diff}}(\varphi_{1},\varphi_{2})=\mathrm{mod}(\varphi_{1}-\varphi_{2},2\pi)-\pi$
(12)

with $\mathrm{mod}$ indicating the modulo operation.

Additionally, the model underlying PSPR (see Eq. 19) maps multiple phases to
the same images. This needs to be taken into account when comparing inputs and
outputs or outputs from two different analysis runs. When comparing
$\Phi_{\mathrm{inp}}(i,x,y)$ and $\Phi_{\mathrm{rec}}(i,x,y)$, there may be a
sign switch which still results in the same images:

$\displaystyle\cos[\Phi_{\mathrm{inp}}(i,x,y)]$ $\displaystyle=\cos[s\
\Phi_{\mathrm{inp}}(i,x,y)]=\cos[\Phi^{\prime}_{\mathrm{inp}}(i,x,y)]\ \
\forall s\in\\{-1,1\\}.$

Comparing the individual outputs $\theta_{\mathrm{rec}}(i)$ and
$\gamma_{\mathrm{rec}}(x,y)$ with the inputs $\theta_{\mathrm{inp}}(i)$ and
$\gamma_{\mathrm{inp}}(x,y)$ adds another degree of freedom $c$ while all
phases $\theta_{\mathrm{inp}}^{\prime}(i)$ and
$\gamma_{\mathrm{inp}}^{\prime}(x,y)$ remain a solution for the same image
set:

$\displaystyle\begin{split}\cos[\Phi_{\mathrm{inp}}(i,x,y)]&=\cos[\theta_{\mathrm{inp}}(i)+\gamma_{\mathrm{inp}}(x,y)]\\\
&=\cos(s\\{[\theta_{\mathrm{inp}}(i)-c]+[\gamma_{\mathrm{inp}}(x,y)+c]\\})\\\
&=\cos[\theta_{\mathrm{inp}}^{\prime}(i)+\gamma_{\mathrm{inp}}^{\prime}(x,y)]\\\
\forall\ \ c&\in[-\pi,\pi),\ s\in\\{-1,1\\}.\end{split}$ (13)

The constants $c$ and $s$ are unique for each run and determined by the PCA.
Both constants are determined by the signs of the eigenvalues when computing
the components and the maximation of variance under orthogonality.

In the following, $c$ and $s$ will be computed from $\theta_{\mathrm{inp}}(i)$
and $\theta_{\mathrm{rec}}(i)$ as the solution to:

$\displaystyle\min_{c,s}$
$\displaystyle\left(\sum_{i}\left\\{\varphi_{\mathrm{diff}}\left[\theta_{\mathrm{rec}}(i),s\,\theta_{\mathrm{inp}}(i)-c\right]\right\\}\right)\textrm{with}\
\ c\in[-\pi,\pi),\ s\in\\{-1,1\\}.$ (14)

This minimization problem can be solved directly without iteration as
described in the appendix.

Since $c$ and $s$ are the same for both $\theta_{\mathrm{rec}}(i)$ and
$\gamma_{\mathrm{rec}}(x,y)$, they are used for the computation of both
$\theta_{\mathrm{diff}}(i)$ and $\gamma_{\mathrm{diff}}(x,y)$:

$\displaystyle\theta_{\mathrm{diff}}(i)$
$\displaystyle\coloneqq\varphi_{\mathrm{diff}}\left[s\,\theta_{\mathrm{inp}}(i)-c,\theta_{\mathrm{rec}}(i)\right]$
(15) $\displaystyle\gamma_{\mathrm{diff}}(x,y)$
$\displaystyle\coloneqq\varphi_{\mathrm{diff}}\left[s\,\gamma_{\mathrm{inp}}(x,y)+c,\gamma_{\mathrm{rec}}(x,y)\right].$
(16)

##  III Noise Free Phase Reconstruction

This section establishes an accuracy baseline for the performance of PSPR by
analyzing computer-generated sample images, inspired by a simplified model of
a two-port atom interferometer. The two ports of the interferometer have a
Gaussian intensity profile $g(x,y,\sigma)$ and the same spatial phase pattern
$\gamma_{\mathrm{inp}}(x,y)$. They are shifted relative to each other in
position by $2\delta x$ in x-direction, do not exhibit any overlap and the
left port has an additional phase shift of $\pi$. This is parametrized by:

$\displaystyle g(x,y,\Delta)$
$\displaystyle=\frac{1}{2\pi\Delta^{2}}\mathrm{exp}\left(\frac{-(x^{2}+y^{2})}{2\Delta^{2}}\right)$
(17) $\displaystyle\Phi_{\mathrm{inp}}(i,x,y)$
$\displaystyle=\theta_{\mathrm{inp}}(i)+\gamma_{\mathrm{inp}}(x,y)$ (18)
$\displaystyle I(i,x,y)$
$\displaystyle=\begin{cases}\frac{1}{2}\,\text{d}x\,\text{d}y\,g(x-\delta
x,y,0.3)\\{1+C\cos[\Phi_{\mathrm{inp}}(i,x-\delta x,y)]\\}&\text{if }x\geq
0\\\ \frac{1}{2}\,\text{d}x\,\text{d}y\,g(x+\delta
x,y,0.3)\\{1+C\cos[\Phi_{\mathrm{inp}}(i,x+\delta x,y)+\pi]\\}&\text{if
}x<0\\\ \end{cases}.$ (19)

$C$ is the contrast and $\text{d}x$ and $\text{d}y$ are the width and height
of each pixel respectively. This image model was used to generate 100 images
with $n_{\mathrm{images}}=100$ each for three choices for
$\gamma_{\mathrm{inp}}(x,y)$ with the parameters:

$\displaystyle\theta_{\mathrm{inp}}(i)$ $\displaystyle=\frac{2\pi
i}{n_{\mathrm{images}}}\ \text{with}\ i\in[0,n_{\mathrm{images}})$ (20)
$\displaystyle\gamma_{\mathrm{inp}}(x,y)$
$\displaystyle=\begin{cases}2x^{2}+4y\\\ 4x+4y\\\ 4(0.2x)^{2}+2y\end{cases}.$
(21)

The input phase patterns $\gamma_{\mathrm{inp}}(x,y)$ and the main result of
this section $\gamma_{\mathrm{diff}}(x,y)$ are plotted in Fig. 3.
$\gamma_{\mathrm{diff}}(x,y)$ is, with the exception of a few outlier pixels,
close to machine precision. The differences vary slightly from run to run and
one can see weak remnants of the patterns $\gamma_{\mathrm{inp}}(x,y)$ in
$\gamma_{\mathrm{diff}}(x,y)$. The PSPR reconstruction yields an excellent
agreement in this case where $\theta_{\mathrm{diff}}(i)$ is close to the
precision of the used floating point numbers, so for double precision about
$\qty{e-16}{}$.

a) | $2x^{2}+4y$ |
---|---|---
|  |
b) | $4x+4y$ |
|  |
c) | $4(0.2x)^{2}+2y$ |
|  |
Figure 3:  Spatial phase patterns $\gamma_{\mathrm{inp}}(x,y)$ used for the
synthetic image generation as defined in Eq. 21. These cases were chosen to
show different scaling behaviors when including noise in Section IV. The two
columns in each image correspond to the two ports of the interferometer which
are modeled as the two cases in Eq. 19. The reconstruction was done from a
data set of 100 images generated with Eq. 19 and the parameters in Eq. 20 and
Eq. 21. Almost all pixels are reconstructed close to machine precision around
$\qty{e-16}{}$, the original phase patterns are weakly visible. Note the
discrete and non-linear color bars in the difference plots to accommodate
outliers.

##  IV Phase Reconstruction in the Presence of Shot Noise

Shot noise, also referred to as quantum projection noise [41], is a type of
noise which is present in all quantum measurements. This type of noise was
chosen to evaluate the reconstruction accuracy of PSPR in the case of
realistic measurement in a quantum system.

The intensity of the model images in the previous section is used as the
probability density to find an atom in a particular pixel. For each atom, a
single pixel is picked randomly according to this density. The brightness of
each pixel is then defined as the number of atoms that populate it. All
intensity values in the following graphs are therefore in units of atoms.

To understand the reconstruction performance of PSPR in the presence of shot
noise, the image generation process is repeated $n_{\mathrm{run}}$ times. All
graphs in this chapter focus on the standard deviation $\mathrm{std}$, because
the differences between input and reconstruction average to zero for the
chosen parameter regimes. Each of the following figures was generated with
$n_{\mathrm{run}}=74000$ simulation runs per input parameter combination. The
default input parameters are $\phi(i)=2\pi i/1000$ with $i\in[0,1000)$,
$N_{\mathrm{atoms}}=10^{5}$, $n_{\mathrm{images}}=1000$ and the contrast
$C=1$.

The reconstruction accuracy
$\sigma_{i}\coloneqq\mathrm{std}[\theta_{\mathrm{diff}}(i)]$ varies as a
function of $\theta_{\mathrm{inp}}(i)$ (Fig. 4b). This strongly correlates
with the radius of the coefficient pair on the ellipse (Fig. 4a).

(a)

(b)

Figure 4:  Results for interferograms defined in Eq. 19 with shot noise for
$n_{\mathrm{images}}=10000$, $N_{\mathrm{atoms}}=10^{5}$, $C=1$. a) The
distance $r$ of the coordinate $(w_{1}(i),w_{2}(i)$ to the center of the
fitted ellipse as a function of $\theta_{\mathrm{inp}}(i)$. b) The
reconstruction accuracy $\sigma_{i}$ as a function of
$\theta_{\mathrm{inp}}(i)$. One can clearly see the strong correlation between
$\sigma_{i}$ and $r$.

This is abstracted away by taking the maximum $\mathrm{max}(\sigma_{i})$, in
order to understand and characterize the shot-noise scaling as a function of
other simulation parameters.
$\sigma_{xy}\coloneqq\mathrm{std}[\gamma_{\mathrm{diff}}(x,y)]$ does not show
any scaling as a function of $\Phi_{\mathrm{inp}}(i,x,y)$.

When systematically changing the number of atoms per image
$N_{\mathrm{atoms}}$, the number of images $n_{\mathrm{images}}$ per analysis,
the contrast $C$, and the spatial phase distribution
$\gamma_{\mathrm{inp}}(x,y)$ one finds the following empirical scaling laws
for $\mathrm{max}(\sigma_{i})$ and $\sigma_{xy}$:

$\displaystyle\mathrm{max}(\sigma_{i})$ $\displaystyle=\frac{a}{C\
N_{\mathrm{atoms}}^{1/2}}$ (22) $\displaystyle\sigma_{xy}$
$\displaystyle=\frac{b}{C\ \left[\bar{I}(x,y)\
n_{\mathrm{images}}\right]^{1/2}}.$ (23)

$a$ and $b$ are fit constants with $a$ depending on the chosen pattern
$\gamma_{\mathrm{inp}}(x,y)$:

$\displaystyle a$ $\displaystyle\approx\begin{cases}1.70\
&\gamma_{\mathrm{inp}}(x,y)=2x^{2}+4y\\\ 1.67\
&\gamma_{\mathrm{inp}}(x,y)=4x+4y\\\ 2.51\
&\gamma_{\mathrm{inp}}(x,y)=4(0.2x^{2})+2y\end{cases}$ $\displaystyle b$
$\displaystyle\approx 1.41.$

Both laws match the intuition that the reconstruction accuracy depends on how
many interfering atoms were observed per reconstructed value and are
compatible with atomic shot noise scaling.

Figure 6 and Fig. 6 confirm the shot-noise scaling of Eq. 22 and Eq. 23 with
$N_{\mathrm{atoms}}$ and $n_{\mathrm{images}}$, respectively. One can see that
for about $N_{\mathrm{atoms}}>100$ atoms per image and for
$n_{\mathrm{images}}>50$ images per sequence the signal quality is high enough
for the scaling law Eq. 22 to hold. Since $\mathrm{max}(\sigma_{i})$ does not
scale as a function of $n_{\mathrm{images}}$ the reconstruction accuracy
converges towards constant values for large enough image sets.

Figure 5:  Worst-case reconstruction accuracy $\mathrm{max}(\sigma_{i})$ as a
function of $N_{\mathrm{atoms}}$ for $C=1$, $n_{\mathrm{images}}=1000$ and
$\theta_{\mathrm{inp}}(i)$ as in Eq. 20. For approximately
$N_{\mathrm{atoms}}>100$ PCA works well enough so that
$\mathrm{max}(\sigma_{i})$ scales with $a/(C\cdot\sqrt{N_{\mathrm{atoms}}})$,
with $a$ being a fit constant which depends on spatial phase pattern
$\gamma_{\mathrm{inp}}(x,y)$.

Figure 6:  Worst-case reconstruction accuracy $\mathrm{max}(\sigma_{i})$ as a
function of $n_{\mathrm{images}}$. $\mathrm{max}(\sigma_{i})$ is independent
of $n_{\mathrm{images}}$ and therefore only improves until
$n_{\mathrm{images}}=50$ which are enough images for PSPR to work well for
$C=1$, $N_{\mathrm{atoms}}=10^{5}$ and $\theta_{\mathrm{inp}}(i)$ as in Eq.
20.

Similarly, $\sigma_{xy}$ solely scales with the total number of interfering
atoms that fell into the reconstructed pixel during the image sequence as
depicted in Fig. 7.

Figure 7: $\sigma_{xy}$ scaling with the atom number per image. This plot
shows all pixels from all three phase patterns $\gamma_{\mathrm{inp}}(x,y)$.
Four different values from the scan of $n_{\mathrm{images}}$ were selected and
are color-coded according to the legend. The order in the plot from top to
bottom is the same as in the legend. $\sigma_{xy}$ scales with $b/\\{C\
[n_{\mathrm{images}}\bar{I}(x,y)]^{\frac{1}{2}}\\}$ with $b=1.41$ as a fitted
constant. One can see that $\sigma_{xy}$ of a pixel only depends on the number
of atoms that populated that pixel over the course of the whole data set.

##  V Example Reconstruction of the Image and Spatial Phases from Experimental
Data

In this chapter, PSPR is used to extract the spatial phase profile of two
interferometry ports in a set of experimental data images. The experimental
data was taken with the apparatus described in reference [42]. It uses an atom
chip and laser beams for the generation and preparation of Bose-Einstein
condensates (BEC). The experimental sequence is the fountain gravimeter shown
in reference [43] where the atom chip is also used as a mirror for the
interrogating laser beam and therefore acts as the inertial reference of the
measurement.

The atom chip has an optical coating to yield a highly reflective optical
surface, but is not an ideal mirror. Its surface has imperfections, which can
lead to spatial wavefront distortions, beam intensity fluctuations and other
aberrations of the employed light fields. These can produce a spatial phase
distortion in the interferometer outputs, which is examined with PSPR in this
article. A first test with a pure PCA-based evaluation on the experimental
data sets already showed that the intrinsic interferometer sensitivity can be
enhanced by enlarged contrast and respectively reduced technical noise [44].

The first measurement consists of 991 images from a Mach-Zehnder
interferometer operated with Bragg diffraction where each beam-splitter
transfers a momentum of $\qty{2}{}$ onto the atoms [45, 46], with $k$ denoting
the wave number of the Bragg laser and $\hbar$ the reduced Planck’s constant.
The second data set is from an interferometer with larger momentum transfer of
$\qty{6}{}$ via higher-order Bragg transitions and consists of 1015 images.
Figure 8 shows the results of step 2 of the PSPR algorithm (Fig. 2). The
average intensity $\bar{I}(x,y)$ is clearly distributed over two ports. This
figure does not show the $w_{j}(i)$ because they are unordered due to
vibrations causing large phase jumps between each image. For the first data
set the first component shows the typical pattern of mainly moving atoms
between the two ports, depending on the value of $w_{1}(i)$.

Component |  $DataSet1$(2) |  Data Set 2 $(\qty{6}{})$
---|---|---
$\bar{I}(x,y)$ |  |
$\mathrm{pc}_{1}(x,y)$ |  |
$\mathrm{pc}_{2}(x,y)$ |  |
$\mathrm{pc}_{3}(x,y)$ |  |
$\mathrm{pc}_{4}(x,y)$ |  |
Figure 8:  Average intensity $\bar{I}(x,y)$ and the first four principal
components for a Bragg diffraction based atom gravimeter featuring momentum
transfers of $\qty{2}{}$ (Data set 1) and $\qty{6}{}$ (Data set 2). Data set 1
contains 991 images, while the results for data set 2 are reconstructed from
1015 images. The two ports are visible in all images. $\bar{I}(x,y)$ has a
linear color map indicating the intensities while the components are plotted
with a diverging color map with blue indicating negative values and yellow
indicating positive values. $\mathrm{pc}_{1}(x,y)$ in the left column displays
the main interferometer function for relatively flat phase profiles, moving
atoms from left to right positive coefficients and from right to left for
negative coefficients. The mean clearly shows the two ports, the first
component is cosine-like and encodes the movement of atoms between the two
ports. The second component is the sine-like component while the components
three and four are most likely due to motion of the camera relative to the
interferometer.

The result of step 3 of PSPR yields the ellipses of Fig. 9. In the second data
set, other effects are stronger than the interferometry so that the $w_{2}(i)$
and $w_{4}(i)$ form the ellipse used for phase reconstruction. The first data
set has a weaker spatial phase profile, which causes the ellipse to be more
stretched compared to the second data set with a phase profile which has a
larger dynamic range of phases.

aaaaaData Set 1 $(\qty{2}{})$ |  aaaaaData Set 2 $(\qty{6}{})$
---|---
|
Figure 9:  These two plots show the correlations between the two ”sine” and ”cosine” components of the first and second data set (right). The coefficients $w_{j}(i)$ are the projections of each image $I(i,x,y)$ onto the components $\mathrm{pc}_{j}(x,y)$ which are shown in Fig. 8. This ellipse pattern is a signature for the sine-like and cosine-like components. When comparing with Fig. 10 one can see that the stronger ellipse-shape in the first data set corresponds to a more flat spatial phase profile compared to the second data set. This is consistent with the simulations.  Data Set 1 $(\qty{2}{})$a |  Data Set 2 $(\qty{6}{})$a
---|---
|
Figure 10:  The reconstructed spatial phase profile
$\gamma_{\mathrm{rec}}(x,y)$ for the two experimental data sets. The first
data set with $\qty{6}{}$ momentum transfers shows much stronger distortions
which is to be expected since there were far more atom-light interactions with
distorted wavefronts during the sequence. In the areas on the sides of the
ports, where there are almost no atoms, the signal strength is too weak to
give a meaningful signal.

The main result of this section can be seen in Fig. 10 which shows the actual
reconstruction of the spatial phase profiles when using the whole algorithm.
The second data set shows a much stronger spatial phase gradient, than the
first data set, which is to be expected due to the larger number of atom-light
interactions. In the second data set with the stronger $\qty{6}{}$ momentum
transfer, the spatial phases are spanning a gamut of $2\pi$ so for any overall
interferometer phase there are always some interfering atoms in each of the
ports. Analysing this data set without taking into account the spatial phase
profile would leave a lot of useful signal unutilized and would lead to
increased systematic effects.

In the second data set there are also two faint clouds between the strong
clouds which are a loss channel formed by atoms that interacted at a different
momentum transfer order. These clouds indicate the possibility of utilising
information from what would normally regarded as a loss channel.

##  VI Conclusion and Outlook

The PSPR algorithm developed in this article allows to reconstruct the spatial
phase profiles in image sets of atom interferometry ports. Its performance was
proven to theoretically reach the machine precision level. The algorithm was
then successfully applied to reconstruct the experimental phase front pattern
seen in the ports of an atom gravimeter. These results point towards a future
use of PSPR in the analysis of atom interferometers embracing and utilizing
non-plane wavefronts to make measurements more accurate and versatile instead
of simply smoothing the inevitable wavefronts of the interrogation lasers.

It is a first step towards measuring wavefront aberrations using the exact
atom-light phase profiles at the positions where the interaction occurs during
experimental sequences. This will improve the understanding of the systematics
in matter wave systems by varying different parameters in the setup and
characterizing the impact on the spatial phase map.

All parts of PSPR can be implemented very efficiently even on hardware with
limited performance and memory. For PCA, randomized out-of-core algorithms
exist, which can achieve the decomposition without ever loading the whole data
set into the main memory [47]. Furthermore, the ellipse fitting algorithm is
stable even with in the presence of significant instrument noise levels and
completely iteration-free [33].

PSPR performs a lossy compression that is well adapted to interferograms of
atom interferometers. It exploits that the spatial phase imprints do not vary
much during each measurement while still allowing arbitrary phase patterns
which can be tracked over the long-term operation of the interferometer. Using
this knowledge it reduces the three-dimensional $I(i,x,y)$ into the much
smaller one-dimension $\theta(i)$ and two-dimensional $\bar{I}(x,y)$ and
$\gamma(x,y)$.

Additionally, PSPR offers two heuristics that can be used to automatically
detect when the atom interferometer deviates too much from its expected
behavior. The first heuristic is how well the two components that are used for
the phase reconstruction describe the whole image set. This can be computed,
for example, by projecting the images into that basis and evaluating their
deviation from the raw data. The second heuristic is how well the ellipse fits
onto the coefficients giving a clear indicator for how accurately the two
selected components actually describe the expected interferogram defined in
Eq. 1.

Having the potential to realize the phase extraction and tracking of key
performance characteristics directly on the specific atom interferometry
device, PSPR helps to make autonomous operation possible. For instance, as
part of the setup of the interferometer one can calibrate expectation
thresholds for how well the two used PCA components and the ellipse fit
explain the data. When these thresholds are violated, the device may,
depending on the use case, launch debugging routines, safeguards, warnings or
special calibration routines to mitigate the problems or explicitly request
user intervention.

For future atom interferometers deployed in real-world environments on ground
and in space [20, 48, 49], it will be necessary to fully automate the data
analysis directly on an embedded processor. For field applications, the number
of measurements will be too large to be analyzed with manual intervention by
non-experts. Moreover, the storage requirements for the raw data also become a
limiting factor. Finally, in space or in environments where communication is
denied, bandwidth to send raw data to be processed is not available or
insufficient. For all these reasons, PSPR will be a well suited on-board tool.

###### Acknowledgements.

We thank Klemens Hammerer for his insights about shot noise models and Víctor
J. Martínez-Lahuerta and Matthew Glaysher for carefully reading the
manuscript. Additionally we thank Christian Struckmann for helpful discussions
about phase differences and uniqueness. This work was funded by the Deutsche
Forschungsgemeinschaft (German Research Foundation) under Germany’s Excellence
Strategy (EXC-2123 QuantumFrontiers Grants No. 390837967), through CRC 1227
(DQ-mat) within Projects No. A05 and through the QuantERA 2021 co-funded
project No. 499225223 (SQUEIS), and the German Space Agency (DLR) with funds
provided by the German Federal Ministry of Economic Affairs and Energy (German
Federal Ministry of Education and Research (BMBF)) due to an enactment of the
German Bundestag under Grants No. DLR 50WM1952 (QUANTUS-V-Fallturm), No.
50WM2250A (QUANTUS plus), No. 50WM2450A (QUANTUS-VI), No. 50WM2245A (CAL-II),
No. 50WM2263A (CARIOQA-GE), No. 50WM2253A (AI-Quadrat), and No. 50NA2106
(QGYRO+).

*

##  Appendix A Computing the Minimal Difference Between Two Phase Results

The phase offset degree of freedom $c$ can be determined using
$\varphi_{\mathrm{diff}}$:

$\displaystyle\vartheta(i)$
$\displaystyle=\varphi_{\mathrm{diff}}[\theta_{\mathrm{inp}}(i),\theta_{\mathrm{rec}}(i)]$
$\displaystyle x$ $\displaystyle=\mathrm{mean}\\{\cos[\vartheta(i)]\\}$
$\displaystyle y$ $\displaystyle=\mathrm{mean}\\{\sin[\vartheta(i)]\\}$
$\displaystyle c$ $\displaystyle=\mathrm{arctan2}(y,x).$

The value of $c$ is the average of all differences between
$\theta_{\mathrm{inp}}(i)$ and $\theta_{\mathrm{rec}}(i)$ while not making
averaging errors due to the phase wrap-around after $2\pi$. That average is
the angle of the average position of the points of all angles on the unit
circle.

The sign $s$ is determined by computing $c$ for both
$\varphi_{\mathrm{diff}}[\theta_{\mathrm{inp}}(i),\theta_{\mathrm{rec}}(i)]$
and
$\varphi_{\mathrm{diff}}[\theta_{\mathrm{inp}}(i),-\theta_{\mathrm{rec}}(i)]$
and choosing the sign for which that difference is smaller.

## References

  * Kasevich and Chu [1991] M. Kasevich and S. Chu, Physical Review Letters 67, 181 (1991).
  * Ménoret  _et al._ [2018] V. Ménoret, P. Vermeulen, N. Le Moigne, S. Bonvalot, P. Bouyer, A. Landragin, and B. Desruelle, Scientific Reports 8, 12300 (2018).
  * Wu  _et al._ [2019] X. Wu, Z. Pagel, B. S. Malek, T. H. Nguyen, F. Zi, D. S. Scheirer, and H. Müller, Science Advances 5, eaax0800 (2019).
  * Stray  _et al._ [2022] B. Stray, A. Lamb, A. Kaushik, J. Vovrosh, A. Rodgers, J. Winch, F. Hayati, D. Boddice, A. Stabrawa, A. Niggebaum, M. Langlois, Y.-H. Lien, S. Lellouch, S. Roshanmanesh, K. Ridley, G. De Villiers, G. Brown, T. Cross, G. Tuckwell, A. Faramarzi, N. Metje, K. Bongs, and M. Holynski, Nature 602, 590 (2022).
  * Geiger  _et al._ [2011] R. Geiger, V. Ménoret, G. Stern, N. Zahzam, P. Cheinet, B. Battelier, A. Villing, F. Moron, M. Lours, Y. Bidel, A. Bresson, A. Landragin, and P. Bouyer, Nature Communications 2, 474 (2011).
  * Cheiney  _et al._ [2018] P. Cheiney, L. Fouché, S. Templier, F. Napolitano, B. Battelier, P. Bouyer, and B. Barrett, Physical Review Applied 10, 034030 (2018).
  * Morel  _et al._ [2020] L. Morel, Z. Yao, P. Cladé, and S. Guellati-Khélifa, Nature 588, 61 (2020).
  * Parker  _et al._ [2018] R. H. Parker, C. Yu, W. Zhong, B. Estey, and H. Müller, Science 360, 191 (2018).
  * Rosi  _et al._ [2014] G. Rosi, F. Sorrentino, L. Cacciapuoti, M. Prevedelli, and G. M. Tino, Nature 510, 518 (2014).
  * Asenbaum  _et al._ [2020] P. Asenbaum, C. Overstreet, M. Kim, J. Curti, and M. A. Kasevich, Physical Review Letters 125, 191101 (2020).
  * Bordé [1989] Ch.J. Bordé, Physics Letters A 140, 10 (1989).
  * Ketterle W.  _et al._ [1999] Ketterle W., Durfee D.S., and Stamper-Kurn D.M., Proceedings of the International School of Physics “Enrico Fermi” 140, 67 (1999).
  * Dickerson  _et al._ [2013] S. M. Dickerson, J. M. Hogan, A. Sugarbaker, D. M. S. Johnson, and M. A. Kasevich, Physical Review Letters 111, 083001 (2013).
  * Sugarbaker  _et al._ [2013] A. Sugarbaker, S. M. Dickerson, J. M. Hogan, D. M. S. Johnson, and M. A. Kasevich, Physical Review Letters 111, 113002 (2013).
  * Wang  _et al._ [2005] Y.-J. Wang, D. Z. Anderson, V. M. Bright, E. A. Cornell, Q. Diot, T. Kishimoto, M. Prentiss, R. A. Saravanan, S. R. Segal, and S. Wu, Physical Review Letters 94, 090405 (2005).
  * Horikoshi and Nakagawa [2006] M. Horikoshi and K. Nakagawa, Physical Review A 74, 031602 (2006).
  * Niebauer  _et al._ [1995] T. M. Niebauer, G. S. Sasagawa, J. E. Faller, R. Hilt, and F. Klopping, Metrologia 32, 159 (1995).
  * Schkolnik  _et al._ [2015] V. Schkolnik, B. Leykauf, M. Hauth, C. Freier, and A. Peters, Applied Physics B 120, 311 (2015).
  * Karcher  _et al._ [2018] R. Karcher, A. Imanaliev, S. Merlet, and F. P. D. Santos, New Journal of Physics 20, 113041 (2018).
  * Geiger  _et al._ [2020] R. Geiger, A. Landragin, S. Merlet, and F. Pereira Dos Santos, AVS Quantum Science 2, 024702 (2020).
  * Bruning  _et al._ [1974] J. H. Bruning, D. R. Herriott, J. E. Gallagher, D. P. Rosenfeld, A. D. White, and D. J. Brangaccio, Applied Optics 13, 2693 (1974).
  * Schreiber and Bruning [2007] H. Schreiber and J. H. Bruning, in  _Optical Shop Testing_, edited by D. Malacara (Wiley, 2007) 1st ed., pp. 547–666.
  * Servin  _et al._ [2014] M. Servin, J. A. Quiroga, and J. M. Padilla,  _Fringe Pattern Analysis for Optical Metrology: Theory, Algorithms, and Applications_ (Wiley-VCH Verlag GmbH & Co. KGaA, Weinheim, Germany, 2014).
  * Pearson [1901] K. Pearson, The London, Edinburgh, and Dublin Philosophical Magazine and Journal of Science 2, 559 (1901).
  * Hotelling [1933] H. Hotelling, Journal of Educational Psychology 24, 417 (1933).
  * Segal  _et al._ [2010] S. R. Segal, Q. Diot, E. A. Cornell, A. A. Zozulya, and D. Z. Anderson, Physical Review A 81, 053601 (2010).
  * Dubessy  _et al._ [2014] R. Dubessy, C. D. Rossi, T. Badr, L. Longchambon, and H. Perrin, New Journal of Physics 16, 122001 (2014).
  * Pedregosa  _et al._ [2011] F. Pedregosa, G. Varoquaux, A. Gramfort, V. Michel, B. Thirion, O. Grisel, M. Blondel, P. Prettenhofer, R. Weiss, V. Dubourg, J. Vanderplas, A. Passos, D. Cournapeau, M. Brucher, M. Perrot, and E. Duchesnay, Journal of Machine Learning Research 12, 2825 (2011).
  * Yatabe  _et al._ [2017] K. Yatabe, K. Ishikawa, and Y. Oikawa, Journal of the Optical Society of America A 34, 87 (2017).
  * Berman [1997] P. R. Berman,  _Atom Interferometry_ (Academic Press, San Diego, CA, 1997).
  * Hogan J. M.  _et al._ [2009] Hogan J. M., Johnson D. M. S., and Kasevich M. A., Proceedings of the International School of Physics “Enrico Fermi” 168, 411 (2009).
  * Merlet  _et al._ [2009] S. Merlet, J. Le Gouët, Q. Bodart, A. Clairon, A. Landragin, F. Pereira Dos Santos, and P. Rouchon, Metrologia 46, 87 (2009).
  * Halir and Flusser [1998] R. Halir and J. Flusser, Proc. of Int. Conf. in Central Europe on Computer Graphics, Visualization and Interactive Digital Media , 125 (1998).
  * van der Walt  _et al._ [2014] S. van der Walt, J. L. Schönberger, J. Nunez-Iglesias, F. Boulogne, J. D. Warner, N. Yager, E. Gouillart, T. Yu, and t. scikit-image contributors, PeerJ 2, e453 (2014).
  * Harris  _et al._ [2020] C. R. Harris, K. J. Millman, S. J. van der Walt, R. Gommers, P. Virtanen, D. Cournapeau, E. Wieser, J. Taylor, S. Berg, N. J. Smith, R. Kern, M. Picus, S. Hoyer, M. H. van Kerkwijk, M. Brett, A. Haldane, J. F. del Río, M. Wiebe, P. Peterson, P. Gérard-Marchant, K. Sheppard, T. Reddy, W. Weckesser, H. Abbasi, C. Gohlke, and T. E. Oliphant, Nature 585, 357 (2020).
  * Frostig  _et al._ [2018] R. Frostig, M. J. Johnson, and C. Leary, in  _SysML_ (SysML, Stanford, CA, 2018).
  * Bradbury  _et al._ [2018] J. Bradbury, R. Frostig, P. Hawkins, M. J. Johnson, C. Leary, D. Maclaurin, G. Necula, A. Paszke, J. VanderPlas, S. Wanderman-Milne, and Q. Zhang, JAX: composable transformations of Python+NumPy programs, https://github.com/google/jax (2018), version 0.4.25.
  * The Bokeh Development Team [2013] The Bokeh Development Team, bokeh, https://github.com/bokeh/bokeh (2013), version 3.4.0.
  * The Colorcet Development Team [2016] The Colorcet Development Team, colorcet, https://github.com/holoviz/colorcet (2016), version 3.0.1.
  * Nichols [2024] D. Nichols, Coloring for colorblindness, https://davidmathlogic.com/colorblind (2024).
  * Itano  _et al._ [1993] W. M. Itano, J. C. Bergquist, J. J. Bollinger, J. M. Gilligan, D. J. Heinzen, F. L. Moore, M. G. Raizen, and D. J. Wineland, Physical Review A 47, 3554 (1993).
  * van Zoest  _et al._ [2010] T. van Zoest, N. Gaaloul, Y. Singh, H. Ahlers, W. Herr, S. T. Seidel, W. Ertmer, E. Rasel, M. Eckart, E. Kajari, S. Arnold, G. Nandi, W. P. Schleich, R. Walser, A. Vogel, K. Sengstock, K. Bongs, W. Lewoczko-Adamczyk, M. Schiemangk, T. Schuldt, A. Peters, T. Könemann, H. Müntinga, C. Lämmerzahl, H. Dittus, T. Steinmetz, T. W. Hänsch, and J. Reichel, Science 328, 1540 (2010).
  * Abend  _et al._ [2016] S. Abend, M. Gebbe, M. Gersemann, H. Ahlers, H. Müntinga, E. Giese, N. Gaaloul, C. Schubert, C. Lämmerzahl, W. Ertmer, W. P. Schleich, and E. M. Rasel, Physical Review Letters 117, 203003 (2016).
  * Abend [2017] S. Abend,  _Atom-Chip-Based Quantum Gravimetry with Bose-Einstein Condensates_ , Ph.D. thesis, Leibniz Universität Hannover (2017).
  * Giltner  _et al._ [1995] D. M. Giltner, R. W. McGowan, and S. A. Lee, Physical Review A 52, 3966 (1995).
  * Martin  _et al._ [1988] P. Martin, B. Oldaker, A. Miklich, and D. Pritchard, Physical Review Letters 60, 515 (1988).
  * Halko  _et al._ [2011] N. Halko, P.-G. Martinsson, Y. Shkolnisky, and M. Tygert, SIAM Journal on Scientific Computing 33, 2580 (2011).
  * Ahlers  _et al._ [2022] H. Ahlers, L. Badurina, A. Bassi, B. Battelier, Q. Beaufils, K. Bongs, P. Bouyer, C. Braxmaier, O. Buchmueller, M. Carlesso, E. Charron, M. L. Chiofalo, R. Corgier, S. Donadi, F. Droz, R. Ecoffet, J. Ellis, F. Estève, N. Gaaloul, D. Gerardi, E. Giese, J. Grosse, A. Hees, T. Hensel, W. Herr, P. Jetzer, G. Kleinsteinberg, C. Klempt, S. Lecomte, L. Lopes, S. Loriani, G. Métris, T. Martin, V. Martín, G. Müller, M. Nofrarias, F. P. D. Santos, E. M. Rasel, A. Robert, N. Saks, M. Salter, D. Schlippert, C. Schubert, T. Schuldt, C. F. Sopuerta, C. Struckmann, G. M. Tino, T. Valenzuela, W. von Klitzing, L. Wörner, P. Wolf, N. Yu, and M. Zelan, STE-QUEST: Space Time Explorer and QUantum Equivalence principle Space Test (2022), arxiv:2211.15412 [gr-qc, physics:hep-ex, physics:hep-ph, physics:physics, physics:quant-ph] .
  * Lévèque  _et al._ [2023] T. Lévèque, C. Fallet, J. Lefebve, A. Piquereau, A. Gauguet, B. Battelier, P. Bouyer, N. Gaaloul, M. Lachmann, B. Piest, E. Rasel, J. Müller, C. Schubert, Q. Beaufils, and F. Pereira Dos Santos, in  _International Conference on Space Optics — ICSO 2022_, edited by K. Minoglou, N. Karafolas, and B. Cugny (SPIE, Dubrovnik, Croatia, 2023) p. 129.


# TokenMixup: Efficient Attention-guided
Token-level Data Augmentation for Transformers

Hyeong Kyu Choi , Joonmyung Choi∗, Hyunwoo J. Kim
Department of Computer Science and Engineering, Korea University
{imhgchoi, pizard<EMAIL_ADDRESS>
The first two authors contributed equallycorresponding author

###### Abstract

Mixup is a commonly adopted data augmentation technique for image
classification. Recent advances in mixup methods primarily focus on mixing
based on saliency. However, many saliency detectors require intense
computation and are especially burdensome for parameter-heavy transformer
models. To this end, we propose TokenMixup, an efficient attention-guided
token-level data augmentation method that aims to maximize the saliency of a
mixed set of tokens. TokenMixup provides $\times 15$ faster saliency-aware
data augmentation compared to gradient-based methods. Moreover, we introduce a
variant of TokenMixup which mixes tokens within a single instance, thereby
enabling multi-scale feature augmentation. Experiments show that our methods
significantly improve the baseline models’ performance on CIFAR and
ImageNet-1K, while being more efficient than previous methods. We also reach
state-of-the-art performance on CIFAR-100 among from-scratch transformer
models. Code is available at https://github.com/mlvlab/TokenMixup.

##  1 Introduction

Various data augmentation methods have been proposed for computer vision
tasks. One of the most successful augmentation methods is mixup [1], which is
commonly applied to image classification. It attempts to augment the input by
taking a convex combination of two random instances, and by reassigning the
ground truth label correspondingly. Building on mixup, many subsequent works
[2, 3, 4, 5, 6, 7, 8, 9] have been introduced. They focus on mixing instances
in a more meaningful way (e.g., saliency-aware mixup, manifold level mixup,
submodular diversity maximization etc.), while a great majority attempts to
perform mixup based on saliency.
These saliency-aware mixup methods often entail gradient computation. To
detect salient regions, the input is forward-propagated once and back-
propagated to extract the gradient map. However, such a mechanism tends to be
computationally heavy. This is especially pertinent to transformer-based
models [10], as they generally retain high parameter volume. Keeping this in
mind, we aim to present a saliency-aware mixup method favorable for
transformers. Specifically, we regard self-attention as the inherent saliency
detector, which serves as an efficient means for augmentation.

We accordingly propose TokenMixup, an efficient attention-guided token-level
mixup method. Its objective is to mix intermediate token sets so that the
saliency level of a batch is maximized. To detect salient tokens, we take
advantage of the attention map to achieve a $\times 15$ speed-up compared to
the gradient-based saliency estimator, without compromising model performance.
The batch saliency is maximized by optimally matching instance pairs,
accomplished by an algorithm that exactly solves the optimization problem. We
also introduce ScoreNet, a simple module that measures the difficulty of an
instance. Based on ScoreNet, TokenMixup is applied selectively with respect to
sample difficulty, which can be viewed as a type of curriculum learning
method. To the best of our knowledge, this is the first attempt in the
literature to discuss mixup from a curriculum learning perspective.
Furthermore, we present a novel mixup approach that combines intermediate
tokens across transformer layers within a single instance. We name this
Vertical TokenMixup (VTM) since it vertically mixes tokens to provide rich
multi-scale information to the attention layer.

Experiments on image classification datasets imply that our methods are both
effective and efficient. Also, an advantage of TokenMixup is that it can be
used in parallel with other mixup methods. In combination with those methods,
we improve our baselines’ performance by significant margins on CIFAR-10,
CIFAR-100, and ImageNet-1K. Especially for CIFAR-100, we achieve a new state-
of-the-art performance among transformer-based models that do not use pre-
trained models.

Then, our contributions are fourfold:

  * •

We propose TokenMixup, an efficient attention-guided token-level mixup for
transformers. By incorporating the self-attention map, we achieve $\times 15$
faster saliency-aware data augmentation compared to gradient-based methods.

  * •

We incorporate curriculum learning into mixup, to adaptively perform
TokenMixup in an intermediate layer based on the confidence score of ScoreNet.

  * •

We also studied a variant, Vertical TokenMixup, which performs mixup with a
single sample and enables multi-scale feature augmentation.

  * •

We achieve state-of-the-art performance on CIFAR-100 among transformer models
trained from scratch. We also gain significant improvements over baseline
models for image classification.

##  2 Related Works

#### Mixup.

Input Mixup [1] is a data augmentation method widely adopted for image
classification. The classification model is trained with a convex combination
of input images and labels. A special case of Input Mixup is CutMix [2], which
can be seen as a pixel-wise mixup method with binary masks. Recent advances in
these mixup methods mainly focus on appropriately utilizing saliency
information, and on mixing in image feature levels.
The core motivation of saliency-based mixup is that salient regions should be
preserved when mixed, to retain a sufficient amount of information and learn
more natural feature representations. SaliencyMix [5] adopts various saliency
detectors to directly extract salient regions. Puzzle Mix [3] attempts to mix
images using saliency information while retaining local statistics. Co-Mixup
[4] maximizes gradient-based saliency while encouraging supermodular diversity
of the mixed images. SuperMix [9] takes advantage of supervised signals to mix
input images based on saliency.
On the other hand, Manifold Mixup [6] has provided theoretical grounds on the
advantages of mixing images in the higher level of features. Manifold mixing
helps the model learn flatter representation and smoother decision boundaries.
Moreover, it is noted that it enhances robustness to adversarial attacks.
Recently, AlignMix [7] was proposed to further align geometric properties of
objects by mixing instances in the feature space. Furthermore, StyleMix [8]
proposed to manipulate content and style information disparately to generate
more abundant and robust samples.

#### Vision Transformers.

Originating from the natural language processing field, the Transformer [10]
model enjoyed tremendous success in numerous computer vision applications,
e.g., image classification, object detection [11, 12], and human-object
interaction detection [13, 14] et cetera. Specifically, image classification
is one of the most fundamental computer vision tasks. Starting from the ViT
[15], many works have attempted to assimilate the convolutional operation with
transformer blocks [16, 17, 18, 19, 20, 21]. Other variants have effectively
incorporated the convolutional neural network into the transformer to achieve
superior performance [22, 23, 24, 25, 26].

##  3 Methods

Here, we introduce TokenMixup, a simple but effective mixup method for
transformer models. The goal of our method is to augment intermediate tokens
while maximizing the saliency level. This can be accomplished in three major
steps: 1) Sample difficulty assessment, 2) Attention-guided saliency detection
& optimal assignment, 3) Token-level mixup. The following subsections provide
detailed descriptions of the steps.

Figure 1: TokenMixup. Batch samples $X$ are first filtered with respect to
difficulty scores evaluated by ScoreNet in $\mathcal{F}$, resulting in easy
samples $\tilde{X}$. Then, images in $\tilde{X}$ are optimally paired with
samples in $X$ via Hungarian matching, so that the overall saliency of the
mini-batch is maximized after mixup. For a matched pair ($\tilde{X}^{(i)}$,
$X^{(j)}$), tokens of $\tilde{X}^{(i)}$ with sufficiently lower saliency
levels than the tokens of $X^{(j)}$ are replaced, resulting in a new token set
$\tilde{X}^{(i)}_{\textrm{mix}}$ containing a greater saliency level.

###  3.1 Sample Difficulty Assessment

We first assess the difficulty of the input to adaptively decide whether to
apply augmentation. To evaluate sample difficulty, a measuring function
$\mathcal{F}$ needs to be defined. We utilize a parameterized module,
ScoreNet, which is a simple MLP that predicts the target value $Y^{(i)}$ based
on the intermediate tokens $X^{(i)}$, where $i$ is the index of the sample in
a batch (See the supplement for figure). Given the ScoreNet output, the
difficulty score is evaluated with the prediction loss computed as

$\mathcal{F}(X^{(i)},Y^{(i)})=\textrm{CrossEntropy}(\textrm{ScoreNet}(X^{(i)}),Y^{(i)}).$
(1)

If the score is greater than a threshold $\tau$, the sample is deemed
sufficiently hard and no mixup is performed. On the other hand, the set of
easy samples with $\mathcal{F}(X^{(i)},Y^{(i)})<\tau$ denoted as $\tilde{X}$
will be augmented by Token-level mixup after pairing with other samples in a
mini-batch, as illustrated in Figure 1. Note that ScoreNet is trained
simultaneously with the main model and training ScoreNet to assess sample
difficulty at the intermediate layer can be viewed as an auxiliary task.

###  3.2 Attention-guided Saliency Detection & Optimal Assignment

Several works [19, 27, 28, 29] have discussed the imbalanced information of
tokens. Due to this imbalance, random token mixing would potentially cause
significant information loss and meaningless token replacements (See the
supplement for relevant analysis). Thus, we aim to mix tokens based on
saliency.

#### Attention-guided saliency detection.

Instead of the computationally heavy gradient-based saliency detectors [3, 4],
we take advantage of the attention map, an inherent saliency approximator
within transformer modules. Specifically, we can infer the saliency of the
tokens in the $i^{th}$ layer via Attention Rollout [30], which computes the
attention imposed on the tokens from layer $i$ to $i+\ell$ ($\ell\geq 0$) as

$A=\Phi^{(i)}\cdot\Phi^{(i+1)}\cdot\cdot\cdot\Phi^{(i+\ell)},$ (2)

such that $\Phi^{(i)}=\frac{1}{H}\sum_{h=1}^{H}\Phi^{(i)}_{h}$ where $H$ is
the number of heads in the multi-head attention layer, and $\Phi^{(i)}_{h}$ is
the $h^{th}$ attention head of layer $i$. In order to retrieve $A$, the
samples need to be propagated for the subsequent $\ell$ layers with stop-
gradient. Since each additional forward propagation is an overhead, we
approximate Attention Rollout by setting $\ell=0$. That is, we only use the
attention map of the following layer as the saliency estimator. We find this a
sufficient approximation of the full Attention Rollout. See section 5.1 for
analysis on the soundness of the approximated saliency map. Then, the saliency
score $S_{t}$ can be computed as

$S_{t}=\displaystyle\frac{1}{n}\sum_{i=1}^{n}A_{i,t},$ (3)

where $t=1,2,\ldots,n$ denotes the token index.

#### Optimal assignment.

Based on the estimated saliency of tokens, we aim to maximize the overall
saliency level by optimally assigning a different mixup target for each
instance. So, we first define $P_{i,j}\in\rm I\\!R^{n}$ as the saliency
difference between a random instance pair $(i,j)$, which is computed as

$P_{i,j}\overset{\triangle}{=}S_{j}-S_{i},$ (4)

where $i,j=1,2,\ldots,b$ and $S_{i}\in\rm I\\!R^{n}$ is the token saliency map
from the $i^{th}$ instance. Then, we represent our objective as an
optimization problem written as

$\max_{\sigma\in\mathcal{M}}\>\max_{r_{i}\in\rm
I\\!B^{n}}\>\displaystyle\sum_{i=1}^{b}(P_{i,\sigma(i)}-\rho)^{\top}r_{i},$
(5)

where $\mathcal{M}$ is the set of all possible batch permutations of the $b$
instances, and $\sigma\in\mathcal{M}$ refers to an arbitrary permutation.
$\rho$ is the threshold hyperparameter that controls the minimum saliency gain
required for a token to be mixed, and $r_{i}$ is a binary decision vector for
the $i^{th}$ instance. i.e., token $t$ of $i$ is replaced with token $t$ of
$\sigma(i)$ when $r_{i,t}=1$, and preserved when
$r_{i,t}=0\>\>(t=1,2,\ldots,n)$.

The above optimization problem can be exactly solved by utilizing the
Hungarian Matching algorithm [31]. The matching algorithm requires the score
matrix, $C\in\rm I\\!R^{b\times b}$, which is computed as

$C_{i,j}=\displaystyle\sum_{t=1}^{n}\textrm{max}(P_{i,j}^{(t)}-\rho,\>0),$ (6)

where each item in $C_{i,j}$ refers to the maximum saliency gain resulting
from mixing $\tilde{X}^{(i)}$ and $X^{(j)}$. The actual maximum gain will be
reached when $\rho=0$, but we set a positive threshold to control the minimum
saliency gain required for each token to be replaced. Then, by applying the
Hungarian algorithm, we can find the optimal batch permutation $\sigma^{*}$
such that

$\sigma^{*}=\displaystyle\underset{\sigma}{\arg\max}\sum_{i=1}^{n}C_{i,\sigma(i)}.$
(7)

Then, by mixing $\tilde{X}^{(i)}$ and $X^{(\sigma^{*}(i))}$, our objective is
optimized.

Algorithm 1 TokenMixup

Input: $X\in\rm I\\!R^{b\times n\times d}$, $Y\in\rm I\\!R^{b\times c}$,
$\tau$, $\rho$
Output: $X\in\rm I\\!R^{b\times n\times d}$, $Y\in\rm I\\!R^{b\times c}$

1:$U^{(i)}\leftarrow\textrm{evaluate difficulty of }X^{(i)}\>\>\>\>$ s.t.
$i=1,2,\ldots,b$ $\triangleright$ Eq. (1)

2:$S_{t}^{(i)}\leftarrow\textrm{evaluate saliency of each token with
$\ell$-step attention rollout}$ $\triangleright$ Eq. (3)

3:$\tilde{X}^{(i)},\tilde{Y}^{(i)},\tilde{S}^{(i)}\leftarrow$ select easy
samples w.r.t $U$ and $\tau$ s.t. $i=1,2,\ldots,b^{\prime}$

4:$C_{ij}\leftarrow\sum_{t}\textrm{max}(S_{t}^{(j)}-\tilde{S}_{t}^{(i)}-\rho,\>0)\>\>\>\>$
s.t. $i=1,2,\ldots,b^{\prime}$ and $j=1,2,\ldots,b$ $\triangleright$ Eq. (6)

5:$\sigma(m)\leftarrow\textrm{HungarianMatching}(C_{ij})\>\>\>\>$ s.t.
$m=1,2,\ldots,b^{\prime}$ and $\sigma(m)\in\\{1,2,\ldots,b\\}$

6:$\tilde{X}^{(i)}_{\textrm{mix}}\leftarrow$ Mix-
token($\tilde{X}^{(i)},X^{(\sigma(i))};\rho)$ s.t. $i=1,2,\ldots,b^{\prime}$
$\triangleright$ Eq. (10)

7:$\tilde{Y}^{(i)}_{\textrm{mix}}\leftarrow$
Relabel($\tilde{Y}^{(i)},Y^{(\sigma(i))})$ $\triangleright$ Eq. (11)

8:restore $\tilde{X}^{(i)}_{\textrm{mix}},\tilde{Y}^{(i)}_{\textrm{mix}}$ to
$X$, $Y$ s.t. $i=1,2,\ldots,b^{\prime}$

9:return  $X$, $Y$

By incorporating sample difficulty assessment (section 3.1) into (5), our
final objective is written as

$\displaystyle\max_{\sigma\in\mathcal{M}}$ $\displaystyle\max_{r_{i}\in{\rm
I\\!B^{n}}}\>\sum_{i=1}^{b}\mathbb{1}_{\\{\mathcal{F}(X^{(i)},Y^{(i)})<\tau\\}}(P_{i,\sigma(i)}-\rho)^{\top}r_{i}$
(8) $\displaystyle\textrm{s.t.}\>\>\mathbf{1}^{\top}r_{i}\leq
n\times\mathbb{1}_{\\{\mathcal{F}(X^{(i)},Y^{(i)})<\tau\\}},$

where $\mathcal{F}$ is the difficulty measurer, and $\tau$ is the difficulty
threshold. $\mathbb{1}$ denotes the indicator function, while $\mathbf{1}$ is
a one vector. If the sample is hard, i.e.,
$\mathcal{F}(X^{(i)},Y^{(i)})\geq\tau$, the constraint becomes
$\mathbf{1}^{\top}r_{i}\leq 0$, enforcing all elements of $r_{i}$ to be 0. In
that case, none of the tokens will be mixed. On the other hand, if the sample
is easy, i.e., $\mathcal{F}(X^{(i)},Y^{(i)})<\tau$, the constraint is
$\mathbf{1}^{\top}r_{i}\leq n$ which is satisfied by any random $r_{i}$
vector. Then, the objective becomes identical to (5).

###  3.3 Token-level Mixup

Now, the remaining question is how the tokens of $\tilde{X}^{(i)}$ and
$X^{(\sigma^{*}(i))}$ are going to be mixed. As defined in objective (5), we
use a hard replacement of tokens with respect to decision vector $r_{i}$.
Then, mixing is performed by replacing token index $t$ from the original
instance $i$, with token $t$ from the paired instance $j=\sigma^{*}(i)$, so
that $P_{i,j}^{(t)}>\rho$ is satisfied. In other words, only the token indices
that increase saliency level greater than $\rho$ are mixed. This is
accomplished by defining an appropriate binary mask $M_{t}\in\rm
I\\!B^{n}\;(t=1,2,\ldots,n)$ containing

$M_{t}=\begin{cases}\hskip
5.69054pt0&,\>\>S_{t}^{(j)}-\tilde{S}_{t}^{(i)}>\rho\\\ \hskip
5.69054pt1&,\>\>\textrm{otherwise}\end{cases}$ (9)

where $\tilde{S}_{t}^{(i)}$ is the saliency score of $\tilde{X}^{(i)}_{t}$.
Then, the mask is employed as

$\tilde{X}^{(i)}_{\textrm{mix}}=M\odot\tilde{X}^{(i)}+(1-M)\odot X^{(j)},$
(10)

where operator $\odot$ denotes element-wise multiplication and
$\tilde{X}^{(i)}$, $X^{(j)}$ refer to the paired instances. We also reassign
labels based on the replaced tokens’ saliency score. The new mixed label is
computed as

$\tilde{Y}_{\textrm{mix}}^{(i)}=\frac{\sum_{t=1}^{n}M_{t}\cdot\tilde{S}_{t}^{(i)}}{\sum_{t=1}^{n}M_{t}\cdot\tilde{S}_{t}^{(i)}+(1-M_{t})\cdot
S_{t}^{(j)}}\cdot\tilde{Y}^{(i)}+\frac{\sum_{t=1}^{n}(1-M_{t})\cdot
S_{t}^{(j)}}{\sum_{t=1}^{n}M_{t}\cdot\tilde{S}_{t}^{(i)}+(1-M_{t})\cdot
S_{t}^{(j)}}\cdot Y^{(j)}.$ (11)

Consequently, relatively unimportant tokens will be replaced with salient
tokens from another instance, and its label will be adjusted with respect to
the change in overall saliency level. See Figure 1 for visualization, and
Algorithm 1 for overall computation process. More detailed pseudocode is
provided in the supplement.

###  3.4 Vertical TokenMixup

Figure 2: Vertical TokenMixup. By applying VTM, the $\kappa$ most salient
tokens from each previous layer are brought up to form the key and value
tokens. By utilizing a cross-attention mechanism, the dimension is preserved
while being able to aggregate multi-scale information for token augmentation.
The figure demonstrates the case when VTM is applied to layer 3.

Mixing multiple instances is not the only way to mixup. We also present a
simple variant of TokenMixup, which provides token-level augmentation within a
single sample by utilizing the tokens from previous layers. Similar to (5), we
can define a new objective for Vertical TokenMixup as

$\max_{l\in\mathcal{L}}\>\max_{r_{i}\in\rm
I\\!B^{n}}\>\displaystyle\sum_{i=1}^{b}(P_{i,l(i)}-\rho)^{\top}\cdot r_{i},$
(12)

where $\mathcal{L}$ is the set of previous layer indices, and $l(i)$ returns
an arbitrary layer index for the $i^{th}$ instance. This objective can be
optimized by a similar scheme as in Algorithm 1. However, considering that the
tokens from each layer with an identical index tend to contain similar
information, saliency difference matrix $P_{i,l(i)}=S_{l(i)}-S_{i}$ (as in
(4)) will most likely be constant across indices. This may lead to meaningless
mixup without much increase in saliency level. Therefore, we take a different
approach to mix tokens vertically.

The simplest method would be to concatenate all tokens from previous layers
and apply self-attention. But due to the quadratic complexity of self-
attention, naive concatenation may induce excessive computation. To reduce
overhead, we selectively pool the $\kappa$ most salient tokens from each
previous layer and concatenate them to the original token set. Also, to
preserve input dimension, we adopt the cross-attention mechanism. If $X_{1}$
is the original token set and $X_{2}$ is the concatenated token set, vertical
cross-attention is expressed as

$Z=\textrm{Softmax}(X_{1}W_{q}(X_{2}W_{k})^{\top})\hskip 2.84526ptX_{2}W_{v},$
(13)

where $W_{q},W_{k},W_{v}$ denote the projection layers for query, key, and
value, respectively. See Figure 2 for visualization, and the supplement for
pseudocode.

###  3.5 Discussions

#### TokenMixup as Curriculum Learning.

A general framework for curriculum learning consists of two main components:
Difficulty Measurer and Training Scheduler [32]. TokenMixup satisfies these
conditions, as ScoreNet measures the difficulty of each input, and augmenting
the tokens based on input difficulty naturally schedules training. In the
early training phase where the model parameter is not optimized, most
instances will be evaluated difficult and no TokenMixup is performed. As the
model converges, on the other hand, many samples will become easy, triggering
augmentation to make the task more challenging overall. Also, TokenMixup is
applied discriminatively based on individual sample difficulty which enables
more intricate curriculum scheduling. See section 5.3 for empirical
demonstration. These properties will be especially useful when other mixup
methods [1, 2, 3, 4] are used in parallel. For instance, if input mixup is
already applied, further augmentation may be unnecessary. In such a case,
ScoreNet will regard the sample sufficiently difficult, and no further mixup
will be performed.

##  4 Experiments

###  4.1 Preliminaries

#### Baseline.

In experiments on CIFAR [33], we used Compact Convolution Transformer (CCT)
[22] as our baseline. For ImageNet-1k [34] experiments, we used the vanilla
ViT-B/16 [15] as baseline. Then, our augmentation methods are evaluated on
three representative image classification datasets: CIFAR-10, CIFAR-100, and
ImageNet-1K.

#### Experimental setup.

We evaluate TokenMixup and its variant, Vertical TokenMixup (VTM), on CIFAR
and ImageNet with different settings. To avoid confusion, we denote the
original TokenMixup as Horizontal TokenMixup (HTM) in the following sections.
For CIFAR experiments, we adopt the 1500-epoch version of CCT. We modified the
learning rate scheduler and positional embedding type to achieve better
performance than original papers (denoted $*$ in Table 1). Other experiment
settings follow [22], and all experiments on CIFAR datasets were conducted on
a single RTX A6000 GPU. For ImageNet-1K experiments, we used the ViT-B/16
model pre-trained and fine-tuned on ImageNet-21k and ImageNet-1k, where we
took advantage of the officially released pre-trained weights. Experiments for
Horizontal TokenMixup were conducted on a single NVIDIA A100 GPU, and 4 RTX
3090 GPUs were used in parallel for Vertical TokenMixup. Other experiment
settings are reported in the supplement.

###  4.2 Experimental Results

Table 1: Experimental results on CIFAR. Experiment results with TokenMixup on
CCT are compared with state-of-the-art models that do not use pretrained
models for initialization. Top-1 validation accuracy is reported for each
model, and baseline models denoted with * reports retrained performance. Also,
the number in the parenthesis indicates the number of epochs used for
training, and the parameter numbers are retrieved from the CIFAR-100 models.
Best performances are highlighted in yellow.

Models # Params MACs CIFAR-10 CIFAR-100 Convolutional Network based ResNet18
11.18 M 0.04 G 90.27 63.41 ResNet50 23.53 M 0.08 G 90.60 61.68 ResNet1001-v2
[35] 10.33 M 1.55 G 95.08 77.29 MobileNetV2/2.0 [36] 8.72 M 0.02 G 91.02 67.44
WRN-28-10 [37] 36.5 M - 96.00 80.75 WRN-40-4 [37] 8.9 M - 95.47 78.82
ResNeXt-29-8×64d [38] 34.4 M - 96.35 82.23 ResNeXt-29-16×64d [38] 68.1 M -
96.42 82.69 Transformer based ViT-Lite-6/4 [22] 3.19 M 0.22 G 93.98 73.33 ViT-
Lite-7/4 [22] 3.72 M 0.26 G 93.57 73.94 NesT-T [39] 17.0 M - 96.04 78.69
NesT-B [39] 68.0 M - 97.20 82.56 CVT-6/4 [22] 3.19 M 0.22 G 93.60 74.23
CCT-7/3x1(1500) [22] 3.76 M 0.95 G 97.48 82.72 CCT-7/3x1(1500) ∗ 3.78 M 0.95 G
97.48 82.87 CCT-7/3x1(1500) + Horizontal TM (ours) 3.81 M 0.95 G 97.57 83.56
CCT-7/3x1(1500) + Vertical TM (ours) 3.78 M 0.95 G 97.78 83.54 CCT-7/3x1(1500)
+ HTM + VTM (ours) 3.81 M 0.95 G 97.75 83.57

Table 2:  Experimental results on ImageNet. Experiment results with TokenMixup
on ViT are compared. We report Top-1 and Top-5 validation accuracy on
ImageNet-1k. For the baseline model, we used ViT-B/16 ($224\times 224$). †
denotes fine-tuned performance officially reported by the authors of ViT.

Models # Params MACs Top-1 Acc. Top-5 Acc. Convolutional Network based
ResNet50 [40] 25.5 M 4.3 G 76.20 - ResNet152 [40] 60.19 M 11.58 G 78.57 -
DenseNet-264 [41] - - 79.20 94.71 WRN-50-2-bottleneck [37] 68.9 M - 78.10
93.97 ReGNetY-4G [42] 21 M 2.0 G 80.0 - ResNeXt-50-2×40d [38] - - 77.00 -
ResNeXt-101-64x4d [38] - - 79.60 94.70 Transformer based ViT-B/16-224† [15]
86.6 M 16.9 G 81.2 - ViT-B/16-224 + Horizontal TM (ours) 87.3 M 16.9 G 82.37
96.29 ViT-B/16-224 + Vertical TM (ours) 86.6 M 16.9 G 82.30 96.33 ViT-B/16-224
+ HTM + VTM (ours) 86.6 M 16.9 G 82.32 96.27

First, we demonstrate the image classification performance on CIFAR-10 and
CIFAR-100 when TokenMixup is applied to CCT. From Table 1, we can see that
applying Horizontal TokenMixup and Vertical TokenMixup simultaneously improves
its strict CIFAR-100 baseline, by a significant margin of 0.70%. HTM and VTM
each surpasses the baseline performance by a margin of 0.69% and 0.67% as
well. Also, we outperform the previous state-of-the-art model (among from-
scratch models) by a huge margin of 0.85%. In the case of CIFAR-10, we
achieved higher accuracy of 97.57% and 97.78% with HTM and VTM, respectively.
Considering the highly saturated performance on this dataset, a 0.30%
improvement over the baseline is not trivial.
Performance on ImageNet-1K is evaluated on the vanilla ViT [15] model. The
officially reported Top-1 accuracy is 81.2%, which is achieved by using
weights pre-trained on ImageNet-21K and fine-tuning on ImageNet-1K. By fine-
tuning the model further by employing our augmentation methods, we achieve
Top-1 accuracy of 82.37% with Horizontal TokenMixup, 82.30% with Vertical
TokenMixup, and 82.32% by using both. Also note that HTM marginally increases
parameter size by only 0.7M, while VTM does not require any additional
parameters. See the supplement for further experiments on other baselines.

##  5 Analysis

Here, we analyze TokenMixup to answer the following research questions: Q1. Is
the Attention an appropriate saliency approximation? Q2. What are the effects
of each component in TokenMixup? Q3. What properties does TokenMixup have?

Figure 3: Qualitative comparison of saliency detectors. The saliency map
derived for each method is provided. We can observe that the methods render
similar outputs. Note, the values are minmax scaled to [0,1]. Table 3:
Saliency detector comparison. Our method (HTM) achieves high accuracy while
being $\times 15$ as fast as the gradient method. Also, we observe only
$\times 3$ latency compared to the random baseline.

Saliency type Gradient-based Attention-based Random baseline Avg. Latency (ms)
236 ($\times$14.8) 16 ($\times$1.0) 5 ($\times$0.3) CIFAR-100 Acc. 83.28 83.56
83.05

###  5.1 Soundness of Saliency Estimation

In TokenMixup, we estimate the saliency using the attention map from the
subsequent layer. The soundness of this approximation would be one of the main
questions regarding saliency detection. So, we provide relevant analyses on
the appropriateness of our saliency approximation method to answer the
research question Q1.

#### Qualitative examples.

Here, we provide a qualitative analysis on the performance of each saliency
detector. In Figure 3, we can observe that there is no significant difference
between the full attention rollout and the 1-step approximation, i.e., the
subsequent attention map. Slight discrepancy in the sharpness of the
distribution can be observed, but it does not greatly affect our method that
utilizes saliency threshold $\rho$. Moreover, attention-based methods seem to
be more accurate than the gradient-based detector in certain cases. See the
supplement for additional qualitative examples, and also the quantified
comparison on the sharpness measures of the two methods.

#### Efficiency comparison.

Table 4: Ablation study on CIFAR-100. Performance on CIFAR-100 is evaluated by
ablating each component one-by-one. We observe that optimal assignment and
label reassignment plays a significant role.

Ablated Component CIFAR-100 Accuracy CCT-7/3x1 + HTM 83.56 (- 0.00) (–) Sample
difficulty assessment (ScoreNet) 82.85 (\- 0.71) (–) Optimal assignment 82.61
(\- 0.95) (–) Token-level mixup 83.05 (\- 0.51) (–) Saliency-based label
reassignment 82.62 (\- 0.94)

We compare the efficiency of three saliency detectors: gradient-based,
attention-based, and random baseline. For the gradient-based detector, we
follow the method used in Co-Mixup [4], and random baseline refers to a random
selection of tokens. In Table 3, we report the average latency of saliency
detection per iteration in our CIFAR-100 experiment setting with a batch size
of 128. Compared to the gradient-based method, attention-based saliency
detection with 1-step rollout achieves approximately $\times 15$ speed-up
while demonstrating even better performance.

###  5.2 Component Analysis

In this section, we analyze the effect of each component in Horizontal
TokenMixup and Vertical TokenMixup so as to answer Q2. In Table 4, we provide
ablation studies on components of HTM: Sample difficulty assessment
(ScoreNet), Optimal assignment, Token-level mixup, and Saliency-based label
reassignment. By ablating ScoreNet, all instances are mixed regardless of
sample difficulty; note, ScoreNet still receives supervision signals. By
getting rid of optimal assignment, mixup pairs are not assigned via Hungarian
matching; pairs are randomly selected. Token-level mixup and label
reassignment is relevant to the method by which tokens are selected. Then, by
ablating mixup, we randomly select tokens to replace, while label reassignment
ablation derives new labels based on the number of tokens that have been
replaced. We also conducted sensitivity tests on key hyperparameters $\tau$,
$\rho$, and $\kappa$, but tables and discussions were moved to the supplement
due to spatial constraints.

###  5.3 What TokenMixup Learns

Table 5: Mixup combinations. CCT performance on CIFAR-100 is compared for
mixup method combinations. ∗ denotes values reported in the original paper.
Best performances among comparable settings are highlighted.

Mixup Combination CIFAR-100 Accuracy CCT-7/3x1 76.80 CCT-7/3x1 + Manifold
Mixup 76.72 CCT-7/3x1 + Horizontal TM 77.34 CCT-7/3x1 + Vertical TM 78.18
CCT-7/3x1 + Input Mixup & Cutmix ∗ 82.87 CCT-7/3x1 + Input Mixup & Cutmix +
Manifold Mixup 80.08 CCT-7/3x1 + Input Mixup & Cutmix + Horizontal TM 83.56
CCT-7/3x1 + Input Mixup & Cutmix + Vertical TM 83.54 CCT-7/3x1 + Input Mixup &
Cutmix + HTM \+ VTM 83.57

Here, we answer our research question Q3 by analyzing what our mixup methods
are learning. We demonstrate a curriculum learning perspective in terms of
HTM, and compare our method with a direct mixup alternative, manifold mixup
[6].

Figure 4: Trend in the number of mixed instances.

#### A Curriculum Learning perspective.

In Figure 4 we plot the trend in the number of samples per batch TokenMixup is
activated based on ScoreNet. In the first few epochs (regime (A)), the number
of mixed samples is relatively low, as there is less need for augmentation in
the early phase of learning. But as the learnable parameters warm up, the
number increases abruptly. In regime (B), the average number of mixed tokens
slowly increases as model performance steadily improves. In the last learning
phase (regime (C)), mixup frequency converges.

#### Comparison with Manifold Mixup.

TokenMixup is applied to the intermediate tokens, similarly to manifold mixup
[6] which is an input mixup [1] method applied to intermediate features. By
using manifold mixup, transformer tokens are mixed via linear interpolation,
instead of the hard replacement of tokens as in TokenMixup. In Table 5, our
methods are compared with manifold mixup, along with various combinations of
original mixup methods (i.e., input mixup and cutmix). The best setting in the
CCT paper adopted input mixup and cutmix. By applying HTM and VTM on top of
that setting, we achieved state-of-the-art performance. On the other hand, by
applying manifold mixup, we observed severe deterioration in accuracy. We
conjecture that soft mixing of transformer tokens leads to over-smoothing of
features, which may be inadequate for self-attention layers. See the
supplement for visual aid.

##  6 Conclusion

We proposed TokenMixup for transformers, which adopts the attention map in
place of computationally heavy gradient-based saliency detectors. By adopting
the attention as the saliency map, we achieve $\times 15$ speed-up compared to
the gradient-based method with even better performance. We also introduce a
novel perspective of curriculum learning to mixup, which enables adaptive
augmentation based on sample difficulty and training schedule. To add
diversity, Vertical TokenMixup is introduced, which mixes tokens from
different layers for multi-scale feature augmentation within a single sample.
With TokenMixup, we achieve state-of-the-art performance on CIFAR-100, and
improve baselines by significant margins on CIFAR-10 and ImageNet-1K.

## Acknowledgments and Disclosure of Funding

This work was partly supported by ICT Creative Consilience program
(IITP-2022-2020-0-01819) supervised by the IITP; the National Supercomputing
Center with supercomputing resources including technical support
(KSC-2022-CRE-0100) and Kakao Brain Corporation.

## References

  * [1] Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  In ICLR, 2018.
  * [2] Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo.  Cutmix: Regularization strategy to train strong classifiers with localizable features.  In ICCV, 2019.
  * [3] Jang-Hyun Kim, Wonho Choo, and Hyun Oh Song.  Puzzle mix: Exploiting saliency and local statistics for optimal mixup.  In ICML, 2020.
  * [4] Jang-Hyun Kim, Wonho Choo, Hosan Jeong, and Hyun Oh Song.  Co-mixup: Saliency guided joint mixup with supermodular diversity.  arXiv preprint arXiv:2102.03065, 2021.
  * [5] AFM Uddin, Mst Monira, Wheemyung Shin, TaeChoong Chung, Sung-Ho Bae, et al.  Saliencymix: A saliency guided data augmentation strategy for better regularization.  arXiv preprint arXiv:2006.01791, 2020.
  * [6] Vikas Verma, Alex Lamb, Christopher Beckham, Amir Najafi, Ioannis Mitliagkas, David Lopez-Paz, and Yoshua Bengio.  Manifold mixup: Better representations by interpolating hidden states.  In ICML, 2019.
  * [7] Shashanka Venkataramanan, Yannis Avrithis, Ewa Kijak, and Laurent Amsaleg.  Alignmix: Improving representation by interpolating aligned features.  In ICLR, 2022.
  * [8] Minui Hong, Jinwoo Choi, and Gunhee Kim.  Stylemix: Separating content and style for enhanced data augmentation.  In CVPR, 2021.
  * [9] Ali Dabouei, Sobhan Soleymani, Fariborz Taherkhani, and Nasser M Nasrabadi.  Supermix: Supervising the mixing data augmentation.  In CVPR, 2021.
  * [10] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In NeurIPS, 2017.
  * [11] Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko.  End-to-end object detection with transformers.  In ECCV, 2020.
  * [12] Xizhou Zhu, Weijie Su, Lewei Lu, Bin Li, Xiaogang Wang, and Jifeng Dai.  Deformable detr: Deformable transformers for end-to-end object detection.  In ICLR, 2020.
  * [13] Bumsoo Kim, Junhyun Lee, Jaewoo Kang, Eun-Sol Kim, and Hyunwoo J Kim.  Hotr: End-to-end human-object interaction detection with transformers.  In CVPR, 2021.
  * [14] Jihwan Park, SeungJun Lee, Hwan Heo, Hyeong Kyu Choi, and Hyunwoo J. Kim.  Consistency learning via decoding path augmentation for transformers in human object interaction detection.  In CVPR, 2022.
  * [15] 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.  In ICLR, 2021.
  * [16] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In ICCV, 2021.
  * [17] Wenhai Wang, Enze Xie, Xiang Li, Deng-Ping Fan, Kaitao Song, Ding Liang, Tong Lu, Ping Luo, and Ling Shao.  Pyramid vision transformer: A versatile backbone for dense prediction without convolutions.  In ICCV, 2021.
  * [18] Kun Yuan, Shaopeng Guo, Ziwei Liu, Aojun Zhou, Fengwei Yu, and Wei Wu.  Incorporating convolution designs into visual transformers.  In ICCV, 2021.
  * [19] Byeongho Heo, Sangdoo Yun, Dongyoon Han, Sanghyuk Chun, Junsuk Choe, and Seong Joon Oh.  Rethinking spatial dimensions of vision transformers.  In ICCV, 2021.
  * [20] Kai Han, An Xiao, Enhua Wu, Jianyuan Guo, Chunjing Xu, and Yunhe Wang.  Transformer in transformer.  NeurIPS, 2021.
  * [21] Stéphane d’Ascoli, Hugo Touvron, Matthew L Leavitt, Ari S Morcos, Giulio Biroli, and Levent Sagun.  Convit: Improving vision transformers with soft convolutional inductive biases.  In ICML, 2021.
  * [22] Ali Hassani, Steven Walton, Nikhil Shah, Abulikemu Abuduweili, Jiachen Li, and Humphrey Shi.  Escaping the big data paradigm with compact transformers.  arXiv preprint arXiv:2104.05704, 2021.
  * [23] Haiping Wu, Bin Xiao, Noel Codella, Mengchen Liu, Xiyang Dai, Lu Yuan, and Lei Zhang.  Cvt: Introducing convolutions to vision transformers.  In ICCV, 2021.
  * [24] Weijian Xu, Yifan Xu, Tyler Chang, and Zhuowen Tu.  Co-scale conv-attentional image transformers.  In ICCV, 2021.
  * [25] Xiangxiang Chu, Zhi Tian, Bo Zhang, Xinlong Wang, Xiaolin Wei, Huaxia Xia, and Chunhua Shen.  Conditional positional encodings for vision transformers.  arXiv preprint arXiv:2102.10882, 2021.
  * [26] Xiangxiang Chu, Zhi Tian, Yuqing Wang, Bo Zhang, Haibing Ren, Xiaolin Wei, Huaxia Xia, and Chunhua Shen.  Twins: Revisiting the design of spatial attention in vision transformers.  NeurIPS, 2021.
  * [27] Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsieh.  Dynamicvit: Efficient vision transformers with dynamic token sparsification.  NeurIPS, 2021.
  * [28] Byungseok Roh, JaeWoong Shin, Wuhyun Shin, and Saehoon Kim.  Sparse detr: Efficient end-to-end object detection with learnable sparsity.  In ICLR, 2021.
  * [29] Tao Wang, Li Yuan, Yunpeng Chen, Jiashi Feng, and Shuicheng Yan.  Pnp-detr: towards efficient visual analysis with transformers.  In ICCV, 2021.
  * [30] Samira Abnar and Willem Zuidema.  Quantifying attention flow in transformers.  arXiv preprint arXiv:2005.00928, 2020.
  * [31] James Munkres.  Algorithms for the assignment and transportation problems.  Journal of the society for industrial and applied mathematics, 5(1):32–38, 1957.
  * [32] Xin Wang, Yudong Chen, and Wenwu Zhu.  A survey on curriculum learning.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2021\.
  * [33] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images, 2009.
  * [34] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei.  Imagenet large scale visual recognition challenge.  IJCV, 2015.
  * [35] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Identity mappings in deep residual networks.  In ECCV, 2016.
  * [36] Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen.  Mobilenetv2: Inverted residuals and linear bottlenecks.  In CVPR, 2018.
  * [37] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  arXiv preprint arXiv:1605.07146, 2016.
  * [38] Saining Xie, Ross Girshick, Piotr Dollár, Zhuowen Tu, and Kaiming He.  Aggregated residual transformations for deep neural networks.  In CVPR, 2017.
  * [39] Zizhao Zhang, Han Zhang, Long Zhao, Ting Chen, Sercan Arik, and Tomas Pfister.  Nested hierarchical transformer: Towards accurate, data-efficient and interpretable visual understanding.  In AAAI, 2022.
  * [40] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In CVPR, 2016.
  * [41] Gao Huang, Zhuang Liu, Laurens Van Der Maaten, and Kilian Q Weinberger.  Densely connected convolutional networks.  In CVPR, 2017.
  * [42] Ilija Radosavovic, Raj Prateek Kosaraju, Ross Girshick, Kaiming He, and Piotr Dollár.  Designing network design spaces.  In CVPR, 2020.

## Checklist

  1. 1.

For all authors…

    1. (a)

Do the main claims made in the abstract and introduction accurately reflect
the paper’s contributions and scope? [Yes] See Section 3.

    2. (b)

Did you describe the limitations of your work? [Yes] See the supplement

    3. (c)

Did you discuss any potential negative societal impacts of your work? [Yes]
See the supplement

    4. (d)

Have you read the ethics review guidelines and ensured that your paper
conforms to them? [Yes]

  2. 2.

If you are including theoretical results…

    1. (a)

Did you state the full set of assumptions of all theoretical results? [N/A]

    2. (b)

Did you include complete proofs of all theoretical results? [N/A]

  3. 3.

If you ran experiments…

    1. (a)

Did you include the code, data, and instructions needed to reproduce the main
experimental results (either in the supplemental material or as a URL)? [Yes]

    2. (b)

Did you specify all the training details (e.g., data splits, hyperparameters,
how they were chosen)? [Yes] See Section 4 and the supplement for technical
details

    3. (c)

Did you report error bars (e.g., with respect to the random seed after running
experiments multiple times)? [N/A] Our experiment settings do not conduct each
experiment multiple times.

    4. (d)

Did you include the total amount of compute and the type of resources used
(e.g., type of GPUs, internal cluster, or cloud provider)? [Yes] See Section 4
and the supplement for technical details.

  4. 4.

If you are using existing assets (e.g., code, data, models) or
curating/releasing new assets…

    1. (a)

If your work uses existing assets, did you cite the creators? [Yes] See the
supplement or Section 4 baseline.

    2. (b)

Did you mention the license of the assets? [Yes] See the supplement

    3. (c)

Did you include any new assets either in the supplemental material or as a
URL? [Yes] See the supplement

    4. (d)

Did you discuss whether and how consent was obtained from people whose data
you’re using/curating? [Yes] See the supplement

    5. (e)

Did you discuss whether the data you are using/curating contains personally
identifiable information or offensive content? [Yes] See the supplement

  5. 5.

If you used crowdsourcing or conducted research with human subjects…

    1. (a)

Did you include the full text of instructions given to participants and
screenshots, if applicable? [N/A]

    2. (b)

Did you describe any potential participant risks, with links to Institutional
Review Board (IRB) approvals, if applicable? [N/A]

    3. (c)

Did you include the estimated hourly wage paid to participants and the total
amount spent on participant compensation? [N/A]