capbtabboxtable[][]

# The Curious Case of Benign Memorization

Sotiris Anagnostidis∗, Gregor Bachmann∗, Lorenzo Noci∗, Thomas Hofmann
Department of Computer Science
ETH Zürich, Switzerland
<EMAIL_ADDRESS>

###### Abstract

Despite the empirical advances of deep learning across a variety of learning
tasks, our theoretical understanding of its success is still very restricted.
One of the key challenges is the overparametrized nature of modern models,
enabling complete overfitting of the data even if the labels are randomized,
i.e. networks can completely memorize all given patterns. While such a
memorization capacity seems worrisome, in this work we show that under
training protocols that include data augmentation, neural networks learn to
memorize entirely random labels in a benign way, i.e. they learn embeddings
that lead to highly non-trivial performance under nearest neighbour probing.
We demonstrate that deep models have the surprising ability to separate noise
from signal by distributing the task of memorization and feature learning to
different layers. As a result, only the very last layers are used for
memorization, while preceding layers encode performant features which remain
largely unaffected by the label noise. We explore the intricate role of the
augmentations used for training and identify a memorization-generalization
trade-off in terms of their diversity, marking a clear distinction to all
previous works. Finally, we give a first explanation for the emergence of
benign memorization by showing that malign memorization under data
augmentation is infeasible due to the insufficient capacity of the model for
the increased sample size. As a consequence, the network is forced to leverage
the correlated nature of the augmentations and as a result learns meaningful
features. To complete the picture, a better theory of feature learning in deep
neural networks is required to fully understand the origins of this
phenomenon.

††∗Equal contribution.

##  1 Introduction

Deep learning has made tremendous advances in the past decade, leading to
state-of-the-art performance on various learning tasks such as computer vision
(He et al., 2016), natural language processing (Devlin et al., 2019) and graph
learning (Kipf & Welling, 2017). While some progress has been made regarding
the theoretical understanding of these deep models (Arora et al., 2018;
Bartlett et al., 2019; 2017; Neyshabur et al., 2015; 2018; Dziugaite & Roy,
2017), the considered settings are unfortunately often very restrictive and
the insights made are only qualitative or very loose. One of the key technical
hurdles hindering progress is the highly overparametrized nature of neural
networks employed in practice, which is in stark contrast with classical
learning theory, according to which simpler hypotheses compatible with the
data should be preferred. The challenge of overparametrization is beautifully
illustrated in the seminal paper of Zhang et al. (2017), showing that deep
networks are able to fit arbitrary labelings of the data, i.e. they can
completely memorize all the patterns. This observation renders tools from
classical learning theory such as VC-dimension or Rademacher complexity
vacuous and commands new avenues to be investigated. The random label
experiment has been applied as a sanity check for new techniques (Arora et
al., 2018; 2019a; Bartlett et al., 2017; Dziugaite & Roy, 2017), where an
approach is evaluated based on its ability to distinguish between networks
that memorize or truly learn the data. From a classical perspective,
memorization is thus considered as a bug, not a feature, and goes hand in hand
with bad generalization.
In this work we challenge this view by revisiting the randomization experiment
of Zhang et al. (2017) with a slight twist: we change the training protocol by
adding data augmentation, a standard practice used in almost all modern deep
learning pipelines. We show that in this more practical setting, the story is
more intricate;

Neural networks trained on random labels with data augmentation learn useful
features!

More precisely, we show that probing the embedding space with the nearest
neighbour algorithm of such a randomly trained network admits highly non-
trivial performance on a variety of standard benchmark datasets. Moreover,
such networks have the surprising ability to separate signal from noise, as
all layers except for the last ones focus on feature learning while not
fitting the random labels at all. On the other hand, the network uses its last
layers to learn the random labeling, at the cost of clean accuracy, which
strongly deteriorates. This is further evidence of a strong, implicit bias
present in modern models, allowing them to learn performant features even in
the setting of complete noise. Inspired by the line of works on benign
overfitting (Bartlett et al., 2020; Sanyal et al., 2021; Frei et al., 2022),
we coin this phenomenon benign memorization. We study our findings through the
lens of capacity and show that under data augmentation, modern networks are
forced to leverage the correlations present in the data to achieve
memorization. As a consequence, the model learns invariant features which have
been identified to have strong discriminatory power in the field of self-
supervised learning (Caron et al., 2021; Grill et al., 2020; Bardes et al.,
2022; Zbontar et al., ; Chen & He, 2021). Specifically, we make the following
contributions:

  * •

We make the surprising observation that learning under complete label noise
still leads to highly useful features (benign memorization), showing that
memorization and generalization are not necessarily at odds.

  * •

We show that deep neural networks exhibit an astonishing capability to
separate noise and signal between different layers, fitting the random labels
only at the very last layers.

  * •

We highlight the intricate role of augmentations and their interplay with the
capacity of the model class, forcing the network to learn the correlation
structure.

  * •

We interpret our findings in terms of invariance learning, an objective that
has instigated large successes in the field of self-supervised learning.

##  2 Related Work

Memorization. Our work builds upon the seminal paper of Zhang et al. (2017)
which showed how neural networks can easily memorize completely random labels.
This observation has inspired a multitude of follow-up works and the
introduced randomization test has become a standard tool to assess the
validity of generalization bounds (Arora et al., 2018; 2019a; Bartlett et al.,
2017; Dziugaite & Roy, 2017). The intriguing capability of neural networks to
simply memorize data has inspired researchers to further dissect the
phenomenon, especially in the setting where only a subset of the targets is
randomized. Arpit et al. (2017) studies how neural networks tend to learn
shared patterns first, before resorting to memorization when given real data,
as opposed to random labels where examples are fitted independently. Feldman &
Zhang (2020) study the setting when real but “long-tailed” data is used and
show how memorization in this case can be beneficial to performance. Maennel
et al. (2020); Pondenkandath et al. (2018) on the other hand show how pre-
training networks on random labels can sometimes lead to faster, subsequent
training on the clean data or novel tasks. Finally, Zhang et al. (2021) show
how training on random labels can be valuable for neural architecture search.
In all these previous works, data augmentation is excluded from the training
pipeline. For partial label noise, it is well-known in the literature that
neural networks exhibit surprising robustness (Rolnick et al., 2017; Song et
al., 2020; Patrini et al., 2017) and generalization is possible. We highlight
however that this setting is distinct from complete label noise, which we
study in this work.

Data augmentation. Being a prominent component of deep learning applications,
the benefits of data augmentation have been investigated theoretically in the
setting of clean labels (Chen et al., 2020b; Dao et al., 2019; Wu et al.,
2020; Hanin & Sun, 2021). The benefits of data augmentation have been verified
empirically when only a subset of the data is corrupted (Nishi et al., 2021).
On the other hand, investigations with pure label noise where no signal
remains in the dataset are absent in the literature. Finally, we want to
highlight the pivotal role of data augmentation in self-supervised learning
frameworks (Caron et al., 2021; Grill et al., 2020; Bardes et al., 2022;
Zbontar et al., ; Chen & He, 2021), where it facilitates learning of invariant
features.

Figure 1: (Left) Encoder-projector pair. (Right) Standard data augmentations.

(a)

(b)

##  3 Background

Setting. In this work we consider the standard classification setting, where
we are given $n\in\mathbb{N}$ i.i.d. samples
$\left(\bm{x}_{1},\bm{y}_{1}\right),\dots,\left(\bm{x}_{n},\bm{y}_{n}\right)\stackrel{{\scriptstyle
i.i.d.}}{{\sim}}\mathcal{D}$ from some data distribution $\mathcal{D}$,
consisting of inputs $\bm{x}\in\mathbb{R}^{d}$ and one-hot targets
$\bm{y}\in\mathbb{R}^{C}$, each encoding one out of $C$ classes. We consider a
family of parametrized functions (in this work, neural networks)
$f_{\bm{\theta}}:\mathbb{R}^{d}\xrightarrow[]{}\mathbb{R}^{C}$, where
$\bm{\theta}\in\Theta$ denotes the (concatenated) weights from some space
$\Theta\subseteq\mathbb{R}^{p}$ where $p$ is the total number of parameters.
Moreover, we specify a loss function
$l:\mathbb{R}^{C}\times\mathbb{R}^{C}\xrightarrow[]{}\mathbb{R}_{+}$ which
measures the discrepancy between a prediction $f_{\bm{\theta}}(\bm{x})$ and
its corresponding target $\bm{y}$, i.e. $l(f_{\bm{\theta}}(\bm{x}),\bm{y})$.
We then perform learning by minimizing the empirical loss $\hat{L}$ as a
function of the parameters $\bm{\theta}$,

$\hat{\bm{\theta}}:=\operatorname{argmin}_{\bm{\theta}\in\Theta}\hat{L}(\bm{\theta}):=\operatorname{argmin}_{\bm{\theta}\in\Theta}\sum_{i=1}^{n}l(f_{\bm{\theta}}(\bm{x}_{i}),\bm{y}_{i})$
(1)

using some form of stochastic gradient descent and measure the resulting
generalization error
$L(\hat{\bm{\theta}})=\mathbb{E}_{(\bm{x},\bm{y})\sim\mathcal{D}}\left[l(f_{\hat{\bm{\theta}}}(\bm{x}),\bm{y})\right]$.
We denote by $p_{\bm{X}}$ the marginal density of the inputs and by
$p_{\bm{Y}|\bm{X}}$ the conditional density of the labels given an input.
$p_{\bm{Y}|\bm{X}}$ encodes the statistical relationship between an input
$\bm{x}$ and the associated label $\bm{y}$. As typical in practice, we assume
that we are in the so-called interpolation regime (Ma et al., 2018), i.e. we
assume that stochastic gradient descent can find a parameter configuration
$\hat{\bm{\theta}}$ that achieves zero training loss (see Eq. 1).

Architecture. Throughout this work, we consider networks composed of an
encoder $h_{\bm{\psi}}:\mathbb{R}^{d}\xrightarrow[]{}\mathbb{R}^{m}$ and a
projector $g_{\bm{\phi}}:\mathbb{R}^{m}\xrightarrow[]{}\mathbb{R}^{C}$, where
both $\bm{\psi}$ and $\bm{\phi}$ denote the parameters of the respective
building block. As an encoder, we typically employ modern convolutional
architectures such as ResNets (He et al., 2016) or VGG (Simonyan & Zisserman,
2014), excluding the final fully-connected layers, while the projector is
usually an MLP with one hidden layer. We illustrate such a network in Fig.
1(a). Such architectures have become very popular in the domain of feature
learning and are extensively used in unsupervised and self-supervised learning
(Caron et al., 2021; Grill et al., 2020; Bardes et al., 2022; Zbontar et al.,
). In this work, we are interested in assessing the quality of the encoder’s
features $h_{\bm{\psi}}$ when the network $f_{\bm{\theta}}=g_{\bm{\phi}}\circ
h_{\bm{\psi}}$ is trained on random labels.

Probing. To evaluate the embeddings $h_{\bm{\psi}}$, we apply nearest-
neighbour-based probing, which refers to performing $K$-nearest-neighbour
classification based on the embeddings
$\\{\left(h_{\bm{\psi}}(\bm{x}_{i}),\bm{y}_{i}\right)\\}_{i=1}^{n}$. We fix
the number of neighbours to $K=20$ throughout this work unless otherwise
specified. Probing measures how useful a given set of features $h_{\bm{\psi}}$
is for a fixed task. A special case we will often consider in this work is
probing the encoder of a network at initialization, which we refer to as
probing at initialization. Due to the very restricted complexity of probing,
the resulting performance is very dependent on the quality of the input
representations $h_{\bm{\psi}}$. Such method has been used to assess the
quality of a given embedding space in various works, for instance Alain &
Bengio (2017); Chen et al. (2020a); Caron et al. (2021); Grill et al. (2020);
Bardes et al. (2022); Zbontar et al. .

