# Toward Learning Robust and Invariant Representations with Alignment
Regularization and Data Augmentation

Haohan Wang  Carnegie Mellon UniversityPittsburghPAUSA<EMAIL_ADDRESS>,
Zeyi Huang  University of Wisconsin-MadisonMadisonWIUSA<EMAIL_ADDRESS>,
Xindi Wu  Carnegie Mellon UniversityPittsburghPAUSA<EMAIL_ADDRESS>and
Eric P. Xing  Carnegie Mellon UniversityPittsburghPAUSA<EMAIL_ADDRESS>
(2022)

###### Abstract.

Data augmentation has been proven to be an effective technique for developing
machine learning models that are robust to known classes of distributional
shifts (e.g., rotations of images), and alignment regularization is a
technique often used together with data augmentation to further help the model
learn representations invariant to the shifts used to augment the data. In
this paper, motivated by a proliferation of options of alignment
regularizations, we seek to evaluate the performances of several popular
design choices along the dimensions of robustness and invariance, for which we
introduce a new test procedure. Our synthetic experiment results speak to the
benefits of squared $\ell_{2}$ norm regularization. Further, we also formally
analyze the behavior of alignment regularization to complement our empirical
study under assumptions we consider realistic. Finally, we test this simple
technique we identify (worst-case data augmentation with squared $\ell_{2}$
norm alignment regularization) and show that the benefits of this method
outrun those of the specially designed methods. We also release a software
package in both TensorFlow and PyTorch for users to use the method with a
couple of lines111 https://github.com/jyanln/AlignReg.

machine learning, data augmentation, robustness, trustworthy

††journalyear: 2022††copyright: rightsretained††conference: Proceedings of the
28th ACM SIGKDD Conference on Knowledge Discovery and Data Mining; August
14–18, 2022; Washington, DC, USA††booktitle: Proceedings of the 28th ACM
SIGKDD Conference on Knowledge Discovery and Data Mining (KDD ’22), August
14–18, 2022, Washington, DC, USA††doi: 10.1145/3534678.3539438††isbn:
978-1-4503-9385-0/22/08††ccs: Computing methodologies Regularization††ccs:
Computing methodologies Computer vision††ccs: Computing methodologies
Supervised learning

##  1\. Introduction

Data augmentation, i.e., to increase the dataset size through generating new
samples by transforming the existing samples with some predefined functions,
is probably one of the most often used techniques to improve a machine
learning model’s performance. It has helped machine learning models achieve
high prediction accuracy over various benchmarks (e.g., Lim et al., 2019; Ho
et al., 2019; Zhong et al., 2020; Buslaev et al., 2020; Ghiasi et al., 2020;
Kostrikov et al., 2020).

In addition to improving prediction accuracy, the community has also leveraged
data augmentation to help the models learn more robust representations that
can generalize to the datasets distributed differently (e.g., Zheng et al.,
2016; Sun et al., 2018; Huang et al., 2020; Min et al., 2020). To improve
robustness, the community usually designed the transformation functions used
to augment the data in correspondence to the transformations we see in the
real world (Hernandez-Garcia, 2020), such as the changes of image texture or
contrast. Thus, models trained with these augmented data are more likely to be
invariant to these designed transformations, such as the texture or contrast
variations of the input images.

To further help a model learn representations invariant to the
transformations, we can regularize the model so that the distance between
representations learned by the model from a pair of data (the original one and
the transformed counterpart) will be small. This regularization has been used
extensively recently to help models learn more robust and invariant
representations (e.g., Liang et al., 2018; Kannan et al., 2018; Zhang et al.,
2019b; Hendrycks et al., 2020). Motivated by this popularity, this paper
mainly studies the behaviors of this regularization, which we refer to as
_alignment regularization_ (AR). In particular, we seek to answer the
question: _how should we use alignment regularization to take advantage of the
augmented data to the fullest extent to learn robust and invariant models?_

To answer this, we first conduct a range of experiments over image
classification benchmarks to evaluate how popular variants of AR contribute to
learning robust and invariant models. We test for accuracy, robustness, and
invariance, for which we propose a new test procedure. Our empirical study
favors the squared $\ell_{2}$ norm. Our contributions of this paper are as
follows.

  * •

With a new invariance test, we show that alignment regularization is important
to help the model learn representations invariant to the transformation
function, and squared $\ell_{2}$ norm is considered the favorable choice as
assessed by a variety of empirical evaluations (Section 4).

  * •

We formalize a generalization error bound for models trained with AR and
augmented data (Section 5).

  * •

We test the method we identified (squared $\ell_{2}$ norm as AR) in multiple
scenarios. We notice that this generic approach can compete with methods
specially designed for different scenarios, which we believe endorses its
empirical strength (Section 6).

##  2\. Related Work and Key Differences

Tracing back to the earliest convolutional neural networks, (Wang and Raj,
2017), we notice that even early models for the MNIST dataset have been
boosted by data augmentation (Abu-Mostafa, 1990; Simard et al., 1991; LeCun et
al., 1998). Later, the rapidly growing machine learning community has seen a
proliferate development of data augmentation techniques that have helped
models climb the state-of-the-art ladder (Shorten and Khoshgoftaar, 2019).
Among the augmentation techniques, the most relevant one to this paper is to
generate the samples (with constraint) that maximize the training loss along
with training (Fawzi et al., 2016).

While the above paragraph mainly discusses how to generate the augmented
samples, we mainly study how to train the models with augmented samples. For
example, instead of directly mixing augmented samples with the original
samples, one can consider regularizing the representations (or outputs) of
original samples and augmented samples to be close under a distance metric
(which we refer to as alignment regularization, AR). Many concrete ideas have
been explored in different contexts. For example, $\ell_{2}$ distance and
cosine similarities between internal representations in speech recognition
(Liang et al., 2018), squared $\ell_{2}$ distance between logits (Kannan et
al., 2018), or KL divergence between softmax outputs (Zhang et al., 2019b) in
adversarially robust vision models, Jensen–Shannon divergence (of three
distributions) between embeddings for texture invariant image classification
(Hendrycks et al., 2020). These are but a few highlights of the concrete and
successful implementations for different applications out of a vast collection
(e.g., (Wu et al., 2019; Guo et al., 2019; Zhang et al., 2019a; Shah et al.,
2019; Asai and Hajishirzi, 2020; Sajjadi et al., 2016; Zheng et al., 2016; Xie
et al., 2015)), and we can expect methods permuting these three elements
(distance metrics, representation or outputs, and applications) to be
discussed in the future. Further, given the popularity of GAN (Goodfellow,
2016) and domain adversarial neural network (Ganin et al., 2016), we can also
expect the distance metric generalizes to a specialized discriminator (i.e., a
classifier), which can be intuitively understood as a calculated (usually
maximized) distance measure, and one example here is the Wasserstein-1 metric
(Arjovsky et al., 2017; Gulrajani et al., 2017).

Key Differences: With this rich collection of regularizing choices, which one
method should we consider in general? More importantly, do we need the
regularization at all? These questions are important for multiple reasons,
especially since sometimes AR may worsen the results (Jeong et al., 2019). In
this paper, we first conduct an empirical study to show that AR (especially
squared $\ell_{2}$ norm) can help learn robust and invariant models, we then
also derive generalization error bounds to complement our empirical
conclusion.

There are also several previous discussions regarding the detailed
understandings of data augmentation (Yang et al., 2019; Chen et al., 2019;
Hernández-García and König, 2018; Rajput et al., 2019; Dao et al., 2019; Ghosh
and Thiery, 2021; Zhang et al., 2021), among which, (Yang et al., 2019) is
probably the most relevant as it also defends the usage of the AR. In addition
to what is reported in (Yang et al., 2019), our work also connects to
invariance and shows that another advantage of AR is to learn invariant
representations.

##  3\. Accuracy, Robustness, and Invariance

This section discusses the three major evaluation metrics we will use to test
AR. We will first recapitulate the background of accuracy and robustness. Then
we will introduce our definition of invariance and our proposed evaluation.

Notations $(\mathbf{X},\mathbf{Y})$ denotes the data, where
$\mathbf{X}\in\mathcal{R}^{n\times p}$ and $\mathbf{Y}\in\\{0,1\\}^{n\times
k}$ (one-hot vectors for $k$ classes). $(\mathbf{x},\mathbf{y})$ denotes a
sample. $f(,\theta)$ denotes the model, which takes in the data and outputs
the softmax (probabilities of the prediction) and $\theta$ denotes the
corresponding parameters. $g()$ completes the prediction (i.e., mapping
softmax to one-hot prediction). $a()$ denotes a function for data
augmentation, i.e., a transformation function. $a\in\mathcal{A}$, which is the
set of transformation functions of interest. $\mathbf{P}$ denotes the
distribution of $(\mathbf{x},\mathbf{y})$. For any sampled
$(\mathbf{x},\mathbf{y})$, we can have $(a(\mathbf{x}),\mathbf{y})$, and we
use $\mathbf{P}_{a}$ to denote the distribution of these transformed samples.
Further, we use $\mathbf{Q}_{a(\mathbf{x}),\widehat{\theta}}$ to denote the
distribution of $f(a(\mathbf{x});\widehat{\theta})$ for
$(\mathbf{x},\mathbf{y})\sim\mathbf{P}$. $D(\cdot,\cdot)$ is a distance
measure over two distributions. $r(\theta)$ denotes the risk of model
$\theta$. $\widehat{\cdot}$ denotes the estimation of the term $\cdot$, e.g.,
$\widehat{r}_{\mathbf{P}}(\widehat{\theta})$ denotes the empirical risk of the
estimated model.

###  3.1. Accuracy

The community studying the statistical property of the error bound usually
focuses on the expected risk defined as

(1) $\displaystyle
r_{\mathbf{P}}(\widehat{\theta})=\mathbb{E}_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}}\mathbb{I}[g(f(\mathbf{x};\widehat{\theta}))\neq\mathbf{y}],$

where $\mathbb{I}[\mathbf{c}]$ is a function that returns $1$ if the condition
$\mathbf{c}$ holds.

In practice, the error is evaluated by replacing $\mathbf{P}$ with a hold-out
test dataset, and the accuracy is $1-r_{\mathbf{P}}(\widehat{\theta})$.

###  3.2. Robustness

We define robustness as the worst-case expected risk when the test data is
allowed to be transformed by functions in $\mathcal{A}$, following (e.g.,
Szegedy et al., 2013; Goodfellow et al., 2014). Formally, we study the worst-
case error as

(2) $\displaystyle
r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})=\mathbb{E}_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}}\max_{a\sim\mathcal{A}}\mathbb{I}[g(f(a(\mathbf{x});\widehat{\theta}))\neq\mathbf{y}],$

where we use $r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})$ to denote the
robust error as it will depend on $\mathcal{A}$. In practice, the robust error
is also evaluated by replacing $\mathbf{P}$ with a hold-out dataset.

###  3.3. Invariance

Further, high robustness performances do not necessarily mean the model is
truly invariant to the transformation functions (Hernández-García et al.,
2019), and we continue to introduce a new test to evaluate the model’s
behavior in learning representations invariant to the transformations.

##### Invariance

