\section{Introduction}

The center of gravity of NLP research has shifted to the development
of language models (LMs) for representation and generation of text,
and most recent high-impact research contributions describe new
LMs.
For some tasks, a model needs to take into account not only a text but
also some non-textual information, and a wide range of multimodal LMs
have been developed that allow the representation of a text jointly
with some external modality. Most of this work focuses
on \emph{visual} tasks where NLP models need to be integrated with
computer vision models; examples of tasks in this area include visual
question answering and caption generation.
A range of combined language-and-vision LMs have been developed using
different approaches for integrating representations of text and of
images or videos.

But can we be sure that a multimodal model actually uses the provided visual information instead of just relying on statistical tendencies in the text corpus?
With the development of multimodal LMs, some recent work has
investigated what information is stored in the representations of the
multiple modalities and how the multiple representations interact.
For instance, \newcite{frank2021vision} carried out a set of controlled tests
to tease apart the effects of the textual and visual modalities.

\begin{figure}[t]
    \centering
    \includegraphics[width=0.98\linewidth]{images/bananas}

    \caption{An example of a controlled test of a masked language
    model for a color stereotype. We compute the output from the MLM
    head when providing an image of an object with a stereotypical
    color (a yellow banana) and compare it to the output when the object has an unusual color (green). If the MLM is strongly affected by a stereotype bias, the predictions change little.}
    \label{fig:overview}
\end{figure}

It has been widely noted that representations of language are
affected by several kinds of \emph{stereotypes}, which we loosely
define as any type of phenomenon that has a highly skewed prior
probability distribution.
In these cases, the skewed distribution may cause a model to
simply go with the default choice and
ignore contextual information that would suggest an unusual analysis.
Most of the discussion in the field has been about stereotypes
relating to various demographic attributes \cite{bolukbasi2016}, but
in this work, we use the term ``stereotype'' in the more general sense
mentioned above.
This issue is likely to affect multimodal LMs as well, although we are
aware of no previous work that investigates this phenomenon
systematically; for instance, if some object is often associated with
some visual property (e.g. a color or shape), this property may be
predicted by the model even in cases where it is not present.
This effect may also have methodological implications in benchmarks
for the evaluation of LMs:
if a model predicted the correct answer, did it do so because of the stereotype or because it actually used the available
visual information?

In this work, we propose a methodology and develop two benchmark sets
for stress-testing multimodal LMs to determine to what
extent they are affected by problems related to stereotypes.
The key idea is to look at predictions of a language/vision LM with
different visual inputs and compare the behavior of the LM in the
presence or absence of stereotypes.
For cases when a stereotype is present, we compare model outputs when
the image \emph{does} correspond to the stereotype to when
it \emph{does not}.

The rest of the paper is organized as
follows. Section~\ref{sec:benchmarks} discusses the design of the
benchmark sets and how we use them to investigate multimodal LMs for
stereotypes. Details about the multimodal LMs we have used
are covered in Section~\ref{sec:models}, and
Section~\ref{sec:method} describes how they are applied for the
benchmarks, while
Section~\ref{sec:results} presents the figures achieved on the
benchmarks and discusses their implications.
In Section~\ref{sec:relwork}, we discuss related research.
Finally, Section~\ref{sec:conclusions} summarizes the main points and
discusses limitations and possible extensions.

\section{Design of Benchmark Datasets}
\label{sec:benchmarks}

We have collected two datasets consisting of textual templates and
corresponding images. These datasets were selected because in these
cases it was relatively easy to
collect images exemplifying some visual property, and where on the one
hand we could find images corresponding to a stereotype, but on the
other hand also control images \emph{not} corresponding to the stereotype.

These datasets also contain subsets we call ``neutral''
where stereotypes are not present.
The purpose of these images is to investigate whether LMs are more
sensitive to the choice of images
in the cases when they cannot rely on stereotypes.

\subsection{The Memory Colors Dataset}

The first dataset is an extension of the \emph{Memory Colors}
dataset \cite{norlund2021transferring}, originally developed for the
purpose of measuring the transfer of information between visual and
textual representations.
The original dataset lists a set of 109 common physical objects, where each
object is listed with a ``memory color'': a stereotypical color we
typically associate with the object. For instance, the dataset lists
tomatoes as stereotypically red although tomatoes frequently have
other colors. The set of objects was annotated by multiple annotators,
and only the objects where there was a perfect or almost perfect
consensus among annotators were included.

The dataset comes with a set of textual templates that can be
used to generate prompts for LMs.
Since the dataset was originally intended for use in LMs where no image was
available, these text templates were intentionally formulated to
elicit stereotypical responses, e.g. \emph{``The typical color of a tomato
is\ldots''}. In our case, we changed the templates to encourage the
model to focus on the image, e.g. \emph{``The color of this tomato
is\ldots''}.

The Memory Colors dataset also includes a set of prototypical images
exemplifying the stereotypical color. For each of the object types, we
collected an additional image where the color was not the stereotypical
one, e.g. a green tomato.
All images were collected by carrying out a Google image
search and picking the first result.
The majority of objects with unusual colors includes examples of natural
images (e.g. unripe tomatoes, orange sky); in a few cases, the
color had been artifically modified

We also extended the Memory Colors dataset with 19 neutral
object types selected so that they were not expected to have a stereotypical
color. This set includes common objects such as cars, houses, etc.
We refer to the combined set, including the images with
non-stereotypical colors and the neutral instances, as
the \emph{Extended Memory Colors} dataset.

\subsection{Gender Stereotypes Dataset}
\label{sec:method:gender}

The effect of gender in neural language representation models has been
widely investigated and it is relevant to consider this in multimodal
representations as well.
We compiled a second dataset we term the \emph{Gender Stereotypes}
dataset.
The aim is to identify how good a multimodal model performs in the
prediction of a person's gender when it is fed two different images,
which will act as visual signals for us, one corresponding to a man
and another one corresponding to a woman.
For each pair, there is a sentence that describes the activity.
As in the color dataset, we include stereotypical cases (male-coded and
female-coded, respectively) as well as cases where no stereotype is
present.