[0.33] [0.64]  Layer Initialization End of training Layer 1 36.9 34.3 Layer 2
36.4 33.5 Layer 3 36.5 30.4 Layer 4 36.2 14.9 Embedding 40.2 12.4 Projector
hidden layer 38.4 10.5 Output 24.2 10.5

Figure 2: Fitting random labels on unaugmented data on CIFAR10 with a ResNet18
is not significantly slower than fitting clean labels. Figure 3: $K$-NN
probing accuracies (percentage) for a ResNet18 at initialization versus at the
end of training on random labels without augmentations on CIFAR10. The four
layers correspond to the intermediate representations after each stage of a
ResNet.

Random labels and memorization. Denote by $\bm{e}_{l}\in\mathbb{R}^{C}$ the
$l$-th unit vector, i.e. the one-hot encoding of class $l$. Zhang et al.
(2017) introduced a randomization test, where any clean label
$\bm{y}_{i}\in\\{\bm{e}_{1},\dots,\bm{e}_{C}\\}$ in the training set is
replaced by a random one-hot encoding, $\tilde{\bm{y}}_{i}=\bm{e}_{l}$ where
$l\sim\mathcal{U}(\\{1,\dots,C\\})$ with $\mathcal{U}$ denoting the uniform
distribution over the discrete set $\\{1,\dots,C\\}$. Throughout this text, we
will denote randomized variables with $\sim$ on top. Notice that such an
intervention destroys any statistical relationship between inputs and targets,
i.e. $p_{\bm{Y}|\bm{X}}=p_{\bm{Y}}$. As a consequence, training on such a
dataset should become very difficult as there is no common pattern left to
exploit and an algorithm has to resort to a pure memorization approach. Zhang
et al. (2017) showed that deep neural networks have an astonishing capacity to
perform such memorization tasks, easily overfitting any random label
assignment even on large-scale datasets such as ImageNet (Deng et al., 2009)
without a large increase in training time (see Fig. 3). Even explicit
regularizers such as weight decay and Dropout (Srivastava et al., 2014) can be
used in a standard manner, data augmentation however is excluded from the
pipeline. We have reproduced a subset of the results of Zhang et al. (2017) in
Table 1, first column. As expected, deep neural networks indeed manage to
achieve zero training error across a variety of datasets while not
generalizing at all, both with respect to a random test labeling as well as
the original, clean test labels. In order to further study the amount of
distortion in the resulting embedding space, we apply nearest-neighbour
probing with respect to the clean data. More precisely, given the features
$h_{{\bm{{\psi}}}}$ learnt from training on the random label task
$\\{(\bm{x}_{i},\tilde{\bm{y}}_{i})\\}_{i=1}^{n}$, we apply probing based on
the clean training data $\\{(\bm{x}_{i},{\bm{y}}_{i})\\}_{i=1}^{n}$ and
evaluate with respect to clean test data. We display the results in Table 3.
In line with previous works, (Cohen et al., 2018; Maennel et al., 2020), we
find that while very early layers might retain their performance at
initialization, a significant drop occurs with increasing depth, further
highlighting the lack of feature learning and the malignant nature of
memorization in this setting.

Data Augmentation. A standard technique present in almost all computer vision
pipelines is data augmentation. We consider transformations
$T:\mathbb{R}^{d}\xrightarrow[]{}\mathbb{R}^{d}$, that take a given input
$\bm{x}$ and produce a new augmentation $\bar{\bm{x}}:=T(\bm{x})$, which by
design, should preserve the associated label, i.e. $\bar{\bm{y}}=\bm{y}$. Such
transformations are usually given as a composition of smaller augmentations
including random crops, flips, color-jittering etc. In Fig. 1(b) we show a set
of different augmentations of the same underlying image $\bm{x}$. Notice how
these transformations indeed leave the associated label invariant. We denote
by $\mathcal{T}$ the set of all possible augmentations. Data augmentation is
usually applied in an online fashion, i.e. at every step of gradient descent,
we uniformly sample a fresh transformation
$T\sim\mathcal{U}\left(\mathcal{T}\right)$ and propagate it through the
network. We highlight that data augmentation is a standard technique necessary
for state-of-the-art performance for a variety of vision tasks. Indeed, the
top five leaders on ImageNet111https://paperswithcode.com/sota/image-
classification-on-imagenet (Yu et al., 2022; Dai et al., 2021; Zhai et al.,
2021; Pham et al., 2021; Liu et al., 2021) all rely on some form of data
augmentation in their training pipeline. If one hence wants to study the
memorization potential of neural networks in practical settings, data
augmentation needs to be considered. We notice that the results of Zhang et
al. (2017) and subsequent studies on the properties of memorization under
random labels (Arpit et al., 2017; Maennel et al., 2020) were obtained without
the use of data augmentation, leaving the memorization picture thus
incomplete.

##  4 Benign Memorization

In this section, we present the curious phenomenon of benign memorization i.e.
how neural networks manage to completely fit random labels under data
augmentation, while at the same time learning predictive features for
downstream tasks. Let us first formally introduce the terms benign and malign
memorization, which are central to the results of this work. In the following,
$\mathcal{S}:=\\{(\bm{x}_{i},\bm{y}_{i})\\}_{i=1}^{n}$ denotes the original
_clean_ dataset and
$\tilde{\mathcal{S}}=\\{(\bm{x}_{i},\tilde{\bm{y}}_{i})\\}_{i=1}^{n}$ its
randomly labeled version.

###### Definition 4.1

We call an encoder-projector pair
$\left(h_{\bm{\phi}_{*}},g_{\bm{\psi}_{*}}\right)$ a memorization of
$\tilde{\mathcal{S}}$, if $f_{*}$ perfectly fits $\tilde{\mathcal{S}}$.
Moreover, we call $\left(h_{\bm{\phi}_{*}},g_{\bm{\psi}_{*}}\right)$ a malign
memorization if additionally, probing of $h_{\bm{\phi}_{*}}$ on $\mathcal{S}$
does not improve over probing at initialization. On the contrary, we call
$\left(h_{\bm{\phi}_{*}},g_{\bm{\psi}_{*}}\right)$ a benign memorization of
$\tilde{\mathcal{S}}$ if probing of $h_{\bm{\phi}_{*}}$ on $\mathcal{S}$
outperforms probing at initialization.

As highlighted in Sec. 3 and shown in Table 1, memorizing solutions found with
stochastic gradient descent without data augmentation on standard vision
benchmarks are of malign nature. This is very intuitive, as randomizing the
targets destroys any signal present in the dataset and thus generalization
seems impossible. We show now how including data augmentation in the training
pipeline completely reverses this observation.

Dataset | Model | Random | Random + DA | Clean | Clean + DA | Init
---|---|---|---|---|---|---
CIFAR10 | ResNet18 | $12.4$ | $76.2$ | $83.6$ | $91.3$ | $40.6$
VGG11 | $15.0$ | $71.9$ | $83.0$ | $89.0$ | $44.3$
CIFAR100 | ResNet18 | $5.0$ | $42.7$ | $50.5$ | $68.2$ | $16.6$
VGG11 | $6.0$ | $46.4$ | $51.2$ | $59.0$ | $18.9$
TinyImageNet | ResNet18 | $2.4$ | $33.4$ | $38.3$ | $46.6$ | $4.6$
VGG11 | $1.2$ | $29.4$ | $34.6$ | $43.8$ | $5.4$
Table 1: $K$-NN probing accuracies (in percentage) of the embeddings for
various datasets under different settings. DA refers to training under data
augmentation. INIT refers to the performance at initialization. Except for
INIT, all models reach perfect (unaugmented) training accuracy.

Training details. In all subsequent experiments involving data augmentations,
we use the standard transformations employed in self-supervised learning
frameworks such as Chen et al. (2020c); Grill et al. (2020); Chen & He (2021).
These consist of a composition of random crops, color-jittering, random
greyscaling, and random horizontal flips, leading to a diverse set of
transformations. Moreover, we rely on mixup augmentations Zhang et al. (2018),
where two images $\bm{x}_{1}$, $\bm{x}_{2}$ are combined into a linear
interpolation, according to some weighting $\alpha\in[0,1]$, i.e.
$\bar{\bm{x}}:=\alpha\bm{x}_{1}+(1-\alpha)\bm{x}_{2}$. The corresponding label
$\bar{\bm{y}}$ is accordingly subject to the same linear combination, i.e.
$\bar{\bm{y}}=\alpha\bm{y}_{1}+(1-\alpha)\bm{y}_{2}$ where labels are
represented as their one-hot encodings. We use the standard vision datasets
CIFAR10 and CIFAR100 (Krizhevsky & Hinton, 2009), as well as TinyImageNet (Le
& Yang, 2015). For more details, we refer the reader to Appendix D.

Benign Memorization. We display the results of training under data
augmentation in Table 1. We observe that, surprisingly, nearest-neighbour
probing of the learnt embeddings leads to clearly non-trivial performance,
strongly improving over the models trained under random labels without data
augmentation. Moreover, we strongly outperform probing at initialization,
showing that indeed rich feature learning is happening. As a consequence, data
augmentation does not simply prevent malign memorization by preserving the
signal at initialization but actually leads to learning from the data. On the
other hand, it holds that the projector achieves perfect training accuracy on
the random labels and the network thus indeed memorizes the training data
perfectly. Under data augmentation, deep models hence exhibit benign
memorization. In Appendix B, Fig. 11, we further underline the utility of the
learnt features by evaluating their performance under transfer learning. We
remark that this is, to the best of our knowledge, the first work showing that
learning under random labels can lead to strong downstream performance,
highlighting that memorization and generalization are not necessarily at odds.
We notice that the training time under randomized targets together with data
augmentation increases significantly, both compared to training under clean
labels as well as fitting random labels without augmentations. In Fig. 4 we
show the evolution of both the training loss as well as nearest-neighbour
probing test accuracy as a function of the number of epochs. Observe that for
as long as $3000$ epochs, neither training loss nor probing accuracy show any
progress but then suddenly start to rapidly improve.
We stress that our goal is not to compete with standard training under clean
labels, it can be seen in Table 1 that there is a large gap between the two
methods. Instead we rather aim for a deeper understanding of how deep networks
generalize. We show that generalization, contrary to prior beliefs, remains
possible even in the most adversarial setting of complete label noise.

Signal-Noise Separation. We now further inspect models trained under random
labels and data augmentation, with an emphasis on how the noise stemming from
the random labeling affects different parts of the network. To gain insights
into this, we perform nearest-neighbour probing of different layers in the
network, with the twist that we fit the $K$-NN classifier both with respect to
the clean training labels, as well as with respect to the random labels. This
way we can assess how much a given feature has adapted to the particular
structure (clean vs. random). We visualize the results of such a clean and
noisy probing strategy in Fig. 4 at different stages of training.
Surprisingly, we can see a striking separation between feature learning and
memorization within the architecture; noisy probing only starts to improve
over random guessing once we reach the projector, whereas the encoder remains
largely unaffected. On the other hand, clean probing outperforms probing at
initialization throughout the entire embedding stage but sharply decays once
we get to the projector. While previous works have identified that the very
first layers provide non-trivial performance even without data augmentation,
including it strengthens this effect significantly. We give an interpretation
of this finding in Sec. 5.2. We further investigate the role of the projector
in Appendix B, Fig. 9, finding that higher widths lead to better performance
in general.