If the model can learn a representation invariant to the transformation
functions, it will map the samples of different transformations to the same
representation. Intuitively, to measure how invariant a model is to the
transformations in $\mathcal{A}$, we can calculate the distances between each
pair of the two transformed samples when a sample is transformed with
functions in $\mathcal{A}$. Thus, we define the following term to measure
invariance:

(3) $\displaystyle
I_{\mathbf{P},\mathcal{A}}(\widehat{\theta})=\sup_{a_{1},a_{2}\in\mathcal{A}}D(\mathbf{Q}_{a_{1}(\mathbf{x}),\widehat{\theta}},\mathbf{Q}_{a_{2}(\mathbf{x}),\widehat{\theta}}),$

We suggest using Wasserstein metric as $D(\cdot,\cdot)$, considering its
favorable properties (e.g., see practical examples in Figure 1 of (Cuturi and
Doucet, 2014) or theoretical discussions in (Villani, 2008)).

In practice, we also need to replace $\mathbf{P}$ with a hold-out dataset so
that the evaluation can be performed. In addition, we notice that
$I_{\mathbf{P},\mathcal{A}}(\widehat{\theta})$, although intuitive, is not
convenient in practice because the evaluated values are not bounded. Thus, we
reformulate it into the following invariance test procedure, whose final score
will be bounded between 0 and 1 (the higher, the better). Therefore, the score
can be conveniently discussed together with accuracy and robust accuracy,
which are also bounded between 0 and 1.

##### Invariance test

Given a family of transformation functions used in data augmentation
$\mathcal{A}=\\{a_{1}(),a_{2}(),\dots,a_{t}()\\}$ of $t$ elements, and a
collection of samples (from the hold-out dataset) of the same label $i$,
denoted as $\mathbf{X}^{(i)}$, the evaluation procedure is as follows. We
first generate the transformed copies of $\mathbf{X}^{(i)}$ with
$\mathcal{A}$, resulting in
$\mathbf{X}^{(i)}_{a_{1}},\mathbf{X}^{(i)}_{a_{2}},\dots,\mathbf{X}^{(i)}_{a_{t}}$.
We combined these copies into a dataset, denoted as $\mathcal{X}^{(i)}$. For
every sample $\mathbf{x}$ in $\mathcal{X}^{(i)}$, we retrieve its $t$ nearest
neighbors of other samples in $\mathcal{X}^{(i)}$, and calculate the overlap
of the retrieved samples with the transformed copies of $\mathbf{x}$ by
$\mathcal{A}$, i.e.,
$\\{a_{1}(\mathbf{x}),a_{2}(\mathbf{x}),\dots,a_{t}(\mathbf{x})\\}$. The
calculated overlap score will be in $[0,1]$ in general, but since the identity
map is usually in $\mathcal{A}$, this score will usually be in $[1/t,1]$.

During the retrieval of nearest neighbors, we consider the distance function
of the two samples (namely $\mathbf{x}$ and $\mathbf{x}^{\prime}$) as

$d(f(\mathbf{x};\widehat{\theta}),f(\mathbf{x}^{\prime};\widehat{\theta}))$

, where $\widehat{\theta}$ is the model we are interested in examining. In the
empirical study later, we consider
$d(\mathbf{u},\mathbf{v})=\|\mathbf{u}-\mathbf{v}\|_{1}$, with $\mathbf{u}$
and $\mathbf{v}$ denoting two vectors. If we use other distance functions, the
reported values may differ, but we notice that the rank of the methods
compared in terms of this test barely changes.

Finally, we iterate through label $i$ and report the averaged score for all
the labels as the final score. A higher score indicates the model
$\widehat{\theta}$ is more invariant to the transformation functions in
$\mathcal{A}$.

This invariance test procedure is formally presented in Algorithm 1 below.

Result: $\widehat{I}(\widehat{\theta})$

Input: a family of transformation functions
$\mathcal{A}=\\{a_{1}(),a_{2}(),\dots,a_{t}()\\}$, a hold-out dataset
$(\mathbf{X},\mathbf{Y})$, the model of interest $\widehat{\theta}$, and a
distance metric $d()$;

for  _every label $i$ _ do

identify all the samples from $(\mathbf{X},\mathbf{Y})$ with label $i$, name
this set of samples $\mathbf{X}^{(i)}$;

for  _every $a()\in\mathcal{A}$_ do

generate $\mathbf{X}_{a}^{(i)}$ by applying $a()$ to every
$\mathbf{x}\in\mathbf{X}^{(i)}$;

end for

$\mathcal{X}^{(i)}=\mathbf{X}^{(i)}_{a_{1}}\cup\mathbf{X}^{(i)}_{a_{2}}\cup,\dots,\cup\mathbf{X}^{(i)}_{a_{t}}$;

for  _every $\mathbf{x}\in\mathbf{X}^{(i)}$_ do

generate set
$T=\\{a_{1}(\mathbf{x}),a_{2}(\mathbf{x}),\dots,a_{t}(\mathbf{x})\\}$;

for  _every $\mathbf{x}^{\prime}\in\mathcal{X}^{(i)}$_ do

calculate the distance between $\mathbf{x}$ and $\mathbf{x}^{\prime}$ with
$d(f(\mathbf{x};\widehat{\theta}),f(\mathbf{x}^{\prime};\widehat{\theta}))$;

end for

retrieve the $t$ nearest neighbors of $\mathbf{x}$ out of $\mathcal{X}^{(i)}$,
and name this set of samples $K$;

calculate the score for $\mathbf{x}$ with $|T\cap K|/|T|$, where $|T|$ denotes
the cardinality of the set $T$;

end for

calculate the score for label $i$ as the average score across all
$\mathbf{x}\in\mathbf{X}^{(i)}$;

end for

calculate the final score $\widehat{I}(\widehat{\theta})$ as the average score
across all the labels;

Algorithm 1 Invariance test

##  4\. Empirical Study

In this section, we conduct experiments to study the relationship between
robustness and invariance, as well as how training with AR can help improve
the invariance score. In short, our empirical study in this section will lead
us to the following three major conclusions:

  * •

High robust accuracy does not necessarily mean a high invariance score and
vice versa.

  * •

AR can help improve the invariance score.

  * •

Squared $\ell_{2}$ norm over logits is considered the empirically most
favorable AR option for learning robust and invariant representations.

Table 1. Test results on MNIST dataset for different ARs over three evaluation metrics and three distribution shifts. B denotes Baseline, i.e., the model does not use any data augmentation; V denotes vanilla augmentation, i.e., the model uses data augmentation but not AR; L denotes $\ell_{1}$ norm; S denotes squared $\ell_{2}$ norm; C denotes cosine similarity; K denotes KL divergence; W denotes Wasserstein-1 metric; D denotes GAN discriminator.  |  | B | V | L | S | C | K | W | D
---|---|---|---|---|---|---|---|---|---
Texture | Accuracy | $99.2_{\pm 0.0}$ | $99.2_{\pm 0.0}$ | $99.0_{\pm 0.1}$ | $99.1_{\pm 0.0}$ | $99.4_{\pm 0.0}$ | $68.7_{\pm 41}$ | $98.7_{\pm 0.1}$ | $99.1_{\pm 0.1}$
Robustness | $98.3_{\pm 0.3}$ | $99.0_{\pm 0.0}$ | $99.0_{\pm 0.0}$ | $99.0_{\pm 0.0}$ | $99.1_{\pm 0.0}$ | $68.7_{\pm 41}$ | $98.4_{\pm 0.1}$ | $98.8_{\pm 0.1}$
Invariance | $92.4_{\pm 0.0}$ | $99.2_{\pm 0.0}$ | $100_{\pm 0.0}$ | $100_{\pm 0.0}$ | $99.0_{\pm 0.0}$ | $76.0_{\pm 34}$ | $60.7_{\pm 2.9}$ | $35.0_{\pm 6.7}$
Rotation | Accuracy | $99.2_{\pm 0.0}$ | $99.0_{\pm 0.1}$ | $99.3_{\pm 0.0}$ | $99.3_{\pm 0.0}$ | $99.0_{\pm 0.0}$ | $98.8_{\pm 0.0}$ | $98.5_{\pm 0.4}$ | $98.9_{\pm 0.0}$
Robustness | $28.9_{\pm 0.6}$ | $93.6_{\pm 0.3}$ | $95.2_{\pm 0.1}$ | $95.1_{\pm 0.1}$ | $93.5_{\pm 0.1}$ | $94.5_{\pm 0.2}$ | $92.3_{\pm 0.8}$ | $93.2_{\pm 0.7}$
Invariance | $20.6_{\pm 0.4}$ | $58.3_{\pm 2.2}$ | $66.0_{\pm 3.8}$ | $65.4_{\pm 3.5}$ | $29.1_{\pm 0.6}$ | $71.9_{\pm 2.8}$ | $48.7_{\pm 1.9}$ | $39.3_{\pm 6.9}$
Contrast | Accuracy | $99.2_{\pm 0.0}$ | $98.9_{\pm 0.3}$ | $99.4_{\pm 0.0}$ | $99.4_{\pm 0.0}$ | $99.2_{\pm 0.0}$ | $98.9_{\pm 0.0}$ | $98.7_{\pm 0.0}$ | $99.1_{\pm 0.0}$
Robustness | $26.0_{\pm 1.0}$ | $95.4_{\pm 2.6}$ | $96.8_{\pm 0.8}$ | $97.4_{\pm 0.6}$ | $97.9_{\pm 0.4}$ | $88.4_{\pm 4.5}$ | $87.2_{\pm 9.6}$ | $97.7_{\pm 0.6}$
Invariance | $20.7_{\pm 1.1}$ | $37.5_{\pm 6.9}$ | $41.4_{\pm 0.3}$ | $41.3_{\pm 0.4}$ | $26.3_{\pm 1.1}$ | $40.3_{\pm 0.9}$ | $28.4_{\pm 1.7}$ | $20.0_{\pm 0.1}$
Table 2. Test results on CIFAR10 dataset for different ARs over three evaluation metrics and three distribution shifts. Notations are the same as in Table 1.  |  | B | V | L | S | C | K | W | D
---|---|---|---|---|---|---|---|---|---
Texture | Accuracy | $88.5_{\pm 1.7}$ | $86.3_{\pm 0.3}$ | $82.8_{\pm 0.4}$ | $82.0_{\pm 0.0}$ | $86.8_{\pm 0.1}$ | $84.6_{\pm 0.4}$ | $86.8_{\pm 0.1}$ | $86.5_{\pm 0.4}$
Robustness | $38.3_{\pm 0.7}$ | $76.5_{\pm 0.0}$ | $79.1_{\pm 0.2}$ | $79.4_{\pm 0.1}$ | $76.8_{\pm 0.1}$ | $75.6_{\pm 0.1}$ | $76.8_{\pm 0.1}$ | $77.3_{\pm 0.2}$
Invariance | $44.7_{\pm 0.5}$ | $94.1_{\pm 0.6}$ | $100_{\pm 0.0}$ | $100_{\pm 0.0}$ | $94.2_{\pm 1.2}$ | $96.7_{\pm 1.0}$ | $93.4_{\pm 0.4}$ | $95.4_{\pm 0.8}$
Rotation | Accuracy | $88.5_{\pm 1.7}$ | $81.1_{\pm 2.3}$ | $80.3_{\pm 4.7}$ | $78.1_{\pm 2.2}$ | $80.7_{\pm 1.8}$ | $80.4_{\pm 6.8}$ | $87.5_{\pm 1.6}$ | $83.7_{\pm 4.0}$
Robustness | $15.3_{\pm 1.1}$ | $49.0_{\pm 0.4}$ | $50.7_{\pm 2.6}$ | $51.1_{\pm 1.0}$ | $47.0_{\pm 0.6}$ | $46.6_{\pm 4.5}$ | $49.4_{\pm 0.9}$ | $47.7_{\pm 1.7}$
Invariance | $47.3_{\pm 0.6}$ | $54.6_{\pm 2.0}$ | $40.7_{\pm 1.4}$ | $55.2_{\pm 1.6}$ | $55.7_{\pm 1.1}$ | $55.5_{\pm 1.0}$ | $55.7_{\pm 1.3}$ | $54.5_{\pm 0.7}$
Contrast | Accuracy | $88.5_{\pm 1.7}$ | $85.2_{\pm 4.5}$ | $88.8_{\pm 1.3}$ | $86.9_{\pm 2.0}$ | $89.6_{\pm 0.9}$ | $83.4_{\pm 3.7}$ | $87.2_{\pm 2.2}$ | $89.7_{\pm 0.9}$
Robustness | $54.5_{\pm 0.8}$ | $77.7_{\pm 0.8}$ | $83.1_{\pm 1.3}$ | $83.4_{\pm 1.1}$ | $80.7_{\pm 2.4}$ | $79.5_{\pm 2.9}$ | $82.8_{\pm 1.0}$ | $80.8_{\pm 5.6}$
Invariance | $53.1_{\pm 1.6}$ | $67.5_{\pm 0.1}$ | $69.8_{\pm 4.8}$ | $71.3_{\pm 2.3}$ | $53.6_{\pm 4.7}$ | $73.0_{\pm 2.2}$ | $74.6_{\pm 0.8}$ | $66.6_{\pm 4.3}$