The dataset contains 50 different text sentences and 100 images with,
where half of the images show male individuals and half show females.
Internally in the dataset, 19 and 21 text templates were created for
the male and female stereotypical activities, respectively.\footnote{Stereotypical activities were selected from \href{
https://careersmart.org.uk/occupations/equality/which-jobs-do-men-and-women-do-occupational-breakdown-gender}{\color{darkblue}this website}.}
Further,
we defined a list of 10 different neutral tasks:
\emph{eating, walking, reading, writing, meditating, talking,
studying, listening to music, clapping, crying}. For these cases, we
assumed that there is no stereotypical gender associated with the
activities.

As we will discuss in more detail in Section~\ref{sec:method}, the
property to be predicted will be represented in the sentence as
a \texttt{[MASK]} token to be substituted by a masked LM.
To include an example from the gender stereotype dataset, the sentence
is as \emph{'My therapist is very good,} \texttt{[MASK]} \emph{helped me get
myself together'}; according to the source where we selected the
stereotypical occupations, therapy professionals are more frequently
female.

For each of the 50 text templates, we selected two images, one for
each of the genders. As for the colors dataset, we used the first
result in an image search judged by an annotator to correspond to the
gender in question. We did not take the self-identified gender into
account.

\section{Multimodal Language Models}
\label{sec:models}

The Transformer \cite{vaswani2017attention} is a sequence-based model
that is now the standard architecture in NLP for devising
representation and generation components in neural models.
Pre-trained language models such as BERT \cite{devlin2019bert}
based on the architecture of Transformers, have proven capable of
learning powerful representations applicable to a wide range of
tasks. They have yielded state-of-the-art performance in many
downstream tasks.

Multimodal models fusing the textual and visual modalities have been
devised by researchers after looking at the huge success of
pre-trained language models. %
In such models, multiple modalities are considered, and data for the
training of the models is in multiple modalities. As our
research problem revolves around the aspect of multimodality, we will
focus on two modalities: a textual and a visual signal. The visual
signal is in the form of images, and the natural language is the
written text accompanying the images, such as captions or descriptions
of the images.
Examples of such visual/textual Transformers include VilBERT \cite{lu2019vilbert}, LXMERT \cite{tan2019lxmert}, VisualBERT \cite{li2020what},
OSCAR \cite{li2020oscar}, ImageBERT \cite{qi2020imagebert}, FLAVA \cite{singh2022flava}, and others.
Most of the earlier models use features extracted from a Faster-RCNN pipeline \cite{ren2015faster}, while later models use visual Transformer architectures \cite{dosovitskiy2021vit}.
These types of models are then trained on datasets that contain text/image pairs such as SBU Captions~\citep{sbu_captions}, MS COCO~\citep{coco},
Conceptual Captions~\citep{sharma2018conceptual}, and Visual Genome
QA~\citep{krishnavisualgenome}, using various pre-training tasks. They
are sometimes trained from scratch on the combined language/vision
data and sometimes warm-started from a unimodal model such as BERT.

For this study, we selected three different multimodal models
to run our experiments on. These image-augmented Transformer models are
\mbox{VisualBERT,} LXMERT, and FLAVA. These three are specifically chosen
to give a certain diversity in the selection of model architecture:
one single-stream CNN-based model, one dual-stream CNN-based model,
and one visual Transformer-based model.

All the models we selected are BERT-like variations that use a
the technique of Masked Language
Modelling (MLM) during pre-training. This idea was presented in the
original BERT paper \cite{devlin2019bert}. In
the task of Masked Language Modelling, we predict a token which has
been masked by us in the sentence, given a set of unmasked tokens. In
our case, unmasked tokens are supplemented by the the visual
signals.
The random masking ratio for the MLM is around 15\%, and for
investigation of our experiments one special \texttt{[MASK]} token is
taken.
As we will discuss in Section~\ref{sec:method}, we rely on the ability
of the MLM to predict missing tokens in our experiments.

\paragraph{VisualBERT} This is a single stream multimodal model, i.e,
the language and vision embeddings are processed via a single
Transformer. It is an extension of BERT, by redefining the
process of how input is processed. The language embeddings are
extracted from BERT's tokenizer, which acts as text encoder.
For the embeddings of the visual signals, Faster-RCNN is used. It
extracts image features in the form of 36 RoI (region of interest)
boxes for each image, and these RoI boxes are used as features.
Each of these 36 ROI boxes are vectors of size 2048. The boxes with
highest probability/confidence are chosen.
The visual representations are
appended at the end of the sequence of word embeddings.

\paragraph{LXMERT}  This model is a dual stream multimodal model,
where the inputs are processed through two Transformers, for natural
language and vision signals respectively. Text is processed in the
same manner as of VisualBERT, based on BERT's tokenizer. The image
features for the LXMERT are extracted by the Faster-RCNN, in the same
way as of VisualBERT, but we also feed the normalized boxes alongside
features, which are locations of these bounding boxes. At last,
the Transformers are fused.

\paragraph{FLAVA}  FLAVA has a text encoder, an image encoder, and a
multimodal encoder. It is a dual stream multimodal model. The text
encoder, has an architecture of ViT (visual Transformers) to extract
single-modal text representations. For the images, an image encoder
based on ViT architecture extracts single-modal image
representations. A separate Transformer, multimodal encoder, is then
applied. The unimodal representations are passed through the fusion
encoder which fuses two modalities, and thus obtaining cross-modal
representations.

\subsection{Model Details}

There is a slight difference in how the two CNN-based models,
VisualBERT and LXMERT, are applied.
In the case of VisualBert, we also input
locations of bounding boxes. For the experiments concerning
VisualBERT, we have used the pretrained BERT
tokenizer,\footnote{bert-base-uncased from the HuggingFace library.} and
VisualBERT with COCO pretraining
checkpoint\footnote{uclanlp/visualbert-vqa-coco-pre from HuggingFace.}
for the model. In the case of LXMERT, the LXMERT base
tokenizer and model\footnote{unc-nlp/lxmert-base-uncased from
HuggingFace library.} were used. For FLAVA, we used the pretrained processor and
model.\footnote{facebook/flava-full from HuggingFace library.}

\section{Methodology of Analysis}
\label{sec:method}

Our benchmarking method uses a cloze-style fill-in-the-blank
approach \cite{petroni2019lmaskb,jiang2020know}, which has previously been applied in experiments investigating the interaction between visual and linguistic representation \cite{norlund2021transferring,hagstrom2022b,hagstrom2022a}.
This approach is easy to apply to BERT-style models that include a
masked language model (MLM) as part of their pre-training pipeline.
When applying the MLM in our experiment,
the model is provided with an image and a text prompt,
where the visual property to be predicted by the model has been
replaced by the mask dummy token.
We then investigate how well the missing token is predicted under
different circumstances.

Since the nature of the two benchmarks
is different, we had to apply different methodologies to get the
results. We discuss these details below.

\subsection{The Memory Colors Dataset}

For the Memory Colors dataset, we compare the image having a
stereotypical color to an image with an unusual color for the
particular object, and to a dummy image containing no meaningful
information.
Following previous work that applied image-augmented LMs to text-only
inputs, we have considered different types of dummy images.
We have used two types of dummy
images: the first one being a completely black image
following \newcite{iki2021effect}, and the second
consisting of white noise.
However, in experiments we did generally not see major differences
between the behavior of the models when using the black dummy images
and when using the noise images, so we limit the discussion to black
dummy images in the rest of this paper.

For a given text prompt and image, we  mark the output as correctly or
incorrectly predicted depending on whether the
token predicted at the \texttt{[MASK]} position matches
the color of the label we have provided in the dataset or not.

In these experiments, we did not restrict the output vocabulary to
color terms. In general, after going through the results, it seems
that all the three models tend to output color at the position
of \texttt{[MASK]} token.

\begin{table*}[t!]
\centering
\begin{tabular}{@{}lcccccc}
\toprule
 & \multicolumn{3}{c}{\textbf{Stereotypes}} & & \multicolumn{2}{c}{\textbf{No stereotypes}} \\ \cmidrule(l){2-4}\cmidrule(l){6-7}
\textbf{Model} & \shortstack{\textbf{Original}\\\textbf{image}} & \shortstack{\textbf{Control}\\\textbf{image}} & \shortstack{\textbf{Black}\\\textbf{image}} & & \shortstack{\textbf{Original}\\\textbf{image}} & \shortstack{\textbf{Black}\\\textbf{image}} \\ \midrule
VisualBERT & 0.23 & 0.08 (0.50) & 0.28 (0.41) & & 0.0 & 0.0 (0.84)\\
LXMERT & 0.72 & 0.11 (0.76) & 0.69 (0.87) & & 0.47 & 0.05 (0.47)\\
FLAVA & 0.74 & 0.69 (0.06) & 0.08 (0.08) & & 0.89 & 0.11 (0.11)\\
\bottomrule
\end{tabular}
\caption{Accuracies on the extended Memory Colors datasets. For
control images with unexpected colors, the accuracies are computed with respect to the \emph{new} color, while for the black images the accuracies are with respect to the \emph{original} color. Figures in brackets show the proportion of predictions that are equal to the original prediction.}
\label{tab:colorsresults}
\end{table*}

\subsection{Gender Stereotypes Dataset}

For the Gender Stereotypes dataset, we also consider the output of the
MLM head at the masked position, but in this case we also need to take
into account that
several words may be applicable in the given context. For this reason,
we create two buckets of male and female words:
\emph{he, male, man, men, boy, his} and \emph{she, female, woman,
women, girl, her}, respectively.
We choose the predicted gender based on the highest probability the elements in the buckets get for the masked token. If the element with the highest probabilty falls in the bucket containing male words, we count this instance as predicted male by the model and vice versa for the female bucket.

\section{Results}
\label{sec:results}

We evaluated the three selected models on the two benchmarks. In both
cases, we compare the predictions when a stereotype is present and the
image corresponds to the stereotype to the case where the
image \emph{does not} correspond to the stereotype.
We also evaluate
cases where there is no stereotype and we carry out similar
comparisons in this case.
Additionally, we look at the model's predictions when provided with a
black dummy image.

\subsection{The Extended Memory Colors Dataset}

Table~\ref{tab:colorsresults} shows the results on the extended Memory
Colors  stereotypes dataset. When using real images, the figures
outside the brackets should be interpreted as predictive accuracies;
for the black dummy images, the figures show the proportions of cases
predicted as the stereotypical color. The figures in brackets show
the proportion of predictions that are identical to the original
prediction.

We note that VisualBERT performs poorly on this dataset, confirming
previously published results that this model is underfitted on visual
data and mostly sticks to the prediction by an equivalent BERT
model.
The effect of the image seems minimal and
its performance is close to the majority-class baseline accuracy
of 0.25.

The LXMERT and FLAVA models achive better scores on the original
Memory Colors dataset: both models have accuracies in the 0.70--0.75
range. However, we see clearly that this similarity of performance is
superficial and that the LXMERT model mostly relies on stereotypes:
when we consider the control images with
unexpected colors, the performance of LXMERT is very poor and it
mostly keeps predicting the stereotypical color.
Its performance is somewhat better for the non-stereotypical cases,
but far from perfect.
FLAVA on
the other hand predicts fairly well on the control set, although
somewhat worse than for the images with stereotypical colors; it also
predicts with a good accuracy for the non-stereotypical cases. It is
clear that FLAVA is much more sensitive to the choice of images in
this task.

For the dummy images that are completely black,
the LXMERT model's prediction are again to a large extent identical to the
original predictions.
Again, the FLAVA model is more receptive to the choice of images: it
predicts the color \emph{black} in 92\% of the cases and there is no
discernible effect of stereotypes; it can be discussed whether this is
a desired behavior in this case, since the image does not include an
object of the kind mentioned in the prompt.

Finally, we note that for the non-stereotypical instances, LXMERT's
predictions seem to shift more between the original images and the
black dummy images. This suggests that in cases where the model cannot
rely on a stereotype, the model is more sensitive to the visual input.

\subsection{Gender Stereotypes Dataset}

\begin{table*}[t!]
\centering
\begin{tabular}{lccccccccc}
\toprule
 & \multicolumn{3}{c}{\textbf{Male stereotypes}} &  \multicolumn{3}{c}{\textbf{Female stereotypes}} &  \multicolumn{3}{c}{\textbf{No stereotypes}} \\ \cmidrule(l){2-4}\cmidrule(l){5-7}\cmidrule(l){8-10}
\textbf{Model} & \shortstack{\textbf{Male}\\\textbf{image}} & \shortstack{\textbf{Female}\\\textbf{image}} & \shortstack{\textbf{Black}\\\textbf{image}} &  \shortstack{\textbf{Male}\\\textbf{image}} & \shortstack{\textbf{Female}\\\textbf{image}} & \shortstack{\textbf{Black}\\\textbf{image}} &  \shortstack{\textbf{Male}\\\textbf{image}} & \shortstack{\textbf{Female}\\\textbf{image}} & \shortstack{\textbf{Black}\\\textbf{image}} \\ \midrule
VisualBERT & 0.89 & 0.89 & 0.89 & 0.71 & 0.81 & 0.86 & 0.60 & 0.70 & 0.60\\
LXMERT & 0.84 & 0.68 & 0.73 & 0.95 & 0.76 & 0.90 & 0.90 & 0.40 & 0.80\\
FLAVA & 0.84 & 0.32 & 0.84 & 0.81 & 0.19 & 0.33 & 0.90 & 0.10 & 0.50\\
\bottomrule
\end{tabular}
\caption{Results on the gender stereotypes datasets. The figures
show the proportion predicted as \emph{male}.}
\label{tab:gendersresults}
\end{table*}

Table~\ref{tab:gendersresults} shows the results on the gender
stereotypes dataset. Note that for consistency, the figures
show the proportion of instances predicted
as \emph{male}, so they should not be interpreted as accuracies when
predicting with an image of a female.

Generally speaking, all models tend to predict the \emph{male} class
when provided with an image showing male individuals. When the input
shows a female individual, the picture is more varied.
As in the previous experiment, FLAVA reacts much more strongly to the
choice of images than VisualBERT and LXMERT, and tends to predict
the \emph{male} class for images with males and vice versa.

Unexpectedly, VisualBERT as well as LXMERT both seem to generally
assign higher probabilities to male-coded words, even when the prompt
is stereotypically female; this is surprising since we had expected
these models to predict the stereotypical classes in these cases.
It seems that FLAVA is the only model that shows signs
of \emph{contextual} gender stereotypes
in this experiment: when provided with a black dummy image, this model
predicts according to what would have been expected stereotypically,
and at 50\% for the non-stereotypical cases.
As we saw in the color experiment, for the non-stereotypical cases
LXMERT seems at least somewhat affected by the choice of images,
although less so than FLAVA.

\section{Related Work}
\label{sec:relwork}

This work falls in the broad category of model
analysis \cite{belinkov2019analysis} of Transformer
models \cite{rogers2020primer}. \newcite{belinkov2019analysis} divide
previous approaches to model analysis into several methodological
categories; in the current work, we use an approach based on
behavioral testing of a specific model behavior.
Specifically, our analysis is based on the outputs of the masked
language model head of BERT-like models, similarly to how
\newcite{petroni2019lmaskb} and \newcite{jiang2020know} tested BERT models for basic encyclopedic
and commonsense knowledge.

The methodology based on targeted behavioral testing has also been
used to investigate a number of research questions in the analysis of
language-and-vision Transformer models.
In particular, a number of investigations look at what type of
generalizations happen between the visual and textual modalities.
\newcite{cao2020revealing} claimed that when considering attention scores,
the effect of the visual modality is limited and that the textual
modality dominates.
\newcite{norlund2021transferring} investigated the effect of multimodal training on textual representations, and concluded that the degree of transfer between the representations of the respective modalities is limited, at least for CNN-based models;  \newcite{hagstrom2022b,hagstrom2022a} drew similar conclusions based on more extensive experiments that also include the FLAVA model.
\newcite{parcalabescu2021seeing} considered the task of predicting
numbers and arrived at a conclusion similar to ours: frequently
occurring numbers are predicted more often by the model.

The previous work that is most closely related to our in terms of
research questions and methodology is that by
\newcite{frank2021vision}. They designed ablation tests where parts of
the image or the text are hidden; as we have discussed, this setup is
comparable to our experiments where black and white-noise images are used.
\newcite{parcalabescu2022valse} introduced the idea of ``foils'': texts
that differs minimally from the one describing the image.
Our use of adversarially selected images can be seen as similar to the idea of foils, but focused on the visual modality.

\section{Conclusions}
\label{sec:conclusions}

In this work, we have proposed a methodological framework based on
controlled tests designed to tease out the influence of stereotypes on
the predictions of visually augmented language models.
The key idea is that we expect common evaluation benchmarks to include
many stereotypical cases that can easily be predicted simply by
relying on language statistics. In order to disentangle the effect of the
stereotype and the contribution of the visual representations
we compare the model's output in cases where the
provided image adheres to the stereotype to cases where it does not.
We also consider the model's behavior in cases where there are no
stereotypes, that is when the prior distribution of outputs is more
evenly distributed.

As an application of this framework, we created two datasets to
facilitate the investigation of stereotypes for two properties: the
color of objects and the gender of people. Each dataset contains a
set of text prompts and corresponding image pairs, where one image in
the pair corresponds to the stereotype and the other is a control
where the stereotypical property is not present. This allows
comparisons to be carried out in a controlled fashion.

Using the two benchmark sets, we evaluated three MLM-based visually
augmented Transformer models: VisualBERT, LXMERT, and FLAVA. There are
clear differences between the models, and in particular some of these
differences emerge much more clearly in the controlled setting.
For instance, the CNN-based LXMERT and Transformer-based FLAVA achieve
similar scores in terms of raw accuracy scores for predicting the
color of objects in images. However, if we consider the control images
where the objects do not have the stereotypical color, the FLAVA
outperforms LXMERT by a wide margin, since LXMERT keeps
predicting the stereotypical color.
This means that we can see clear differences among the models with
respect to how sensitive they are to the choice of images.

For the gender stereotypes experiments, the results were somewhat
unexpected since it turned out that the older CNN-based models
almost consistently assigned higher probabilities to male-related
words, where we had expected at least the LXMERT model to be somewhat
affected by stereotypes suggested by the textual prompt.
The newer FLAVA model on the other hand again predicts more
consistently with the input image in this experiment, and only falls
back on stereotypes when the input images are uninformative.

\subsection{Limitations and Possible Extensions}

As discussed in \S\ref{sec:method:gender}, we have intentionally used a
simplistic operationalization of the notion of gender in this work and selected
images returned by the image search engine when queried for
`male' or `female' respectively, and that the annotator then decided
were prototypical representatives of the male or the female genders.
The self-identified gender of the people in the images was not taken
into account in this experiment and since our goal was to investigate
the sensitivity of visually augmented LMs to the choice of images, it
was a priority to carry out such an evaluation using clear-cut cases.
In a more thorough investigation, it could potentially be useful to
also consider how e.g. the FLAVA model, which seems to be more
affected by the visual input, reacts when presented with images that
do not fall into such clear-cut categories.

The most obvious way that this work could be improved would be to
improve the robustness of the conclusions by scaling up the
investigations along all dimensions: instead of
considering just the two properties of color and gender, we would like
to investigate a wider selection of properties that would be
meaningful to test in language and vision models. Shape, size, and
orientation are a few possible examples. For each
scenario, it would also be useful to collect more examples than what
we have included here, in order to improve the statistical
robustness. Furthermore, since LMs are sensitive to the choice of a prompt \cite{jiang2020know}, our conclusions would be on firmer ground if we would evaluate on several text prompts for each image.
Naturally, it would be interesting to consider a more extensive selection of models as well.

In this work, we treated the property of being stereotypical as binary
and divided the test cases into groups based on this
property. However, as discussed in the introduction, in reality the
notion of stereotypicality is related to prior probability
distributions.
For this reason, a natural generalization of the experiments we have
carried out here would be to consider stereotypicality on a continuous
scale, e.g. by computing the entropy of the prior distribution and
then to see how this correlates with the probability of incorrect predictions
when encountering an unusual case.

The experiments in this work have been limited to evaluations of the
model's behavior for selected visual-linguistic properties. It remains
to see whether the same idea can be extended beyond evaluation to
devise new \emph{training} methods as well, in order to inject a bias
into the training process aimed at reducing the effects of stereotypes
and encouraging the model to rely on the visual information. This type
of training would typically involve more work in data collection,
unless methods can be devised to adversarially generate images with
unusual properties.

We finally note that the proposed methodology is not limited to the
evaluation of visually augmented LMs, but could be relevant when
considering any extra-linguistic extension of LMs. For instance,
similar pitfalls may occur in the evaluation of LMs augmented with
structural knowledge representations. If a knowledge-augmented LM
correctly predicted some encyclopedic fact \cite{petroni2019lmaskb,jiang2020know}, was
this because of what the knowledge resource contained or because
of text statistics?

\section*{Acknowledgements}

Richard Johansson was supported by the projects \emph{Interpreting and Grounding Pre-trained Representations for NLP} and \emph{Representation Learning for Conversational AI}, both funded by Wallenberg  AI,  Autonomous  Systems  and  Software Program (WASP) funded  by  the  Knut  and  Alice Wallenberg Foundation.


\section{Additional CIFAR100 Results}
\label{sec:add-cifar}
This section includes additional experiment results on applying ResMem to CIFAR100 dataset.
\subsection{Additional robustness results}
In addition to the results already presented in Section \ref{sec:exp-rob}, we also evaluate ResMem performance for each architecture in CIFAR-ResNet\{8, 14, 20, 32, 44, 56\} and each subset (10\%, 20\%, ..., 100\%) of CIFAR100 training data.
We use the same training hyperparameter and the ResMem hyperparameter as described in Section \ref{sec:exp-setup}.
Generally, we see that ResMem yields larger improvement over the baseline DeepNet when the network is small and dataset is large.
\begin{figure*}[ht]
    \centering
    \subfigure[\label{fig:cifar-sample-resnet8-appendix}CIFAR-ResNet-8]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet8_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet8_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet14}CIFAR-ResNet-14]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet14_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet14_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet20}CIFAR-ResNet-20]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet20_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet20_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet32}CIFAR-ResNet-32]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet32_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet32_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet44}CIFAR-ResNet-44]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet44_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet44_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet56}CIFAR-ResNet-56]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet56_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet56_train.pdf}}
    \caption{Test(left)/Training (right) accuracy for different sample sizes.}
    \label{fig:cifar-all}