Figure 4: $K$-NN probing of a ResNet18 on CIFAR10, based on (flattened)
representation of different layers. Dashed blue lines indicate $K$-NN
accuracies fitted on the clean training data and evaluated w.r.t. clean test
data. Solid lines indicate $K$-NN training accuracies based on the random
labels. Notice the sharp phase transition when reaching the projector
layers“P” and “O”, i.e. the projector reaches strong performance on the noisy
labels (dashed) but randomly guesses on the clean data (solid). The opposite
happens for the embedding layer, reaching $76\%$ $K$-NN accuracy.

##  5 The Role of Augmentations

The empirical evidence in the previous section highlights the crucial role of
data augmentation for benign memorization. In this section, we aim to gain
insights into the phenomenon by investigating more closely what properties of
data augmentation are essential for benign memorization and how “ideal”
augmentations might strongly differ when learning with clean or random labels.

Label Preservation. The first important characteristic of augmentations is
their label-preserving nature, i.e. by augmenting an image $(\bm{x},\bm{y})$
twice to produce
$\left(T_{1}(\bm{x}),\bm{y}\right),\left(T_{2}(\bm{x}),\bm{y}\right)$, we have
effectively added information, as indeed $T_{1}(\bm{x})$ and $T_{2}(\bm{x})$
share the same label. Unfortunately, this reasoning is flawed in the setting
of random labels, or at the least only forms part of a larger picture, as we
show in the following thought experiment. Consider two “original” samples
$(\bm{x}_{1},\tilde{\bm{y}})$ and $(\bm{x}_{2},\tilde{\bm{y}})$ which happen
to share the same random label $\tilde{\bm{y}}$ but in truth have distinct
labels $\bm{y}_{1}\not=\bm{y}_{2}$. In this case, forming augmentations
$\left(T_{1}(\bm{x}_{1}),\tilde{\bm{y}}\right)$ and
$\left(T_{2}(\bm{x}_{1}),\tilde{\bm{y}}\right)$ might lead to some correct
signal as $T_{1}(\bm{x}_{1})$ and $T_{2}(\bm{x}_{1})$ share the same, true
label, but at the same time leads to more distortion as $\bm{x}_{2}$ has the
same, random label, reinforcing the wrong correlation even more. As a
consequence, separating noise from signal remains equally challenging as
before. We illustrate the argument in Fig. 6. To check this hypothesis, we
consider the extreme case where augmentations $T\in\mathcal{T}$ produce a new,
i.i.d. sample $T(\bm{x})$ that shares the same, true label with $\bm{x}$. In a
sense, this is the ideal augmentation and leads to the highest information
gain. We implement such augmentations by only using a subset of the training
data while using the remainder to assign to each training point $B$ potential,
independent examples with the same, true label. We choose the subset size as
$s=1000$ and the number of augmentations as $B=50$ for CIFAR10 and train the
same models as in Sec. 4, both under random and clean labels. We display the
results in Table 6. Notice that for clean label training, such augmentations
increase the training set size from $1000$ to $1000\times 50=50000$, thus
leading to almost identical performance as training on the full dataset.
Counter-intuitively, but for the reasons outlined above, training under random
labels severely suffers under such ideal, independent augmentations and leads
to malign memorization. In fact, this experiment is equivalent to fitting
random labels without augmentations on the full training set where malign
memorization occurs, as seen in Sec. 4.

This thought experiment demonstrates that under random labels, augmentations
seem to play a very distinct role compared to the clean setting and label
preservation in itself is not enough to guarantee benign memorization. This
raises the following question:

What other properties of augmentations, besides label preservation, cause
benign memorization?

Figure 5: Illustration of the thought experiment in Sec. 5. The true cat
$\bm{x}_{1}$ labeled as “giraffe” is augmented, leading to a positive signal
between the augmentations but to more noise due to $\bm{x}_{2}$, that has as
true label “dog”.

Dataset Size | $1000$ | $50000$
---|---|---
Random | $11.41$ | $12.41$
Random + i.i.d. DA | $12.41$ | -
Clean | $36.75$ | 83.66
Clean + i.i.d. DA | $82.73$ | -

Figure 6: $K$-NN probing accuracies of the embeddings of a ResNet18 for
CIFAR10 with and without i.i.d. augmentations under clean and random labels.
$50000$ refers to full CIFAR10 for reference.

We hypothesize that the origins of the phenomenon lie at the interplay between
the highly correlated nature of augmentations and the inflated effective
sample size that exceeds the model capacity (Sec 5.1), forcing the model to
learn meaningful features (Sec 5.2).

###  5.1 Going Beyond the Model Capacity

We now study the phenomenon of benign memorization from the view point of
capacity of a model class. Intuitively speaking, the capacity of a model
captures how many distinct datapoints we can fit, even if the corresponding
targets are completely randomized. If a model has enough capacity, it can
potentially memorize all the patterns in the dataset. As seen in Zhang et al.
(2017), deep models used in practice in conjunction with standard datasets,
operate below the capacity threshold but nevertheless, they do not “abuse”
their power and fit clean data in a meaningful way. On the other hand, by
using data augmentation we inflate the number of samples and consequently
operate above the capacity level. As a result, a model needs to efficiently
encode augmentations (pure memorization is not possible) and hence meaningful
features need to be learnt. We start by mathematically introducing the concept
of capacity of a model class trained with Gradient Descent (GD). We remark
that similar definitions have been explored in prior work (Arpit et al.,
2017).

###### Definition 5.1

Consider $\bm{x}_{1},\dots,\bm{x}_{n}$ for $n\in\mathbb{N}$ in general
position and a fixed number of classes $C$. Given a labeling
$\tilde{\mathcal{S}}=\\{\left(\bm{x}_{i},\tilde{\bm{y}}_{i}\right)\\}_{i=1}^{n}$,
let $\mathcal{A}_{t}(\tilde{\mathcal{S}})\subset\Theta$ denote the set of
solutions reachable by GD with a computational budget $t$. We define the
capacity $\mathcal{C}_{t}$ of the model class
$\\{f_{\bm{\theta}}:\bm{\theta}\in\Theta\\}$ as

$\mathcal{C}_{t}:=\text{sup}\\{n\in\mathbb{N}|~{}\forall\tilde{\mathcal{S}}~{}\exists\bm{\theta}\in\mathcal{A}_{t}(\tilde{\mathcal{S}})\text{
s.t. }f_{\bm{\theta}}\text{ memorizes }\tilde{\mathcal{S}}\\}.$

While similar to standard measures such as the VC-dimension (Vapnik &
Chervonenkis, 1971), the capacity defined here depends on the learning
algorithm (including the computational budget $t$) and is thus always smaller
than the corresponding VC-dimension. As seen in Sec. 3, standard datasets such
as CIFAR10 have size $n$ below the capacity $\mathcal{C}_{t}$ of modern
networks such as ResNets. When using data augmentation however, we now show
that the resulting dataset exceeds the capacity.

Inflated Sample Size. Consider a set of augmentations $\mathcal{T}$ and assume
that it has a finite size, $B:=|\mathcal{T}|<\infty$. Augmenting thus leads to
a larger effective dataset $\mathcal{S}_{\text{aug}}$ of size $Bn$. We can now
study whether the capacity of standard networks trained by GD for a fixed
number of epochs exceeds such an augmented dataset by varying $B$ and randomly
labeling each sample. We pre-compute a fixed set of $B$ random augmentations
for each sample for varying $B$ and attempt to memorize a random labeling,
where labels of augmentations are not preserved. We use the same setup as in
Sec. 4 and augment CIFAR10 using the standard augmentations in the way
described above. We display the results in Fig. 7. We see that indeed,
overfitting the random labels becomes more difficult as we increase $B$ and
actually infeasible for a fixed number of gradient steps $t=5000\times 196$
(fixed batch size on CIFAR10). Moreover, notice that in the standard setting
of online augmentations, this observation becomes even more drastic as
effectively $B\xrightarrow[]{}\infty$ if typical, continuous augmentations
such as color jittering are included in the pipeline. We hypothesize that
malign memorization under such an augmentation strategy thus becomes
infeasible.

Figure 7: Illustration of nearest-neighbour probing accuracy (left) and
training loss (right) of a ResNet18 on CIFAR10 with increasing number of
augmentations. We consider both label-preserving (blue) and completely random
labeling of augmentations (red).

Learn if you must. We now investigate the influence of capacity on the
resulting probing accuracy, in the case where augmentations are label-
preserving and thus offer valuable signal. We thus consider the same setup as
in the previous paragraph for varying number of augmentations $B$ while
assigning the same label to augmentations of the same image. We show the
results in Fig. 7 on the left. We observe that as the number of augmentations
$B$ increases, probing accuracy improves and eventually surpasses probing at
initialization. We hypothesize that as we approach and eventually surpass
capacity, the model is increasingly forced to leverage the signal in
augmentations and thus learns more and more relevant features. As we increase
B, we saturate the information gain provided by the augmentations, the signal
becomes redundant and performance starts to plateau.

###  5.2 What can you Learn from Augmentations?

While we have seen that the model is forced to leverage the signal in
augmentations, it remains unclear why this leads to high-quality embeddings.
We now show how augmentations encourage features to become invariant, a
property that has been identified in SSL to be strongly predictive.

Normalized Invariance. To measure the invariance of a model $f_{\bm{\theta}}$,
we introduce the following quantity:

$I(\bm{x};f_{\bm{\theta}},\mathcal{T}):=\frac{\mathbb{E}_{T_{1},T_{2}\sim\mathcal{U}(\mathcal{T})}\lVert
f_{\bm{\theta}}(T_{1}(\bm{x}))-f_{\bm{\theta}}(T_{2}(\bm{x}))\rVert}{\mathbb{E}_{\bm{x}^{\prime}\neq\bm{x}}\lVert
f_{\bm{\theta}}(\bm{x})-f_{\bm{\theta}}(\bm{x^{\prime}})\rVert}.$ (2)

Intuitively, $I(\bm{x};f_{\bm{\theta}},\mathcal{T})$ captures the features’
similarity of augmentations of the same datapoint $\bm{x}$ compared to the
representations of different datapoints $\bm{x}^{\prime}\neq\bm{x}$. Model’s
invariance implies $I(\bm{x};f_{\bm{\theta}},\mathcal{T})=0$, hence different
augmentations are mapped to the same datapoint, while the model is still able
to meaningfully distinguish between the representations of different
datapoints (i.e. $\mathbb{E}_{\bm{x}^{\prime}\neq\bm{x}}\lVert
f_{\bm{\theta}}(\bm{x})-f_{\bm{\theta}}(\bm{x^{\prime}})\rVert\neq 0$). In
Fig. 8, we show how $I(\bm{x};f,\mathcal{T})$ correlates with probing
performance and indeed decreases over time. Due to its better implicit bias
(ResNet vs MLP), invariance is largely learnt in the encoder, leading to the
striking signal-noise separation identified in Sec. 4. These results suggest
that when the model has insufficient capacity to memorize the (augmented)
samples, it learns to be more invariant with respect to the label-preserving
augmentations. Consequently, this mechanism reduces the “effective sample
size” and allows the model to fit the data in a _benign_ (i.e. augmentation-
invariant) way. But why do invariant features imply a better clustering in
embedding space as evidenced by a high $K$-NN accuracy? This is a heavily
researched topic with several plausible theories in the area of self-
supervised learning (SSL) (Saunshi et al., 2022; Arora et al., 2019b; Wen &
Li, 2021). In Appendix A, we derive a more formal connection between SSL loss
and training with random labels. Others have looked at the improved sample
complexity caused by incorporating invariances into the model (Bietti et al.,
2021; Xiao & Pennington, 2022).