###  4.1. Experiment Setup

Our empirical investigation is conducted over two benchmark datasets (MNIST
dataset with LeNet architecture and CIFAR10 dataset with ResNet18
architecture) and three sets of the transformations.

##### Transformation Functions

We consider three sets of transformation functions:

  * •

Texture: we use Fourier transform to perturb the texture of the data by
discarding the high-frequency components cut-off by a radius $r$, following
(Wang et al., 2020). The smaller $r$ is, the fewer high-frequency components
the image has. We consider

$\mathcal{A}=\\{a_{0}(),a_{12}(),a_{10}(),a_{8}(),a_{6}()\\}$

, where the subscript denotes the radius $r$ except that $a_{0}()$ is the
identity map. We consider $\mathcal{A}$ during test time, but only $a_{0}()$
and $a_{6}()$ during training.

  * •

Rotation: we rotate the images clockwise $r$ degrees with
$\mathcal{A}=\\{a_{0}(),a_{15}(),a_{30}(),$ $a_{45}(),a_{60}()\\}$, where the
subscript denotes the degree of rotation and $a_{0}()$ is the identity map. We
consider $\mathcal{A}$ during test time, but only $a()_{0}$ and $a_{60}()$
during training.

  * •

Contrast: we create the images depicting the same visual information, but with
different scales of the pixels, including the negative color representation.
Therefore, we have
$\mathcal{A}=\\{a_{0}(\mathbf{x})=\mathbf{x},a_{1}(\mathbf{x})=\mathbf{x}/2,a_{2}(\mathbf{x})=\mathbf{x}/4,a_{3}(\mathbf{x})=1-\mathbf{x},a_{4}(\mathbf{x})=(1-\mathbf{x})/2,a_{5}(\mathbf{x})=(1-\mathbf{x})/4$,
where $\mathbf{x}$ stands for the image whose pixel values have been set to be
between 0 and 1. We consider $\mathcal{A}$ during test time, but only
$a_{0}()$ and $a_{3}()$ during training.

##### Alignment Regularizations

We consider the following popular choices of AR (with $\mathbf{u}$ and
$\mathbf{v}$ denoting two vector embeddings):

  * L:

$\ell_{1}$ norm of the vector differences, i.e.,
$\|\mathbf{u}-\mathbf{v}\|_{1}$

  * S:

squared $\ell_{2}$ norm of the vector differences, i.e.,
$\|\mathbf{u}-\mathbf{v}\|_{2}^{2}$

  * C:

cosine similarity, i.e.,
$\mathbf{u}^{T}\mathbf{v}/\|\mathbf{u}\|\cdot\|\mathbf{v}\|$

  * K:

KL divergence over a batch of paired embeddings; the second argument are
augmented samples.

  * W:

Wasserstein-1 metric over a batch of paired embeddings, with implementation
following Wasserstein GAN (Arjovsky et al., 2017; Gulrajani et al., 2017)

  * D:

a vanilla GAN discriminator over a batch of paired embeddings, the one-layer
discriminator is trained to classify samples vs. augmented samples.

We mainly discuss applying the AR to logits (embeddings prior to the final
softmax function). We have also experimented with applying to the final
softmax output and the embeddings one layer prior to logits. Both cases lead
to substantially worse results, so we skip the discussion.

##### Hyperparameters

We use standard data splits. We first train the baseline models to get
reasonably high performances (for MNIST, we train 100 epochs with learning
rate set to be $10^{-4}$; for CIFAR10, we train 200 epochs with learning rate
initialized to be $10^{-1}$ and reduced one magnitude every 50 epochs; batch
sizes in both cases are set to be 128). Then we train other augmented models
with the same learning rate and batch size etc. The regularization weight is
searched with 8 choices evenly split in the logspace from $10^{-7}$ to $1$.
For each method, the reported score is from the weight resulting in the
highest robust accuracy. We test with three random seeds.

##### Evaluation Metrics:

We consider the three major evaluation metrics as we discussed in Section 3:

  * •

Accuracy: test accuracy on the original test data.

  * •

Robustness: the worst-case accuracy when each sample can be transformed with
$a\in\mathcal{A}$.

  * •

Invariance: the metric to test whether the learned representations are
invariant to the transformation functions, as introduced in Section 3.3.

###  4.2. Results and Discussion

Tables 1 and 2 show the empirical results across the three distribution shifts
and the three evaluation metrics. First of all, no method can dominate across
all these evaluations, probably because of the tradeoff between accuracy and
robustness (Tsipras et al., 2018; Zhang et al., 2019b; Wang et al., 2020).
Similarly, the tradeoff between accuracy and invariance can be expected from
the role of the regularization weight: when the weight is small, AR has no
effect, and the model is primarily optimized for improving accuracy; when the
weight is considerable, the model is pushed toward a trivial solution that
maps every sample to the same embedding, ignoring other patterns of the data.
This is also the reason that results in Tables 1 and 2 are selected according
to the robust accuracy.

Due to the tradeoffs, it may not be strategic if we only focus on the highest
number of each row. Instead, we suggest studying the three rows of each test
case together and compare the tradeoffs, which is also a reason we reformulate
the invariance test in Section 3.3 so that we can have bounded invariance
scores directly comparable to accuracy and robust accuracy.

For example, in the texture rows of Table 1, while cosine similarity can
outperform squared $\ell_{2}$ norm in accuracy and robustness with 0.3 and 0.1
margins, respectively, it is disadvantageous in invariance with a larger
margin (1.0). Similarly, for rotation rows of Table 1, KL-divergence shows the
overall highest scores, followed by $\ell_{1}$ norm and squared $\ell_{2}$
norm. For contrast rows, both $\ell_{1}$ norm and squared $\ell_{2}$ norm
stand out. Overall, experiments in MNIST suggest the desired choice to be
$\ell_{1}$ norm or squared $\ell_{2}$ norm, and we believe squared $\ell_{2}$
norm is marginally better.

On the other hand, the experiments in CIFAR10 in Table 2 mostly favor
$\ell_{1}$ norm and squared $\ell_{2}$ norm. Good performances can also be
observed from Wasserstein-1 metric for rotation. Also, we notice that squared
$\ell_{2}$ norm, in general, outperforms $\ell_{1}$ norm.

Thus, our empirical study recommends squared $\ell_{2}$ norm for AR to learn
robust and invariant models, with $\ell_{1}$ norm as a runner-up.

##  5\. Analytical Support

According to our discussions in Section 3.3, Wasserstein metric is supposed to
be a favorable option as AR. However, its empirical performance does not stand
out. We believe this disparity is mainly due to the difficulty in calculating
the Wasserstein metric in practice.

On the other hand, norm-based consistency regularizations stand out. We are
interested in studying the properties of these metrics. In particular, this
section aims to complement the empirical study by showing that norm-based AR
can lead to bounded robust generalization error under certain assumptions.

Also, our experiments only use two transformation functions from $\mathcal{A}$
during training but are tested with all the functions in $\mathcal{A}$. We
also discuss the properties of these two functions and argue that training
with only these two functions can be a good strategy when certain assumptions
are met.

###  5.1. Overview of Analytical Results

All together, we need six assumptions (namely, A1-A6). Out of these
assumptions, A1 is necessary for using data augmentation, and A4 is necessary
for deriving machine learning generalization error bound. A2 and A3 are
properties of data transformation functions, and A5 and A6 are technical
assumptions. Along with introducing these assumptions, we will also offer
empirical evidence showing that these assumptions are likely to hold in
practice.

In particular, we will show that

  * •

With A2 holds, $\ell_{1}$ norm can replace the empirical Wasserstein metric to
regularize invariance (Proposition A.2) and then we can derive a bounded
robust error if all the functions in $\mathcal{A}$ are available (Theorem
5.1).

  * •

With the above result and A3, we can derive a bounded robust error if only two
special functions in $\mathcal{A}$ (which we refer to as _vertices_) are
available (Lemma 5.2).

###  5.2. Assumptions Setup

####  5.2.1. Assumptions on Data Augmentation Functions

Our first three assumptions are for the basic properties of the data
transformation functions used. These properties are formally introduced as
assumptions below.

  * A1:

Dependence-preservation: the transformation function will not alter the
dependency regarding the label (i.e., for any $a()\in\mathcal{A}$,
$a(\mathbf{x})$ will have the same label as $\mathbf{x}$) or the features
(i.e., $a_{1}(\mathbf{x}_{1})$ and $a_{2}(\mathbf{x}_{2})$ are independent if
$\mathbf{x}_{1}$ and $\mathbf{x}_{2}$ are independent).

Intuitively, “Dependence-preservation” has two perspectives: Label-wise, the
transformation cannot alter the label of the data, which is a central
requirement of almost all the data augmentation functions in practice.
Feature-wise, the transformation will not introduce new dependencies between
the samples.

We consider the label-wise half of this argument as a fundamental property of
any data augmentations. It has to be always true for data augmentation to be a
useful technique. On the other hand, the feature-wise half of this argument is
a fundamental property required to derive the generalization error bounds.
Intuitively, we believe this property holds for most data augmentation
techniques in practice.

  * A2:

Efficiency: for $\widehat{\theta}$ and any $a()\in\mathcal{A}$,
$f(a(\mathbf{x});\widehat{\theta})$ is closer to $\mathbf{x}$ than any other
samples under a distance metric $d_{e}(\cdot,\cdot)$, i.e.,

$\displaystyle
d_{e}(f(a(\mathbf{x});\widehat{\theta}),f(\mathbf{x};\widehat{\theta}))\leq\min_{\mathbf{x}^{\prime}\in\mathbf{X}_{-\mathbf{x}}}d_{e}(f(a(\mathbf{x});\widehat{\theta}),f(\mathbf{x}^{\prime};\widehat{\theta})).$

We define $d_{e}(\cdot,\cdot)$ to be $\ell_{1}$ norm.

Intuitively, the efficiency property means the augmentation should only
generate new samples of the same label as minor perturbations of the original
one. If a transformation violates this property, there should exist other
simpler transformations that can generate the same target sample.

  * A3:

Vertices: For a model $\widehat{\theta}$ and a transformation $a()$, we use
$\mathbf{P}_{a,\widehat{\theta}}$ to denote the distribution of
$f(a(\mathbf{x});\widehat{\theta})$ for
$(\mathbf{x},\mathbf{y})\sim\mathbf{P}$. “Vertices” argues that exists two
extreme elements in $\mathcal{A}$, namely $a^{+}$ and $a^{-}$, with certain
metric $d_{x}(\cdot,\cdot)$, we have

$\displaystyle
d_{x}(\mathbf{P}_{a^{+},\widehat{\theta}},\mathbf{P}_{a^{-},\widehat{\theta}})=\sup_{a_{1},a_{2}\in\mathcal{A}}d_{x}(\mathbf{P}_{a_{1},\widehat{\theta}},\mathbf{P}_{a_{2},\widehat{\theta}})$

We define $d_{x}(\cdot,\cdot)$ to be Wasserstein-1 metric.

Intuitively, “Vertices” suggests that there are extreme cases of the
transformations. For example, if one needs the model to be invariant to
rotations from $0^{\circ}$ to $60^{\circ}$, we consider the vertices to be
$0^{\circ}$ rotation function (thus identity map) and $60^{\circ}$ rotation
function. In practice, one usually selects the transformation vertices with
intuitions or domain knowledge.

Notice that we do not need to argue that A3 always holds. All we need is that
A3 can sometimes hold, and when it holds, we can directly train with the
regularized vertex augmentation. Thus, anytime RVA empirically performs well
is a favorable argument for A3. To show that RVA can sometimes perform well,
we compare the RVA with vanilla (non-regularized) worst-case data augmentation
(VWA) method across our synthetic experiment setup. We notice that out of six
total scenarios ({texture, rotation, contrast} $\times$ {MNIST, CIFAR10}), RVA
outperforms VWA frequently. This suggests that the domain-knowledge of
vertices can help in most cases, although not guaranteed in every case.

####  5.2.2. Assumptions on Background and Generalization Error Bound

In an abstract manner, when the test data and train data are from the same
distribution, several previous analyses on the generalization error can be
sketched as (see examples in A4):

(4) $\displaystyle
r_{\mathbf{P}}(\widehat{\theta})\leq\widehat{r}_{\mathbf{P}}(\widehat{\theta})+\phi(|\Theta|,n,\delta)$

which suggests that the expected risk $r_{\mathbf{P}}(\widehat{\theta})$ can
be bounded by the empirical risk $\widehat{r}_{\mathbf{P}}(\widehat{\theta})$
and a function of hypothesis space $|\Theta|$ and number of samples $n$;
$\delta$ accounts for the probability when the bound holds. $\phi()$ is a
function of these three terms. Dependent on the details of different analyses,
different concrete examples of this generic term will need different
assumptions. We use a generic assumption A4 to denote the assumptions required
for each example.

Following our main goal to study how alignment regularization and data
augmentation help in accuracy, robustness, and invariance, our strategy in
theoretical analysis is to derive error bounds for accuracy and robustness,
and the error bound directly contains terms to regularize the invariance.
Further, as robustness naturally bounds accuracy (i.e.,
$r_{\mathbf{P}}(\widehat{\theta})\leq
r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})$ following the definitions in (1)
and (2) respectively), we only need to study the robust error.