\end{figure*}

\subsection{Sensitivity analysis for CIFAR100}
\label{sec:sensitivity-cifar}
\begin{figure*}[ht]
    \centering
    \subfigure[\label{fig:cifar-k}\# of neighbours $k$.]{\includegraphics[width=.32\textwidth]{figures/cifar/cifar_k.pdf}}
    \subfigure[\label{fig:cifar-sigma}Radius parameter $\sigma$.]{\includegraphics[width=.32\textwidth]{figures/cifar/cifar_sigma.pdf}}
    \subfigure[\label{fig:cifar-temp}Temperature $T$.]{\includegraphics[width=.32\textwidth]{figures/cifar/cifar_temp.pdf}}
    \caption{
    Sensitivity analysis of ResMem hyperparameters.
    The $y$-axis represents the CIFAR100 test accuracies,
    and the
    $x$-axis represents the sweeping of respective hyperparameters.}
    \label{fig:robustness-hyper}
\end{figure*}

\paragraph{Varying locality parameter $k$ and $\sigma$.}
We vary the number of neighbours from $k=27$ to $k=500$.
We find that  ResMem test accuracy is relatively stable across the choice of the number of neighbours (cf.~Figure \ref{fig:cifar-k}).
The trend of the curve suggests that as $k\rightarrow\infty$, the ResMem test accuracy seems to be converging to a constant level.
For $\sigma$, we explored different values of $\sigma \in (0.1, 2.0)$. We observe that the test accuracy has a unimodal shape as a function of $\sigma$, suggesting that there is an optimal choice of $\sigma$ (cf.~Figure \ref{fig:cifar-sigma}).

\paragraph{Varying temperature $T$ and connection to distillation.}
We tried $T=0.1$ to $T=5$, and also identified an unimodal shape for the test accuracy (Figure \ref{fig:cifar-temp}).
The fact that we can use different temperatures for (a) training the network and (b) constructing the $k$-NN predictor reminds us of the well-established knowledge distillation  procedure \citep{distillation}.
In knowledge distillation, we first use one model (the teacher network) to generate targets at a higher temperature, and then train a second model (the student network) using the \emph{combination} of the true labels and the output of the first network.

ResMem operates in a reversed direction:
Here we have a second model ($\knn$) that learns the \emph{difference} between true labels and the output of the first model.
In both cases, we can tune the temperature of the first model to control how much information is retained.
This connection offers an alternative perspective that regards ResMem as a ``dual procedure'' to knowledge distillation.

\section{Additional ImageNet results}
\label{sec:add-imagenet}
This section includes additional experiment results on applying ResMem to ImageNet dataset.

\subsection{Sensitivity analysis for ImageNet}
\label{sec:sensitivity-imagenet}
\begin{figure*}[ht]
    \centering
    \subfigure[\label{fig:imagenet-k}\# of neighbours $k$.]{\includegraphics[width=.32\textwidth]{figures/imagenet/imagenet_k.pdf}}
    \subfigure[\label{fig:imagenet-sigma}Radius parameter $\sigma$.]{\includegraphics[width=.32\textwidth]{figures/imagenet/imagenet_sigma.pdf}}
    \subfigure[\label{fig:imagenet-temp}Temperature $T$.]{\includegraphics[width=.32\textwidth]{figures/imagenet/imagenet_temp.pdf}}
    \caption{
    Sensitivity analysis of ResMem hyperparameters.
    The $y$-axis represents the imagenet100 test accuracies,
    and the
    $x$-axis represents the sweeping of respective hyperparameters.}
    \label{fig:imagenet-robustness-hyper}
\end{figure*}

We vary the number of neighbours from $k=15$ to $k=90$.
Similar to CIFAR100 finding,  ResMem test accuracy is relatively stable across the choice of the number of neighbours (cf.~Figure \ref{fig:imagenet-k}).
For $\sigma$, we explored different values of $\sigma \in (0.3, 2.1)$. We observe that the test accuracy has a unimodal shape as a function of $\sigma$, suggesting that there is an optimal choice of $\sigma$ (cf.~Figure \ref{fig:imagenet-sigma}).
We tried $T=0.3$ to $T=2.1$, and also identified an unimodal shape for the test accuracy (Figure \ref{fig:imagenet-temp}).

\subsection{ImageNet-LT}
\label{sec:imagenet-lt}
We also evaluated the proposed ResMem on ImageNet-LT~\citep{liu2019large} -- a long-tail version of ImageNet~\citep{ILSVRC15}.
We again utilize ResNet~\citep{resnet} architecture.
In particular, we studied the performance of ResMem with both ResNet-18 and ResNet-50 as the underlying DeepNet.
We trained the DeepNet for 90 epochs by using SGD with Nesterov momentum.
We employ cosine learning rate schedule with the base learning rate of 0.4, and use 0.9 as momentum parameter.
Based on a coarse hyperparameter search, we utilized $\sigma = 2.0$ and temperature $T=.8$ in ResMem for both ResNet18 and ResNet50 as the underlying DeepNet.
As for the number of neighbors, we utlized $k = 70$ and $k = 90$ for ResNet18 and ResNet50, respectively.

Table~\ref{table:imagenet-lt} showcases the performance of both DeepNet an ResMem on ImageNet-LT. As expected the (balanced) accuracy on ImageNet-LT is worse that that on ImageNet.
However, it is also evident from the table that our proposed ResMem leads to performance improvement even in the a large-scale long-tail setting.

\begin{table}[ht]
\centering
\caption{Test accuracy for ResMem and baseline deep network.}
\label{table:imagenet-lt}
\vskip 0.1in
\scalebox{1}{
\begin{tabular}{@{}cccc@{}}
\toprule
\multirow{2}{*}{Dataset}  & \multirow{2}{*}{Architecture}  & \multicolumn{2}{c}{Test accuracy} \\
\cmidrule(l){3-4}
& & \nn  & \alg \\
\toprule
ImageNet & ResNet18 & 66.86\%  & \textbf{67.44\%} \\
ImageNet-LT & ResNet18 & 39.56\%  & \textbf{39.91\%} \\
ImageNet-LT & ResNet50 & 46.18\%  & \textbf{46.39\%} \\
\toprule
\end{tabular}
}
\end{table}
\section{Additional details of NLP experiments}
\label{sec:detail-nlp}

The Decoder-Only model used in our experiments is essentially the normal Encoder-Decoder architecture with Encoder and Cross-Attention removed.
We pretrained both the T5-small and T5-base model on C4~\citep{2020t5} dataset with auto-regressive language modeling task for 1,000,000 steps, with dropout rate of 0.1 and batch size of 128.
The learning rate for the first 10,000 steps is fixed to 0.01 and the rest steps follow a square root decay schedule.

During the inference for retrieval key, query embeddings and residuals, we ensured every token has at least 64 preceding context by adopting a sliding window strategy, where a window of 256 token slides from the beginning to the end on each of the articles, with a stride of $256 - 64 = 192$.

For residuals, we only stored the top 128 residuals measured by the absolute magnitude, as the residual vector is as large as T5 vocabulary size (i.e., 32128), and storing all 32128 residuals for each token is too demanding for storage. However, when weight-combining the residuals, we zero filled the missing residuals so that all the residual vectors have 32128 elements.

\newpage
\section{Some concentration results for uniform random variables}
In this section, we state some concentration results that are useful for the theoretical analysis in Section \ref{sec:thy}.
Let $\bxt, \bx_1, \dots, \bx_n\overset{\rm i.i.d.}{\sim}$Unif$(\cB_{\bzero, \sqrt{d+2}})$ be i.i.d. samples from the uniform distribution over the Euclidean norm ball of radius $\sqrt{d+2}$ in $\R^d$.
Let
\begin{equation}
  Z_n = \min_{\bx\in \{\bx_1, \bx_2, \dots, \bx_n\}}\|\bxt-\bx\|^2.
\end{equation}
If $n=1$, $\E Z_1$ is the sum of the variance of each coordinate of Unif$(\cB_{\bzero, \sqrt{d+2}})$.
Therefore, $\E Z_n$ provides a generalized measure of concentration.
Intuitively, $\E Z_n \rightarrow 0$ as $n\rightarrow \infty$.
The proposition below provides a upper bound on the rate of convergence.

\begin{lemma}[Nearest Neighbor concentration]
\label{lem:nearestneighbor}
Given the assumptions above
\begin{equation}
\E Z_n \lesssim  d^2 \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},
\end{equation}
where $\lesssim $ means inequality up to an universal constant independent of $d$ and $n$.
\end{lemma}

\begin{proof}
Define
\begin{equation}
\begin{aligned}
\cE_1 &= \{Z_n\leq \delta^2\},\\
\cE_2 &= \{\delta\leq \sqrt{d+2}-\|\bxt\|\}.
\end{aligned}
\end{equation}

We will compute two probabilities $\P(\cE_1|\cE_2)$ and $\P(\cE_2)$ that will be useful latter.
\begin{equation}
\begin{aligned}
\P(\cE_1^c|\cE_2) &= \P(Z_n\geq\delta^2 | \cE_2) = \P(\|\bxt-\bx_i\|\geq\delta,~\forall i|\cE_2),\\
&= \E_{\bxt} \P(\|\bxt-\bx_i\|\geq\delta|\cE_2,\bxt)^n = \E_{\bxt} (1-\P(\|\bxt-\bx_i\|\leq\delta|\cE_2, \bxt))^n,\\
&= \E_{\bxt} \left[1- \frac{\text{Vol}(\cB_{\bxt, \delta})}{\text{Vol}(\cB_{\bzero, \sqrt{d+2}})} \right]^n = \left[1 - \left(\frac{\delta}{\sqrt{d+2}}\right)^d\right]^n,\\
&\leq \exp\left[-n\left(\frac{\delta}{\sqrt{d+2}}\right)^d\right].
\end{aligned}
\end{equation}

Next, we compute $\P(\cE_2)$
\begin{equation}
\begin{aligned}
\P(\cE_2) = \P(\|\bxt\|\leq \sqrt{d+2}-\delta) = \left(\frac{\sqrt{d+2}-\delta}{\sqrt{d+2}}\right)^d = \left(1-\frac{\delta}{\sqrt{d+2}}\right)^d.
\end{aligned}
\end{equation}

We use $\cE_1$ and $\cE_2$ to compute the following upper bound
\begin{equation}
\begin{aligned}
\E Z_n &= \E(Z_n | \cE_1\cap\cE_2) \P(\cE_1\cap\cE_2) + \E (Z_n|(\cE_1\cap\cE_2)^c) P((\cE_1\cap\cE_2)^c), \\
&\leq \delta^2 + (2\sqrt{d+2})^2 \left(1 - \P(\cE_1\cap\cE_2)\right), \\
&= \delta^2 + 4(d+2) \left[1-\P(\cE_1|\cE_2)\P(\cE_2)\right].
\end{aligned}
\end{equation}

To find an upper bound for $\E Z_n$, we need to find an upper bound for $1-\P(\cE_1|\cE_2)\P(\cE_2)$.
\begin{equation}
\begin{aligned}
1-\P(\cE_1|\cE_2)\P(\cE_2) &= 1-\left[1-\P(\cE_1^c|\cE_2)\right]\P(\cE_2), \\
&= 1 -\P(\cE_2) + \P(\cE_1^c|\cE_2)\P(\cE_2),\\
&\leq 1-\P(\cE_2) + \P(\cE_1^c|\cE_2).
\end{aligned}
\end{equation}
Now choose $\delta=\sqrt{d+2}n^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d}$.
\begin{equation}
\P(\cE_1^c|\cE_2)\leq \exp\left[-n\left(\frac{\delta}{\sqrt{d+2}}\right)^d\right] = \exp\left[-nn^{-1}\log\left(n^{1/d}\right)\right] = n^{-1/d},
\end{equation}