Figure 8: Normalized invariance of a ResNet18 on CIFAR10, as a function of
epochs, lower value means more invariant. The dashed line represents $K$-NN
probing error (scaled for comparability).

Method | Normalized invariance
---|---
Initialization | 0.867
Clean | 0.912
Clean + DA | 0.691
Random | 0.989
Random + DA | 0.680

##  6 Discussion and Conclusion

In this work we have identified the surprising phenomenon of benign
memorization, demonstrating that generalization and memorization are not
necessarily at odds. We put forward an interpretation through the lens of
model capacity, where the inflation in sample size forces the model to exploit
the correlation structure by learning the invariance with respect to the
augmentations. Furthermore, we have shown that invariance learning happens
largely in the encoder, while the projector performs the noisy memorization
task.

Our findings underline the complicated and mysterious inner workings of neural
networks, showing that generalization can be found where we least expect it.
To describe benign memorization, a complete generalization theory needs to
capture the strong implicit bias built into deep models, which enables a clean
separation of noise and signal. Secondly, such a theory needs to incorporate
feature learning, as benign memorization only emerges in the encoder. Both
those goals remain very challenging. In particular, we speculate that the line
of work based on the neural tangent kernel (Jacot et al., 2018) which connects
SGD-optimized neural networks in the infinite width regime and the realm of
kernels cannot explain benign memorization. In fact, in this regime the
weights do not move from initialization, thus preventing feature learning. On
the other hand, approaches that aim at computing perturbative finite width
corrections to these kernels (Hanin & Nica, 2019; Zavatone-Veth et al., 2021)
might be a promising approach to explain the phenomenon of benign
memorization, as they incorporate feature learning. We leave exploring benign
memorization in such a mathematical framework as future work.

##  7 Reproducibility Statement

We refer the reader to Appendix D for a complete description of the training
protocol.