To study the robust error, we need two additional technical assumptions. A5
connects the distribution of expected robust risk and the distribution of
empirical robust risk, and A6 connects the 0-1 classification error and cross-
entropy error.

  * A4:

We list two examples here:

    * –

when A4 is “$\Theta$ is finite, $l(\cdot,\cdot)$ is a zero-one loss, samples
are i.i.d”,
$\phi(|\Theta|,n,\delta)=\sqrt{(\log(|\Theta|)+\log(1/\delta))/2n}$

    * –

when A4 is “samples are i.i.d”,
$\phi(|\Theta|,n,\delta)=2\mathcal{R}(\mathcal{L})+\sqrt{(\log{1/\delta})/2n}$,
where $\mathcal{R}(\mathcal{L})$ stands for Rademacher complexity and
$\mathcal{L}=\\{l_{\theta}\,|\,\theta\in\Theta\\}$, where $l_{\theta}$ is the
loss function corresponding to $\theta$.

For more information or more concrete examples of the generic term, one can
refer to relevant textbooks such as (Bousquet et al., 2003).

A4 stands for the fundamental assumptions used to derive standard
generalization bounds. We rely on this assumption as how previous theoretical
works rely on them.

  * A5:

the distribution for expected robust risk equals the distribution for
empirical robust risk, i.e.,

$\displaystyle\operatorname*{arg\,max}_{\mathcal{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}r_{\mathbf{P}^{\prime}}(\widehat{\theta})=\operatorname*{arg\,max}_{\mathcal{P^{\prime}}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})$

where $T(\mathbf{P},\mathcal{A})$ is the collection of distributions created
by elements in $\mathcal{A}$ over samples from $\mathbf{P}$.

Eq. (2) can be written equivalently into the expected risk over a pseudo
distribution $\mathbf{P}^{\prime}$ (see Lemma 1 in (Tu et al., 2019)), which
is the distribution that can sample the data leading to the expected robust
risk. Thus, equivalently, we can consider $\sup_{\mathcal{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}r_{\mathbf{P}^{\prime}}(\widehat{\theta})$ as a
surrogate of $r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})$, where
$T(\mathbf{P},\mathcal{A})$ denotes the set of possible resulting
distributions. Following the empirical strength of techniques such as
adversarial training (Madry et al., 2018), we introduce an assumption relating
the distribution of expected robust risk and the distribution of empirical
robust risk (namely, A5, in Appendix 5.2). Thus, the bound of our interest
(i.e., $\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}r_{\mathbf{P}^{\prime}}(\widehat{\theta})$) can be
analogously analyzed through $\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})$.

A5 is likely to hold in practice: Assumption A5 appears very strong, however,
the successes of methods like adversarial training (Madry et al., 2018)
suggest that, in practice, A5 might be much weaker than it appears.

  * A6:

With $(\mathbf{x},\mathbf{y})\in(\mathbf{X},\mathbf{Y})$, the sample
maximizing cross-entropy loss and the sample maximizing classification error
for model $\widehat{\theta}$ follows:

(5)
$\displaystyle\forall\mathbf{x},\quad\dfrac{\mathbf{y}^{\top}f(\mathbf{x};\widehat{\theta})}{\inf_{a\in\mathcal{A}}\mathbf{y}^{\top}f(a(\mathbf{x});\widehat{\theta})}\geq\exp\big{(}\mathbb{I}(g(f(\mathbf{x};\widehat{\theta}))\neq
g(f(\mathbf{x}^{\prime};\widehat{\theta})))\big{)}$

where $\mathbf{x}^{\prime}$ stands for the sample maximizing classification
error, i.e.,

$\displaystyle\mathbf{x}^{\prime}=\operatorname*{arg\,min}_{\mathbf{x}}\mathbf{y}^{\top}g(f(\mathbf{x};\widehat{\theta}))$

Also,

(6)
$\displaystyle\forall\mathbf{x},\quad|\inf_{a\in\mathcal{A}}\mathbf{y}^{\top}f(a(\mathbf{x});\widehat{\theta})|\geq
1$

Intuitively, although Assumption A6 appears complicated, it describes the
situations of two scenarios:

If
$g(f(\mathbf{x};\widehat{\theta}))=g(f(\mathbf{x}^{\prime};\widehat{\theta}))$,
which means either the sample is misclassified by $\widehat{\theta}$ or
$\mathcal{A}$ is not rich enough for a transformation function to alter the
prediction, the RHS of Eq. 5 is 1, thus Eq. 5 always holds (because
$\mathcal{A}$ has the identity map as one of its elements).

If $g(f(\mathbf{x};\widehat{\theta}))\neq
g(f(\mathbf{x}^{\prime};\widehat{\theta}))$, which means a transformation
alters the prediction. In this case, A6 intuitively states that the
$\mathcal{A}$ is reasonably rich and the transformation is reasonably powerful
to create a gap of the probability for the correct class between the original
sample and the transformed sample. The ratio is described as the ratio of the
prediction confidence from the original sample over the prediction confidence
from the transformed sample is greater than $e$.

###  5.3. Analytical Support

##### Regularized Worst-case Augmentation

To have a model with a small invariance score, we should probably directly
regularize the empirical counterpart of Eq. (3). However, Wasserstein distance
is difficult to calculate in practice. Fortunately, Proposition A.2
conveniently allows us to use $\ell_{1}$ norm to replace Wasserstein metric.
With Proposition A.2, now we can offer our main technical result to study the
robust error $r_{\mathbf{P},\mathcal{A}}{\widehat{\theta}}$ (as defined in Eq.
(2)).

###### Theorem 5.1.

With Assumptions A1, A2, A4, A5, and A6, with probability at least $1-\delta$,
we have

$\displaystyle
r_{\mathbf{P},\mathcal{A}}{\widehat{\theta}}\leq\widehat{r}_{\mathbf{P}}(\widehat{\theta})+\sum_{i}||f(\mathbf{x}_{i};\widehat{\theta})-f(\mathbf{x}^{\prime}_{i};\widehat{\theta})||_{1}+\phi(|\Theta|,n,\delta)$

and $\mathbf{x}^{\prime}=a(\mathbf{x})$, where
$a=\operatorname*{arg\,min}_{a\in\mathcal{A}}\mathbf{y}^{\top}f(a(\mathbf{x});\widehat{\theta})$.
$\phi(|\Theta|,n,\delta)$ is defined in A4.

This technical result immediately inspires the method to guarantee worst case
performance, as well as to explicitly enforce the concept of invariance. The
method
$a=\operatorname*{arg\,min}_{a\in\mathcal{A}}\mathbf{y}^{\top}f(a(\mathbf{x});\widehat{\theta})$
is selecting the transformation function maximizing the cross-entropy loss
(notice the sign difference between here and the cross-entropy loss), which we
refer to as worst-case data augmentation. This method is also closely
connected to adversarial training (e.g., Madry et al., 2018).

##### Regularized Vertex Augmentation

As $\mathcal{A}$ in practice usually has a large number of (and possibly
infinite) elements, we may not always be able to identify the worst-case
transformation function with reasonable computational efforts. We further
leverage the vertex property (boundary cases of transformation functions,
discussed as Assumption A3 in the appendix) of the transformation function to
bound the worst-case generalization error:

###### Lemma 5.2.