and
\begin{equation}
\P(\cE_2) = \left(1-\frac{\delta}{\sqrt{d+2}}\right)^d\geq 1-d\frac{\delta}{\sqrt{d+2}} = 1-dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d}.
\end{equation}
Thus
\begin{equation}
1-\P(\cE_1|\cE_2)\P(\cE_2) \leq 1 - 1 + dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d} + n^{-1/d} \lesssim dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d}.
\end{equation}
Combining everything together, we get
\begin{equation}
\begin{aligned}
\E Z_n &\leq (d+2) n^{-2/d} \left[\log\left(n^{1/d}\right)\right]^{2/d} + 4(d+2)\times dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d},\\
&\lesssim d^2 n^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d},\\
&= d^2 \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.
\end{aligned}
\end{equation}
This completes the proof.
\end{proof}

\begin{proposition}[\cite{wainwright2019high} Corollary 6.20]
    \label{prop:matcon}
    Let $\bx_i\overset{\rm i.i.d.}{\sim}$Unif$~(\cB_{\bzero, \sqrt{d+2}})$ for $i=1, \dots, n$ be uniformly distributed over a ball of radius $B$ in $\R^d$ centered at $\bzero$.
    Let
    \[
        \bSigma_n = \frac{1}{n} \sum_{i=1}^n \bx_i\bx_i^\sT
    \]
    be the sample covariance matrix. Then
    \[
        \P(\|\bSigma_n - \bI\|_{\rm op} > \eps) \leq 2d\exp\left[ - \frac{n\eps^2}{2(d+2)(1+\eps)}\right].
    \]
\end{proposition}

\newpage
\section{Proof of Theorem \ref{thm:test-loss-ltm}}
\label{sec:lin-proof}
In this section, we present the proof of Theorem \ref{thm:test-loss-ltm}.
In Section \ref{sec:lin-decomp-detail}, we provide the detail of the decomposition of the risk into $T_1$ and $T_2$.
Then in Section \ref{sec:upb-t1} we compute an upper bound for $T_1$, and compute an upper bound for $T_2$ in Section \ref{sec:upd-t2}.
Finally, we combine everything together in Section \ref{sec:final-thm} and completes the proof.

\subsection{Decomposition of the test risk}
\label{sec:lin-decomp-detail}
\begin{equation}
\begin{aligned}
&~~~~~\E  \left[f^{\alg}(\bxt)-f_\star(\bxt)\right]^2
    = \E  \left[f_n(\bxt)+r_n(\bxt)-f_\star(\bxt)\right]^2,\\
    &= \E  \left[ f_n(\bxt) - f_\star(\bxt)
                      -f_n(\bxnn) + f_\star(\bxnn)\right]^2,\\
    &= \E  \left[ f_n(\bxt)
                      -f_\infty(\bxt) + f_\infty(\bxt)
                      -f_\star(\bxt)
                      -f_n(\bxnn)
                      +f_\infty(\bxnn) - f_\infty(\bxnn)
                      +f_\star(\bxnn)
                 \right]^2,\\
    &\leq 3\times [ \underbrace{\E  (f_n(\bxt)-f_\infty(\bxt))^2
                        +\E  (f_n(\bxnn)-f_\infty(\bxnn))^2}_{T_1}

                    +\underbrace{\E  (f_\infty(\bxt)
                       -f_\star(\bxt)- f_\infty(\bxnn)
                       +f_\star(\bxnn))^2}_{T_2}
                  ],
\end{aligned}
\end{equation}
where in the last inequality, we used the fact that $(a+b+c)^2<3(a^2+b^2+c^2)$ for any $a, b, c\in\R$.

\ActivateWarningFilters[pdftoc]
\subsection{Upper bound on $T_1$.}
\DeactivateWarningFilters[pdftoc]

\label{sec:upb-t1}
\noindent Since $\P_\bx = \text{Unif}(\cB_{\bzero, B}~ )$, we apply the bound $\|\bxt\|, \|\bxnn\|\leq B$ to obtain
\begin{equation}
\begin{aligned}
T_1
&= \E [f_n(\bxt) - f_\infty (\bxt)]^2 + \E [f_n(\bxnn) - f_\infty(\bxnn)]^2, \\
&= \E \<\btheta_n-\btheta_\infty, \bxt\>^2 + \E \<\btheta_n - \btheta_\infty, \bxnn\>^2, \\
&\leq \E\|\btheta_n-\btheta_\infty\|^2 \|\bxt\|^2 + \E \|\btheta_n-\btheta_\infty\|^2 \|\bxnn\|^2,\\
&\leq 2B^2\E\|\btheta_n - \btheta_\infty\|^2.
\end{aligned}
\end{equation}

As $n$ gets large, the empirical covariance matrix $\bSigma_n = \bX^\sT \bX/n$ is concentrated around its mean $\bI$.
Let $\bDelta_n = \bI - \bSigma_n$ denote this deviation.
For some $\eps \in (0, 1)$, define the following ``good event'' over the randomness in $\bSigma_n$
\begin{equation}
\cA = \{\|\bDelta_n\|_{\rm op} <\eps\},
\end{equation}