## References

  * Alain & Bengio (2017) Guillaume Alain and Yoshua Bengio.  Understanding intermediate layers using linear classifier probes.  _International Conference on Learning Representations (ICLR)_ , abs/1610.01644, 2017.
  * Arora et al. (2018) Sanjeev Arora, Rong Ge, Behnam Neyshabur, and Yi Zhang.  Stronger generalization bounds for deep nets via a compression approach.  _Proceedings of the 35th International Conference on Machine Learning (ICML)_ , 2018.
  * Arora et al. (2019a) Sanjeev Arora, Simon S. Du, Wei Hu, Zhiyuan Li, and Ruosong Wang.  Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks.  _Proceedings of the 36th International Conference on Machine Learning (ICML_ , 2019a.
  * Arora et al. (2019b) Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi.  A theoretical analysis of contrastive unsupervised representation learning.  _arXiv preprint arXiv:1902.09229_ , 2019b.
  * Arpit et al. (2017) Devansh Arpit, Stanisław Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, et al.  A closer look at memorization in deep networks.  In _International conference on machine learning_ , pp. 233–242. PMLR, 2017.
  * Bardes et al. (2022) Adrien Bardes, Jean Ponce, and Yann LeCun.  VICReg: Variance-invariance-covariance regularization for self-supervised learning.  In _International Conference on Learning Representations (ICLR)_ , 2022.
  * Bartlett et al. (2017) Peter Bartlett, Dylan J. Foster, and Matus Telgarsky.  Spectrally-normalized margin bounds for neural networks.  _31st Conference on Neural Information Processing Systems (Neurips)_ , 2017.
  * Bartlett et al. (2019) Peter L. Bartlett, Nick Harvey, Chris Liaw, and Abbas Mehrabian.  Nearly-tight vc-dimension and pseudodimension bounds for piecewise linear neural networks.  _Journal of Machine Learning Research 20_ , pp. 1–17, 2019.
  * Bartlett et al. (2020) Peter L. Bartlett, Philip M. Long, Gá bor Lugosi, and Alexander Tsigler.  Benign overfitting in linear regression.  _Proceedings of the National Academy of Sciences_ , 117(48):30063–30070, apr 2020.  doi: 10.1073/pnas.1907378117.
  * Bietti et al. (2021) Alberto Bietti, Luca Venturi, and Joan Bruna.  On the sample complexity of learning under geometric stability.  _Advances in Neural Information Processing Systems_ , 34:18673–18684, 2021.
  * Caron et al. (2021) Mathilde Caron, Hugo Touvron, Ishan Misra, Herv’e J’egou, Julien Mairal, Piotr Bojanowski, and Armand Joulin.  Emerging properties in self-supervised vision transformers.  _2021 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pp. 9630–9640, 2021.
  * Chen et al. (2020a) Mark Chen, Alec Radford, Jeff Wu, Heewoo Jun, Prafulla Dhariwal, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In _Proceedings of the 37th International Conference on Machine Learning (ICML)_ , 2020a.
  * Chen et al. (2020b) Shuxiao Chen, Edgar Dobriban, and Jane Lee.  A group-theoretic framework for data augmentation.  In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 21321–21333. Curran Associates, Inc., 2020b.
  * Chen et al. (2020c) Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton.  A simple framework for contrastive learning of visual representations.  In Hal Daumé III and Aarti Singh (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 1597–1607. PMLR, 13–18 Jul 2020c.
  * Chen & He (2021) Xinlei Chen and Kaiming He.  Exploring simple siamese representation learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 15750–15758, 2021.
  * Coates et al. (2011) Adam Coates, Andrew Ng, and Honglak Lee.  An analysis of single-layer networks in unsupervised feature learning.  In Geoffrey Gordon, David Dunson, and Miroslav Dudík (eds.), _Proceedings of the Fourteenth International Conference on Artificial Intelligence and Statistics (AISTATS)_ , volume 15 of _Proceedings of Machine Learning Research_ , pp. 215–223, Fort Lauderdale, FL, USA, 11–13 Apr 2011. PMLR.
  * Cohen et al. (2018) Gilad Cohen, Guillermo Sapiro, and Raja Giryes.  Dnn or k-nn: That is the generalize vs. memorize question.  _ArXiv_ , abs/1805.06822, 2018.
  * Cubuk et al. (2020) Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le.  Randaugment: Practical automated data augmentation with a reduced search space.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops_ , pp. 702–703, 2020.
  * Dai et al. (2021) Zihang Dai, Hanxiao Liu, Quoc V Le, and Mingxing Tan.  Coatnet: Marrying convolution and attention for all data sizes.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , 2021.
  * Dao et al. (2019) Tri Dao, Albert Gu, Alexander Ratner, Virginia Smith, Chris De Sa, and Christopher Re.  A kernel theory of modern data augmentation.  In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 1528–1537. PMLR, 09–15 Jun 2019.
  * Deng et al. (2009) Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _2009 IEEE Conference on Computer Vision and Pattern Recognition_ , pp. 248–255, 2009.  doi: 10.1109/CVPR.2009.5206848.
  * Devlin et al. (2019) Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _Proceedings of NAACL-HLT_ , 2019.
  * Dosovitskiy et al. (2014) Alexey Dosovitskiy, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox.  Discriminative unsupervised feature learning with convolutional neural networks.  _Advances in neural information processing systems_ , 27, 2014.
  * Dziugaite & Roy (2017) Gintare Karolina Dziugaite and Daniel M. Roy.  Computing nonvacuous generalization bounds for deep (stochastic) neural networks with many more parameters than training data.  _Proceedings of the Thirty-Third Conference on Uncertainty in Artificial Intelligence_ , 2017.
  * Feldman & Zhang (2020) Vitaly Feldman and Chiyuan Zhang.  What neural networks memorize and why: Discovering the long tail via influence estimation.  _Advances in Neural Information Processing Systems_ , 33:2881–2891, 2020.
  * Frei et al. (2022) Spencer Frei, Niladri S Chatterji, and Peter Bartlett.  Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data.  In Po-Ling Loh and Maxim Raginsky (eds.), _Proceedings of Thirty Fifth Conference on Learning Theory_ , volume 178 of _Proceedings of Machine Learning Research_ , pp. 2668–2703. PMLR, 02–05 Jul 2022.
  * Grill et al. (2020) Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukcuoglu, Remi Munos, and Michal Valko.  Bootstrap your own latent - a new approach to self-supervised learning.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 21271–21284. Curran Associates, Inc., 2020.
  * Hanin & Nica (2019) Boris Hanin and Mihai Nica.  Finite depth and width corrections to the neural tangent kernel.  _arXiv preprint arXiv:1909.05989_ , 2019.
  * Hanin & Sun (2021) Boris Hanin and Yi Sun.  How data augmentation affects optimization for linear regression.  In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , 2021.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  _IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , 2016.
  * Hua et al. (2021a) Tianyu Hua, Wenxiao Wang, Zihui Xue, Sucheng Ren, Yue Wang, and Hang Zhao.  On feature decorrelation in self-supervised learning.  In _Proceedings of the IEEE/CVF International Conference on Computer Vision_ , pp. 9598–9608, 2021a.
  * Hua et al. (2021b) Tianyu Hua, Wenxiao Wang, Zihui Xue, Yue Wang, Sucheng Ren, and Hang Zhao.  On feature decorrelation in self-supervised learning.  _2021 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pp. 9578–9588, 2021b.
  * Jacot et al. (2018) Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  _Advances in neural information processing systems_ , 31, 2018.
  * Kipf & Welling (2017) Thomas N. Kipf and Max Welling.  Semi-supervised classification with graph convolutional networks.  _Proceedings of the 5th International Conference on Learning Representations_ , 2017.
  * Krizhevsky & Hinton (2009) Alex Krizhevsky and Geoffrey Hinton.  Learning multiple layers of features from tiny images.  Technical Report 0, University of Toronto, Toronto, Ontario, 2009.
  * Le & Yang (2015) Ya Le and Xuan S. Yang.  Tiny imagenet visual recognition challenge.  2015\.
  * Liu et al. (2021) Ze Liu, Han Hu, Yutong Lin, Zhuliang Yao, Zhenda Xie, Yixuan Wei, Jia Ning, Yue Cao, Zheng Zhang, Li Dong, Furu Wei, and Baining Guo.  Swin transformer v2: Scaling up capacity and resolution, 2021.
  * Ma et al. (2018) Siyuan Ma, Raef Bassily, and Mikhail Belkin.  The power of interpolation: Understanding the effectiveness of SGD in modern over-parametrized learning.  In Jennifer Dy and Andreas Krause (eds.), _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 3325–3334. PMLR, 10–15 Jul 2018.
  * Maennel et al. (2020) Hartmut Maennel, Ibrahim M Alabdulmohsin, Ilya O Tolstikhin, Robert Baldock, Olivier Bousquet, Sylvain Gelly, and Daniel Keysers.  What do neural networks learn when trained with random labels?  _Advances in Neural Information Processing Systems_ , 33:19693–19704, 2020.
  * Neyshabur et al. (2015) Behnam Neyshabur, Ryota Tomioka, and Nathan Srebro.  Norm-based capacity control in neural networks.  _Proceedings of The 28th Conference on Learning Theory (PMLR)_ , 2015\.
  * Neyshabur et al. (2018) Behnam Neyshabur, Srinadh Bhojanapalli, and Nathan Srebro.  A pac-bayesian approach to spectrally-normalized margin bounds for neural networks.  _International Conference on Learning Representations (ICLR)_ , 2018\.
  * Nishi et al. (2021) Kento Nishi, Yi Ding, Alex Rich, and Tobias Höllerer.  Augmentation strategies for learning with noisy labels.  _2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 8018–8027, 2021.
  * Patrini et al. (2017) Giorgio Patrini, Alessandro Rozza, Aditya Krishna Menon, Richard Nock, and Lizhen Qu.  Making deep neural networks robust to label noise: A loss correction approach.  _2017 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 2233–2241, 2017.
  * Pham et al. (2021) Hieu Pham, Zihang Dai, Qizhe Xie, Minh-Thang Luong, and Quoc V. Le.  Meta pseudo labels.  In _IEEE Conference on Computer Vision and Pattern Recognition_ , 2021\.
  * Pondenkandath et al. (2018) Vinaychandran Pondenkandath, Michele Alberti, Sammer Puran, Rolf Ingold, and Marcus Liwicki.  Leveraging random label memorization for unsupervised pre-training.  _Workshop of Integration of Deep Learning Theories at Conference on Neural Information Processing Systems (NIPS)_ , abs/1811.01640, 2018.
  * Rolnick et al. (2017) David Rolnick, Andreas Veit, Serge Belongie, and Nir Shavit.  Deep learning is robust to massive label noise.  _arXiv preprint arXiv:1705.10694_ , 2017.
  * Sanyal et al. (2021) Amartya Sanyal, Puneet K. Dokania, Varun Kanade, and Philip Torr.  How benign is benign overfitting?  In _International Conference on Learning Representations_ , 2021.
  * Saunshi et al. (2022) Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy.  Understanding contrastive learning requires incorporating inductive biases.  _Proceedings of the 39th International Conference on Machine Learning, (ICML)_ , 2022.
  * Simonyan & Zisserman (2014) Karen Simonyan and Andrew Zisserman.  Very deep convolutional networks for large-scale image recognition.  _International Conference on Learning Representations (ICLR)_ , 2014\.
  * Song et al. (2020) Hwanjun Song, Minseok Kim, Dongmin Park, Yooju Shin, and Jae-Gil Lee.  Learning from noisy labels with deep neural networks: A survey, 2020.
  * Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov.  Dropout: A simple way to prevent neural networks from overfitting.  _Journal of Machine Learning Research_ , 15(56):1929–1958, 2014.
  * Vapnik & Chervonenkis (1971) V. N. Vapnik and A. Ya. Chervonenkis.  On the uniform convergence of relative frequencies of events to their probabilities.  _Theory of Probability and its Applications_ , 16(2):264–280, 1971.  doi: 10.1137/1116025.
  * Wen & Li (2021) Zixin Wen and Yuanzhi Li.  Toward understanding the feature learning process of self-supervised contrastive learning.  In _International Conference on Machine Learning_ , pp. 11112–11122. PMLR, 2021.
  * Wu et al. (2020) Sen Wu, Hongyang R. Zhang, Gregory Valiant, and Christopher Ré.  On the generalization effects of linear transformations in data augmentation.  In _Proceedings of the 37th International Conference on Machine Learning (ICML)_ , 2020.
  * Xiao & Pennington (2022) Lechao Xiao and Jeffrey Pennington.  Synergy and symmetry in deep learning: Interactions between the data, model, and inference algorithm.  _arXiv preprint arXiv:2207.04612_ , 2022.
  * Yu et al. (2022) Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu.  Coca: Contrastive captioners are image-text foundation models.  _Transactions on Machine Learning Research_ , 2022.
  * Zavatone-Veth et al. (2021) Jacob Zavatone-Veth, Abdulkadir Canatar, Ben Ruben, and Cengiz Pehlevan.  Asymptotics of representation learning in finite bayesian neural networks.  _Advances in neural information processing systems_ , 34:24765–24777, 2021.
  * (58) Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny.  Barlow twins: Self-supervised learning via redundancy reduction.  In _Proceedings of the 38th International Conference on Machine Learning (ICML_.
  * Zhai et al. (2021) Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer.  Scaling vision transformers, 2021.
  * Zhang et al. (2017) Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning requires rethinking generalization.  _International Conference on Learning Representations (ICLR)_ , 2017\.
  * Zhang et al. (2018) Hongyi Zhang, Moustapha Cisse, Yann N. Dauphin, and David Lopez-Paz.  mixup: Beyond empirical risk minimization.  In _International Conference on Learning Representations_ , 2018.
  * Zhang et al. (2021) Xuanyang Zhang, Pengfei Hou, X. Zhang, and Jian Sun.  Neural architecture search with random labels.  _2021 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)_ , pp. 10902–10911, 2021.

##  Appendix A Connection to Self-Supervised Learning

In this section, we investigate the connection between non-contrastive SSL
(Hua et al., 2021b; Grill et al., 2020; Chen & He, 2021) and training with
random labels on the mean squared error (MSE) loss. We start with the
following result:

###### Lemma A.1

Fix $B\in\mathbb{N}$ vectors $\bm{x}_{1},\dots,\bm{x}_{B}\in\mathbb{R}^{d}$
and $\bm{a}\in\mathbb{R}^{d}$. Then it holds that

$\frac{1}{B}\sum_{i=1}^{B}\|\bm{x}_{i}-\bm{a}\|^{2}=\frac{1}{2B^{2}}\sum_{i,j=1}^{B}\|\bm{x}_{i}-\bm{x}_{j}\|^{2}+\|\bm{a}-\frac{1}{B}\sum_{i=1}^{B}\bm{x}_{i}\|^{2}$

Proof: We simply expand the first term on the right-hand-side:

$\begin{split}\frac{1}{2B^{2}}\sum_{i,j=1}^{B}\|\bm{x}_{i}-\bm{x}_{j}\|^{2}&=\frac{1}{2B^{2}}\sum_{i,j=1}^{B}\|\bm{x}_{i}-\bm{a}+\bm{a}-\bm{x}_{j}\|^{2}\\\
&=\frac{1}{2B^{2}}\left(\sum_{i,j=1}^{B}\|\bm{x}_{i}-\bm{a}\|^{2}+\|\bm{x}_{j}-\bm{a}\|^{2}-2(\bm{x}_{i}-\bm{a})^{T}(\bm{x}_{j}-\bm{a})\right)\\\
&=\frac{1}{B}\sum_{i}^{B}\|\bm{x}_{i}-\bm{a}\|^{2}-\frac{1}{B^{2}}\sum_{i,j=1}^{B}(\bm{x}_{i}-\bm{a})^{T}(\bm{x}_{j}-\bm{a})\\\
&=\frac{1}{B}\sum_{i}^{B}\|\bm{x}_{i}-\bm{a}\|^{2}-\|\bm{a}\|^{2}-\frac{1}{B^{2}}\sum_{i,j=1}^{n}\bm{x}_{i}^{T}\bm{x}_{j}+\frac{2}{B}\sum_{i=1}^{n}\bm{a}^{T}\bm{x}_{i}\end{split}$

On the other hand, we have that

$\begin{split}\|\bm{a}-\frac{1}{B}\sum_{i=1}^{B}\bm{x}_{i}\|^{2}&=\|\bm{a}\|^{2}+\frac{1}{B^{2}}\sum_{i,j=1}^{n}\bm{x}_{i}^{T}\bm{x}_{j}-\frac{2}{B}\sum_{i=1}^{n}\bm{a}^{T}\bm{x}_{i}\end{split}$

Hence we see that

$\begin{split}\frac{1}{2B^{2}}\sum_{i,j=1}^{B}\|\bm{x}_{i}-\bm{x}_{j}\|^{2}&=\frac{1}{B}\sum_{i}^{B}\|\bm{x}_{i}-\bm{a}\|^{2}-\|\bm{a}-\frac{1}{B}\sum_{i=1}^{B}\bm{x}_{i}\|^{2}\end{split}$

and re-arranging terms concludes the proof.
We now apply this result in the case where the label plays the role of
$\bm{a}$ and $\bm{x}_{i}$’s play the role of different augmentations. Let us
thus consider $B$ augmentations $T_{1},\dots,T_{B}$ and $n\in\mathbb{N}$
samples $\bm{x}_{1},\dots,\bm{x}_{n}\in\mathbb{R}^{d}$. Moreover we have some
labels $\bm{y}_{1},\dots,\bm{y}_{n}\in\mathbb{R}^{K}$ that could be completely
random. Using the previous result we can write the (random) supervised loss
(assuming mean-squared error) as follows:

###### Theorem A.2

Denote the supervised loss under data augmentation as

$\hat{L}^{\text{super}}(\bm{\theta})=\frac{1}{nB}\sum_{i=1}^{n}\sum_{a=1}^{B}\big{\|}f_{\bm{\theta}}(T_{a}(\bm{x}_{i}))-\bm{y}_{i}\big{\|}^{2}$

Then we can decompose it into the following two terms,

$\hat{L}^{\text{super}}(\bm{\theta})=\underbrace{\frac{1}{2nB^{2}}\sum_{i=1}^{n}\sum_{a,b=1}^{B}\Big{\|}f_{\bm{\theta}}(T_{a}(\bm{x}_{i}))-f_{\bm{\theta}}(T_{b}(\bm{x}_{i}))\Big{\|}^{2}}_{=:\operatorname{Inv}(\bm{\theta})\geq
0}+\underbrace{\frac{1}{n}\sum_{i=1}^{n}\Big{\|}\bm{y}_{i}-\frac{1}{B}\sum_{a=1}^{B}f_{\bm{\theta}}(T_{a}(\bm{x}_{i}))\Big{\|}^{2}}_{=:\operatorname{Bias}(\bm{\theta})\geq
0}$

Notice that in $\operatorname{Inv}(\bm{\theta})$, we group augmentations of
the same input $\bm{x}_{i}$ together, measuring thus how invariant a given
model $f_{\bm{\theta}}$ is w.r.t. to the augmentations. This is the positive
signal illustrated in the thought experiment in Fig. 6. Interestingly,
minimizing $\operatorname{Inv}(\bm{\theta})$ is a core ingredient for so-
called non-contrastive self-supervised learning methods (Hua et al., 2021b;
Grill et al., 2020; Chen & He, 2021).

The bias term on the other hand is influenced by the random labeling. To
better understand the bias term, define
$N_{c}:=\\{i\in\\{1,\dots,n\\}:\bm{y}_{i}=\bm{e}_{c}\\}$, i.e. all samples
that have label $c$. Furthermore, let
$\bar{f}_{\bm{\theta}}(\bm{x}_{i})=\frac{1}{B}\sum_{a=1}^{B}f_{\bm{\theta}}(T_{a}(\bm{x}_{i}))$
be the model average over all the $B$ augmentations. We can show that we can
express the bias as

$\operatorname{Bias}(\bm{\theta})=\frac{1}{n}\sum_{c=1}^{C}\sum_{i\in
N_{c}}\Big{\|}\bm{e}_{c}-\bar{f}_{\bm{\theta}}(\bm{x}_{i})\Big{\|}^{2}.$

Notice that as in a random labeling experiment, the number of classes $C$ can
be considered as a hyperparameter, hence if we choose $C>>n$, then each sample
$\bm{x}_{1},\dots,\bm{x}_{n}$ is assigned a different class with high
probability. In this case, the “bad bias” given by assigning the same label to
samples of different classes vanishes, and one has:

$\operatorname{Bias}(\bm{\theta})=\frac{1}{n}\sum_{i=1}^{n}\Big{\|}\bm{e}_{i}-\bar{f}_{\bm{\theta}}(\bm{x}_{i})\Big{\|}^{2},$

where w.l.o.g we have re-ordered the labels such that the $i$-th label is
assigned to sample $\bm{x}_{i}$. In this limiting case, the bias term
represents a contrastive term that encourages the average model
$\bar{f}_{\bm{\theta}}(\bm{x}_{i})$ to be different from the other samples
$\bm{x}_{j}\neq\bm{x}_{i}$ and their augmentations. In this sense, this term
can be related to contrastive SSL. On the other hand — unlike contrastive SSL
— the distance between the average models of two datapoints cannot be
arbitrarily large, and once the loss is driven to zero,
$\operatorname{Bias(\bm{\theta})}=0$ and it is trivial to see that
$\Big{\|}\bar{f}_{\bm{\theta}}(\bm{x}_{i})-\bar{f}_{\bm{\theta}}(\bm{x}_{j})\Big{\|}^{2}=2$.
We refer the reader to Appendix B for experiments with varying number of
classes. We observe that, as hypothesized, the influence of the bias
diminishes and performance increases as a function of the number of classes
$C$.

##  Appendix B More Experiments

In the following we present more empirical evidence that may help better
interpret our findings.

Role of the Projector. We visualize in Fig. 9 the effect of varying the size
of the hidden layer in the projector. There is a threshold under which
learning is impossible (or at least very hard given a fixed computational
budget). Surprisingly, using a larger projector seems beneficial beyond the
point where we manage to decrease the loss significantly and achieve
memorization of the original dataset.

Figure 9: Nearest-neighbour probing accuracy as function of the size of the
projector.

The role of the number of random classes. So far we assigned to each sample a
random label in $\mathbb{R}^{C}$, where $C$ is the number of classes in the
original dataset. In principle, increasing the number of classes decreases the
strength of the noise provided to the model. We verify this in Fig. 10, where
we vary the number of possible classes that each sample is randomly assigned
to. In the extreme case, every sample is assigned its unique class label
$\bm{y}_{i}=\bm{e}_{i}$, where $\bm{e}_{i}\in\mathbb{R}^{n}$, recovering the
result of Dosovitskiy et al. (2014) which we refer to as ‘Per sample’. We
observe that while increasing number of classes, performance improves (as less
noisy assignments are made). We refer the reader to Appendix A for a more
formal connection.

Figure 10: Nearest-neighbour probing accuracy as a function of the number of
classes.

Transferability. Features learnt with random labels are also useful when
evaluated on different datasets. In Fig. 11 on the right-hand-side, we
transfer the features learnt on CIFAR10, to CIFAR100 and STL10 (Coates et al.,
2011). We observe that we can achieve strong downstream performance,
highlighting the strength of the features learnable under random labels.
Moreover, we see that using data augmentation is crucial when label noise is
high as without data augmentation

Figure 11: Nearest-neighbour probing accuracy for transfer learning to
different datasets as a function of label noise. A ResNet18 is trained based
on random labels on CIFAR10 with (right) and without (left) data augmentation.

Dependence on noise and speed of convergence. Augmentations have been shown to
be especially useful under the presence of heavy label noise. In Fig. 12 we
visualize the performance achieved during the first stage of training under
varying level of label noise. Label noise specifies the percentage of labels
that are randomly permuted versus the percentage of labels that are kept the
same.

(a)

(b)

Figure 12: Nearest-neighbour probing during training for the CIFAR10 and
TinyImageNet datasets during training.

Effect on the strength of the augmentation. Although malign memorization gets
impossible as the number of augmentations $B\to\infty$, the strength of the
augmentations themselves plays a role on the generalization achieved. The
strength of the augmentation basically dictates the strength of the invariance
that the model has to learn, which also correlates with the quality of the
features learned as seen in Fig. 13. To control strength we use strength level
of RandAugment (Cubuk et al., 2020). We clearly see that stronger
augmentations are beneficial for the quality of the embeddings, underlining
the fact that invariance indeed plays a key role in benign memorization.

Figure 13: Nearest-neighbour probing accuracy when trained on random labels
with varying strength of (infinite) augmentations.

Varying levels of label noise. Previous works have highlighted the importance
of augmentations in the presence of high label noise. We examine the
relationship between label noise and generalization both with and without
augmentations in Fig. 14. We see that while data augmentation is not so
crucial for clean labels, its role becomes more and more critical as we
increase the amount of label noise. Without augmentations, we suffer from a
strong decrease in performance, eventually ending with random guessing as we
reach complete label noise. On the other hand, using augmentations stabilizes
this decay and we still achieve strongly non-trivial performance even in the
setting of complete label noise.

Figure 14: Nearest-neighbour probing accuracy when trained on varying levels
of label noise with and without augmentations.

##  Appendix C Toy example

To better understand the role of augmentations with respect to the capacity of
the model we devise a simple toy setting.

We consider samples $\bm{x_{i}},\dots,\bm{x_{n}}\in\mathbb{R}^{d}$, that
belong to one of the $C$ underlying classes. We select the first $d_{1}<d$
coordinates to denote the true class assignment $z_{i}$, uniformly selected
from the set $\\{1,\dots,C\\}$, by sampling from a mixture of Gaussian
distributions and let the rest $d-d_{1}$ dimensions correspond to noise. More
specifically:

$\displaystyle\bm{x_{i[:d_{1}]}}|_{z_{i}=k}$
$\displaystyle\sim\mathcal{N}(\bm{\mu}_{i},\bm{\Sigma_{1}}),$
$\displaystyle\bm{x_{i[d_{1}:]}}$
$\displaystyle\sim\mathcal{N}(\bm{0},\bm{\Sigma_{2}}).$

With $\bm{\mu}_{i}$ we denote the cluster centers that are sampled randomly.
In this case, we consider augmentations of the samples:

$\displaystyle\bar{\bm{x}}=\bm{x}+\bm{a},$ where
$\displaystyle\bm{a}_{[:d_{1}]}=\bm{0}$ and
$\displaystyle\bm{a}_{[d_{1}:]}\sim\mathcal{N}(\bm{0},\bm{\Sigma_{3}}).$

As set the covariance matrices as
$\Sigma_{3}=\Sigma_{2}=10\,\Sigma_{1}=\sigma^{2}\bm{I}$, where $\bm{I}$ is the
identity matrix. It is obvious that learning invariance for this task under
these augmentations leads to embeddings that ignore the noise subspace. We
employ in this case a simple linear encoder $h_{\bm{W}}(\bm{x})=\bm{W}\bm{x}$,
with $\bm{W}\in\mathbb{R}^{d\times d}$.

As a projector we use
$g_{\bm{V}}(\bm{x})=\sum_{i=1}^{K}\cfrac{\cfrac{1}{\|\bm{x}-\bm{v_{i}}\|}}{\sum_{j=1}^{K}\cfrac{1}{\|\bm{x}-\bm{v_{j}}\|}}\,\bm{l}_{i}$.
Here $\bm{l_{i}}$ is a one hot encoding sampled randomly from the set
$\\{\bm{e}_{1},\dots,\bm{e}_{C}\\}$. We choose such a projector as it is not
difficult to verify that for $K>=d$ the capacity of the model, in this case,
will be approximately fixed to the number of $K$ patterns that we compare the
embedding space against. Additionally, its non-linear nature guarantees that
any possible invariance learning will occur at the encoder instead. In Fig. 15
we visualize the $K$-NN probing for the downstream task of correctly
predicting the clean label of unseen test samples, after training on random
labels. By varying the number of samples $n$ and augmentations $B$ available,
we directly control whether full memorization can take place or not. When the
number of samples $n$ is smaller than $\frac{\mathcal{C}}{B}$, where
$\mathcal{C}$ is the capacity of the model, full memorization is possible,
which discourages any feature learning, leading to bad generalization.

Figure 15: $K$-NN probing for the downstream class assignment on unseen test
points for the toy example, after training on random labels.

##  Appendix D Experimental Setup

As commonly done for smaller size images (Chen & He, 2021; Hua et al., 2021a),
we use a variant of the ResNet18 architecture, where the max-pooling layer is
removed, and the first convolution is modified to have a kernel size of $3$
and a stride of $1$. We also remove the last fully connected layer, as we
replace it with our projector. We provide more details about the
hyperparameters used in Table 2. For the experiments on TinyImageNet, we
rescale images to a size of $64\times 64$ and additionally restore the stride
of the first convolution to the value $2$. We apply the same modification to
VGG when trained on TinyImageNet.

| Hyperparameters | Value
---|---|---
Augmentations | Random cropping scale | (0.08, 1)
Horizontal flip probability | 0.5
Color-Jittering | (0.8, 0.8, 0.8, 0.2)
Grayscale probability | 0.2
Mixup | yes
Parameters | image-size | 32/64
clip-norm | ‘none’
dropout | ‘none’
Projector size | 65536
Projector MLP normalization | ‘none’
Training | Learning rate | $4\epsilon^{-3}$
Learning rate scheduler | ‘none’
Adam $(\beta_{1},\beta_{2})$ | (0.9, 0.999)
Batch size | 256
Weight decay | 0.0
Loss | ‘MSE-loss’
Table 2: Hyperparameters for the random labels experiments.


# Turn Down the Noise: Leveraging Diffusion Models for Test-time Adaptation
via Pseudo-label Ensembling

Mrigank Raman
Machine Learning Department
Carnegie Mellon University
<EMAIL_ADDRESS>
&Rohan Shah
Machine Learning Department
Carnegie Mellon University
<EMAIL_ADDRESS>
&Akash Kannan
Machine Learning Department
Carnegie Mellon University
<EMAIL_ADDRESS>
&Pranit Chawla
Machine Learning Department
Carnegie Mellon University
<EMAIL_ADDRESS>

###### Abstract

The goal of test-time adaptation is to adapt a source-pretrained model to a
continuously changing target domain without relying on any source data.
Typically, this is either done by updating the parameters of the model (model
adaptation) using inputs from the target domain or by modifying the inputs
themselves (input adaptation). However, methods that modify the model suffer
from the issue of compounding noisy updates whereas methods that modify the
input need to adapt to every new data point from scratch while also struggling
with certain domain shifts. We introduce an approach that leverages a pre-
trained diffusion model to project the target domain images closer to the
source domain and iteratively updates the model via pseudo-label ensembling.
Our method combines the advantages of model and input adaptations while
mitigating their shortcomings. Our experiments on CIFAR-10C demonstrate the
superiority of our approach, outperforming the strongest baseline by an
average of $1.7\%$ across 15 diverse corruptions and surpassing the strongest
input adaptation baseline by an average of $18\%$.

##  1 Introduction

Supervised deep learning algorithms often assume that the data used for
training (source) and testing (target) follow the same ideal pattern, where
they are independent and identically distributed (IID). However, in real-world
situations, this ideal rarely holds true. Even small differences or changes in
data distribution can seriously affect the performance of these models
[Quionero-Candela et al., 2009]. This becomes a major challenge when using
supervised learning for tasks that involve data with different distributions.
To deal with shifts in data distribution, adaptation methods typically use the
target data to continuously improve and update their predictions.

Typically, adaptation techniques combat the problem of distribution shift by
jointly optimizing on the source and the target distributions during training
[Ganin et al., 2016, Hoffman et al., 2018, Saenko et al., 2010]. These methods
exhibit strong performance when the shifts in the target domain are well-
understood and predefined. However, they tend to falter when faced with the
arrival of an entirely new and previously unseen target domain. This
underscores the imperative need for development of test-time adaptation, a
technique that adapts the model at inference time without relying on the
source data and without interrupting the inference process. Broadly speaking,
there are two orthogonal ways of performing test-time adaptation. First, model
adaptation [Wang et al., 2022, 2020, Chen et al., 2022, Liang et al., 2020]
wherein the model is updated iteratively during inference without using test
labels. Second, Input adaptation [Gao et al., 2022] wherein the input is
modified to match the source distribution. Model adaptation helps to
continually adapt to changing distributions but may suffer from noisy updates
depending on the amount of distribution shift. Moreover, input adaptation must
adapt to each example from scratch and sometimes struggle with certain
distribution shifts.

Our proposed method, D-TAPE (Diffusion-infused Test-time Adaptation via
Pseudo-label Ensembling), seamlessly integrates the strengths of model and
input adaptation while minimizing their weaknesses. Utilizing a source domain-
trained diffusion model, we align target images more closely with the source
domain. However, a direct application of a pre-trained diffusion model risks
losing class information. To rectify this, we implement a low-pass filtering
technique, akin to [Choi et al., 2021], enabling class-preserving projections.
Our model then iteratively updates by ensembling pseudo-labels from both the
projected and original test images, offering dynamic adaptability to sudden
domain shifts. D-TAPE outperforms existing state-of-the-art approaches like
CoTTA [Wang et al., 2022], TENT [Wang et al., 2020], and DDA [Gao et al.,
2022]. We also illustrate how D-TAPE effectively brings the input data closer
to the source domain compared to the corrupted images, as evidenced by a
reduction in the $\mathcal{A}$-distance [Ben-David et al., 2006]. We
anticipate that this study will serve as a catalyst for the broader adoption
and exploration of diffusion models in future test-time adaptation research.

##  2 Prior Work

#### Test Time Adaptation:

The goal of test-time adaptation is to adapt a source-pretrained model to a
continuously changing target domain without relying on any source data. To
that end, the landscape of Test Time Adaptation can broadly be segmented into
two paradigms namely Model Adaptation and Input Adaptation. In the realm of
Model Adaptation, particularly in source-free and test-time settings, various
methodologies have emerged. For instance, models like TENT [Wang et al.,
2020], SHOT [Liang et al., 2020], and MEMO [Zhang et al., 2021] employ
unsupervised loss minimization techniques specifically tailored to target data
distributions. Alternatively, some strategies resort to the generation of
pseudo labels via a range of innovative techniques such as a conjugate pseudo-
label function [Goyal et al., 2022], weight-averaged teacher models [Wang et
al., 2022], and nearest neighbor soft-voting mechanisms [Chen et al., 2022].
These pseudo-labels serve as scaffolding for self-training procedures,
effectively amplifying the model’s generalization capabilities. In contrast,
Input Adaptation strategies predominantly leave the model untouched and only
update the input examples from the target domain. Among these, Diffusion
Driven Adaptation [Gao et al., 2022] stands out for its application of a
diffusion model trained on source data, thereby enabling more effective test-
time adaptation.

#### Diffusion Modelling:

Diffusion modeling serves as a pivotal approach in generative techniques,
relying on a two-step process that initially introduces noise into data,
followed by a denoising phase. Seminal works such as DDPM [Ho et al., 2020]
and DDIM [Song et al., 2020] have been instrumental in establishing the
framework of this field. An advanced extension is Guided Diffusion [Dhariwal
and Nichol, 2021], which leverages class labels to enhance the quality of
generated images. Since we do not have access to test labels, we employ DDPM
models that have been specifically trained on the CIFAR10 dataset for
projecting target images closer to the source domain.

##  3 Methodology

#### Diffusion with Interpolated Latent Refinement:

Initially, a diffusion model is pre-trained to generate images from the source
distribution. Then, at test-time, we run the forward process on the domain-
shifted input image by iteratively adding Gaussian noise $N$ times, resulting
in a sequence $x_{0},\ldots,x_{N}$. We subsequently apply a modified reverse
process to project the image closer to the source domain, using the diffusion
model trained to maximize the likelihood of the source data. Initially, we
sample $\hat{x}^{g}_{N-1}\sim p_{\theta}(x^{g}_{N-1}|x^{g}_{N}=x_{N})$ from
the reverse denoising process of the diffusion model. Since the diffusion
model is not explicitly conditioned to preserve class information, we apply
the low-pass filter $h()$ from ILVR [Choi et al., 2021] to the sampled
$\hat{x}^{g}_{N-1}$ to constrain the low-frequency features of
$\hat{x}^{g}_{N-1}$ to be equal to that of $x_{N-1}$ in the same manner as Gao
et al. [2022]. This is a sequence of downsampling and upsampling operations
that helps preserve the high-level image structure and hence the class
information. However, downsampling results in blurring for low-resolution
images. To mitigate this effect, we interpolate the filtered image
$h(\hat{x}^{g}_{N-1})$ with the original sample $\hat{x}^{g}_{N-1}$ to
generate $x^{g}_{N-1}$. This is fed back into the diffusion model and the
process is repeated to generate $x^{g}_{N-2}$ and so on. We provide the exact
algorithm below.

#### Model Adaptation via Pseudo-label ensembling:

Under the self-training framework, the model weights are adapted using some
form of the model predictions themselves. Given a pre-trained model,
$f_{\theta}()$ we follow a strategy similar to that used in CoTTA [Wang et
al., 2022] to adapt the model weights. We store two copies of the original
model - a student $f_{\theta_{s}}()$ and a teacher $f_{\theta_{t}}()$. The
teacher generates pseudo-labels that provide supervision to update the student
model. The weights of the teacher themselves are adapted by exponential moving
average using the student weights. We incorporate the diffusion model output
$x^{g}_{0}$ in both student predictions
$\hat{y}_{S}=g(x_{0},x^{g}_{0},\theta_{S})$ and pseudo-labels
$\hat{y}_{T}=g(x_{0},x^{g}_{0},\theta_{T})$ where the function $g$ is a
combination of a conditional ensembling scheme and logit averaging. We provide
ablations for these choices in Table 3. Mathematically,

$g(x_{0},x^{g}_{0},\theta)=\begin{cases}f_{\theta}(x_{0})&\text{if
}\max_{c}f_{\theta_{T}}(x_{0})>\max_{c}f_{\theta_{T}}(x_{0}^{g})\\\ \\\
\frac{f_{\theta}(x_{0})+f_{\theta_{T}}(x_{0}^{g})}{2}&\text{otherwise}\end{cases}$

We additionally generate student predictions on augmented versions of the
inputs $x^{\prime}_{0}$ and $x^{\prime g}_{0}$, and denote
$\hat{y}_{S}^{\prime}=g(x^{\prime}_{0},x^{\prime g}_{0},\theta_{S})$. The loss
used to update the student is
$\mathcal{L}=-\sum\limits_{c}0.5{y_{T}^{c}}(\log\hat{y}_{S}^{c}+\log\hat{y}_{S}^{{}^{\prime}c})$.

Figure 1: Illustration of our complete adaptation mechanism. Given an image
from the target distribution, we first apply our interpolated latent
refinement procedure to project the image closer to the source distribution.
Then we use a student-teacher framework and update the models by using an
ensemble of pseudo-labels generated by the teacher using both the original and
projected test image. Algorithm 1 Generation using Diffusion

1:Input: Reference image $x_{0}$

2:Output: Generated image $x^{g}_{0}$

3:$N$: diffusion range, $\phi_{D}(\cdot)$: low-pass filter of scale $D$

4:Sample $x_{N}\sim q(x_{N}|x_{0})$ $\triangleright$ perturb input

5:$x^{g}_{N}\leftarrow x_{N}$

6:for $t\leftarrow N$ to $1$ do

7: $\hat{x}^{g}_{t-1}\sim p_{\theta}(x^{g}_{t-1}|x^{g}_{t})$ $\triangleright$
unconditional proposal

8:
$\hat{x}^{g}_{0}\leftarrow\frac{1}{\bar{\alpha}_{t}}x^{g}_{t}-\frac{1}{\bar{\alpha}_{t}-1}\epsilon_{\theta}(x^{g}_{t},t)$

9:
$x^{g}_{t-1}\leftarrow\hat{x}^{g}_{t-1}-w\nabla_{xt}\lVert\phi_{D}(x_{0})-\phi_{D}(\hat{x}^{g}_{0})\rVert^{2}$

10: $x^{g}_{t-1}\leftarrow\alpha x^{g}_{t-1}+(1-\alpha)\hat{x}^{g}_{t-1}$

11:end for

12:return $x_{0}^{g}$

##  4 Experimental Setup and Results

#### Datasets and Evaluation Methodology:

We train our models on the CIFAR10 [Krizhevsky et al., 2009] train set and
evaluated on CIFAR10-C [Hendrycks and Dietterich, 2019]. CIFAR10-C comprises
of 15 different corruptions at 5 different severity levels. We evaluate the
continuous adaptability of the models using the approach used in [Wang et al.,
2022, Gan et al., 2022] which can briefly described as follows: We arrange the
corruption types in the same order as demonstrated in Table 1 and perform
inference on the model using batches of test images with a severity level of 5
for each corruption type. This process results in inference on 15 times the
size of the clean test data. We exclusively use images with a severity level
of 5 to introduce sudden distribution shifts in the batches. We report error
rates for each corruption type and calculate the mean error rate across all 15
corruptions. We also provide details and results on a different setting namely
gradual TTA in 2

#### Baselines:

We compare D-TAPE to the following test-time adaptation methods.
TENT [Wang et al., 2020]: Minimizes model prediction entropy on test data
using the model’s own predictions as pseudo labels, updating only
normalization parameters.
CoTTA [Wang et al., 2022]: Utilizes a weight-averaged teacher model for
pseudo-label generation and updates a student network. It also incorporates
augmentation-averaged pseudo labels for low-confidence test inputs and
restores a fraction of model weights to their original values.
AdaContrast [Chen et al., 2022]: Maintains a memory queue of target features
and predictions with a momentum model. Generates pseudo-labels for using a
combination of cross-entropy loss and contrastive loss to update the model.
DDA [Gao et al., 2022]: Employs a diffusion process to transform test images
back to the source domain, using the output as predictions. No parameter
updates occur during testing, and a low-frequency filter retains class
information during diffusion. Following the approach of Wang et al. [2022], we
employ a pre-trained WideResNet-28-10 [Zagoruyko and Komodakis, 2016] from
RobustBench [Croce et al., 2020] as the backbone for all baselines on
CIFAR10-C.

#### Results Discussion:

Our method outperforms the strongest baseline which is CoTTA on $14$ out of
the $15$ corruptions while beating it by an average of $1.7\%$ across all $15$
corruptions (Table 1). Additionally, our method also beats DDA which only
modifies the input by an average of $18\%$ across all corruptions. This
underscores the effectiveness of our method.

Method |  gauss |  shot |  impulse |  defocus |  glass |  motion |  zoom |  snow |  frost |  fog |  brightness |  contrast |  elastic |  pixelate |  jpeg |  mean
---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---
Source | 72.33 | 65.71 | 72.92 | 46.95 | 54.33 | 34.76 | 42.01 | 25.09 | 41.30 | 26.01 | 9.31 | 46.71 | 26.59 | 58.45 | 30.30 | 43.52
DDA | 33.09 | 30.38 | 38.18 | 49.51 | 33.61 | 37.94 | 41.28 | 19.70 | 24.18 | 39.10 | 11.22 | 61.24 | 25.38 | 30.95 | 19.47 | 33.02
TENT | 24.79 | 20.48 | 28.54 | 14.54 | 32.27 | 16.05 | 13.88 | 20.21 | 20.45 | 17.77 | 11.17 | 15.20 | 24.06 | 19.96 | 25.38 | 20.32
AdaContrast | 29.19 | 22.50 | 30.09 | 13.85 | 32.82 | 14.00 | 12.10 | 16.47 | 14.73 | 14.31 | 8.04 | 9.83 | 22.11 | 17.76 | 19.95 | 18.52
COTTA | 24.30 | 21.43 | 25.91 | 11.72 | 27.90 | 12.33 | 10.66 | 14.84 | 13.82 | 12.36 | 7.54 | 10.76 | 18.23 | 13.61 | 17.68 | 16.21
D-TAPE (Ours) | 19.29 | 17.19 | 22.58 | 11.54 | 22.73 | 11.94 | 10.05 | 14.13 | 12.98 | 12.15 | 8.17 | 10.58 | 16.87 | 12.70 | 14.54 | 14.50

Table 1: Classification error rate (%) for the standard CIFAR10-to-CIFAR10-C
time adaptation task in the sudden setting. Evaluated on WideResNet-28 with
the severity level 5.

#### Gradual TTA:

Unlike the sudden setting, here we make use of all the severity levels. Like
the sudden setting, we generate a permutation of all the different types of
corruptions and for each type we sample batches with severity in the following
order $1\rightarrow 2\rightarrow 3\rightarrow 4\rightarrow 5\rightarrow
4\rightarrow 3\rightarrow 2\rightarrow 1$ on which we do inference, update
weights and then go to the next type of corruption in the same ordering of
severity. Note that we do not reset the weights to the pre-trained values
anywhere in between. For each type of corruption, the average error on its
images is the indicator of performance. D-TAPE outperforms the strongest
baseline by $\sim 1\%$ on average across $15$ corruptions and improves
massively over the best input adaptation technique which is DDA (Table 2).

Model | Gradual CIFAR-10C
---|---
Source | 43.52
DDA | 33.02
TENT | 24.56
AdaContrast | 15.78
CoTTA | 13.99
D-TAPE | 13.09
Table 2: Mean results over different types of corruption over the gradual
CIFAR10-C at level 5 corruption, we report error percentage for each method.

#### Ablations

We perform ablations to demonstrate the importance of using conditional
ensembling and logit averaging in the pseudo-label ensembling phase of D-TAPE
and report the results in Table 3. We notice that without conditional
ensembling, our performance on fog and contrast drastically decreases. This is
expected since diffusion fails to project the images closer to the source for
these corruptions as seen in Figure 2. We also observe that without logit
averaging our performance drops slightly in the case of noises as well where
we expect the best performance. Logit averaging helps us to defend against the
rare cases where diffusion fails to retain the same object in the image.

Conditional Ensembling | Logit Averaging | gauss | shot | impulse | fog | contrast | mean
---|---|---|---|---|---|---|---
✗ | ✗ | 26.19 | 24.9 | 33.76 | 33.69 | 41.19 | 23.68
✗ | ✓ | 19.70 | 18.10 | 23.79 | 16.79 | 17.37 | 16.02
✓ | ✗ | 20.21 | 18.57 | 24.26 | 14.73 | 14.14 | 16.03
✓ | ✓ | 19.29 | 17.19 | 22.58 | 12.15 | 10.58 | 14.50
Table 3: Ablation on two components of our model, conditional ensembling and
logit averaging. Mean is over 15 noises

#### Does Diffusion help?

To answer the question of whether using diffusion helps or not, we use
$\mathcal{A}$-distance [Ben-David et al., 2006] to quantify shifts in
distribution. We follow Ben-David et al. [2006] and train a linear classifier
to separate the two domains for computing the $\mathcal{A}$-distance. We
observe that, on average, the generated images closely match the clean image
distribution, with particularly close alignment in the case of gaussian, shot,
and impulse noises. However, for fog and contrast, the $\mathcal{A}$-distance
slightly increases for generated images.(Figure 2). We believe this happens
because our interpolated latent refinement scheme retains low frequency
noises. We report all the $\mathcal{A}$-distance values in Table 4 which
indicate that diffusion is beneficial for almost all corruptions.

Figure 2: $\mathcal{A}$-distance between the CIFAR-10C dataset and the
original CIFAR-10 dataset, and $\mathcal{A}$-distance between generated images
and original CIFAR-10 dataset. Mean is over 15 noises

## References

  * Quionero-Candela et al. [2009] Joaquin Quionero-Candela, Masashi Sugiyama, Anton Schwaighofer, and Neil D. Lawrence.  Dataset shift in machine learning.  2009\.  URL https://api.semanticscholar.org/CorpusID:61294087.
  * Ganin et al. [2016] Yaroslav Ganin, Evgeniya Ustinova, Hana Ajakan, Pascal Germain, Hugo Larochelle, François Laviolette, Mario Marchand, and Victor Lempitsky.  Domain-adversarial training of neural networks.  _The journal of machine learning research_ , 17(1):2096–2030, 2016.
  * Hoffman et al. [2018] Judy Hoffman, Eric Tzeng, Taesung Park, Jun-Yan Zhu, Phillip Isola, Kate Saenko, Alexei Efros, and Trevor Darrell.  Cycada: Cycle-consistent adversarial domain adaptation.  In _International conference on machine learning_ , pages 1989–1998. Pmlr, 2018.
  * Saenko et al. [2010] Kate Saenko, Brian Kulis, Mario Fritz, and Trevor Darrell.  Adapting visual category models to new domains.  In Kostas Daniilidis, Petros Maragos, and Nikos Paragios, editors, _Computer Vision – ECCV 2010_ , pages 213–226, Berlin, Heidelberg, 2010. Springer Berlin Heidelberg.  ISBN 978-3-642-15561-1.
  * Wang et al. [2022] Qin Wang, Olga Fink, Luc Van Gool, and Dengxin Dai.  Continual test-time domain adaptation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 7201–7211, 2022.
  * Wang et al. [2020] Dequan Wang, Evan Shelhamer, Shaoteng Liu, Bruno Olshausen, and Trevor Darrell.  Tent: Fully test-time adaptation by entropy minimization, 2020.  URL https://arxiv.org/abs/2006.10726.
  * Chen et al. [2022] Dian Chen, Dequan Wang, Trevor Darrell, and Sayna Ebrahimi.  Contrastive test-time adaptation.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 295–305, 2022.
  * Liang et al. [2020] Jian Liang, Dapeng Hu, and Jiashi Feng.  Do we really need to access the source data? source hypothesis transfer for unsupervised domain adaptation, 2020.  URL https://arxiv.org/abs/2002.08546.
  * Gao et al. [2022] Jin Gao, Jialing Zhang, Xihui Liu, Trevor Darrell, Evan Shelhamer, and Dequan Wang.  Back to the source: Diffusion-driven test-time adaptation.  _arXiv preprint arXiv:2207.03442_ , 2022.
  * Choi et al. [2021] Jooyoung Choi, Sungwon Kim, Yonghyun Jeong, Youngjune Gwon, and Sungroh Yoon.  Ilvr: Conditioning method for denoising diffusion probabilistic models, 2021.
  * Ben-David et al. [2006] Shai Ben-David, John Blitzer, Koby Crammer, and Fernando Pereira.  Analysis of representations for domain adaptation.  In B. Schölkopf, J. Platt, and T. Hoffman, editors, _Advances in Neural Information Processing Systems_ , volume 19. MIT Press, 2006.  URL https://proceedings.neurips.cc/paper_files/paper/2006/file/b1b0432ceafb0ce714426e9114852ac7-Paper.pdf.
  * Zhang et al. [2021] Marvin Zhang, Sergey Levine, and Chelsea Finn.  Memo: Test time robustness via adaptation and augmentation, 2021.  URL https://arxiv.org/abs/2110.09506.
  * Goyal et al. [2022] Sachin Goyal, Mingjie Sun, Aditi Raghunathan, and J Zico Kolter.  Test time adaptation via conjugate pseudo-labels.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=2yvUYc-YNUH.
  * Ho et al. [2020] Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  _Advances in Neural Information Processing Systems_ , 33:6840–6851, 2020.
  * Song et al. [2020] Jiaming Song, Chenlin Meng, and Stefano Ermon.  Denoising diffusion implicit models.  _arXiv preprint arXiv:2010.02502_ , 2020.
  * Dhariwal and Nichol [2021] Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in Neural Information Processing Systems_ , 34:8780–8794, 2021.
  * Krizhevsky et al. [2009] Alex Krizhevsky, Geoffrey Hinton, et al.  Learning multiple layers of features from tiny images.  2009\.
  * Hendrycks and Dietterich [2019] Dan Hendrycks and Thomas Dietterich.  Benchmarking neural network robustness to common corruptions and perturbations.  _Proceedings of the International Conference on Learning Representations_ , 2019.
  * Gan et al. [2022] Yulu Gan, Xianzheng Ma, Yihang Lou, Yan Bai, Renrui Zhang, Nian Shi, and Lin Luo.  Decorate the newcomers: Visual domain prompt for continual test time adaptation.  _arXiv preprint arXiv:2212.04145_ , 2022.
  * Zagoruyko and Komodakis [2016] Sergey Zagoruyko and Nikos Komodakis.  Wide residual networks.  _arXiv preprint arXiv:1605.07146_ , 2016.
  * Croce et al. [2020] Francesco Croce, Maksym Andriushchenko, Vikash Sehwag, Edoardo Debenedetti, Nicolas Flammarion, Mung Chiang, Prateek Mittal, and Matthias Hein.  Robustbench: a standardized adversarial robustness benchmark.  _arXiv preprint arXiv:2010.09670_ , 2020.

###  .1 $\mathcal{A}$-distance

We report the $\mathcal{A}$-distance values for all the $15$ corruptions in
Table 4

| CIFAR10-C | Generated
---|---|---
gauss | 1.17 | 0.42
shot | 1.12 | 0.43
impulse | 1.36 | 0.55
defocus | 0.95 | 0.92
glass | 0.60 | 0.50
motion | 0.75 | 0.65
zoom | 0.79 | 0.73
snow | 1.51 | 1.48
frost | 1.39 | 1.36
fog | 0.99 | 1.00
brightness | 1.65 | 1.62
contrast | 1.29 | 1.31
elastic | 0.70 | 0.60
pixelate | 1.21 | 0.85
jpeg | 0.58 | 0.50
mean | 1.07 | 0.86
Table 4: We report the $\mathcal{A}$-distance between CIFAR10-C and CIFAR10
images and also between the images generated using our interpolated latent
refinement scheme and CIFAR10 images

###  .2 Hyperparameter Details

We perform $N=1000$ steps of the diffusion process, use $\alpha=0.9$ which is
decided using the FID score of the generated images. During the model
adaptation stage, we use a learning rate of $1e-3$, batch size of $200$, Adam
optimizer with a $\beta=0.9$ and we use $1$ optimizer step. We use a momentum
of $0.999$ to update the teacher via exponential averaging.