With Assumptions A1-A6, assuming there is a $a^{\prime}()\in\mathcal{A}$ where
$\widehat{r}_{\mathbf{P}_{a^{\prime}}}(\widehat{\theta})=\frac{1}{2}\big{(}\widehat{r}_{\mathbf{P}_{a^{+}}}(\widehat{\theta})+\widehat{r}_{\mathbf{P}_{a^{-}}}(\widehat{\theta})\big{)}$,
with probability at least $1-\delta$, we have:

$\displaystyle r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})\leq$
$\displaystyle\dfrac{1}{2}\big{(}\widehat{r}_{\mathbf{P}_{a^{+}}}(\widehat{\theta})+\widehat{r}_{\mathbf{P}_{a^{-}}}(\widehat{\theta})\big{)}$
$\displaystyle+\sum_{i}||f(a^{+}(\mathbf{x}_{i});\widehat{\theta})-f(a^{-}(\mathbf{x}^{\prime});\widehat{\theta})||_{1}+\phi(|\Theta|,n,\delta),$

where $a^{+}()$ and $a^{-}()$ are defined in A3, and $\phi(|\Theta|,n,\delta)$
in A4.

This result corresponds to the method that can be optimized conveniently
without searching for the worst-case transformations. However, the method
requires good domain knowledge of the vertices (i.e., boundary cases) of the
transformation functions.

Table 3. Comparison to advanced rotation-invariant models. We report the accuracy on the test sets rotated. “main” means the resulting images are highly likely to be semantically the same as the original ones. “all” means the average accuracy of all rotations. The underlined scores show that data augmentation and AR can help a vanilla model to compete with advanced methods. The bold scores (highest at each row) show that data augmentation and AR can further improve the advanced methods.  | ResNet | GC | ST | ETN
---|---|---|---|---
| Base | RVA | RWA | Base | RVA | RWA | Base | RVA | RWA | Base | RVA | RWA
main | 45.4 | 66.5 |  71.1 | 38.5 | 72.2 | 73.8 | 45.9 | 58.3 | 62.9 | 56.9 | 65.1 | 57.7
all | 31.2 | 48.1 |  52.8 | 26.7 | 54.4 | 55.0 | 32.1 | 40.2 | 42.7 | 39.5 | 52.6 | 46.1
Table 4. Comparison to advanced methods on 9 super-class ImageNet classification with different distribution shifts. | Acc. | WAcc. | ImageNet-A | ImageNet-S
---|---|---|---|---
Base | 90.8 | 88.8 | 24.9 | 41.1
SIN | 88.4 | 86.6 | 24.6 | 40.5
LM | 67.9 | 65.9 | 18.8 | 36.8
RUBi | 90.5 | 88.6 | 27.7 | 42.3
RB | 91.9 | 90.5 | 29.6 | 41.8
RVA | 92.2 | 91.2 | 28.0 | 42.5
RWA | 92.8 | 91.6 | 28.8 | 43.2
Table 5. The generic methods can also improve standard accuracy.  | ResNet18 | ResNet50 | ResNet101
---|---|---|---
B | RVA | RWA | B | RVA | RWA | B | RVA | RWA
Top-1 | 75.6 | 100 | 77.2 | 77.4 | 100 | 78.2 | 77.8 | 100 | 78.7
Top-5 | 93.1 | 100 | 93.8 | 93.9 | 100 | 94.4 | 94.4 | 100 | 94.9

Thus, our theoretical discussions have complemented our empirical findings in
Section 4 by showing that norm-based regularizations can lead to bounded
robust error. There is a disparity that our analytical result is about
$\ell_{1}$ norm while our empirical study suggests squared $\ell_{2}$ norm. We
conjecture the disparity is mainly caused by the difficulty in passing the
gradient of $\ell_{1}$ norm in practice.

##  6\. Experiments with Advanced Methods

We continue to test the methods we identified in comparison to more advanced
methods. Although we argued for the value of invariance, for a fair
comparison, we will test the performances evaluated by the metrics the
previous methods are designed for. Our method will use the same generic
approach and the same transformation functions as in the previous empirical
study, although these functions are not necessarily part of the distribution
shift we test now. In summary, our method can outperform (or be on par with)
these SOTA techniques in the robustness metric they are designed for (Section
6.2). In addition, we run a side test to show that our method can also improve
accuracy (Section 6.3).

###  6.1. Methods

Section 4 and Section 5 lead us to test the following two methods:

  * •

RVA (regularized vertex augmentation): using squared $\ell_{2}$ norm as AR
over logits between the original samples and the augmented samples of a fixed
vertex transformation function (original samples are considered as from
another vertex).

  * •

RWA (regularized worst-case augmentation): using squared $\ell_{2}$ norm as AR
over logits between the original samples and the worst-case augmented samples
identified at each iteration. Worst-case samples are generated by the function
with the maximum loss when we iterate through all the transformation
functions.

###  6.2. Robustness

##### Rotation

We compare our results with rotation-invariant models, mainly Spatial
Transformer (ST) (Jaderberg et al., 2015), Group Convolution (GC) (Cohen and
Welling, 2016), and Equivariant Transformer Network (ETN) (Tai et al., 2019).
We also tried to run CGNet (Kondor et al., 2018), but the method does not seem
to scale to the CIFAR10 and ResNet level. All these methods are tested with
ResNet34 following popular settings in the community. The results are in Table
3. We test the models every $15^{\circ}$ rotation from $0^{\circ}$ rotation to
$345^{\circ}$ rotation. Augmentation-related methods use the $\mathcal{A}$ of
“rotation” in synthetic experiments, so the testing scenario goes beyond what
the augmentation methods have seen during training.

We report two summary results in Table 3. “main” means the average prediction
accuracy from images rotated from $300^{\circ}$ to $60^{\circ}$ (passing
$0^{\circ}$), when the resulting images are highly likely to preserve the
class label. “all” means the average accuracy of all rotations.

Our results can be interpreted from two perspectives. First, by comparing all
the columns in the first panel to the first column of the other three panels,
data augmentation and AR can boost a vanilla model to outperform other
advanced techniques. On the other hand, by comparing the columns within each
panel, data augmentation and AR can further improve the performances of these
techniques.

Interestingly, the baseline model with our generic approach (RWA in the first
panel) can almost compete with the advanced methods even when these methods
also use augmentation and AR (RWA in GC panel). We believe this result
strongly indicates the potential of this simple augmentation and
regularization method to match the advanced methods.

In summary, RWA can boost the vanilla model to outperform advanced methods.
Data augmentation and squared $\ell_{2}$ AR can further improve the
performances when plugged onto advanced methods.

##### Texture & Contrast

We follow (Bahng et al., 2020) and compare the models for a nine super-class
ImageNet classification (Ilyas et al., 2019) with class-balanced strategies.
Also, we follow (Bahng et al., 2020) to report standard accuracy (Acc.),
weighted accuracy (WAcc.), a scenario where samples with unusual texture are
weighted more, and accuracy over ImageNet-A (Hendrycks et al., 2019), a
collection of failure cases for most ImageNet trained models. Additionally, we
also report the performance over ImageNet-Sketch (Wang et al., 2019a), an
independently collected ImageNet test set with only sketch images. As (Bahng
et al., 2020) mainly aims to overcome the texture bias, we also use our
texture-wise functions in Section 4 for augmentation. However, there are no
direct connections between these functions and the distribution shift of the
test samples. Also, we believe the distribution shifts here, especially the
one introduced by our newly added ImageNet-Sketch, are more than texture, and
also correspond to the contrast case of our study.

Following (Bahng et al., 2020), the base network is ResNet, and we compare
with the vanilla network (Base), and several methods designed for this task:
including StylisedIN (SIN) (Geirhos et al., 2019), LearnedMixin (LM) (Clark et
al., 2019), RUBi (RUBi) (Cadene et al., 2019) and ReBias (RB) (Bahng et al.,
2020). Results are in Table 4.

The results favor our generic method in most cases. RVA outperforms other
methods in standard accuracy, weighted accuracy, and ImageNet-Sketch, and is
shy from ReBias on ImageNet-A. RWA shows the same pattern as that of RVA and
further outperforms RVA. Overall, these results validate the empirical
strength of data augmentation (even when the augmentation is not designed for
the task) and squared $\ell_{2}$ norm AR for learning robust models.

###  6.3. Accuracy

Further, these experiments help us notice that the generic technique can also
help improve the accuracy, although the technique is motivated by robustness
and invariance. Therefore, we follow the widely accepted CIFAR100 test
pipeline and test the performances of different architectures of the ResNet
family. The results are reported in Table 5, where Base stands for the
baseline model with the default accuracy boosting configurations.

For both top-1 and top-5 accuracies and across the three ResNet architectures,
our techniques can help improve the accuracy. In addition, we notice that our
techniques can help bridge the gap of different architectures within the
ResNet family: for example, RWA helps ResNet50 to outperform the vanilla
ResNet101.

##  7\. Conclusion

In this paper, we seek to answer how to train with augmented data so that
augmentation can be taken to the fullest extent. We first defined a new
evaluation metric called invariance and conducted a line of empirical studies
to show that norm-based alignment regularization can help learn robust and
invariant models. Further, we complement our observations with formal
derivations of bounded generalization errors. We notice that regularizing
squared $\ell_{2}$ norm between the logits of the originals samples and those
of the augmented samples is favorable: the trained model tends to have the
most favorable performances in robust accuracy and invariance. In general, the
method we recommend is “regularized worst-case augmentation” with squared
$\ell_{2}$ norm as the alignment regularization. One can also consider
“regularized vertex augmentation” when extra assumptions on the vertex
properties of the transformation functions are met. Lastly, we would like to
remind a potential limitation of alignment regularization: it may not always
help improve the i.i.d accuracy due to the tradeoff between accuracy and
robustness or invariance. In addition, to simplify the procedure of users in
leveraging our contribution, we also release a software package in both
TensorFlow and PyTorch for users to use our identified methods with a couple
lines of code.

###### Acknowledgements.

This work was supported by NIH R01GM114311, NIH P30DA035778, and NSF
IIS1617583; NSF CAREER IIS-2150012 and IIS-2204808. The authors would like to
thank Hanru Yan for the implementation of the software package.