where $\|\bDelta_n\|_{\rm op}$ denotes the operator norm of the deviation matrix.
The high level idea of the proof is to condition on the event $\cA$ and deduce and upper bound of $\|\btheta_n-\btheta_\infty\|$ in terms of $\eps$.
Then, we use the fact that $\cA$ happens with high probability.

Recall that $\btheta_\infty = L \btheta_\star$, and
\begin{equation}
\btheta_n = \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \, \frac{1}{n} \|\bX\btheta-\by\|^2.
\end{equation}

Since $\by = \bX\btheta_\star$ by definition,
the Lagrangian of the convex program above is
\begin{equation}
\mathcal{L}(\btheta, \lambda) = \frac{1}{n} \|\bX\btheta-\bX\btheta_\star\|^2 + \lambda (\|\btheta\|^2 - L).
\end{equation}

The KKT condition suggests that the primal-dual optimal pair $(\btheta_n, \lambda_n)$ is given by
\begin{equation}
\begin{aligned}
\|\btheta_n\|&\leq L ,\\
\lambda_n    &\geq 0 ,\\
\lambda_n(\|\btheta_n\|-L)&=0 ,\\
\end{aligned}
\end{equation}

and at optimality
\begin{equation}
\begin{aligned}
\grad_\btheta \mathcal{L}(\btheta_n, \lambda_n) = 0 &\iff \frac{2}{n}\bX^\sT\bX(\btheta-\btheta_\star) + 2\lambda_n \btheta = 0,  \\
&\iff \btheta_n = (\bSigma_n+\lambda_n \bI)^{-1}\bSigma_n\btheta_\star.\\
\end{aligned}
\end{equation}

The complementary slackness condition
$\lambda_n(\|\btheta_n\|-L)=0$ suggests that either $\lambda_n=0$ or $\|\btheta_n\|=L$.
But if $\lambda_n = 0$, the stationary condition $\grad_\btheta \mathcal{L}(\btheta, \lambda) = 0$ would suggest that
$\btheta_n = \bSigma_n^{-1}\bSigma_n\btheta_\star = \btheta_\star \Rightarrow \|\btheta_n\|=1>L,$ a contradiction.
(Note that here $\bSigma_n$ is invertible condition on the event $\cA$.)
Therefore, we must have $\|\btheta_n\|=L$.
As a result, the primal and dual pair $(\btheta_n, \lambda_n)$ is determined by the system of equations
\begin{equation}
\begin{cases}
\btheta_n &= (\bSigma_n + \lambda_n \bI)^{-1} \bSigma_n \btheta_\star, \\
\|\btheta_n\| &= L, \\
\lambda_n &> 0.
\end{cases}
\end{equation}
Next, we proceed to compute the deviation $\|\btheta_n - \btheta_\infty\|$.
\begin{equation}
\begin{aligned}
\btheta_n
&= \left[ (\lambda_n+1)\bI - \bDelta_n\right]^{-1} \bSigma_n \btheta_\star, \\
&= (\lambda_n+1)^{-1} \left[ \bI - \frac{\bDelta_n}{\lambda_n+1}\right]^{-1}\bSigma_n \btheta_\star, \\
&= (\lambda_n+1)^{-1} \left[ \bI + \sum_{k=1}^{\infty} \frac{\bDelta_n^k}{(\lambda_n+1)^k}\right](\bI-\bDelta_n) \btheta_\star,\\
&=(\lambda_n+1)^{-1} \left[ \bI + \sum_{k=1}^\infty \frac{\bDelta_n^k}{(\lambda_n+1)^k} - \bDelta_n - \sum_{k=1}^\infty \frac{\bDelta_n^{k+1}}{(\lambda_n+1)^k} \right]\btheta_\star,\\
&=(\lambda_n+1)^{-1}\btheta_\star + (\lambda_n+1)^{-1} \bDelta_n \left[\sum_{k=1}^\infty \frac{\bDelta_n^{k-1}}{(\lambda_n+1)^k} - \bI - \sum_{k=1}^\infty \frac{\bDelta_n^{k}}{(\lambda_n+1)^k} \right]\btheta_\star,\\
&= (\lambda_n+1)^{-1}\btheta_\star + (\lambda_n+1)^{-1} \bDelta_n\left[\sum_{k=1}^{\infty} \frac{\bDelta_n^{k-1}-\bDelta_n^k}{(\lambda_n+1)^k} - \bI\right]\btheta_\star.
\end{aligned}
\end{equation}

Define
\begin{equation}
\bD_n = \bDelta_n\left[\sum_{k=1}^{\infty} \frac{\bDelta_n^{k-1}-\bDelta_n^k}{(\lambda_n+1)^k} - \bI\right].
\end{equation}

Then $\btheta_n = (\lambda_n+1)^{-1}\btheta_\star + (\lambda_n+1)^{-1}\bD_n\btheta_\star$, and
\begin{equation}
\begin{aligned}
\|\bD_n\| &\leq \|\bDelta_n\| \left[1+\sum_{k=1}^{\infty} \frac{\|\bDelta_n\|^{k-1}+\|\bDelta_n\|^k}{(\lambda_n+1)^k}\right] ,\\
&\leq \eps \left[1+2(1+\lambda_n)^{-1}\sum_{k=1}^\infty \left(\frac{\eps}{1+\lambda_n}\right)^k\right],\\
&= \eps\left(1 + \frac{2}{1+\lambda_n} \frac{1}{1 - \frac{\eps}{1+\lambda_n}}\right) \leq 3\eps.
\end{aligned}
\end{equation}

Therefore
\begin{equation}
\begin{aligned}
&L=\|\btheta_n\|^2 = (\lambda_n+1)^{-2} + (\lambda_n+1)^{-2} \btheta_\star^\sT \bD_n^\sT \bD_n \btheta_\star + 2(\lambda_n+1)^{-2} \btheta_\star \bD_n \btheta_\star, \\
\Rightarrow & (\lambda_n+1)^2L^2 = 1 + \delta_n,~ \delta_n = \btheta_\star^\sT\bD_n^\sT\bD_n\btheta_\star + 2\btheta_\star^\sT\bD_n\btheta_\star.
\end{aligned}
\end{equation}

We can obtain the following bound for $\delta_n$:
\begin{equation}
|\delta_n| \leq \|\btheta_\star\|^2 \|\bD_n\|^2 + 2\|\btheta_\star\|^2 \|\bD_n\| \leq 9\eps^2 + 6\eps \leq 15\eps.
\end{equation}

Since $1-\delta_n/2\leq \sqrt{1+\delta_n}\leq 1+\delta_n/2$, and $|\delta_n|\leq 15\eps$, we obtain
\begin{equation}
|(\lambda_n+1)L-1|\leq \frac{15\eps}{2} \Rightarrow \left|L-(\lambda_n+1)^{-1}\right|\leq \frac{15\eps}{2}(\lambda_n+1)^{-1} \leq \frac{15\eps}{2},
\end{equation}

where the last inequality follows as we have $\lambda_n > 0$.
Finally,
\begin{equation}
\begin{aligned}
\btheta_n -\btheta_\infty~~~ &= (\lambda_n+1)^{-1}\btheta_\star - L \btheta_\star + (\lambda_n+1)^{-1}\bD_n\btheta_\star,\\
\Rightarrow \|\btheta_n-\btheta_\infty\|^2 &= [(1+\lambda_n)^{-1}-L]^2 + (1+\lambda_n)^{-2} \btheta_\star \bD_n^\sT\bD_n\btheta_\star + 2(\lambda_n+1)^{-1}[(1+\lambda_n)^{-1}-L] \btheta_\star \bD_n \btheta_\star,\\
&\leq 64\eps^2 + 9\eps^2 + 45\eps^2 = 118\eps^2, \\
\Rightarrow \|\btheta_n - \btheta_\infty\|^2 &\lesssim \eps^2.
\end{aligned}
\end{equation}

Combine the above result with Proposition \ref{prop:matcon}, we get that

\begin{equation}
\begin{aligned}
\E\|\btheta_n - \btheta_\infty\|^2 &= \E (\|\btheta_n - \btheta_\infty\|^2|\cA)\P(\cA) + \E(\|\btheta_n - \btheta_\infty\|^2|\cA^c)\P(\cA^c), \\
&\leq \eps^2 + 4L^2 \times 4d\exp\left[ -\frac{n\eps^2}{2(d+2)(1+\eps)} \right], \\
\end{aligned}
\end{equation}

If we choose $\eps = n^{-1/3}$, we get
\begin{equation}
\E\|\btheta_n - \btheta_\infty\|^2 \lesssim dL^2n^{-2/3},
\end{equation}

which implies that
\begin{equation}
T_1 \lesssim d^2L^2n^{-2/3}.
\end{equation}

\ActivateWarningFilters[pdftoc]
\subsection{Upper bound on $T_2$.}
\DeactivateWarningFilters[pdftoc]

\label{sec:upd-t2}
Plugging in the formula for $f_\perp(\bxt) = f_\star(\bxt)-f_\infty(\bxt) = \<\bxt, \btheta_\perp\>$, we get
\begin{equation}
\begin{aligned}
  T_2 &= \E [f_\perp(\bxnn)-f_\perp(\bxt)]^2,\\
  &= \E\<\btheta_\perp, \bxnn-\bxt\>^2,\\
  &\leq (1-L)^2\|\btheta_\star\|^2 \E \|\bxt-\bxnn\|^2,\\
  &= (1-L)^2 \E \|\bxt-\bxnn\|^2,
\end{aligned}
\end{equation}

where in the last inequality, we used the relation that $\btheta_\perp = (1-L)\btheta_\star$.
Proposition \ref{lem:nearestneighbor} suggests that
\begin{equation}
\E \|\bxt-\bxnn\|^2 \lesssim d^2 \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},
\end{equation}

which implies
\begin{equation}
T_2 \lesssim d^2(1-L)^2  \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.
\end{equation}

\begin{remark}[Comparison with pure nearest neighbor and ERM]
If we rely solely on nearest neighbor method, the prediction error is
\begin{equation}
  \E[f_\star(\bxt)-f_\star(\bxnn)]^2 = \E \<\bxt-\bxnn, \btheta_\star\>^2 \leq \E \|\bxt-\bxnn\|^2.
\end{equation}

On the other hand, if we solely rely on ERM, even with infinite sample, we get
\begin{equation}
  \E[f_\star(\bxt)-f_\infty(\bxt)]^2 = \E \<\bxt, \btheta_\star-\btheta_\infty\>^2 \leq (1-L)^2\E \|\bxt\|^2.
\end{equation}

We can see from the upper bound that $\alg$ takes advantage of both
\begin{itemize}
    \item Projecting $f_\star$ onto $f_\infty$, so that the dependence on the prediction function is reduced from $1$ to $(1-L)^2$.
    \item Memorizing the residuals using nearest neighbor, so that the variance is reduced from $\E\|\bxt\|^2$ to $\E\|\bxnn-\bxt\|^2$.
\end{itemize}
\end{remark}

\ActivateWarningFilters[pdftoc]
\subsection{Test loss for $\alg$.}
\DeactivateWarningFilters[pdftoc]
\label{sec:final-thm}

If we combine the previous two parts together, we get
\begin{equation}
\begin{aligned}
\E \left[\hat{f}(\bxt) - f_\star(\bxt)\right]^2 &\lesssim d^2L^2n^{-2/3} + d^2(1-L)^2  \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.
\end{aligned}
\end{equation}