## References

  * (1)
  * Abu-Mostafa (1990) Yaser S Abu-Mostafa. 1990\.  Learning from hints in neural networks.  _Journal of complexity_ 6, 2 (1990), 192–198.
  * Achille and Soatto (2018) Alessandro Achille and Stefano Soatto. 2018.  Information dropout: Learning optimal representations through noisy computation.  _IEEE transactions on pattern analysis and machine intelligence_ 40, 12 (2018), 2897–2905.
  * Arjovsky et al. (2017) Martin Arjovsky, Soumith Chintala, and Léon Bottou. 2017\.  Wasserstein GAN.  arXiv:1701.07875 [stat.ML]
  * Asai and Hajishirzi (2020) Akari Asai and Hannaneh Hajishirzi. 2020.  Logic-Guided Data Augmentation and Regularization for Consistent Question Answering.  arXiv:2004.10157 [cs.CL]
  * Bahng et al. (2020) Hyojin Bahng, Sanghyuk Chun, Sangdoo Yun, Jaegul Choo, and Seong Joon Oh. 2020.  Learning De-biased Representations with Biased Representations.  119 (2020), 528–539.
  * Bobkov and Ledoux (2019) Sergey Bobkov and Michel Ledoux. 2019.  _One-dimensional empirical measures, order statistics, and Kantorovich transport distances_. Vol. 261.  American Mathematical Society.
  * Bousquet et al. (2003) Olivier Bousquet, Stéphane Boucheron, and Gábor Lugosi. 2003.  Introduction to statistical learning theory. In _Summer School on Machine Learning_. Springer, 169–207.
  * Buslaev et al. (2020) Alexander Buslaev, Vladimir I. Iglovikov, Eugene Khvedchenya, Alex Parinov, Mikhail Druzhinin, and Alexandr A. Kalinin. 2020\.  Albumentations: Fast and Flexible Image Augmentations.  _Inf._ 11, 2 (2020), 125.
  * Cadene et al. (2019) Remi Cadene, Corentin Dancette, Matthieu Cord, Devi Parikh, et al. 2019\.  Rubi: Reducing unimodal biases for visual question answering. In _Advances in neural information processing systems_. 841–852.
  * Chen et al. (2019) Shuxiao Chen, Edgar Dobriban, and Jane H Lee. 2019\.  A Group-Theoretic Framework for Data Augmentation.  arXiv:1907.10905 [stat.ML]
  * Clark et al. (2019) Christopher Clark, Mark Yatskar, and Luke Zettlemoyer. 2019\.  Don’t Take the Easy Way Out: Ensemble Based Methods for Avoiding Known Dataset Biases.  _arXiv preprint arXiv:1909.03683_ (2019).
  * Cohen and Welling (2016) Taco Cohen and Max Welling. 2016.  Group equivariant convolutional networks. In _International conference on machine learning_. 2990–2999.
  * Cubuk et al. (2019) Ekin D Cubuk, Barret Zoph, Dandelion Mane, Vijay Vasudevan, and Quoc V Le. 2019.  Autoaugment: Learning augmentation strategies from data. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 113–123.
  * Cuturi and Doucet (2014) Marco Cuturi and Arnaud Doucet. 2014.  Fast computation of Wasserstein barycenters.  (2014).
  * Dao et al. (2019) Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Ré. 2019\.  A Kernel Theory of Modern Data Augmentation. In _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ _(Proceedings of Machine Learning Research, Vol. 97)_ , Kamalika Chaudhuri and Ruslan Salakhutdinov (Eds.). PMLR, 1528–1537.
  * Fawzi et al. (2016) Alhussein Fawzi, Horst Samulowitz, Deepak S. Turaga, and Pascal Frossard. 2016.  Adaptive data augmentation for image classification. In _2016 IEEE International Conference on Image Processing, ICIP 2016, Phoenix, AZ, USA, September 25-28, 2016_. IEEE, 3688–3692.
  * Ganin et al. (2016) Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky. 2016.  Domain-adversarial training of neural networks.  _The Journal of Machine Learning Research_ 17, 1 (2016), 2096–2030.
  * Geirhos et al. (2019) Robert Geirhos, Patricia Rubisch, Claudio Michaelis, Matthias Bethge, Felix A. Wichmann, and Wieland Brendel. 2019.  ImageNet-trained CNNs are biased towards texture; increasing shape bias improves accuracy and robustness.. In _International Conference on Learning Representations_.
  * Ghiasi et al. (2020) Golnaz Ghiasi, Yin Cui, Aravind Srinivas, Rui Qian, Tsung-Yi Lin, Ekin D. Cubuk, Quoc V. Le, and Barret Zoph. 2020\.  Simple Copy-Paste is a Strong Data Augmentation Method for Instance Segmentation.  _CoRR_ abs/2012.07177 (2020).  arXiv:2012.07177
  * Ghosh and Thiery (2021) Atin Ghosh and Alexandre H. Thiery. 2021.  On Data-Augmentation and Consistency-Based Semi-Supervised Learning. In _International Conference on Learning Representations_.
  * Goodfellow (2016) Ian Goodfellow. 2016\.  NIPS 2016 tutorial: Generative adversarial networks.  _arXiv preprint arXiv:1701.00160_ (2016).
  * Goodfellow et al. (2014) Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. 2014\.  Explaining and harnessing adversarial examples.  _arXiv preprint arXiv:1412.6572_ (2014).
  * Gulrajani et al. (2017) Ishaan Gulrajani, Faruk Ahmed, Martin Arjovsky, Vincent Dumoulin, and Aaron Courville. 2017\.  Improved Training of Wasserstein GANs.  arXiv:1704.00028 [cs.LG]
  * Guo et al. (2019) Hao Guo, Kang Zheng, Xiaochuan Fan, Hongkai Yu, and Song Wang. 2019.  Visual Attention Consistency Under Image Transforms for Multi-Label Image Classification. In _IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2019, Long Beach, CA, USA, June 16-20, 2019_. Computer Vision Foundation / IEEE, 729–739.
  * Hendrycks and Dietterich (2019) Dan Hendrycks and Thomas Dietterich. 2019.  Benchmarking Neural Network Robustness to Common Corruptions and Perturbations.  _Proceedings of the International Conference on Learning Representations_ (2019).
  * Hendrycks et al. (2020) Dan Hendrycks, Norman Mu, Ekin Dogus Cubuk, Barret Zoph, Justin Gilmer, and Balaji Lakshminarayanan. 2020.  AugMix: A Simple Data Processing Method to Improve Robustness and Uncertainty. In _8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020_. OpenReview.net.
  * Hendrycks et al. (2019) Dan Hendrycks, Kevin Zhao, Steven Basart, Jacob Steinhardt, and Dawn Song. 2019.  Natural adversarial examples.  _arXiv preprint arXiv:1907.07174_ (2019).
  * Hernandez-Garcia (2020) Alex Hernandez-Garcia. 2020\.  Data augmentation and image understanding.  _arXiv preprint arXiv:2012.14185_ (2020).
  * Hernández-García et al. (2019) Alex Hernández-García, Peter König, and Tim C Kietzmann. 2019\.  Learning robust visual representations using data augmentation invariance.  _arXiv preprint arXiv:1906.04547_ (2019).
  * Hernández-García and König (2018) Alex Hernández-García and Peter König. 2018.  Data augmentation instead of explicit regularization.  arXiv:1806.03852 [cs.CV]
  * Ho et al. (2019) Daniel Ho, Eric Liang, Xi Chen, Ion Stoica, and Pieter Abbeel. 2019.  Population Based Augmentation: Efficient Learning of Augmentation Policy Schedules. In _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ _(Proceedings of Machine Learning Research, Vol. 97)_ , Kamalika Chaudhuri and Ruslan Salakhutdinov (Eds.). PMLR, 2731–2741.
  * Huang et al. (2020) Zeyi Huang, Haohan Wang, Eric P. Xing, and Dong Huang. 2020\.  Self-challenging Improves Cross-Domain Generalization. In _Computer Vision - ECCV 2020 - 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part II_ _(Lecture Notes in Computer Science, Vol. 12347)_ , Andrea Vedaldi, Horst Bischof, Thomas Brox, and Jan-Michael Frahm (Eds.). Springer, 124–140.
  * Ilyas et al. (2019) Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Logan Engstrom, Brandon Tran, and Aleksander Madry. 2019.  Adversarial examples are not bugs, they are features. In _Advances in Neural Information Processing Systems_. 125–136.
  * Jaderberg et al. (2015) Max Jaderberg, Karen Simonyan, Andrew Zisserman, et al. 2015\.  Spatial transformer networks. In _Advances in neural information processing systems_. 2017–2025.
  * Jeong et al. (2019) Jisoo Jeong, Seungeui Lee, Jeesoo Kim, and Nojun Kwak. 2019\.  Consistency-based Semi-supervised Learning for Object detection. In _Advances in Neural Information Processing Systems_. 10758–10767.
  * Kannan et al. (2018) Harini Kannan, Alexey Kurakin, and Ian Goodfellow. 2018\.  Adversarial Logit Pairing.  arXiv:1803.06373 [cs.LG]
  * Kondor et al. (2018) Risi Kondor, Zhen Lin, and Shubhendu Trivedi. 2018.  Clebsch–gordan nets: a fully fourier space spherical convolutional neural network. In _Advances in Neural Information Processing Systems_. 10117–10126.
  * Kostrikov et al. (2020) Ilya Kostrikov, Denis Yarats, and Rob Fergus. 2020\.  Image Augmentation Is All You Need: Regularizing Deep Reinforcement Learning from Pixels.  _CoRR_ abs/2004.13649 (2020).
  * LeCun et al. (1998) Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. 1998.  Gradient-based learning applied to document recognition.  _Proc. IEEE_ 86, 11 (1998), 2278–2324.
  * Liang et al. (2018) Davis Liang, Zhiheng Huang, and Zachary C Lipton. 2018\.  Learning noise-invariant representations for robust speech recognition. In _2018 IEEE Spoken Language Technology Workshop (SLT)_. IEEE, 56–63.
  * Liang (2016) Percy Liang. 2016\.  CS229T/STAT231: Statistical Learning Theory (Winter 2016).
  * Lim et al. (2019) Sungbin Lim, Ildoo Kim, Taesup Kim, Chiheon Kim, and Sungwoong Kim. 2019.  Fast AutoAugment. In _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché-Buc, Emily B. Fox, and Roman Garnett (Eds.). 6662–6672.
  * Lopes et al. (2019) Raphael Gontijo Lopes, Dong Yin, Ben Poole, Justin Gilmer, and Ekin D Cubuk. 2019\.  Improving robustness without sacrificing accuracy with patch gaussian augmentation.  _arXiv preprint arXiv:1906.02611_ (2019).
  * Madry et al. (2018) Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. 2018\.  Towards Deep Learning Models Resistant to Adversarial Attacks. In _6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Conference Track Proceedings_. OpenReview.net.
  * Min et al. (2020) Junghyun Min, R. Thomas McCoy, Dipanjan Das, Emily Pitler, and Tal Linzen. 2020.  Syntactic Data Augmentation Increases Robustness to Inference Heuristics. In _ACL_.
  * Rajput et al. (2019) Shashank Rajput, Zhili Feng, Zachary Charles, Po-Ling Loh, and Dimitris Papailiopoulos. 2019\.  Does Data Augmentation Lead to Positive Margin?. In _International Conference on Machine Learning_. 5321–5330.
  * Sajjadi et al. (2016) Mehdi Sajjadi, Mehran Javanmardi, and Tolga Tasdizen. 2016\.  Regularization with stochastic transformations and perturbations for deep semi-supervised learning. In _Advances in neural information processing systems_. 1163–1171.
  * Shah et al. (2019) Meet Shah, Xinlei Chen, Marcus Rohrbach, and Devi Parikh. 2019\.  Cycle-Consistency for Robust Visual Question Answering. In _The IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_.
  * Shorten and Khoshgoftaar (2019) Connor Shorten and Taghi M Khoshgoftaar. 2019.  A survey on image data augmentation for deep learning.  _Journal of Big Data_ 6, 1 (2019), 60.
  * Simard et al. (1991) Patrice Simard, Bernard Victorri, Yann LeCun, and John S Denker. 1991.  Tangent prop-a formalism for specifying selected invariances in an adaptive network. In _NIPS_.
  * Sun et al. (2018) Sining Sun, Ching-Feng Yeh, Mari Ostendorf, Mei-Yuh Hwang, and Lei Xie. 2018.  Training Augmentation with Adversarial Examples for Robust Speech Recognition. In _ISCA_.
  * Szegedy et al. (2013) Christian Szegedy, Wojciech Zaremba, Ilya Sutskever, Joan Bruna, Dumitru Erhan, Ian Goodfellow, and Rob Fergus. 2013\.  Intriguing properties of neural networks.  _arXiv preprint arXiv:1312.6199_ (2013).
  * Tai et al. (2019) Kai Sheng Tai, Peter Bailis, and Gregory Valiant. 2019\.  Equivariant Transformer Networks.  _arXiv preprint arXiv:1901.11399_ (2019).
  * Tsipras et al. (2018) Dimitris Tsipras, Shibani Santurkar, Logan Engstrom, Alexander Turner, and Aleksander Madry. 2018\.  Robustness May Be at Odds with Accuracy. In _International Conference on Learning Representations_.
  * Tu et al. (2019) Zhuozhuo Tu, Jingwei Zhang, and Dacheng Tao. 2019\.  Theoretical analysis of adversarial learning: A minimax approach. In _Advances in Neural Information Processing Systems_. 12259–12269.
  * Villani (2003) Cédric Villani. 2003\.  _Topics in optimal transportation_.  Number 58. American Mathematical Soc.
  * Villani (2008) Cédric Villani. 2008\.  _Optimal transport: old and new_. Vol. 338.  Springer Science & Business Media.
  * Wang et al. (2019a) Haohan Wang, Songwei Ge, Zachary C. Lipton, and Eric P. Xing. 2019a.  Learning Robust Global Representations by Penalizing Local Predictive Power.  , 10506–10518 pages.
  * Wang et al. (2019b) Haohan Wang, Zexue He, Zachary C. Lipton, and Eric P. Xing. 2019b.  Learning Robust Representations by Projecting Superficial Statistics Out. In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net.
  * Wang and Raj (2017) Haohan Wang and Bhiksha Raj. 2017.  On the origin of deep learning.  _arXiv preprint arXiv:1702.07800_ (2017).
  * Wang et al. (2020) Haohan Wang, Xindi Wu, Zeyi Huang, and Eric P. Xing. 2020\.  High Frequency Component Helps Explain the Generalization of Convolutional Neural Networks. In _Computer Vision and Pattern Recognition (CVPR)_.
  * Wu et al. (2019) Xindi Wu, Yijun Mao, Haohan Wang, Xiangrui Zeng, Xin Gao, Eric P. Xing, and Min Xu. 2019.  Regularized Adversarial Training (RAT) for Robust Cellular Electron Cryo Tomograms Classification. In _2019 IEEE International Conference on Bioinformatics and Biomedicine, BIBM 2019, San Diego, CA, USA, November 18-21, 2019_ , Illhoi Yoo, Jinbo Bi, and Xiaohua Hu (Eds.). IEEE, 1–6.
  * Xie et al. (2015) Saining Xie, Tianbao Yang, Xiaoyu Wang, and Yuanqing Lin. 2015\.  Hyper-class augmented and regularized deep learning for fine-grained image classification. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 2645–2654.
  * Yang et al. (2019) Fanny Yang, Zuowen Wang, and Christina Heinze-Deml. 2019\.  Invariance-inducing regularization using worst-case transformations suffices to boost accuracy and spatial robustness. In _Advances in Neural Information Processing Systems_. 14757–14768.
  * Zhang et al. (2019b) Hongyang Zhang, Yaodong Yu, Jiantao Jiao, Eric P. Xing, Laurent El Ghaoui, and Michael I. Jordan. 2019b.  Theoretically Principled Trade-off between Robustness and Accuracy. In _Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA_ _(Proceedings of Machine Learning Research, Vol. 97)_ , Kamalika Chaudhuri and Ruslan Salakhutdinov (Eds.). PMLR, 7472–7482.
  * Zhang et al. (2021) Linjun Zhang, Zhun Deng, Kenji Kawaguchi, Amirata Ghorbani, and James Zou. 2021.  How Does Mixup Help With Robustness and Generalization?. In _International Conference on Learning Representations_.
  * Zhang (2019) Richard Zhang. 2019\.  Making convolutional networks shift-invariant again.  _arXiv preprint arXiv:1904.11486_ (2019).
  * Zhang et al. (2019a) Zhirui Zhang, Shuangzhi Wu, Shujie Liu, Mu Li, Ming Zhou, and Tong Xu. 2019a.  Regularizing neural machine translation by target-bidirectional agreement. In _Proceedings of the AAAI Conference on Artificial Intelligence_ , Vol. 33. 443–450.
  * Zheng et al. (2016) Stephan Zheng, Yang Song, Thomas Leung, and Ian Goodfellow. 2016\.  Improving the robustness of deep neural networks via stability training. In _Proceedings of the ieee conference on computer vision and pattern recognition_. 4480–4488.
  * Zhong et al. (2020) Zhun Zhong, Liang Zheng, Guoliang Kang, Shaozi Li, and Yi Yang. 2020.  Random Erasing Data Augmentation. In _AAAI_.