\noindent This completes the proof of Theorem~\ref{thm:test-loss-ltm}.
\section{Introduction}
\input{intro}

\input{method}

\section{Empirical results}
\label{sec:exp}
\input{experiments}

\section{Theoretical results}
\label{sec:thy}
\input{theory}

\section{Discussion and future work}
\label{sec:diss}
\input{discussion}

\section*{Acknowledgements}
\input{ack}

\subsection{Experimental setup}
\label{sec:exp-setup}

\paragraph{CIFAR100.}
For the CIFAR100~\citep{cifar}, we use CIFAR-ResNet8 \citep{resnet} as the base DeepNet.
For the DeepNet training, we use SGD with batch size 128, trained for 256 epochs.
We use a peak learning rate 0.4, and momentum 0.9.
We warm up the learning rate linearly for the first 15 epochs, and decay the learning rate by $0.1$ after epochs $\{ 96, 192, 224 \}$.
For ResMem, we use the pre-logit layer as the image embedding, which has dimension 64.
For the nearest neighbor search (Step 3, Section \ref{sec:method}), we define the distance between two images to be the $\ell_2$ distance between their embeddings.
We use $\sigma=0.7$, $k=53$, and $T=1.4$ to compute the weights for the nearest neighbor regressor.
We provide the sensitivity analysis of test accuracy against ResMem parameters in Appendix \ref{sec:add-cifar} (cf. Figure \ref{fig:robustness-hyper}).
We can see that ResMem boosts the test accuracy of CIFAR-ResNet8 from 56.46\% to 59.66\%, which is between the base DeepNet test accuracy for CIFAR-ResNet8 and CIFAR-ResNet14 (cf. Figure \ref{fig:cifar-arch}).

\paragraph{ImageNet.}
The experimental setting for ImageNet~\citep{ILSVRC15} is similar to CIFAR100.
We employ a ResNet-18 model \citep{resnet} and employ the following training procedure.
We train the DeepNet for 120 epochs with the batch size of 90. We use an initial learning rate equal $0.8$ and an annealing schedule, decreasing the learning rate 10 times after 30, 60 and 80 epochs.
We also use warm-up for the first 5 epochs of training, linearly increasing the learning rate from 0.1 up to its initial learning rate.
For ResMem, we again use the second last layer of DeepNet as a $512$-dimensional embedding of an image and rely on the $\ell_2$ distance between the embeddings for nearest neighbor search (Step 3, Section \ref{sec:method}).
We use $\sigma=1$, temperature $T=1$ and $k=80$ to compute the weights for the nearest neighbor regressor.
We provide a similar sensitivity analysis of ResMem hyperparameter for ImageNet in Appendix \ref{sec:sensitivity-imagenet} (cf. Figure \ref{fig:imagenet-robustness-hyper}).
The ResMem improvement on ImageNet is relatively small compared with CIFAR100, since the base ResNet18 test accuracy at 66.86\% is already high compared with other base DeepNet accuracies.
In addition to the standard ImageNet dataset, we also evaluate ResMem in ImageNet-LT~\citep{liu2019large}, a long-tailed version of ImageNet, and observed similar improvement (cf. Appendix \ref{sec:imagenet-lt}).

\paragraph{Natural language experiments}
For the language experiment, we use a Decoder-Only T5-small~\citep{2020t5} model and C4~\citep{2020t5} dataset.
C4 is generated from scraping the internet and commonly used as a pretraining dataset or part of the pretraining mix.
We pre-trained the DeepNet on C4 training split with auto-regressive language modeling task.
For experimental efficiency, we used 1\% of the C4 training split (which corresponds to 1,639 million tokens) as the retrieval database, and extracted last transformer layer's pre-MLP, post-LayerNorm representations as the key embeddings for $k$NN search, and we created the query embeddings using the whole validation split and the same representation location.
For each query, we retrieved 50 neighbors with $L_2$ distance using ScaNN~\citep{avq_2020}.
We used the temperature $T = 1$ for the residual computation and $\sigma = 1$ for computing the neighbor weights.
The predicted token is the one with highest probability, similar to greedy decoding, and we measured the prediction accuracy to match the vision experiments.
The result in Table~\ref{table:main-table} shows that ResMem improved test accuracy by 2.86\% which is around the accuracy (40.08\%) of a T5-base Decoder-Only model trained in a similar way. Please see Appendix~\ref{sec:detail-nlp} for details on base model training, data processing etc.

\subsection{Where does the improvement come from?}
\label{sec:emp-ana}

\begin{figure*}[t]
    \centering
    \includegraphics[width=\textwidth]{figures/tpr.pdf}
    \caption{Examples from CIFAR100 and C4 test set with the property that
    \textbf{(i)} $y^\textsf{ResMem}$ is correct;
    \textbf{(ii)} $y^\textsf{DeepNet}$ is wrong but close in meaning.
    We use red to denote the ResMem prediction and blue to denote DeepNet prediction.
    The DeepNet predictions capture \emph{coarse} structure (e.g., predicting \texttt{poppy} for a sample whose true label is \texttt{rose}),
    which can be refined by ResMem capturing the remaining \emph{fine-grained} structure.}
    \label{fig:cifar-nlp-analysis}
\end{figure*}

In this section, we discuss the accuracy improvement in Table \ref{table:main-table}.
Let ${\sf Gain}_{\sf ResMem}$ be the difference between the test accuracy of ResMem and baseline DeepNet:
\[
{\sf Gain}_{\sf ResMem} = L_{01}(f_\alg) - L_{01}(f_\nn),
\]
where $L_{01}$ is the misclassification error as defined in equation \eqref{eqn:L01}.
We offer a decomposition of ${\sf Gain}_{\sf ResMem}$ that sheds light into the mechanism behind ResMem.
For a test set $\{(x_i, y_i)\}_{i=1}^m$, let $y_i^\alg$ be the ResMem prediction on instance $x_i$ and let $y_i^\nn$ be the baseline neural network prediction on $x_i$.
When $y^\alg_i = y^\nn_i$, sample $x_i$ does not contribute to ${\sf Gain}_{\sf ResMem}$.
When $y^\alg_i\neq y^\nn_i$, this could arise
either from the desirable event where the deep network misclassifies while ResMem classifies correctly;
or from the undesirable event where the ResMem misclassifies, while the deep network classifies correctly.
These can be summarized by the \textsf{\small TPR} (\textit{true positive rate})
and \textsf{\small FPR} (\textit{false positive rate})
respectively:
\begin{equation}
    \textsf{\small TPR} = \frac{1}{m}\sum_{i=1}^m \mathds{1}\{y^\nn_i\neq y_i ~\text{and}~ y^\alg_i = y_i\}.
\end{equation}
\begin{equation}
    \textsf{\small FPR} = \frac{1}{m}\sum_{i=1}^m \mathds{1}\{y^\nn_i= y_i ~\text{and}~ y^\alg_i \neq y_i\}.
\end{equation}

Note that~${\sf Gain}_{\sf ResMem} = \textsf{\small TPR} - \textsf{\small FPR}.$
Table~\ref{tab:TPR_FPR} summarizes $\textsf{\small TPR}$, $\textsf{\small FPR}$, and ${\sf Gain}_{\sf ResMem}$ for the three main experiments discussed in Table~\ref{table:main-table}.
The decomposition of ${\sf Gain}_{\sf ResMem}$ says that the gain of ResMem came from the \textsf{\small TPR} samples, provided they outweigh the \textsf{\small FPR} samples.

\begin{table}[ht]
\centering
\caption{True positive rate (TPR) and false positive rate (FPR) for vision and language experiments in Table~\ref{table:main-table}.}
\label{tab:TPR_FPR}
\vskip 0.1in
\scalebox{1}{
\begin{tabular}{@{}ccccc@{}}
\toprule
Dataset & Architecture & \textsf{\small TPR} & \textsf{\small FPR} & ${\sf Gain}_{\sf ResMem}$ \\
\toprule
CIFAR100 & ResNet8  & 5.89\%  & 2.70\%  & \textbf{3.19\%} \\
ImageNet & ResNet18 & 3.15\%  & 2.57\%  & \textbf{0.58\%} \\
C4-subset & T5-small & 5.37\%  & 2.44\%  & \textbf{2.93\%} \\
\toprule
\end{tabular}
}
\vskip -0.1in
\end{table}

\paragraph{Analysis of \textsf{\small TPR} samples}
Focusing on the test samples where ResMem indeed helps ($y^\nn_i\neq y_i$ and $y^\alg_i = y_i$), we identify a common underlying pattern:
while the DeepNet makes an incorrect prediction, it still captures some coarse structure.
For example, in the CIFAR100,
one sample has correct label
$
y_i = y^\alg_i = \texttt{rose},
$
but baseline DeepNet predicts
$
y^\nn_i = \texttt{poppy}
$
(cf. Figure~\ref{fig:cifar-nlp-analysis}).
In this case, both \texttt{rose} and \texttt{poppy} are flowers.
We find similar behavior for the language modeling task (cf. Figure~\ref{fig:cifar-nlp-analysis}).

The empirical analysis suggests the DeepNet in isolation can already learn some large scale structures, but is unable to make fine-grained distinctions.
This is where ResMem helps: \emph{ResMem helps memorize information in the training label that the DeepNet cannot learn.}

\paragraph{Additional insights from the decomposition.}
We discuss an alternative way of choosing hyperparameters inspired by the decomposition of ${\sf Gain}_{\sf ResMem}$ into \textsf{\small TPR} and \textsf{\small FPR}.
In this paper, we choose the ResMem hyperparameters that minimizes the test error on the validation set or, equivalently, maximize ${\sf Gain}_{\sf ResMem}$.
With the decomposition of ${\sf Gain}_{\sf ResMem}$, an alternative hyperparameter selection procedure is to use
the following optimization problem:
\[
{\rm maximize}_{\textsf{\small FPR}(\texttt{hyperparam.})<0.05} \textsf{\small TPR}(\texttt{hyperparam.}),
\]
which ensures that ResMem modifies the predictions of DeepNet in a more conservative manner.
In particular, bounding \textsf{\small FPR} implies that ResMem has minimal impact on the examples where DeepNet already makes correct predictions.
At the same time, a higher value of \textsf{\small TPR} corresponds to maximization of the desirable occurrences where ResMem can correct a wrong prediction by DeepNet.

\subsection{When does ResMem help?}
\label{sec:exp-rob}

\begin{figure*}[t]
    \centering
    \subfigure[\label{fig:cifar-sample}Test(left)/Training (right) accuracy for different sample sizes.]{
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_subset.pdf}
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_subset_train.pdf}}
    \subfigure[\label{fig:cifar-arch}Test(left)/Training (right) accuracy for different architectures.]{
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_arch.pdf}
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_arch_train.pdf}}
    \caption{
    Robustness of ResMem improvement with respect to training sample size and deep network architecture.
    \textbf{(a):} Using $10\%, 20\%, \dots, 100\%$ of training data.
    \textbf{(b):} Using progressively larger CIFAR-ResNet architecture.
    }
    \label{fig:robustness-sample-arch}
\end{figure*}

In this subsection, we present several experiments on CIFAR100 that demonstrates the robustness of the ResMem improvement.

\paragraph{Varying sample size.}
The main experiment~(cf.~Table~\ref{table:main-table}) uses the entire CIFAR100 training data.
We repeat the experiment with subsets ($10\%, 20\%, \dots, 100\%$) of CIFAR100 training data (subsampled uniformly across different classes).
On the left (right) of Figure \ref{fig:cifar-sample}, we report the test (training) accuracy of ResMem and baseline DeepNet.
As a sanity check, we can see that ResMem always achieves perfect training accuracy, and the DeepNet training accuracy decreases as samples increase (since it's harder to fit larger dataset).
We can see that ResMem yields \emph{progressively larger margin of improvement when more data is used}.
This trend suggests a desirable property of ResMem: in real problems where the dataset is extremely large, ResMem is expected to bring even greater benefit.

\paragraph{Varying the DeepNet size.}
We again consider the main CIFAR100 experiment, which was only for a CIFAR-ResNet8 model.
We repeat the ResMem evaluation for different models in the CIFAR-ResNet family, i.e., CIFAR-ResNet-$\{8, 14, 20, 32, 44, 56\}$.
We use the same choice of training hyperparameter for all six CIFAR-ResNet models, so that only the choice of model architecture changes.
DeepNet has progressively better training accuracy as model gets larger, which is intuitive (cf. Figure \ref{fig:cifar-arch}).
We find that ResMem yields the \emph{largest margin of improvement when the smallest architecture is applied}, i.e., when it is difficult for the DeepNet itself to interpolate/memorize the data.

\section{Background and Notation}

In this section, we fix notation and discuss the relevant literature that motivates us to propose the ResMem algorithm.

\subsection{Notation}
We consider multi-class classification problems over instances $\mathscr{X}$ and labels $\mathscr{Y} \stackrel{.}{=} \{1, 2, \dots, L\} = [L]$.
Given training examples $S = \{ ( x_i, y_i ) \}_{i \in [n]} \in ( \mathscr{X} \times \mathscr{Y} )^n$,
the goal is to learn a scorer $f \colon \mathscr{X} \to \mathbb{R}^L$ that,
given an instance,
assigns an affinity score for each label.
Such an $f$ should minimize the \emph{misclassification error} on test samples:
\begin{equation}
\label{eqn:L01}
L_{01}( f ) \stackrel{.}{=} \Pr_{(x, y)}( y \neq {\tt pred}( f( x ) ) ),
\end{equation}
where ${\tt pred}( z ) \stackrel{.}{=} \argmax_{y' \in [L]} z_{y'}$,
and $\Pr$ is the distribution over labelled instances.
To achieve this, one typically minimizes the \emph{empirical loss}
\[
\hat{L}_\ell( f )\stackrel{.}{=} \frac{1}{n} \sum_{i \in [n]} \ell( y_i, f( x_i ) ),
\]
where $\ell \colon [L] \times \mathbb{R}^L \to \mathbb{R}_+$ is a loss function.
Ideally, one would like to use
$\ell_{01}( y, f( x ) ) \stackrel{.}{=} 1( y \neq {\tt pred}( f( x ) ) )$;
for computational tractability, it is popular to instead use a \emph{surrogate loss}, such as the softmax cross-entropy:
$$ \ell( y, f( x ) ) = -\log \frac{\exp( f_y( x ) )}{\sum_{y' \in [L]} \exp( f_{y'}( x ) )}. $$

\subsection{Memorization for generalization: prior work}

\paragraph{Memorization is sufficient for generalization.}
Overparameterised neural models with many more parameters than training samples have the capacity to perfectly fit (or \emph{interpolate}) even random training labels~\citep{Zhang:2017};
i.e., they can drive the empirical loss $\hat{L}_\ell( f ) \to 0$ for \emph{any} training set.
At the same time, when trained on real-world datasets,
increasing model complexity tends to \emph{improve} model performance~\citep{Neyshabur:2019, pmlr-v119-yang20j};
that is, the models do not simply memorize the training sample, but rather learn generalizable patterns.

Several works have sought to understand the reasons behind this behaviour, both empirically~\citep{Arpit:2017} and theoretically~\citep{Bartlett:2017,Dziugaite:2017,Brutzkus:2018,Belkin:2018,Neyshabur:2019,Liang:2020,Montanari:2020,Bartlett:2020,Vapnik:2021,Wang:2021b, pmlr-v139-yang21a}.
One recurring message from the theory is that memorization (in the form of interpolation) can be sufficient for generalization.

\paragraph{Memorization can be necessary for generalization.}
Some recent works~\citep{Feldman2020,CJK2022Mem} showed that memorization
---
either in the sense of interpolation, or in a more general sense of stability~\citep{Feldman2020_DiscoveringLT}
---
may be necessary for generalization.
\citet{Feldman:2019} considered a setting where the label distribution exhibits a \emph{long-tailed} distribution, and showed that to prevent incurring a large error on the large number of under-represented classes, it may be necessary to memorize many of their associated training samples.
\citet{CJK2022Mem} considered a linear regression setting where it is beneficial to fit the training targets to error lower than the Bayes-error (i.e., the inherent noise in the targets).

\section{Our method: ResMem}
\label{sec:method}

Prior works suggests that \emph{implicit} memorization through increase in model size leads to better generalization.
When such implicit memorization is infeasible
---
e.g., because of the increased computational cost of using a large model
---
is it possible to perform \emph{explicit} memorization and get generalization benefits?\\

\noindent This motivates us to the propose the ResMem algorithm (c.f. Figure \ref{fig:cartoon}), which operates as follows:

\begin{tcolorbox}[colback=black!10]
\begin{center}
    {\large \textbf{ResMem} --- Residual Memorization}
\end{center}
\begin{enumerate}
    \item \textbf{Train the base DeepNet.}
    Train a neural network $f_\nn$ on the training samples $S$ as usual.

    %
    \item \textbf{Prepare the residual data.} Compute the \emph{residual} \footnote{For an overparameterized network that perfectly fits the training sample, the residuals will all be $0$. However,
    we are interested in employing
    smaller networks, which will likely have non-zero residuals.} prediction of each training example as
    \[
    r_i = \textsf{onehot}(y_i) - \textsf{softmax}(f_\nn(x_i)/T),~\forall~i\in [n],
    \]
    where $\textsf{onehot} \colon \mathscr{Y} \to \R^L$ is the standard encoding that maps the label to a probability vector.
    Here, $T$ is a hyperparameter corresponding to the ``temperature'' scaling of the softmax operation.
    Then, we employ the output of an intermediate layer of the base network $f_\nn$, denoted by $z_i = \phi(x_i)$, as an embedding for the training instance $x_i$.
    These embeddings are utilized for the nearest neighbor search in the next step.
    %
    \item \textbf{Predict via memorized residuals.}
    To make a prediction on a test sample $\tilde{x}\in\mathscr{X}$, first compute its embedding $\tilde{z} = \phi(\tilde{x})$.
    Then, use soft $k$-nearest neighbor method to build a function $r_{\knn}$ defined by weights $\overline{w}_i(\tilde{x})$:
    \begin{equation}
        r_{\knn}(\tilde{x}) = \sum_{i=1}^n \overline{w}_i(\tilde{x}) \cdot r_i.
    \end{equation}
    The weights $\overline{w}_i(\tilde{x})$ satisfy $\sum_{i}\overline{w}_i(\tilde{x}) = 1$, and
    are computed from raw weights $w_i$ as follows:
    \[
    w_i = \exp(-\|\tilde{z}-z_i\|_2/\sigma),
    \quad
    \overline{w}_i(\tilde{x}) \propto \mathds{1}\left(w_i\geq w_{(k)}\right) w_i.
    \]
    Here $w_{(k)}$ represents the $k$-th largest entry of $w_i$'s, and $k, \sigma$ are two hyperparameters that collectively controls the locality of nearest neighbor search.
\end{enumerate}

\noindent We make the final prediction based on the following scorer:
\begin{equation}
\label{eqn:resmem-pred-classification}
f_\alg(\tilde{x}) = \textsf{softmax}(f_\nn(\tilde{x})/T) + r_{\knn}(\tilde{x}).
\end{equation}
\end{tcolorbox}

\begin{remark}[Classification versus regression]
The above introduces ResMem in a classification setting,
which serves as our primary focus in the experiments in Section \ref{sec:exp}.
That said, we emphasize that ResMem can easily be extended to other supervised learning problems beyond classification with neural networks.
For example, in Section \ref{sec:thy}, we analyze ResMem in a regression setting.
\end{remark}

\begin{remark}[Explicit memorization]
Smaller $k$ or $\sigma$ corresponds to putting higher weight on residuals of the closest neighboring training examples.
For sufficiently small $k$ and $\sigma$, $f_\alg$ achieves exact memorization of the training sample, i.e., ${\tt pred}(f_\alg(x_i))=y_i$ for all $i \in [n]$.
\end{remark}

\subsection{Applicable scenarios of ResMem}
We posit that $f_\alg$ yields the largest margin of improvement over $f_\nn$ when the learning task is ``hard'' (and thus \emph{implicit} memorization is infeasible).
We discuss three such scenarios below.
Each of our main empirical or theoretical results roughly corresponds to one of these settings.

\begin{itemize}[leftmargin=5mm, itemsep=1mm, partopsep=0pt,parsep=0pt]
    %
    \item \textbf{Hard generalization.} In this scenario, the Bayes-optimal decision boundary is very complex, and is beyond the capability of the neural network itself.
    %
    To demonstrate this, we analyze a theoretical linear regression problem where the target regression function is not contained in the hypothesis class (cf.~Section \ref{sec:thy}).

    %
    \item \textbf{Hard interpolation.} Here, the number of samples is large enough to make training set interpolation (achieving zero training error) infeasible for a given neural network model.
    In contrast, ResMem can circumvent the large sample size by explicitly memorizing the training samples.
    For example, current large language models (LLMs) are typically trained for at most a single epoch over trillions of examples~\citep{Chowdhery:2022}.
    %
    We emulate this scenario by considering a causal language modeling task (cf.~Section \ref{sec:exp}).

    %
    \item \textbf{Efficiency constraint.} In many practical settings, one may prefer a smaller model over a state-of-the-art model due to the training and deployment cost constraints.
    %
    We emulate such a setting through an image classification task where it is indeed feasible to memorize the training data perfectly using state-of-the-art neural networks, but instead, we use smaller neural networks for computational efficiency (cf. Section \ref{sec:exp}).
\end{itemize}

\subsection{Relation to existing algorithms}

\paragraph{Nearest neighbor method.} The nearest neighbor~\citep{NNCover, KnuthAOP, Hastie:2001, locallearning} method assigns label to a test sample based on the label of its nearest neighbor(s) in the training set.
Owing to its \emph{simplicity}, \emph{flexibility} in defining input similarity, and \emph{computational efficiency} with various approximation schemes~\citep{Gionis1999SimilaritySI, Muja2009FastAN}, this method remains popular.
However, the performance of nearest neighbor drops as data becomes high dimensional~\citep{nnrates, Muja2009FastAN}.
Therefore, to apply it to high dimensional data such as image and text~\citep{knnsvm}, one approach is to learn a representation of data using neural networks~\citep{pmlr-v2-salakhutdinov07a}.
Following this approach, \citep{arXivDnnKnn} finds that applying nearest neighbor methods directly to memorize the training labels $y_i$ yields similar performance with the original softmax based neural network classification.
In contrast, ResMem applies nearest neighbor to memorize the \emph{residual} $r_i$ of the training samples.