## Appendices

##  Appendix A Proof of Theoretical Results

###  A.1. Lemma A.1 and its proof

###### Lemma A.1.

With Assumptions A1, A4, and A5, with probability at least $1-\delta$, we have

$\displaystyle
r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})\leq\dfrac{1}{n}\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}}\sup_{a\in\mathcal{A}}\mathbb{I}(g(f(a(\mathbf{x});\widehat{\theta}))\neq\mathbf{y})+\phi(|\Theta|,n,\delta)$

###### Proof.

With Assumption A5, we have

$\displaystyle\operatorname*{arg\,max}_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}r_{\mathbf{P}^{\prime}}(\widehat{\theta})=\operatorname*{arg\,max}_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})=\mathbf{P}_{w}$

we can analyze the expected risk following the standard classical techniques
since both expected risk and empirical risk are studied over distribution
$\mathbf{P}_{w}$.

Now we only need to make sure the classical analyses (as discussed in A4) are
still valid over distribution $\mathbf{P}_{w}$:

  * •

when A4 is “$\Theta$ is finite, $l(\cdot,\cdot)$ is a zero-one loss, samples
are i.i.d”,
$\phi(|\Theta|,n,\delta)=\sqrt{\dfrac{\log(|\Theta|)+\log(1/\delta)}{2n}}$.
The proof of this result uses Hoeffding’s inequality, which only requires
independence of random variables. One can refer to Section 3.6 in (Liang,
2016) for the detailed proof.

  * •

when A4 is “samples are i.i.d”,
$\phi(|\Theta|,n,\delta)=2\mathcal{R}(\mathcal{L})+\sqrt{\dfrac{\log{1/\delta}}{2n}}$.
The proof of this result relies on McDiarmid’s inequality, which also only
requires independence of random variables. One can refer to Section 3.8 in
(Liang, 2016) for the detailed proof.

Assumption A1 guarantees the samples from distribution $\mathbf{P}_{w}$ are
still independent, thus the generic term holds for at least these two concrete
examples, thus the claim is proved.

∎

###  A.2. Proposition A.2 and Proof

###### Proposition A.2.

With A2, for any $a\in\mathcal{A}$, we have

$\displaystyle
W_{1}(\widehat{\mathbf{Q}}_{\mathbf{x},\widehat{\theta}},\widehat{\mathbf{Q}}_{a(\mathbf{x}),\widehat{\theta}})=\sum_{i}^{|(\mathbf{X},\mathbf{Y})|}||f(\mathbf{x}_{i};\widehat{\theta})-f(a(\mathbf{x}_{i});\widehat{\theta})||_{1},$

where $\widehat{\mathbf{Q}}_{\mathbf{x},\widehat{\theta}}$ denotes the
empirical distribution of $f(\mathbf{x};\widehat{\theta})$ for
$(\mathbf{x},\mathbf{y})\in(\mathbf{X},\mathbf{Y})$.

###### Proof.

We use the order statistics representation of Wasserstein metric over
empirical distributions (e.g., see Section 4 in (Bobkov and Ledoux, 2019))

$\displaystyle
W_{1}(\widehat{\mathbf{Q}}_{\mathbf{x},\widehat{\theta}},\widehat{\mathbf{Q}}_{a(\mathbf{x}),\widehat{\theta}}))=\inf_{\sigma}\sum_{i}^{|(\mathbf{X},\mathbf{Y})|}||f(\mathbf{x}_{i};\widehat{\theta})-f(a(\mathbf{x}_{\sigma(i)}),\widehat{\theta})||_{1}$

where $\sigma$ stands for a permutation of the index, thus the infimum is
taken over all possible permutations. With Assumption A2, when
$d_{e}(\cdot,\cdot)$ in A2 chosen to be $\ell_{1}$ norm, we have:

$\displaystyle||f(\mathbf{x}_{i};\widehat{\theta})-f(a(\mathbf{x}_{i});\widehat{\theta})||_{1}\leq\min_{j\neq
i}||f(\mathbf{x}_{i};\widehat{\theta})-f(a(\mathbf{x}_{j}),\widehat{\theta})||_{1}$

Thus, the infimum is taken when $\sigma$ is the natural order of the samples,
which leads to the claim. ∎

###  A.3. Proof of Theorem 5.1

###### Theorem 0.

With Assumptions A1, A2, A4, A5, and A6, with probability at least $1-\delta$,
we have

(7) $\displaystyle
r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})\leq\widehat{r}_{\mathbf{P}}(\widehat{\theta})+\sum_{i}||f(\mathbf{x}_{i};\widehat{\theta})-f(\mathbf{x}^{\prime}_{i};\widehat{\theta})||_{1}+\phi(|\Theta|,n,\delta)$

and $\mathbf{x}^{\prime}=a(\mathbf{x})$, where
$a=\operatorname*{arg\,min}_{a\in\mathcal{A}}\mathbf{y}^{\top}f(a(\mathbf{x});\widehat{\theta})$.

###### Proof.

First of all, in the context of multiclass classification, where
$g(f(\mathbf{x},;\theta))$ predicts a label with one-hot representation, and
$\mathbf{y}$ is also represented with one-hot representation, we can have the
empirical risk written as:

$\displaystyle\widehat{r}_{\mathbf{P}}(\widehat{\theta})=1-\dfrac{1}{n}\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}}\mathbf{y}^{\top}g(f(\mathbf{x};\widehat{\theta}))$

Thus,

$\displaystyle\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})=$
$\displaystyle\widehat{r}_{\mathbf{P}}(\widehat{\theta})+\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})-\widehat{r}_{\mathbf{P}}(\widehat{\theta})$
$\displaystyle=$
$\displaystyle\widehat{r}_{\mathbf{P}}(\widehat{\theta})+\dfrac{1}{n}\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\big{(}\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}}\mathbf{y}^{\top}g(f(\mathbf{x};\widehat{\theta}))$
$\displaystyle-\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}^{\prime}}\mathbf{y}^{\top}g(f(\mathbf{x};\widehat{\theta}))\big{)}$

With A6, we can continue with:

$\displaystyle\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})\leq$
$\displaystyle\widehat{r}_{\mathbf{P}}(\widehat{\theta})+\dfrac{1}{n}\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\big{(}\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}}\mathbf{y}^{\top}\log(f(\mathbf{x};\widehat{\theta}))$
$\displaystyle-\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}^{\prime}}\mathbf{y}^{\top}\log(f(\mathbf{x};\widehat{\theta}))\big{)}$