\paragraph{Boosting and residual fitting.}
Boosting algorithms such as AdaBoost~\citep{Freund:1997} seek to construct an ensemble of ``weak learner'' models with good generalization.
AdaBoost achieves this in an iterative manner, and can be interpreted as a particular instantiation of \emph{forward stage-wise regression}~\citep{Friedman:2000}, a classical procedure from statistics~\citep{Goldberger:1961,Alley:1987,Tibshirani:2015}.
Intuitively, at each round, one fits the residual of the ensemble of weak learners constructed thus far so as to minimize the exponential loss $\ell_{\rm exp}( y, f( x ) ) = e^{-y \cdot f( x ) }$.
Concretely, for binary classification problems,
AdaBoost fits a scorer
$f( x ) = \sum_{t = 1}^{T} \alpha_t \cdot h_t( x )$,
where $\alpha_t \in \mathbb{R}$,
$h_t \colon \mathscr{X} \to \{ \pm 1 \}$ is a weak classifier (e.g., a decision stump),
and
$T$ is the number of rounds of fitting.
This fitting is performed iteratively, with $(\alpha_t, h_t)$ learned during the $t$-th iteration.

\paragraph{Memory-augmented language models.}
In the language modelling literature, several works explore combining neural models with an external database or memory, which can be queried to retrieve additional context~\citep{Lample:2019,Guu:2020,Borgeaud:2021,li2022decoupled}.
Closer to our work,~\citet{Khandelwal:2020} employ a linear combination of neural network and $k$-nearest neighbour classifier components.
However, our work is quite different in that our $k$-nearest neighbor components memorizes the \emph{residual} of the DeepNet prediction, whereas \citet{Khandelwal:2020} memorizes the target label directly.
Various forms of memory have also been considered in generic classification problems~\citep{Panigrahy:2021,Vapnik:2021,Wang:2022}.
This line of literature again differs from ResMem in that their memory tries to memorize labels directly, whereas ResMem memorizes the \emph{residuals}, leading to a natural combination of the neural network and the memory component.

\paragraph{Model compression for large neural networks.}
Since ResMem boosts the test accuracy of a small, non-memorizing neural network, we can also view it as a technique that allows a small network to match the performance of a larger one.
This relates to the model compression literature.
Distillation~\citep{Hinton:2015, Bucilua:2006} is a popular strategy for compressing a large neural model to a smaller one.
For a survey of other effective strategies, including pruning, see~\citet{Menghani:2021}.
In Appendix \ref{sec:sensitivity-cifar} we discuss that ResMem can be regarded as a ``dual procedure'' of distillation.
\subsection{Assumptions and setting}

In this section, we present the setup and assumptions for the stylized linear regression problem.
We consider a setting where the function class that we minimize over does \emph{not} include the ground-truth function that relates the covariates to the response.
Therefore, even with infinite samples, the test loss will decay to a positive constant.
We exactly characterize the rate of decay, and show that it converges to $0$ under ResMem. Our analysis rests on the following assumptions.

\begin{assumption}[Distribution of covariates]
\label{ass:lin-dist}
The distribution of covariate $\bx\in\R^d$, denoted by $\P_\bx$, is the uniform distribution over a Euclidean norm ball centered at the origin of radius $\sqrt{d+2}$. The choice of radius ensures that $\E_{\bx\sim\P_\bx} \bx\bx^\sT = \bI$.
\end{assumption}

\begin{assumption}[Linear regression over norm ball]
\label{ass:lin-prob-def}
Consider the problem of learning a linear function $f_\star(\bx) = \<\bx, \btheta_\star\>$ with $\|\btheta_\star\|=1$ from training data $\{(\bx_i, y_i)\}_{i=1:n}$ where $\bx_i\overset{\rm i.i.d.}{\sim} \P_\bx$ and $y_i=f_\star(\bx_i)$ using the function class
\begin{equation}\label{eqn:linfunc}
  \cF = \{\bx\mapsto\<\bx,\btheta\>, \|\btheta\|<L\}.
\end{equation}
We assume $L<1$ so that the problem belongs to the ``hard generalization'' scenario discussed in Section \ref{sec:method}, where the hypothesis space is inadequate to fit the function on its own.
\end{assumption}

ResMem proceeds by first learning a linear function $f_n(\bx) = \<\btheta_n, \bx\>$ from $\cF$ through empirical risk minimization (ERM):
\begin{equation}
    \btheta_n = \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \, \frac{1}{n} \sum_{i=1}^n \left[\<\bx_i, \btheta\>-y_i\right]^2.
\end{equation}
The empirical risk minimizer $f_n$ should be thought of as the analog of $f_\nn$ in the deep learning context.
It defines a ground-truth residual function $r_\star(\bx)=f_\star(\bx)-f_n(\bx)$.
Now we fix a test covariate $\bxt\sim\P_x$.
ResMem ``memorizes'' the residual function through the $1$-nearest neighbor to $\bxt$
\begin{equation}
    r_n(\bxt) = r_\star(\bxnn) = f_\star(\bxnn) - f_n(\bxnn),
\end{equation}
where $\bxnn$ is the nearest neighbor to $\bxt$ among the training covariates $\bx_1,\dots,\bx_n$:
\[
 \bxnn = \underset{\bx\in\{\bx_1, \dots, \bx_n\}}{\operatorname{argmin}} \, \|\bx - \bxt\|.
\]
The final prediction is
\begin{equation}
\label{eqn:lin-pred-rule}
f_{n}^{\alg}(\bxt) = f_n(\bxt) + r_n(\bxt).
\end{equation}
Observe that if $\bxt$ coincides with any training sample,
$f_{n}^{\alg}(\bxt) = f_\star(\bxt)$, i.e., we have explicit memorization.

Note that we worked with $1$-nearest neighbor regressor for simplicity instead of the general $k$-nearest neighbor algorithm.
The effect of choosing different $k$ is not the main focus of this theoretical analysis.

\subsection{A decomposition of the target function}

Next, we introduce a decomposition of $f_\star$, which will help us analyze various components that make up the risk.
Define
\[
\begin{aligned}
\btheta_\infty
    &= \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \,
       \E_{\bx\sim\P_x}\left[\<\btheta, \bx\> - \<\btheta_\star, \bx\>\right]^2,\\
    &= \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \,
       \|\btheta-\btheta_\star\| = L \btheta_\star,
\end{aligned}
\]
which is what ERM learns in the limit of $n\rightarrow\infty$.
We can think of  $\btheta_\infty$ as the best function that ERM can learn.
Then, we can decompose $\btheta_\star$ into $\btheta_\star = \btheta_\infty + \btheta_\perp$, where $\btheta_\perp=\btheta_\star - \btheta_\infty$.
This decomposition can be generalized beyond linear regression.
Since $\btheta_\infty$ defines a function $f_\infty(\bx) = \<\bx, \btheta_\infty\>$, for general non-linear functions, the argument above can be generalized to the decomposition of $f_\star$ to an learnable and non-learnable part
\[f_\star =  f_\infty + f_\perp.\]
Intuitively, $f_\infty$ is the best function in $\mathcal{F}$ that ERM can learn, and $f_\perp$ is beyond the capacity of ERM due to the particular choice of function class.
ResMem approximates $f_\perp$ using the non-parametric nearest neighbor method.

\subsection{A decomposition of the prediction error}
\label{sec:lin-risk-decomp}

We now introduce a decomposition of the prediction risk that reveals how $\alg$ algorithm boosts generalization.
Note that the prediction error of $\alg$ is
\begin{equation}
\label{eqn:lin-test-loss}
\E  \left[\left(f_{n}^{\alg}(\bxt)-f_\star(\bxt)\right)^2\right].
\end{equation}
It can be decomposed into two components:
\[
\begin{aligned}
&\E  \left[f_{n}^{\alg}(\bxt)-f_\star(\bxt)\right]^2  \leq\\
&3\times[~ \underbrace{\E  (f_n(\bxt)-f_\infty(\bxt))^2
                    +\E  (f_n(\bxnn)-f_\infty(\bxnn))^2}_{T_1}
+\underbrace{\E  (f_\infty(\bxt)
                   -f_\star(\bxt)- f_\infty(\bxnn)
                   +f_\star(\bxnn))^2}_{T_2}
              ~].
\end{aligned}
\]
We provide the detail of the decomposition in Section \ref{sec:lin-decomp-detail}.
We can see that $T_1$ arises due to the difference between $f_n$ and $f_\infty$ (i.e., the estimation error), which, as we will show later, goes to $0$ as $n$ goes to infinity:
\[T_1\rightarrow 0 ~\text{as}~ n\rightarrow\infty.\]

On the other hand, $T_2$ arises due to the limited capacity of $\cF$.
It captures an irreducible error of the risk, which in general is \textbf{not} asymptotically zero.
However, because of the explicit memorization $\alg$ algorithm introduces ($\bxnn\rightarrow\bxt$ as $n\rightarrow\infty$), we also have
\[T_2\rightarrow 0 ~\text{as}~ n\rightarrow\infty.\]
This decomposition provides a statistical perspective on ResMem: it preserves the asymptotic consistency of $T_1$ as in classical learning problems while enforcing the asymptotic consistency of $T_2$ through the nearest-neighbor method.

\subsection{Main theoretical result}
Given the set up above, we are ready to state the main theoretical result of the paper, which characterizes the rate at which test risk of ResMem approaches 0. The proof of the result below can be found in Appendix \ref{sec:lin-proof}.

\begin{theorem}[Risk for $\alg$ algorithm]
\label{thm:test-loss-ltm}
For the problem defined in Assumption \ref{ass:lin-prob-def} with covariates distribution in Assumption \ref{ass:lin-dist}, the $\alg$ prediction rule $f_{n}^{\alg}(\bxt)$ defined in equation \eqref{eqn:lin-pred-rule} achieves risk \eqref{eqn:lin-test-loss}
\[
\begin{aligned}
\E \left[f_{n}^{\alg}(\bxt) - f_\star(\bxt)\right]^2
\lesssim d^2L^2n^{-2/3} + d^2(1-L)^2  \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},
\end{aligned}
\]
where $\lesssim$ denotes inequality up to a universal constant independent of $d, n$ and $L$.
\end{theorem}

The result includes contribution from two terms introduced in Section \ref{sec:lin-risk-decomp}:
\begin{itemize}
    \item $T_1\lesssim d^2L^2n^{-2/3}$ that arises due to the difference between $f_n$ and $f_\infty$.
    \item $T_2\lesssim \left[\log\left(n^{1/d}\right)/n\right]^{1/d}$ that vanishes as the nearest neighbor of the test point approaches the test point itself $\bxnn\rightarrow\bxt$.
\end{itemize}

The two terms $T_1$ and $T_2$ can be viewed as ``two stages of learning''.
Without the ResMem memorization component, we have the usual story of machine learning: $T_1\rightarrow 0$ at the usual parametric rate, and $T_2$ stays as an irreducible error, so the overall test error diminishes to a constant at a very fast rate.
With the introduction of nearest neighbor memorization procedure, $T_2$ can also be reduced to $0$ at a slower rate, whereas the fast decay of $T_1$ is still preserved.

This theoretical result shows why it is not favorable to use the $k$-nearest neighbor component to memorize the response directly:
As a corollary of setting $L=0$ in Theorem \ref{thm:test-loss-ltm}, pure nearest neighbor would result in an overall slow rate of $\approx n^{-1/d}$.
However, with ResMem, we can enjoy benefit of having the test loss being asymptotically $0$, while also enjoy the fast rate of $n^{-2/3}$ for smaller sample sizes.