If we use $e(\cdot)=-\mathbf{y}^{\top}\log(\cdot)$ to replace the cross-
entropy loss, we simply have:

$\displaystyle\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})\leq$
$\displaystyle\widehat{r}_{\mathbf{P}}(\widehat{\theta})+\dfrac{1}{n}\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\big{(}\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}^{\prime}}e(f(\mathbf{x};\widehat{\theta}))$
$\displaystyle-\sum_{(\mathbf{x},\mathbf{y})\sim\mathbf{P}}e((f(\mathbf{x};\widehat{\theta}))\big{)}$

Since $e(\cdot)$ is a Lipschitz function with constant $\leq 1$ (because of
A6, Eq.(6)) and together with the dual representation of Wasserstein metric
(See e.g., (Villani, 2003)), we have

$\displaystyle\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})$
$\displaystyle\leq\widehat{r}_{\mathbf{P}}(\widehat{\theta})+W_{1}(\widehat{\mathbf{Q}}_{\mathbf{x},\widehat{\theta}},\widehat{\mathbf{Q}}_{a(\mathbf{x}),\widehat{\theta}}))$

where $\mathbf{x}^{\prime}=a(\mathbf{x})$, where
$a=\operatorname*{arg\,min}_{a\in\mathcal{A}}\mathbf{y}^{\top}f(a(\mathbf{x});\widehat{\theta})$;
$\widehat{\mathbf{Q}}_{\mathbf{x},\widehat{\theta}}$ denotes the empirical
distribution of $f(a(\mathbf{x});\widehat{\theta})$ for
$(\mathbf{x},\mathbf{y})\in(\mathbf{X},\mathbf{Y})$. Note that
$r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})$, by definition, is a shorthand
notation for

$\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}r_{\mathbf{P}^{\prime}}(\widehat{\theta})$

.

Further, we can use the help of Proposition B.2 to replace Wassertein metric
with $\ell_{1}$ distance. Finally, we can conclude the proof with Assumption
A5 as how we did in the proof of Lemma B.1. ∎

| Clean | Noise | Blur | Weather | Digital | mCE
---|---|---|---|---|---|---
Gauss | Shot | Impulse | Defocus | Glass | Motion | Zoom | Snow | Frost | Fog | Bright | Contrast | Elastic | Pixel | JPEG
B | 23.9 | 79 | 80 | 82 | 82 | 90 | 84 | 80 | 86 | 81 | 75 | 65 | 79 | 91 | 77 | 80 | 80.6
RVA | 23.6 | 78 | 78 | 79 | 74 | 87 | 79 | 76 | 78 | 75 | 69 | 58 | 68 | 85 | 75 | 75 | 75.6
RWA | 22.4 | 61 | 63 | 63 | 68 | 75 | 65 | 66 | 70 | 69 | 64 | 56 | 55 | 70 | 61 | 63 | 64.6
SU | 24.5 | 67 | 68 | 70 | 74 | 83 | 81 | 77 | 80 | 74 | 75 | 62 | 77 | 84 | 71 | 71 | 74.3
AA | 22.8 | 69 | 68 | 72 | 77 | 83 | 80 | 81 | 79 | 75 | 64 | 56 | 70 | 88 | 57 | 71 | 72.7
MBP | 23 | 73 | 74 | 76 | 74 | 86 | 78 | 77 | 77 | 72 | 63 | 56 | 68 | 86 | 71 | 71 | 73.4
SIN | 27.2 | 69 | 70 | 70 | 77 | 84 | 76 | 82 | 74 | 75 | 69 | 65 | 69 | 80 | 64 | 77 | 73.3
AM | 22.4 | 65 | 66 | 67 | 70 | 80 | 66 | 66 | 75 | 72 | 67 | 58 | 58 | 79 | 69 | 69 | 68.4
AMS | 25.2 | 61 | 62 | 61 | 69 | 77 | 63 | 72 | 66 | 68 | 63 | 59 | 52 | 74 | 60 | 67 | 64.9
Table 6. Comparison to advanced models over ImageNet-C data. Performance reported (mCE) follows the standard in ImageNet-C data: mCE is the smaller the better. | B | InfoDrop | HEX | PAR | VA | RVA | RSC | VWA | RWA
---|---|---|---|---|---|---|---|---|---
Top-1 | 0.1204 | 0.1224 | 0.1292 | 0.1306 | 0.1362 | 0.1405 | 0.1612 | 0.1432 | 0.1486
Top-5 | 0.2408 | 0.256 | 0.2564 | 0.2627 | 0.2715 | 0.2793 | 0.3078 | 0.2846 | 0.2933
Table 7. Comparison to advanced cross-domain image classification models, over
ImageNet-Sketch dataset. We report top-1 and top-5 accuracy following
standards on ImageNet related experiments.

###  A.4. Proof of Lemma 5.2

Lemma. With Assumptions A1-A6, assuming there is a
$a^{\prime}()\in\mathcal{A}$ where
$\widehat{r}_{\mathbf{P}_{a^{\prime}}}(\widehat{\theta})=\frac{1}{2}\big{(}\widehat{r}_{\mathbf{P}_{a^{+}}}(\widehat{\theta})+\widehat{r}_{\mathbf{P}_{a^{-}}}(\widehat{\theta})\big{)}$,
with probability at least $1-\delta$, we have:

(8) $\displaystyle r_{\mathbf{P},\mathcal{A}}(\widehat{\theta})\leq$
$\displaystyle\dfrac{1}{2}\big{(}\widehat{r}_{\mathbf{P}_{a^{+}}}(\widehat{\theta})$
(9)
$\displaystyle+\widehat{r}_{\mathbf{P}_{a^{-}}}(\widehat{\theta})\big{)}+\sum_{i}||f(a^{+}(\mathbf{x}_{i});\widehat{\theta})-f(a^{-}(\mathbf{x}^{\prime});\widehat{\theta})||_{1}+\phi(|\Theta|,n,\delta)$

###### Proof.

We can continue with

$\displaystyle\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})$
$\displaystyle\leq\widehat{r}_{\mathbf{P}}(\widehat{\theta})+W_{1}(\widehat{\mathbf{Q}}_{\mathbf{x},\widehat{\theta}},\widehat{\mathbf{Q}}_{a(\mathbf{x}),\widehat{\theta}})),$

where $\widehat{\mathbf{Q}}_{\mathbf{x},\widehat{\theta}}$ denotes the
empirical distribution of $f(a(\mathbf{x});\widehat{\theta})$ for
$(\mathbf{x},\mathbf{y})\in(\mathbf{X},\mathbf{Y})$. from the proof of Theorem
5.2. With the help of Assumption A3, we have:

$\displaystyle
d_{x}(f(a^{+}(\mathbf{x}),\widehat{\theta}),f(a^{-}(\mathbf{x}),\widehat{\theta}))\geq
d_{x}(f(\mathbf{x},\widehat{\theta}),f(\mathbf{x}^{\prime},\widehat{\theta}))$

When $d_{x}(\cdot,\cdot)$ is chosen as Wasserstein-1 metric, we have:

$\displaystyle\sup_{\mathbf{P}^{\prime}\in
T(\mathbf{P},\mathcal{A})}\widehat{r}_{\mathbf{P}^{\prime}}(\widehat{\theta})$
$\displaystyle\leq\widehat{r}_{\mathbf{P}}(\widehat{\theta})+W_{1}(\widehat{\mathbf{Q}}_{a^{+}(\mathbf{x}),\widehat{\theta}},\widehat{\mathbf{Q}}_{a^{-}(\mathbf{x}),\widehat{\theta}}))$

Further, as the LHS is the robust risk generated by the transformation
functions within $\mathcal{A}$, and
$\widehat{r}_{\mathbf{P}}(\widehat{\theta})$ is independent of the term
$W_{1}(\widehat{\mathbf{Q}}_{a^{+}(\mathbf{x}),\widehat{\theta}},\widehat{\mathbf{Q}}_{a^{-}(\mathbf{x}),\widehat{\theta}}))$,
WLOG, we can replace $\widehat{r}_{\mathbf{P}}(\widehat{\theta})$ with the
risk of an arbitrary distribution generated by the transformation function in
$\mathcal{A}$. If we choose to use
$\widehat{r}_{\mathbf{P}_{a^{\prime}}}(\widehat{\theta})=\frac{1}{2}\big{(}\widehat{r}_{\mathbf{P}_{a^{+}}}(\widehat{\theta})+\widehat{r}_{\mathbf{P}_{a^{-}}}(\widehat{\theta})\big{)}$,
we can conclude the proof with help from Proposition B.2 and Assumption A5 as
how we did in the proof of Theorem 5.2. ∎

##  Appendix B Additional Results for Comparisons with Advanced Methods

We have also conducted two full ImageNet level experiments. However, due to
the limitation of resources, we cannot tune the models substantially. Our
current trial suggest that our techniques can improve the vanilla model to
compete with SOTA models, limited by our resources, we cannot do wide-range
hyperparameters search to outperform them. Also, considering the fact that
many of these methods are significantly more complicated than us and also uses
data augmentation specially designed for the tasks, we consider our
experiments a success indication of the empirical strength of our methods.

##### Texture-perturbed ImageNet classification

We also test the performance on the image classification over multiple
perturbations. We train the model over standard ImageNet training set and test
the model with ImageNet-C data (Hendrycks and Dietterich, 2019), which is a
perturbed version of ImageNet by corrupting the original ImageNet validation
set with a collection of noises. Following the standard, the reported
performance is mCE, which is the smaller the better. We compare with several
methods tested on this dataset, including Patch Uniform (PU) (Lopes et al.,
2019), AutoAugment (AA) (Cubuk et al., 2019), MaxBlur pool (MBP) (Zhang,
2019), Stylized ImageNet (SIN) (Hendrycks and Dietterich, 2019), AugMix (AM)
(Hendrycks et al., 2020), AugMix w. SIN (AMS) (Hendrycks et al., 2020). We use
the performance reported in (Hendrycks et al., 2020). Again, our augmention
only uses the generic texture with perturbation (the $\mathcal{A}$ in our
texture synthetic experiments with radius changed to $20,25,30,35,40$). The
results are reported in Table 6, which shows that our generic method
outperform the current SOTA methods after a continued finetuning process with
reducing learning rates.

##### Cross-domain ImageNet-Sketch Classification

We also compare to the methods used for cross-domain evaluation. We follow the
set-up advocated by (Wang et al., 2019b) for domain-agnostic cross-domain
prediction, which is training the model on one or multiple domains without
domain identifiers and test the model on an unseen domain. We use the most
challenging setup in this scenario: train the models with standard ImageNet
training data, and test the model over ImageNet-Sketch data (Wang et al.,
2019a), which is a collection of sketches following the structure ImageNet
validation set. We compare with previous methods with reported performance on
this dataset, such as InfoDrop (Achille and Soatto, 2018), HEX (Wang et al.,
2019b), PAR (Wang et al., 2019a), RSC (Huang et al., 2020) and report the
performances in Table 7. Notice that, our data augmentation also follows the
requirement that the characteristics of the test domain cannot be utilized
during training. Thus, we only augment the samples with a generic augmentation
set ($\mathcal{A}$ of “contrast” in synthetic experiments). The results again
support the usage of data augmentation and alignment regularization.