[pdftoc]hyperrefToken not allowed in a PDF string

ResMem: Learn what you can and memorize the rest

Zitong Yang⋄**footnotemark: *
---
<EMAIL_ADDRESS>
⋄Department of Statistics, Stanford University
Michal Lukasik† Vaishnavh Nagarajan† Zonglin Li† Ankit Singh Rawat†
---
Manzil Zaheer† Aditya Krishna Menon† Sanjiv Kumar†
{mlukasik, vaishnavh, lizonglin, ankitsrawat,
manzilzaheer, adityakmenon<EMAIL_ADDRESS>
†Google Research, New York

11footnotetext: Work partly done at Google Research, New York.

###### Abstract

The impressive generalization performance of modern neural networks is
attributed in part to their ability to _implicitly_ memorize complex training
patterns. Inspired by this, we explore a novel mechanism to improve model
generalization via _explicit_ memorization. Specifically, we propose the
_residual-memorization_ (_ResMem_) algorithm, a new method that augments an
existing prediction model (e.g. a neural network) by fitting the model’s
residuals with a $k$-nearest neighbor based regressor. The final prediction is
then the sum of the original model and the fitted residual regressor. By
construction, ResMem can explicitly memorize the training labels. Empirically,
we show that ResMem consistently improves the test set generalization of the
original prediction model across various standard vision and natural language
processing benchmarks. Theoretically, we formulate a stylized linear
regression problem and rigorously show that ResMem results in a more favorable
test risk over the base predictor.

##  1 Introduction

(a) Step 1: learn the training set. (b) Step 2: compute the residual. (c) Step
3: memorize the residual.

Figure 1: Illustration of the _residual memorization_ (ResMem) algorithm. In a
nutshell, we first fit a small deep network $f_{\sf DeepNet}$ on the training
sample (Figure 1(a)). When this network is non-memorizing, it incurs non-zero
_residual_ errors in its predictions (Figure 1(b)). We then fit a $k$-nearest
neighbor based regressor on these residuals (Figure 1(c)). The final
prediction is given by the sum of the initial network and $k$-NN regressor
predictions. In all three figures, the $x$-axis represents the features in a
supervised learning problem. In Figure 1(a), the $y$-axis represents the
targets of prediction. In Figure 1(b) and 1(c), the $y$-axis represents the
residual of the initial fitting from Step 1.

Large neural networks achieve remarkable _generalization_ on test samples
despite _memorization_ (achieving zero training error) of training samples
(Zhang et al., 2017). Several recent analyses have established that, under
certain settings, memorization is _sufficient_ to achieve generalization
(Bartlett et al., 2017; Dziugaite and Roy, 2017; Belkin et al., 2018;
Neyshabur et al., 2019; Bartlett et al., 2020), and, more surprisingly, can
even be _necessary_ (Feldman, 2020; Feldman and Zhang, 2020b; Cheng et al.,
2022).

The above suggests that suitable memorization can be a valuable desiderata for
learning. While increasing model size is a conceptually simple strategy to
enable memorization, this has the obvious downside of significantly increasing
the cost of model training and serving. This raises a natural question: are
there alternate mechanisms to improve the memorization (and thus
generalization) of a relatively small model?

In this paper, we propose _residual memorization_ (_ResMem_), a simple yet
effective mechanism that achieves this goal (Section 3). Compared to the
_implicit_ memorization performed by large neural models, the key idea of
ResMem is to perform _explicit_ memorization via a separate $k$-nearest
neighbor component. Specifically, ResMem involves first training a standard
neural network $f_{\small\textsf{DeepNet}}$, and then explicitly _memorizing
the model’s residuals_ with a $k$-nearest neighbor based regressor
$r_{{\text{k}\textsf{NN}}}$ (cf. Figure 1). Memorization through $k$-nearest
neighbor can be efficiently computed with various approximation schemes (e.g.
Guo et al. (2020)). Subsequently, the ResMem prediction on an instance $x$ is
given by the sum of the two components, i.e.,
$f_{\small\textsf{DeepNet}}(x)+r_{\text{k}\textsf{NN}}(x)$.

We establish the efficacy of ResMem both theoretically and empirically,
particularly when implicit memorization is infeasible, e.g., when the training
set size is very large, or when the data patterns are too complex to learn
perfectly. In Section 4, we show that ResMem consistently improves test
accuracy (compared to a baseline DeepNet) on image classification tasks with
ImageNet (Russakovsky et al., 2015) and CIFAR100 (Krizhevsky, 2009), and
autoregressive language modeling on C4 (Raffel et al., 2020).

Towards understanding this improved performance, we hypothesize that ResMem
works by learning in two-stages (cf. Section 4.2). Specifically, we posit that
the initial DeepNet $f_{\small\textsf{DeepNet}}$ learns some _coarse_
structure, and ResMem $r_{\text{k}\textsf{NN}}$ supplements the DeepNet
prediction with _fine-grained_ details (cf. Figure 2). We verify our
hypothesis via qualitative analysis on CIFAR100 and C4 (Section 4.2).

To formalize the efficacy of ResMem, we formulate and analyze a stylized high
dimensional linear regression problem that captures the essence behind ResMem
(cf. Section 5). Our analysis (Theorem 5.3) shows that as the number of
samples goes to infinity, the test risk of the base predictor decreases to an
irreducible constant, whereas the risk of ResMem decreases to zero.

To summarize, our contributions are:

  1. 1.

We propose _residual-memorization_ (_ResMem_), a two-stage learning algorithm
that combines a base prediction model with a nearest neighbor regressor (cf.
Section 3);

  2. 2.

We empirically demonstrate that ResMem improves test performance of neural
networks (cf. Section 4), particularly when the training set is extremely
large (cf. Section 4.3);

  3. 3.

We theoretically analyze the rate of convergence of ResMem on a stylized
linear regression problem, and show that it can improve upon the base
prediction model (Section 5).

##  2 Background and Notation

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

###  2.1 Notation

We consider multi-class classification problems over instances $\mathscr{X}$
and labels $\mathscr{Y}\stackrel{{\scriptstyle.}}{{=}}\\{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 _misclassification error_ on test samples:

$L_{01}(f)\stackrel{{\scriptstyle.}}{{=}}\mathbb{P}_{(x,y)}(y\neq{\tt
pred}(f(x))),$ (1)

where ${\tt
pred}(z)\stackrel{{\scriptstyle.}}{{=}}\operatorname*{arg\,max}_{y^{\prime}\in[L]}z_{y^{\prime}}$,
and $\mathbb{P}$ is the distribution over labelled instances. To achieve this,
one typically minimizes the _empirical loss_

$\hat{L}_{\ell}(f)\stackrel{{\scriptstyle.}}{{=}}\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{{\scriptstyle.}}{{=}}1(y\neq{\tt pred}(f(x)))$;
for computational tractability, it is popular to instead use a _surrogate
loss_ , such as the softmax cross-entropy:

$\ell(y,f(x))=-\log\frac{\exp(f_{y}(x))}{\sum_{y^{\prime}\in[L]}\exp(f_{y^{\prime}}(x))}.$

###  2.2 Memorization for generalization: prior work

#### Memorization is sufficient for generalization.

Overparameterised neural models with many more parameters than training
samples have the capacity to perfectly fit (or _interpolate_) even random
training labels (Zhang et al., 2017); i.e., they can drive the empirical loss
$\hat{L}_{\ell}(f)\to 0$ for _any_ training set. At the same time, when
trained on real-world datasets, increasing model complexity tends to _improve_
model performance (Neyshabur et al., 2019; Yang et al., 2020); 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 (Arpit et al., 2017) and theoretically (Bartlett et al.,
2017; Dziugaite and Roy, 2017; Brutzkus et al., 2018; Belkin et al., 2018;
Neyshabur et al., 2019; Liang and Rakhlin, 2020; Montanari and Zhong, 2020;
Bartlett et al., 2020; Vapnik and Izmailov, 2021; Wang et al., 2021; Yang et
al., 2021). One recurring message from the theory is that memorization (in the
form of interpolation) can be sufficient for generalization.

#### Memorization can be necessary for generalization.

Some recent works (Feldman, 2020; Cheng et al., 2022) showed that memorization
— either in the sense of interpolation, or in a more general sense of
stability (Feldman and Zhang, 2020a) — may be necessary for generalization.
Feldman (2019) considered a setting where the label distribution exhibits a
_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. Cheng et al. (2022)
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).

##  3 Our method: ResMem

Prior works suggests that _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 _explicit_ memorization and get
generalization benefits?

This motivates us to the propose the ResMem algorithm (c.f. Figure 1), which
operates as follows:

ResMem — Residual Memorization 1. Train the base DeepNet. Train a neural
network $f_{\small\textsf{DeepNet}}$ on the training samples $S$ as usual. 2.
Prepare the residual data. Compute the _residual_ †††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_{\small\textsf{DeepNet}}(x_{i})/T),~{}\forall~{}i\in[n],$
where $\textsf{onehot}\colon\mathscr{Y}\to\mathbb{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_{\small\textsf{DeepNet}}$, 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. 3. 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_{{\text{k}\textsf{NN}}}$
defined by weights $\overline{w}_{i}(\tilde{x})$:
$r_{{\text{k}\textsf{NN}}}(\tilde{x})=\sum_{i=1}^{n}\overline{w}_{i}(\tilde{x})\cdot
r_{i}.$ (2) 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. We make the final prediction based on
the following scorer:
$f_{\small\textsf{ResMem}}(\tilde{x})=\textsf{softmax}(f_{\small\textsf{DeepNet}}(\tilde{x})/T)+r_{{\text{k}\textsf{NN}}}(\tilde{x}).$
(3)

######  Remark 3.1 (Classification versus regression).

The above introduces ResMem in a classification setting, which serves as our
primary focus in the experiments in Section 4. That said, we emphasize that
ResMem can easily be extended to other supervised learning problems beyond
classification with neural networks. For example, in Section 5, we analyze
ResMem in a regression setting.

######  Remark 3.2 (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_{\small\textsf{ResMem}}$ achieves exact memorization of the
training sample, i.e., ${\tt pred}(f_{\small\textsf{ResMem}}(x_{i}))=y_{i}$
for all $i\in[n]$.

###  3.1 Applicable scenarios of ResMem

We posit that $f_{\small\textsf{ResMem}}$ yields the largest margin of
improvement over $f_{\small\textsf{DeepNet}}$ when the learning task is “hard”
(and thus _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.

  * •

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 5).

  * •

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 (Chowdhery et al., 2022). We emulate
this scenario by considering a causal language modeling task (cf. Section 4).

  * •

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 4).

###  3.2 Relation to existing algorithms

#### Nearest neighbor method.

The nearest neighbor (Cover and Hart, 1967; Knuth, 1973; Hastie et al., 2001;
Bottou and Vapnik, 1992) method assigns label to a test sample based on the
label of its nearest neighbor(s) in the training set. Owing to its
_simplicity_ , _flexibility_ in defining input similarity, and _computational
efficiency_ with various approximation schemes (Gionis et al., 1999; Muja and
Lowe, 2009), this method remains popular. However, the performance of nearest
neighbor drops as data becomes high dimensional (Chaudhuri and Dasgupta, 2014;
Muja and Lowe, 2009). Therefore, to apply it to high dimensional data such as
image and text (Zhang et al., 2006), one approach is to learn a representation
of data using neural networks (Salakhutdinov and Hinton, 2007). Following this
approach, (Cohen et al., 2018) 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 _residual_ $r_{i}$ of the
training samples.

#### Boosting and residual fitting.

Boosting algorithms such as AdaBoost (Freund and Schapire, 1995) 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 _forward stage-wise regression_ (Friedman et al.,
2000), a classical procedure from statistics (Goldberger and Jochems, 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.

#### 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 (Lample et al., 2019; Guu et al., 2020; Borgeaud et al.,
2021; Li et al., 2022). Closer to our work, Khandelwal et al. (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 _residual_ of the DeepNet prediction,
whereas Khandelwal et al. (2020) memorizes the target label directly. Various
forms of memory have also been considered in generic classification problems
(Panigrahy et al., 2021; Vapnik and Izmailov, 2021; Wang and Shao, 2022). This
line of literature again differs from ResMem in that their memory tries to
memorize labels directly, whereas ResMem memorizes the _residuals_ , leading
to a natural combination of the neural network and the memory component.

#### 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 (Hinton et al., 2015b; Bucilǎ et al., 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 Menghani (2021).
In Appendix A.2 we discuss that ResMem can be regarded as a “dual procedure”
of distillation.

##  4 Empirical results

In this section, we present three experiments that showcase the efficacy of
ResMem. Table 1 provides an overview of our results.

Table 1: Test accuracy for ResMem and baseline deep network.

Dataset | Architecture | Test accuracy
---|---|---
DeepNet | ResMem
CIFAR100 | CIFAR-ResNet8 | 56.46% | 59.66%
ImageNet | ResNet18 | 66.86% | 67.44%
C4-subset | T5-small | 38.01% | 40.87%

In Section 4.1, we present the setup for vision (image classification on
CIFAR100 and ImageNet) and natural language experiments (autoregressive
language modeling). In Section 4.2 we conduct an empirical analysis to explain
where the improved accuracy of ResMem comes from. In Section 4.3, we show that
ResMem is particularly helpful with large-scale dataset and relatively simple
network.

###  4.1 Experimental setup

#### CIFAR100.

For the CIFAR100 (Krizhevsky, 2009), we use CIFAR-ResNet8 (He et al., 2016) 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 3), 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 A (cf. Figure 5). 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 3(b)).

#### ImageNet.

The experimental setting for ImageNet (Russakovsky et al., 2015) is similar to
CIFAR100. We employ a ResNet-18 model (He et al., 2016) 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 3). 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
B.1 (cf. Figure 6). 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 (Liu et al.,
2019), a long-tailed version of ImageNet, and observed similar improvement
(cf. Appendix B.2).

#### Natural language experiments

For the language experiment, we use a Decoder-Only T5-small (Raffel et al.,
2020) model and C4 (Raffel et al., 2020) 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 (Guo et al., 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 1 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 C for details on base model
training, data processing etc.

###  4.2 Where does the improvement come from?

Figure 2: Examples from CIFAR100 and C4 test set with the property that (i)
$y^{\textsf{ResMem}}$ is correct; (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 _coarse_ structure
(e.g., predicting poppy for a sample whose true label is rose), which can be
refined by ResMem capturing the remaining _fine-grained_ structure.

In this section, we discuss the accuracy improvement in Table 1. Let ${\sf
Gain}_{\sf ResMem}$ be the difference between the test accuracy of ResMem and
baseline DeepNet:

${\sf Gain}_{\sf
ResMem}=L_{01}(f_{\small\textsf{ResMem}})-L_{01}(f_{\small\textsf{DeepNet}}),$

where $L_{01}$ is the misclassification error as defined in equation (1). 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}^{\small\textsf{ResMem}}$ be the ResMem prediction on instance $x_{i}$
and let $y_{i}^{\small\textsf{DeepNet}}$ be the baseline neural network
prediction on $x_{i}$. When
$y^{\small\textsf{ResMem}}_{i}=y^{\small\textsf{DeepNet}}_{i}$, sample $x_{i}$
does not contribute to ${\sf Gain}_{\sf ResMem}$. When
$y^{\small\textsf{ResMem}}_{i}\neq y^{\small\textsf{DeepNet}}_{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 TPR (true positive rate) and FPR (false positive rate)
respectively:

$\textsf{\small
TPR}=\frac{1}{m}\sum_{i=1}^{m}\mathds{1}\\{y^{\small\textsf{DeepNet}}_{i}\neq
y_{i}~{}\text{and}~{}y^{\small\textsf{ResMem}}_{i}=y_{i}\\}.$ (4)
$\textsf{\small
FPR}=\frac{1}{m}\sum_{i=1}^{m}\mathds{1}\\{y^{\small\textsf{DeepNet}}_{i}=y_{i}~{}\text{and}~{}y^{\small\textsf{ResMem}}_{i}\neq
y_{i}\\}.$ (5)

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

Table 2: True positive rate (TPR) and false positive rate (FPR) for vision and
language experiments in Table 1.

Dataset | Architecture | TPR | FPR | ${\sf Gain}_{\sf ResMem}$
---|---|---|---|---
CIFAR100 | ResNet8 | 5.89% | 2.70% | 3.19%
ImageNet | ResNet18 | 3.15% | 2.57% | 0.58%
C4-subset | T5-small | 5.37% | 2.44% | 2.93%

#### Analysis of TPR samples

Focusing on the test samples where ResMem indeed helps
($y^{\small\textsf{DeepNet}}_{i}\neq y_{i}$ and
$y^{\small\textsf{ResMem}}_{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^{\small\textsf{ResMem}}_{i}=\texttt{rose},$ but baseline
DeepNet predicts $y^{\small\textsf{DeepNet}}_{i}=\texttt{poppy}$ (cf. Figure
2). In this case, both rose and poppy are flowers. We find similar behavior
for the language modeling task (cf. Figure 2).

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: _ResMem helps memorize information in the training
label that the DeepNet cannot learn._

#### Additional insights from the decomposition.

We discuss an alternative way of choosing hyperparameters inspired by the
decomposition of ${\sf Gain}_{\sf ResMem}$ into TPR and 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 FPR implies that ResMem has
minimal impact on the examples where DeepNet already makes correct
predictions. At the same time, a higher value of TPR corresponds to
maximization of the desirable occurrences where ResMem can correct a wrong
prediction by DeepNet.

###  4.3 When does ResMem help?

(a) Test(left)/Training (right) accuracy for different sample sizes.

(b) Test(left)/Training (right) accuracy for different architectures.

Figure 3:  Robustness of ResMem improvement with respect to training sample
size and deep network architecture. (a): Using $10\%,20\%,\dots,100\%$ of
training data. (b): Using progressively larger CIFAR-ResNet architecture.

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

#### Varying sample size.

The main experiment (cf. Table 1) 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 3(a), 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 _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.

#### 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
3(b)). We find that ResMem yields the _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.

##  5 Theoretical results

As discussed in Section 3, ResMem yields the largest improvement when implicit
memorization is infeasible. We provide a concrete result (cf. Theorem 5.3)
that characterizes the benefit of ResMem in a stylized linear regression
setting.

###  5.1 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 _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.

######  Assumption 5.1 (Distribution of covariates).

The distribution of covariate ${\bm{x}}\in\mathbb{R}^{d}$, denoted by
$\mathbb{P}_{\bm{x}}$, is the uniform distribution over a Euclidean norm ball
centered at the origin of radius $\sqrt{d+2}$. The choice of radius ensures
that
$\mathbb{E}_{{\bm{x}}\sim\mathbb{P}_{\bm{x}}}{\bm{x}}{\bm{x}}^{\mathsf{T}}={\bm{I}}$.

######  Assumption 5.2 (Linear regression over norm ball).

Consider the problem of learning a linear function
$f_{\star}({\bm{x}})=\langle{\bm{x}},{\bm{\theta}}_{\star}\rangle$ with
$\|{\bm{\theta}}_{\star}\|=1$ from training data
$\\{({\bm{x}}_{i},y_{i})\\}_{i=1:n}$ where ${\bm{x}}_{i}\overset{\rm
i.i.d.}{\sim}\mathbb{P}_{\bm{x}}$ and $y_{i}=f_{\star}({\bm{x}}_{i})$ using
the function class

${\mathcal{F}}=\\{{\bm{x}}\mapsto\langle{\bm{x}},{\bm{\theta}}\rangle,\|{\bm{\theta}}\|<L\\}.$
(6)

We assume $L<1$ so that the problem belongs to the “hard generalization”
scenario discussed in Section 3, where the hypothesis space is inadequate to
fit the function on its own.

ResMem proceeds by first learning a linear function
$f_{n}({\bm{x}})=\langle{\bm{\theta}}_{n},{\bm{x}}\rangle$ from
${\mathcal{F}}$ through empirical risk minimization (ERM):

${\bm{\theta}}_{n}=\underset{\|{\bm{\theta}}\|\leq
L}{\operatorname{argmin}}\,\frac{1}{n}\sum_{i=1}^{n}\left[\langle{\bm{x}}_{i},{\bm{\theta}}\rangle-
y_{i}\right]^{2}.$ (7)

The empirical risk minimizer $f_{n}$ should be thought of as the analog of
$f_{\small\textsf{DeepNet}}$ in the deep learning context. It defines a
ground-truth residual function
$r_{\star}({\bm{x}})=f_{\star}({\bm{x}})-f_{n}({\bm{x}})$. Now we fix a test
covariate $\widetilde{{\bm{x}}}\sim\mathbb{P}_{x}$. ResMem “memorizes” the
residual function through the $1$-nearest neighbor to $\widetilde{{\bm{x}}}$

$r_{n}(\widetilde{{\bm{x}}})=r_{\star}(\widetilde{{\bm{x}}}_{(1)})=f_{\star}(\widetilde{{\bm{x}}}_{(1)})-f_{n}(\widetilde{{\bm{x}}}_{(1)}),$
(8)

where $\widetilde{{\bm{x}}}_{(1)}$ is the nearest neighbor to
$\widetilde{{\bm{x}}}$ among the training covariates
${\bm{x}}_{1},\dots,{\bm{x}}_{n}$:

$\widetilde{{\bm{x}}}_{(1)}=\underset{{\bm{x}}\in\\{{\bm{x}}_{1},\dots,{\bm{x}}_{n}\\}}{\operatorname{argmin}}\,\|{\bm{x}}-\widetilde{{\bm{x}}}\|.$

The final prediction is

$f_{n}^{{\small\textsf{ResMem}}}(\widetilde{{\bm{x}}})=f_{n}(\widetilde{{\bm{x}}})+r_{n}(\widetilde{{\bm{x}}}).$
(9)

Observe that if $\widetilde{{\bm{x}}}$ coincides with any training sample,
$f_{n}^{{\small\textsf{ResMem}}}(\widetilde{{\bm{x}}})=f_{\star}(\widetilde{{\bm{x}}})$,
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.

###  5.2 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

$\displaystyle{\bm{\theta}}_{\infty}$
$\displaystyle=\underset{\|{\bm{\theta}}\|\leq
L}{\operatorname{argmin}}\,\mathbb{E}_{{\bm{x}}\sim\mathbb{P}_{x}}\left[\langle{\bm{\theta}},{\bm{x}}\rangle-\langle{\bm{\theta}}_{\star},{\bm{x}}\rangle\right]^{2},$
$\displaystyle=\underset{\|{\bm{\theta}}\|\leq
L}{\operatorname{argmin}}\,\|{\bm{\theta}}-{\bm{\theta}}_{\star}\|=L{\bm{\theta}}_{\star},$

which is what ERM learns in the limit of $n\rightarrow\infty$. We can think of
${\bm{\theta}}_{\infty}$ as the best function that ERM can learn. Then, we can
decompose ${\bm{\theta}}_{\star}$ into
${\bm{\theta}}_{\star}={\bm{\theta}}_{\infty}+{\bm{\theta}}_{\perp}$, where
${\bm{\theta}}_{\perp}={\bm{\theta}}_{\star}-{\bm{\theta}}_{\infty}$. This
decomposition can be generalized beyond linear regression. Since
${\bm{\theta}}_{\infty}$ defines a function
$f_{\infty}({\bm{x}})=\langle{\bm{x}},{\bm{\theta}}_{\infty}\rangle$, 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.

###  5.3 A decomposition of the prediction error

We now introduce a decomposition of the prediction risk that reveals how
ResMem algorithm boosts generalization. Note that the prediction error of
ResMem is

$\mathbb{E}\left[\left(f_{n}^{{\small\textsf{ResMem}}}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})\right)^{2}\right].$
(10)

It can be decomposed into two components:

$\displaystyle\mathbb{E}\left[f_{n}^{{\small\textsf{ResMem}}}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})\right]^{2}\leq$
$\displaystyle
3\times[~{}\underbrace{\mathbb{E}(f_{n}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}}))^{2}+\mathbb{E}(f_{n}(\widetilde{{\bm{x}}}_{(1)})-f_{\infty}(\widetilde{{\bm{x}}}_{(1)}))^{2}}_{T_{1}}+\underbrace{\mathbb{E}(f_{\infty}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}}_{(1)})+f_{\star}(\widetilde{{\bm{x}}}_{(1)}))^{2}}_{T_{2}}~{}].$

We provide the detail of the decomposition in Section E.1. 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
${\mathcal{F}}$. It captures an irreducible error of the risk, which in
general is not asymptotically zero. However, because of the explicit
memorization ResMem algorithm introduces
($\widetilde{{\bm{x}}}_{(1)}\rightarrow\widetilde{{\bm{x}}}$ 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.

###  5.4 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 E.

######  Theorem 5.3 (Risk for ResMem algorithm).

For the problem defined in Assumption 5.2 with covariates distribution in
Assumption 5.1, the ResMem prediction rule
$f_{n}^{{\small\textsf{ResMem}}}(\widetilde{{\bm{x}}})$ defined in equation
(9) achieves risk (10)

$\displaystyle\mathbb{E}\left[f_{n}^{{\small\textsf{ResMem}}}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})\right]^{2}\lesssim
d^{2}L^{2}n^{-2/3}+d^{2}(1-L)^{2}\left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},$

where $\lesssim$ denotes inequality up to a universal constant independent of
$d,n$ and $L$.

The result includes contribution from two terms introduced in Section 5.3:

  * •

$T_{1}\lesssim d^{2}L^{2}n^{-2/3}$ that arises due to the difference between
$f_{n}$ and $f_{\infty}$.

  * •

$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
$\widetilde{{\bm{x}}}_{(1)}\rightarrow\widetilde{{\bm{x}}}$.

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 5.3, 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.

##  6 Discussion and future work

#### Theoretical results.

While Theorem 5.3 shows that ResMem achieves asymptotically zero test risk,
another viable strategy is to simply increase the capacity of $\mathcal{F}$.
For example, one may set $\mathcal{F}$ to be a class of suitably wide and deep
neural networks, and appeal to standard universal approximation results
(Cybenko, 1989). However, this strategy of scaling the model size comes at the
price of increased the training and inference cost. By contrast, ResMem
involves augmenting the base predictor with a relatively simple $k$-nearest
neighbor component.

The theory considers the use of the _exact_ minimizer ${\bm{\theta}}_{n}$ of
the empirical risk. For settings involving more complex predictors and losses,
(e.g., non-convex neural network objectives), it may be infeasible to obtain
this exact minimizer: rather, one obtains an _approximate_ minimizer via an
iterative optimization procedure, such as SGD. One may augment the analysis to
additionally consider the effect of _optimization error_ (Bottou and Bousquet,
2007) in such cases. We leave such study for future work.

#### Empirical procedures.

The current formulation of ResMem builds the base DeepNet and $k$NN components
sequentially. Consequently, the DeepNet is trained completely oblivious to the
fact that there is a subsequent $k$NN model that will memorize its residuals.
A natural direction of future work is to consider the _joint_ training of
DeepNet and $k$NN, so that the models can dynamically interact during training
to determine which portion of label is for DeepNet to learn, and the remaining
is for $k$NN to memorize.

A potential problem with applying ResMem to classification is the _mis-
calibration_ of the scorer. The output of the ResMem prediction vector
$f_{\small\textsf{ResMem}}(x)$ (3) is not guaranteed to lie on the probability
simplex. This is not an issue when we only care about the predicted class
membership, since we take the argmax of $f_{\small\textsf{ResMem}}(x)$.
However, this limitation hinders us to access the _confidence_ of the ResMem
prediction. To remedy this, a possible future work is to consider alternative
notions of residual. For example, we can do memorization in the logit space
instead of the probability space. Then, the one-hot encoding of the true label
may be replaced by class mean when defining the residual.

Finally, ResMem can be a promising approach to tackle test-time covariate
shift. The nearest neighbor modifies the prediction of DeepNet based on the
training covariate that are closer to the test covariate, making the algorithm
more _adaptive_ to the specific test covariate (Sun et al., 2020).

## Acknowledgements

Part of the work is done while Zitong Yang is at Google Research, New York. We
would like to thank Chong You, Yu Sun, and Yaodong Yu for their feedback on
the final draft. Zitong Yang would like to thank Shuangping Li for discussion
regarding the proof of Lemma D.1.

## References

  * Alley [1987] William M. Alley.  A note on stagewise regression.  _The American Statistician_ , 41(2):132–134, 1987\.  doi: 10.1080/00031305.1987.10475461.  URL https://www.tandfonline.com/doi/abs/10.1080/00031305.1987.10475461.
  * 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, and Simon Lacoste-Julien.  A closer look at memorization in deep networks.  In _Proceedings of the 34th International Conference on Machine Learning - Volume 70_ , ICML’17, pages 233–242. JMLR.org, 2017.
  * Bartlett et al. [2017] Peter L. Bartlett, Dylan J. Foster, and Matus Telgarsky.  Spectrally-normalized margin bounds for neural networks.  In Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, _Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA_ , pages 6240–6249, 2017.
  * 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, 2020.  doi: 10.1073/pnas.1907378117.  URL https://www.pnas.org/doi/abs/10.1073/pnas.1907378117.
  * Belkin et al. [2018] Mikhail Belkin, Daniel Hsu, and Partha P. Mitra.  Overfitting or perfect fitting? risk bounds for classification and regression rules that interpolate.  In _Proceedings of the 32nd International Conference on Neural Information Processing Systems_ , NIPS’18, pages 2306–2317, Red Hook, NY, USA, 2018. Curran Associates Inc.
  * Borgeaud et al. [2021] Sebastian Borgeaud, Arthur Mensch, Jordan Hoffmann, Trevor Cai, Eliza Rutherford, Katie Millican, George van den Driessche, Jean-Baptiste Lespiau, Bogdan Damoc, Aidan Clark, Diego de Las Casas, Aurelia Guy, Jacob Menick, Roman Ring, Tom Hennigan, Saffron Huang, Loren Maggiore, Chris Jones, Albin Cassirer, Andy Brock, Michela Paganini, Geoffrey Irving, Oriol Vinyals, Simon Osindero, Karen Simonyan, Jack W. Rae, Erich Elsen, and Laurent Sifre.  Improving language models by retrieving from trillions of tokens.  _CoRR_ , abs/2112.04426, 2021.  URL https://arxiv.org/abs/2112.04426.
  * Bottou and Bousquet [2007] Léon Bottou and Olivier Bousquet.  The tradeoffs of large scale learning.  In J. Platt, D. Koller, Y. Singer, and S. Roweis, editors, _Advances in Neural Information Processing Systems_ , volume 20. Curran Associates, Inc., 2007.  URL https://proceedings.neurips.cc/paper/2007/file/0d3180d672e08b4c5312dcdafdf6ef36-Paper.pdf.
  * Bottou and Vapnik [1992] Léon Bottou and Vladimir Vapnik.  Local Learning Algorithms.  _Neural Computation_ , 4(6):888–900, 11 1992\.  ISSN 0899-7667.  doi: 10.1162/neco.1992.4.6.888.  URL https://doi.org/10.1162/neco.1992.4.6.888.
  * Brutzkus et al. [2018] Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz.  SGD learns over-parameterized networks that provably generalize on linearly separable data.  In _International Conference on Learning Representations_ , 2018.  URL https://openreview.net/forum?id=rJ33wwxRb.
  * Bucilǎ et al. [2006] Cristian Bucilǎ, Rich Caruana, and Alexandru Niculescu-Mizil.  Model compression.  In _Proceedings of the 12th ACM SIGKDD International Conference on Knowledge Discovery and Data Mining_ , KDD ’06, pages 535–541, New York, NY, USA, 2006. ACM.
  * Chaudhuri and Dasgupta [2014] Kamalika Chaudhuri and Sanjoy Dasgupta.  Rates of convergence for nearest neighbor classification.  In Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Weinberger, editors, _Advances in Neural Information Processing Systems_ , volume 27. Curran Associates, Inc., 2014.  URL https://proceedings.neurips.cc/paper/2014/file/db957c626a8cd7a27231adfbf51e20eb-Paper.pdf.
  * Cheng et al. [2022] Chen Cheng, John Duchi, and Rohith Kuditipudi.  Memorize to generalize: on the necessity of interpolation in high dimensional linear regression, 2022.  URL https://arxiv.org/abs/2202.09889.
  * Chowdhery et al. [2022] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, Parker Schuh, Kensen Shi, Sasha Tsvyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson, Reiner Pope, James Bradbury, Jacob Austin, Michael Isard, Guy Gur-Ari, Pengcheng Yin, Toju Duke, Anselm Levskaya, Sanjay Ghemawat, Sunipa Dev, Henryk Michalewski, Xavier Garcia, Vedant Misra, Kevin Robinson, Liam Fedus, Denny Zhou, Daphne Ippolito, David Luan, Hyeontaek Lim, Barret Zoph, Alexander Spiridonov, Ryan Sepassi, David Dohan, Shivani Agrawal, Mark Omernick, Andrew M. Dai, Thanumalayan Sankaranarayana Pillai, Marie Pellat, Aitor Lewkowycz, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang, Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstern, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel.  Palm: Scaling language modeling with pathways, 2022.  URL https://arxiv.org/abs/2204.02311.
  * Cohen et al. [2018] Gilad Cohen, Guillermo Sapiro, and Raja Giryes.  Dnn or k-nn: That is the generalize vs. memorize question, 2018.  URL https://arxiv.org/abs/1805.06822.
  * Cover and Hart [1967] T. Cover and P. Hart.  Nearest neighbor pattern classification.  _IEEE Transactions on Information Theory_ , 13(1):21–27, 1967.  doi: 10.1109/TIT.1967.1053964.
  * Cybenko [1989] G Cybenko.  Approximation by superpositions of a sigmoidal function.  _Mathematics of Control, Signals and Systems_ , 2(4):303–314, 1989.
  * Dziugaite and 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.  In _Proceedings of the 33rd Annual Conference on Uncertainty in Artificial Intelligence (UAI)_ , 2017.
  * Feldman [2019] Vitaly Feldman.  Does learning require memorization? A short tale about a long tail.  _CoRR_ , abs/1906.05271, 2019.  URL http://arxiv.org/abs/1906.05271.
  * Feldman [2020] Vitaly Feldman.  _Does Learning Require Memorization? A Short Tale about a Long Tail_ , page 954–959.  Association for Computing Machinery, New York, NY, USA, 2020.  ISBN 9781450369794.  URL https://doi.org/10.1145/3357713.3384290.
  * Feldman and Zhang [2020a] Vitaly Feldman and Chiyuan Zhang.  What neural networks memorize and why: Discovering the long tail via influence estimation.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, _Advances in Neural Information Processing Systems_ , volume 33, pages 2881–2891. Curran Associates, Inc., 2020a.  URL https://proceedings.neurips.cc/paper/2020/file/1e14bfe2714193e7af5abc64ecbd6b46-Paper.pdf.
  * Feldman and Zhang [2020b] Vitaly Feldman and Chiyuan Zhang.  What neural networks memorize and why: Discovering the long tail via influence estimation.  In _Proceedings of the 34th International Conference on Neural Information Processing Systems_ , NIPS’20, Red Hook, NY, USA, 2020b. Curran Associates Inc.  ISBN 9781713829546.
  * Freund and Schapire [1995] Yoav Freund and Robert E. Schapire.  A desicion-theoretic generalization of on-line learning and an application to boosting.  In Paul Vitányi, editor, _Computational Learning Theory_ , pages 23–37, Berlin, Heidelberg, 1995. Springer Berlin Heidelberg.  ISBN 978-3-540-49195-8.
  * Friedman et al. [2000] Jerome Friedman, Trevor Hastie, and Robert Tibshirani.  Additive logistic regression: a statistical view of boosting (With discussion and a rejoinder by the authors).  _The Annals of Statistics_ , 28(2):337 – 407, 2000.  doi: 10.1214/aos/1016218223.  URL https://doi.org/10.1214/aos/1016218223.
  * Gionis et al. [1999] A. Gionis, Piotr Indyk, and Rajeev Motwani.  Similarity search in high dimensions via hashing.  In _Very Large Data Bases Conference_ , 1999.
  * Goldberger and Jochems [1961] Arthur S. Goldberger and D. B. Jochems.  Note on stepwise least squares.  _Journal of the American Statistical Association_ , 56(293):105–110, 1961.  doi: 10.1080/01621459.1961.10482095.  URL https://www.tandfonline.com/doi/abs/10.1080/01621459.1961.10482095.
  * Guo et al. [2020] Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiv Kumar.  Accelerating large-scale inference with anisotropic vector quantization.  In _International Conference on Machine Learning_ , 2020.  URL https://arxiv.org/abs/1908.10396.
  * Guu et al. [2020] Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming-Wei Chang.  Realm: Retrieval-augmented language model pre-training.  In _Proceedings of the 37th International Conference on Machine Learning_ , ICML’20. JMLR.org, 2020.
  * Hastie et al. [2001] Trevor Hastie, Robert Tibshirani, and Jerome Friedman.  _The Elements of Statistical Learning_.  Springer Series in Statistics. Springer New York Inc., New York, NY, USA, 2001.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR)_ , pages 770–778, 2016.  doi: 10.1109/CVPR.2016.90.
  * Hinton et al. [2015a] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network, 2015a.  URL https://arxiv.org/abs/1503.02531.
  * Hinton et al. [2015b] Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean.  Distilling the knowledge in a neural network.  _CoRR_ , abs/1503.02531, 2015b.
  * Khandelwal et al. [2020] Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis.  Generalization through memorization: Nearest neighbor language models.  In _International Conference on Learning Representations_ , 2020.  URL https://openreview.net/forum?id=HklBjCEKvH.
  * Knuth [1973] Donald Knuth.  _The Art Of Computer Programming, vol. 3: Sorting And Searching_.  Addison-Wesley, 1973.
  * Krizhevsky [2009] Alex Krizhevsky.  Learning multiple layers of features from tiny images.  Technical report, CIFAR, 2009.
  * Lample et al. [2019] Guillaume Lample, Alexandre Sablayrolles, Marc’Aurelio Ranzato, Ludovic Denoyer, and Hervé Jégou.  Large memory layers with product keys.  In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Florence d’Alché-Buc, Emily B. Fox, and Roman Garnett, editors, _Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019, December 8-14, 2019, Vancouver, BC, Canada_ , pages 8546–8557, 2019.  URL https://proceedings.neurips.cc/paper/2019/hash/9d8df73a3cfbf3c5b47bc9b50f214aff-Abstract.html.
  * Li et al. [2022] Zonglin Li, Ruiqi Guo, and Sanjiv Kumar.  Decoupled context processing for context augmented language modeling.  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=02dbnEbEFn.
  * Liang and Rakhlin [2020] Tengyuan Liang and Alexander Rakhlin.  Just interpolate: Kernel “Ridgeless” regression can generalize.  _The Annals of Statistics_ , 48(3):1329 – 1347, 2020.  doi: 10.1214/19-AOS1849.  URL https://doi.org/10.1214/19-AOS1849.
  * Liu et al. [2019] Ziwei Liu, Zhongqi Miao, Xiaohang Zhan, Jiayun Wang, Boqing Gong, and Stella X Yu.  Large-scale long-tailed recognition in an open world.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 2537–2546, 2019.
  * Menghani [2021] Gaurav Menghani.  Efficient deep learning: A survey on making deep learning models smaller, faster, and better.  _CoRR_ , abs/2106.08962, 2021.  URL https://arxiv.org/abs/2106.08962.
  * Montanari and Zhong [2020] Andrea Montanari and Yiqiao Zhong.  The interpolation phase transition in neural networks: Memorization and generalization under lazy training, 2020.  URL https://arxiv.org/abs/2007.12826.
  * Muja and Lowe [2009] Marius Muja and David G. Lowe.  Fast approximate nearest neighbors with automatic algorithm configuration.  In _International Conference on Computer Vision Theory and Applications_ , 2009.
  * Neyshabur et al. [2019] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro.  The role of over-parametrization in generalization of neural networks.  In _7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019_. OpenReview.net, 2019.
  * Panigrahy et al. [2021] Rina Panigrahy, Xin Wang, and Manzil Zaheer.  Sketch based memory for neural networks.  In Arindam Banerjee and Kenji Fukumizu, editors, _Proceedings of The 24th International Conference on Artificial Intelligence and Statistics_ , volume 130 of _Proceedings of Machine Learning Research_ , pages 3169–3177. PMLR, 13–15 Apr 2021.  URL https://proceedings.mlr.press/v130/panigrahy21a.html.
  * Raffel et al. [2020] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  _Journal of Machine Learning Research_ , 21(140):1–67, 2020.  URL http://jmlr.org/papers/v21/20-074.html.
  * Russakovsky et al. [2015] Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei.  ImageNet Large Scale Visual Recognition Challenge.  _International Journal of Computer Vision (IJCV)_ , 115(3):211–252, 2015.  doi: 10.1007/s11263-015-0816-y.
  * Salakhutdinov and Hinton [2007] Ruslan Salakhutdinov and Geoff Hinton.  Learning a nonlinear embedding by preserving class neighbourhood structure.  In Marina Meila and Xiaotong Shen, editors, _Proceedings of the Eleventh International Conference on Artificial Intelligence and Statistics_ , volume 2 of _Proceedings of Machine Learning Research_ , pages 412–419, San Juan, Puerto Rico, 21–24 Mar 2007. PMLR.  URL https://proceedings.mlr.press/v2/salakhutdinov07a.html.
  * Sun et al. [2020] Yu Sun, Xiaolong Wang, Liu Zhuang, John Miller, Moritz Hardt, and Alexei A. Efros.  Test-time training with self-supervision for generalization under distribution shifts.  In _ICML_ , 2020.
  * Tibshirani [2015] Ryan J. Tibshirani.  A general framework for fast stagewise algorithms.  _J. Mach. Learn. Res._ , 16(1):2543–2588, jan 2015.  ISSN 1532-4435.
  * Vapnik and Izmailov [2021] Vladimir Vapnik and Rauf Izmailov.  Reinforced SVM method and memorization mechanisms.  _Pattern Recognition_ , 119:108018, 2021.  ISSN 0031-3203.  doi: https://doi.org/10.1016/j.patcog.2021.108018.  URL https://www.sciencedirect.com/science/article/pii/S0031320321002053.
  * Wainwright [2019] Martin J Wainwright.  _High-dimensional statistics: A non-asymptotic viewpoint_ , volume 48.  Cambridge University Press, 2019.
  * Wang et al. [2021] Ke Wang, Vidya Muthukumar, and Christos Thrampoulidis.  Benign overfitting in multiclass classification: All roads lead to interpolation, 2021.  URL https://arxiv.org/abs/2106.10865.
  * Wang and Shao [2022] Zhen Wang and Yuan-Hai Shao.  Generalization-memorization machines, 2022.  URL https://arxiv.org/abs/2207.03976.
  * Yang et al. [2020] Zitong Yang, Yaodong Yu, Chong You, Jacob Steinhardt, and Yi Ma.  Rethinking bias-variance trade-off for generalization of neural networks.  In Hal Daumé III and Aarti Singh, editors, _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pages 10767–10777. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/yang20j.html.
  * Yang et al. [2021] Zitong Yang, Yu Bai, and Song Mei.  Exact gap between generalization error and uniform convergence in random feature models.  In Marina Meila and Tong Zhang, editors, _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 11704–11715. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/yang21a.html.
  * Zhang et al. [2017] Chiyuan Zhang, Samy Bengio, Moritz Hardt, Benjamin Recht, and Oriol Vinyals.  Understanding deep learning requires rethinking generalization.  In _5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings_. OpenReview.net, 2017.
  * Zhang et al. [2006] Hao Zhang, Alexander C. Berg, Michael Maire, and Jitendra Malik.  Svm-knn: Discriminative nearest neighbor classification for visual category recognition.  In _CVPR (2)_ , pages 2126–2136, 2006.  URL https://doi.org/10.1109/CVPR.2006.301.

##  Appendix A Additional CIFAR100 Results

This section includes additional experiment results on applying ResMem to
CIFAR100 dataset.

###  A.1 Additional robustness results

In addition to the results already presented in Section 4.3, 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 4.1. Generally, we see that ResMem yields larger improvement over the
baseline DeepNet when the network is small and dataset is large.

(a) CIFAR-ResNet-8

(b) CIFAR-ResNet-14

(c) CIFAR-ResNet-20

(d) CIFAR-ResNet-32

(e) CIFAR-ResNet-44

(f) CIFAR-ResNet-56

Figure 4: Test(left)/Training (right) accuracy for different sample sizes.

###  A.2 Sensitivity analysis for CIFAR100

(a) # of neighbours $k$. (b) Radius parameter $\sigma$. (c) Temperature $T$.

Figure 5:  Sensitivity analysis of ResMem hyperparameters. The $y$-axis
represents the CIFAR100 test accuracies, and the $x$-axis represents the
sweeping of respective hyperparameters.

#### 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 5(a)). 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 5(b)).

#### 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 5(c)). 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 [Hinton et al.,
2015a]. 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 _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 _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.

##  Appendix B Additional ImageNet results

This section includes additional experiment results on applying ResMem to
ImageNet dataset.

###  B.1 Sensitivity analysis for ImageNet

(a) # of neighbours $k$. (b) Radius parameter $\sigma$. (c) Temperature $T$.

Figure 6:  Sensitivity analysis of ResMem hyperparameters. The $y$-axis
represents the imagenet100 test accuracies, and the $x$-axis represents the
sweeping of respective hyperparameters.

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 6(a)). 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 6(b)). We tried $T=0.3$ to $T=2.1$, and also
identified an unimodal shape for the test accuracy (Figure 6(c)).

###  B.2 ImageNet-LT

We also evaluated the proposed ResMem on ImageNet-LT [Liu et al., 2019] – a
long-tail version of ImageNet [Russakovsky et al., 2015]. We again utilize
ResNet [He et al., 2016] 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 3 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.

Table 3: Test accuracy for ResMem and baseline deep network.

Dataset | Architecture | Test accuracy
---|---|---
DeepNet | ResMem
ImageNet | ResNet18 | 66.86% | 67.44%
ImageNet-LT | ResNet18 | 39.56% | 39.91%
ImageNet-LT | ResNet50 | 46.18% | 46.39%

##  Appendix C Additional details of NLP experiments

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 [Raffel et al., 2020]
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.

##  Appendix D Some concentration results for uniform random variables

In this section, we state some concentration results that are useful for the
theoretical analysis in Section 5. Let
$\widetilde{{\bm{x}}},{\bm{x}}_{1},\dots,{\bm{x}}_{n}\overset{\rm
i.i.d.}{\sim}$Unif$(\mathcal{B}_{{\bm{0}},\sqrt{d+2}})$ be i.i.d. samples from
the uniform distribution over the Euclidean norm ball of radius $\sqrt{d+2}$
in $\mathbb{R}^{d}$. Let

$Z_{n}=\min_{{\bm{x}}\in\\{{\bm{x}}_{1},{\bm{x}}_{2},\dots,{\bm{x}}_{n}\\}}\|\widetilde{{\bm{x}}}-{\bm{x}}\|^{2}.$
(11)

If $n=1$, $\mathbb{E}Z_{1}$ is the sum of the variance of each coordinate of
Unif$(\mathcal{B}_{{\bm{0}},\sqrt{d+2}})$. Therefore, $\mathbb{E}Z_{n}$
provides a generalized measure of concentration. Intuitively,
$\mathbb{E}Z_{n}\rightarrow 0$ as $n\rightarrow\infty$. The proposition below
provides a upper bound on the rate of convergence.

######  Lemma D.1 (Nearest Neighbor concentration).

Given the assumptions above

$\mathbb{E}Z_{n}\lesssim
d^{2}\left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},$ (12)

where $\lesssim$ means inequality up to an universal constant independent of
$d$ and $n$.

###### Proof.

Define

$\displaystyle{\mathcal{E}}_{1}$ $\displaystyle=\\{Z_{n}\leq\delta^{2}\\},$
(13) $\displaystyle{\mathcal{E}}_{2}$
$\displaystyle=\\{\delta\leq\sqrt{d+2}-\|\widetilde{{\bm{x}}}\|\\}.$

We will compute two probabilities
$\mathbb{P}({\mathcal{E}}_{1}|{\mathcal{E}}_{2})$ and
$\mathbb{P}({\mathcal{E}}_{2})$ that will be useful latter.

$\displaystyle\mathbb{P}({\mathcal{E}}_{1}^{c}|{\mathcal{E}}_{2})$
$\displaystyle=\mathbb{P}(Z_{n}\geq\delta^{2}|{\mathcal{E}}_{2})=\mathbb{P}(\|\widetilde{{\bm{x}}}-{\bm{x}}_{i}\|\geq\delta,~{}\forall
i|{\mathcal{E}}_{2}),$ (14)
$\displaystyle=\mathbb{E}_{\widetilde{{\bm{x}}}}\mathbb{P}(\|\widetilde{{\bm{x}}}-{\bm{x}}_{i}\|\geq\delta|{\mathcal{E}}_{2},\widetilde{{\bm{x}}})^{n}=\mathbb{E}_{\widetilde{{\bm{x}}}}(1-\mathbb{P}(\|\widetilde{{\bm{x}}}-{\bm{x}}_{i}\|\leq\delta|{\mathcal{E}}_{2},\widetilde{{\bm{x}}}))^{n},$
$\displaystyle=\mathbb{E}_{\widetilde{{\bm{x}}}}\left[1-\frac{\text{Vol}(\mathcal{B}_{\widetilde{{\bm{x}}},\delta})}{\text{Vol}(\mathcal{B}_{{\bm{0}},\sqrt{d+2}})}\right]^{n}=\left[1-\left(\frac{\delta}{\sqrt{d+2}}\right)^{d}\right]^{n},$
$\displaystyle\leq\exp\left[-n\left(\frac{\delta}{\sqrt{d+2}}\right)^{d}\right].$

Next, we compute $\mathbb{P}({\mathcal{E}}_{2})$

$\displaystyle\mathbb{P}({\mathcal{E}}_{2})=\mathbb{P}(\|\widetilde{{\bm{x}}}\|\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}.$
(15)

We use ${\mathcal{E}}_{1}$ and ${\mathcal{E}}_{2}$ to compute the following
upper bound

$\displaystyle\mathbb{E}Z_{n}$
$\displaystyle=\mathbb{E}(Z_{n}|{\mathcal{E}}_{1}\cap{\mathcal{E}}_{2})\mathbb{P}({\mathcal{E}}_{1}\cap{\mathcal{E}}_{2})+\mathbb{E}(Z_{n}|({\mathcal{E}}_{1}\cap{\mathcal{E}}_{2})^{c})P(({\mathcal{E}}_{1}\cap{\mathcal{E}}_{2})^{c}),$
(16)
$\displaystyle\leq\delta^{2}+(2\sqrt{d+2})^{2}\left(1-\mathbb{P}({\mathcal{E}}_{1}\cap{\mathcal{E}}_{2})\right),$
$\displaystyle=\delta^{2}+4(d+2)\left[1-\mathbb{P}({\mathcal{E}}_{1}|{\mathcal{E}}_{2})\mathbb{P}({\mathcal{E}}_{2})\right].$

To find an upper bound for $\mathbb{E}Z_{n}$, we need to find an upper bound
for
$1-\mathbb{P}({\mathcal{E}}_{1}|{\mathcal{E}}_{2})\mathbb{P}({\mathcal{E}}_{2})$.

$\displaystyle
1-\mathbb{P}({\mathcal{E}}_{1}|{\mathcal{E}}_{2})\mathbb{P}({\mathcal{E}}_{2})$
$\displaystyle=1-\left[1-\mathbb{P}({\mathcal{E}}_{1}^{c}|{\mathcal{E}}_{2})\right]\mathbb{P}({\mathcal{E}}_{2}),$
(17)
$\displaystyle=1-\mathbb{P}({\mathcal{E}}_{2})+\mathbb{P}({\mathcal{E}}_{1}^{c}|{\mathcal{E}}_{2})\mathbb{P}({\mathcal{E}}_{2}),$
$\displaystyle\leq
1-\mathbb{P}({\mathcal{E}}_{2})+\mathbb{P}({\mathcal{E}}_{1}^{c}|{\mathcal{E}}_{2}).$

Now choose
$\delta=\sqrt{d+2}n^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d}$.

$\mathbb{P}({\mathcal{E}}_{1}^{c}|{\mathcal{E}}_{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},$
(18)

and

$\mathbb{P}({\mathcal{E}}_{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}.$
(19)

Thus

$1-\mathbb{P}({\mathcal{E}}_{1}|{\mathcal{E}}_{2})\mathbb{P}({\mathcal{E}}_{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}.$ (20)

Combining everything together, we get

$\displaystyle\mathbb{E}Z_{n}$
$\displaystyle\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},$ (21)
$\displaystyle\lesssim
d^{2}n^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d},$
$\displaystyle=d^{2}\left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.$

This completes the proof. ∎

######  Proposition D.2 (Wainwright [2019] Corollary 6.20).

Let ${\bm{x}}_{i}\overset{\rm
i.i.d.}{\sim}$Unif$~{}(\mathcal{B}_{{\bm{0}},\sqrt{d+2}})$ for $i=1,\dots,n$
be uniformly distributed over a ball of radius $B$ in $\mathbb{R}^{d}$
centered at ${\bm{0}}$. Let

${\bm{\Sigma}}_{n}=\frac{1}{n}\sum_{i=1}^{n}{\bm{x}}_{i}{\bm{x}}_{i}^{\mathsf{T}}$

be the sample covariance matrix. Then

$\mathbb{P}(\|{\bm{\Sigma}}_{n}-{\bm{I}}\|_{\rm op}>\varepsilon)\leq
2d\exp\left[-\frac{n\varepsilon^{2}}{2(d+2)(1+\varepsilon)}\right].$

##  Appendix E Proof of Theorem 5.3

In this section, we present the proof of Theorem 5.3. In Section E.1, we
provide the detail of the decomposition of the risk into $T_{1}$ and $T_{2}$.
Then in Section E.2 we compute an upper bound for $T_{1}$, and compute an
upper bound for $T_{2}$ in Section E.3. Finally, we combine everything
together in Section E.4 and completes the proof.

###  E.1 Decomposition of the test risk

$\displaystyle~{}~{}~{}~{}~{}\mathbb{E}\left[f^{{\small\textsf{ResMem}}}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})\right]^{2}=\mathbb{E}\left[f_{n}(\widetilde{{\bm{x}}})+r_{n}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})\right]^{2},$
(22)
$\displaystyle=\mathbb{E}\left[f_{n}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})-f_{n}(\widetilde{{\bm{x}}}_{(1)})+f_{\star}(\widetilde{{\bm{x}}}_{(1)})\right]^{2},$
$\displaystyle=\mathbb{E}\left[f_{n}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}})+f_{\infty}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})-f_{n}(\widetilde{{\bm{x}}}_{(1)})+f_{\infty}(\widetilde{{\bm{x}}}_{(1)})-f_{\infty}(\widetilde{{\bm{x}}}_{(1)})+f_{\star}(\widetilde{{\bm{x}}}_{(1)})\right]^{2},$
$\displaystyle\leq
3\times[\underbrace{\mathbb{E}(f_{n}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}}))^{2}+\mathbb{E}(f_{n}(\widetilde{{\bm{x}}}_{(1)})-f_{\infty}(\widetilde{{\bm{x}}}_{(1)}))^{2}}_{T_{1}}+\underbrace{\mathbb{E}(f_{\infty}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}}_{(1)})+f_{\star}(\widetilde{{\bm{x}}}_{(1)}))^{2}}_{T_{2}}],$

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\mathbb{R}$.

###  E.2 Upper bound on $T_{1}$.

Since $\mathbb{P}_{\bm{x}}=\text{Unif}(\mathcal{B}_{{\bm{0}},B}~{})$, we apply
the bound $\|\widetilde{{\bm{x}}}\|,\|\widetilde{{\bm{x}}}_{(1)}\|\leq B$ to
obtain

$\displaystyle T_{1}$
$\displaystyle=\mathbb{E}[f_{n}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}})]^{2}+\mathbb{E}[f_{n}(\widetilde{{\bm{x}}}_{(1)})-f_{\infty}(\widetilde{{\bm{x}}}_{(1)})]^{2},$
(23)
$\displaystyle=\mathbb{E}\langle{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty},\widetilde{{\bm{x}}}\rangle^{2}+\mathbb{E}\langle{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty},\widetilde{{\bm{x}}}_{(1)}\rangle^{2},$
$\displaystyle\leq\mathbb{E}\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}\|\widetilde{{\bm{x}}}\|^{2}+\mathbb{E}\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}\|\widetilde{{\bm{x}}}_{(1)}\|^{2},$
$\displaystyle\leq
2B^{2}\mathbb{E}\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}.$

As $n$ gets large, the empirical covariance matrix
${\bm{\Sigma}}_{n}={\bm{X}}^{\mathsf{T}}{\bm{X}}/n$ is concentrated around its
mean ${\bm{I}}$. Let ${\bm{\Delta}}_{n}={\bm{I}}-{\bm{\Sigma}}_{n}$ denote
this deviation. For some $\varepsilon\in(0,1)$, define the following “good
event” over the randomness in ${\bm{\Sigma}}_{n}$

${\mathcal{A}}=\\{\|{\bm{\Delta}}_{n}\|_{\rm op}<\varepsilon\\},$ (24)

where $\|{\bm{\Delta}}_{n}\|_{\rm op}$ denotes the operator norm of the
deviation matrix. The high level idea of the proof is to condition on the
event ${\mathcal{A}}$ and deduce and upper bound of
$\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|$ in terms of $\varepsilon$.
Then, we use the fact that ${\mathcal{A}}$ happens with high probability.

Recall that ${\bm{\theta}}_{\infty}=L{\bm{\theta}}_{\star}$, and

${\bm{\theta}}_{n}=\underset{\|{\bm{\theta}}\|\leq
L}{\operatorname{argmin}}\,\frac{1}{n}\|{\bm{X}}{\bm{\theta}}-{\bm{y}}\|^{2}.$
(25)

Since ${\bm{y}}={\bm{X}}{\bm{\theta}}_{\star}$ by definition, the Lagrangian
of the convex program above is

$\mathcal{L}({\bm{\theta}},\lambda)=\frac{1}{n}\|{\bm{X}}{\bm{\theta}}-{\bm{X}}{\bm{\theta}}_{\star}\|^{2}+\lambda(\|{\bm{\theta}}\|^{2}-L).$
(26)

The KKT condition suggests that the primal-dual optimal pair
$({\bm{\theta}}_{n},\lambda_{n})$ is given by

$\displaystyle\|{\bm{\theta}}_{n}\|$ $\displaystyle\leq L,$ (27)
$\displaystyle\lambda_{n}$ $\displaystyle\geq 0,$
$\displaystyle\lambda_{n}(\|{\bm{\theta}}_{n}\|-L)$ $\displaystyle=0,$

and at optimality

$\displaystyle\nabla_{\bm{\theta}}\mathcal{L}({\bm{\theta}}_{n},\lambda_{n})=0$
$\displaystyle\iff\frac{2}{n}{\bm{X}}^{\mathsf{T}}{\bm{X}}({\bm{\theta}}-{\bm{\theta}}_{\star})+2\lambda_{n}{\bm{\theta}}=0,$
(28)
$\displaystyle\iff{\bm{\theta}}_{n}=({\bm{\Sigma}}_{n}+\lambda_{n}{\bm{I}})^{-1}{\bm{\Sigma}}_{n}{\bm{\theta}}_{\star}.$

The complementary slackness condition $\lambda_{n}(\|{\bm{\theta}}_{n}\|-L)=0$
suggests that either $\lambda_{n}=0$ or $\|{\bm{\theta}}_{n}\|=L$. But if
$\lambda_{n}=0$, the stationary condition
$\nabla_{\bm{\theta}}\mathcal{L}({\bm{\theta}},\lambda)=0$ would suggest that
${\bm{\theta}}_{n}={\bm{\Sigma}}_{n}^{-1}{\bm{\Sigma}}_{n}{\bm{\theta}}_{\star}={\bm{\theta}}_{\star}\Rightarrow\|{\bm{\theta}}_{n}\|=1>L,$
a contradiction. (Note that here ${\bm{\Sigma}}_{n}$ is invertible condition
on the event ${\mathcal{A}}$.) Therefore, we must have
$\|{\bm{\theta}}_{n}\|=L$. As a result, the primal and dual pair
$({\bm{\theta}}_{n},\lambda_{n})$ is determined by the system of equations

$\begin{cases}{\bm{\theta}}_{n}&=({\bm{\Sigma}}_{n}+\lambda_{n}{\bm{I}})^{-1}{\bm{\Sigma}}_{n}{\bm{\theta}}_{\star},\\\
\|{\bm{\theta}}_{n}\|&=L,\\\ \lambda_{n}&>0.\end{cases}$ (29)

Next, we proceed to compute the deviation
$\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|$.

$\displaystyle{\bm{\theta}}_{n}$
$\displaystyle=\left[(\lambda_{n}+1){\bm{I}}-{\bm{\Delta}}_{n}\right]^{-1}{\bm{\Sigma}}_{n}{\bm{\theta}}_{\star},$
(30)
$\displaystyle=(\lambda_{n}+1)^{-1}\left[{\bm{I}}-\frac{{\bm{\Delta}}_{n}}{\lambda_{n}+1}\right]^{-1}{\bm{\Sigma}}_{n}{\bm{\theta}}_{\star},$
$\displaystyle=(\lambda_{n}+1)^{-1}\left[{\bm{I}}+\sum_{k=1}^{\infty}\frac{{\bm{\Delta}}_{n}^{k}}{(\lambda_{n}+1)^{k}}\right]({\bm{I}}-{\bm{\Delta}}_{n}){\bm{\theta}}_{\star},$
$\displaystyle=(\lambda_{n}+1)^{-1}\left[{\bm{I}}+\sum_{k=1}^{\infty}\frac{{\bm{\Delta}}_{n}^{k}}{(\lambda_{n}+1)^{k}}-{\bm{\Delta}}_{n}-\sum_{k=1}^{\infty}\frac{{\bm{\Delta}}_{n}^{k+1}}{(\lambda_{n}+1)^{k}}\right]{\bm{\theta}}_{\star},$
$\displaystyle=(\lambda_{n}+1)^{-1}{\bm{\theta}}_{\star}+(\lambda_{n}+1)^{-1}{\bm{\Delta}}_{n}\left[\sum_{k=1}^{\infty}\frac{{\bm{\Delta}}_{n}^{k-1}}{(\lambda_{n}+1)^{k}}-{\bm{I}}-\sum_{k=1}^{\infty}\frac{{\bm{\Delta}}_{n}^{k}}{(\lambda_{n}+1)^{k}}\right]{\bm{\theta}}_{\star},$
$\displaystyle=(\lambda_{n}+1)^{-1}{\bm{\theta}}_{\star}+(\lambda_{n}+1)^{-1}{\bm{\Delta}}_{n}\left[\sum_{k=1}^{\infty}\frac{{\bm{\Delta}}_{n}^{k-1}-{\bm{\Delta}}_{n}^{k}}{(\lambda_{n}+1)^{k}}-{\bm{I}}\right]{\bm{\theta}}_{\star}.$

Define

${\bm{D}}_{n}={\bm{\Delta}}_{n}\left[\sum_{k=1}^{\infty}\frac{{\bm{\Delta}}_{n}^{k-1}-{\bm{\Delta}}_{n}^{k}}{(\lambda_{n}+1)^{k}}-{\bm{I}}\right].$
(31)

Then
${\bm{\theta}}_{n}=(\lambda_{n}+1)^{-1}{\bm{\theta}}_{\star}+(\lambda_{n}+1)^{-1}{\bm{D}}_{n}{\bm{\theta}}_{\star}$,
and

$\displaystyle\|{\bm{D}}_{n}\|$
$\displaystyle\leq\|{\bm{\Delta}}_{n}\|\left[1+\sum_{k=1}^{\infty}\frac{\|{\bm{\Delta}}_{n}\|^{k-1}+\|{\bm{\Delta}}_{n}\|^{k}}{(\lambda_{n}+1)^{k}}\right],$
(32)
$\displaystyle\leq\varepsilon\left[1+2(1+\lambda_{n})^{-1}\sum_{k=1}^{\infty}\left(\frac{\varepsilon}{1+\lambda_{n}}\right)^{k}\right],$
$\displaystyle=\varepsilon\left(1+\frac{2}{1+\lambda_{n}}\frac{1}{1-\frac{\varepsilon}{1+\lambda_{n}}}\right)\leq
3\varepsilon.$

Therefore

$\displaystyle
L=\|{\bm{\theta}}_{n}\|^{2}=(\lambda_{n}+1)^{-2}+(\lambda_{n}+1)^{-2}{\bm{\theta}}_{\star}^{\mathsf{T}}{\bm{D}}_{n}^{\mathsf{T}}{\bm{D}}_{n}{\bm{\theta}}_{\star}+2(\lambda_{n}+1)^{-2}{\bm{\theta}}_{\star}{\bm{D}}_{n}{\bm{\theta}}_{\star},$
(33) $\displaystyle\Rightarrow$
$\displaystyle(\lambda_{n}+1)^{2}L^{2}=1+\delta_{n},~{}\delta_{n}={\bm{\theta}}_{\star}^{\mathsf{T}}{\bm{D}}_{n}^{\mathsf{T}}{\bm{D}}_{n}{\bm{\theta}}_{\star}+2{\bm{\theta}}_{\star}^{\mathsf{T}}{\bm{D}}_{n}{\bm{\theta}}_{\star}.$

We can obtain the following bound for $\delta_{n}$:

$|\delta_{n}|\leq\|{\bm{\theta}}_{\star}\|^{2}\|{\bm{D}}_{n}\|^{2}+2\|{\bm{\theta}}_{\star}\|^{2}\|{\bm{D}}_{n}\|\leq
9\varepsilon^{2}+6\varepsilon\leq 15\varepsilon.$ (34)

Since $1-\delta_{n}/2\leq\sqrt{1+\delta_{n}}\leq 1+\delta_{n}/2$, and
$|\delta_{n}|\leq 15\varepsilon$, we obtain

$|(\lambda_{n}+1)L-1|\leq\frac{15\varepsilon}{2}\Rightarrow\left|L-(\lambda_{n}+1)^{-1}\right|\leq\frac{15\varepsilon}{2}(\lambda_{n}+1)^{-1}\leq\frac{15\varepsilon}{2},$
(35)

where the last inequality follows as we have $\lambda_{n}>0$. Finally,

$\displaystyle{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}{}{}{}$
$\displaystyle=(\lambda_{n}+1)^{-1}{\bm{\theta}}_{\star}-L{\bm{\theta}}_{\star}+(\lambda_{n}+1)^{-1}{\bm{D}}_{n}{\bm{\theta}}_{\star},$
(36)
$\displaystyle\Rightarrow\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}$
$\displaystyle=[(1+\lambda_{n})^{-1}-L]^{2}+(1+\lambda_{n})^{-2}{\bm{\theta}}_{\star}{\bm{D}}_{n}^{\mathsf{T}}{\bm{D}}_{n}{\bm{\theta}}_{\star}+2(\lambda_{n}+1)^{-1}[(1+\lambda_{n})^{-1}-L]{\bm{\theta}}_{\star}{\bm{D}}_{n}{\bm{\theta}}_{\star},$
$\displaystyle\leq
64\varepsilon^{2}+9\varepsilon^{2}+45\varepsilon^{2}=118\varepsilon^{2},$
$\displaystyle\Rightarrow\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}$
$\displaystyle\lesssim\varepsilon^{2}.$

Combine the above result with Proposition D.2, we get that

$\displaystyle\mathbb{E}\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}$
$\displaystyle=\mathbb{E}(\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}|{\mathcal{A}})\mathbb{P}({\mathcal{A}})+\mathbb{E}(\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}|{\mathcal{A}}^{c})\mathbb{P}({\mathcal{A}}^{c}),$
(37) $\displaystyle\leq\varepsilon^{2}+4L^{2}\times
4d\exp\left[-\frac{n\varepsilon^{2}}{2(d+2)(1+\varepsilon)}\right],$

If we choose $\varepsilon=n^{-1/3}$, we get

$\mathbb{E}\|{\bm{\theta}}_{n}-{\bm{\theta}}_{\infty}\|^{2}\lesssim
dL^{2}n^{-2/3},$ (38)

which implies that

$T_{1}\lesssim d^{2}L^{2}n^{-2/3}.$ (39)

###  E.3 Upper bound on $T_{2}$.

Plugging in the formula for
$f_{\perp}(\widetilde{{\bm{x}}})=f_{\star}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}})=\langle\widetilde{{\bm{x}}},{\bm{\theta}}_{\perp}\rangle$,
we get

$\displaystyle T_{2}$
$\displaystyle=\mathbb{E}[f_{\perp}(\widetilde{{\bm{x}}}_{(1)})-f_{\perp}(\widetilde{{\bm{x}}})]^{2},$
(40)
$\displaystyle=\mathbb{E}\langle{\bm{\theta}}_{\perp},\widetilde{{\bm{x}}}_{(1)}-\widetilde{{\bm{x}}}\rangle^{2},$
$\displaystyle\leq(1-L)^{2}\|{\bm{\theta}}_{\star}\|^{2}\mathbb{E}\|\widetilde{{\bm{x}}}-\widetilde{{\bm{x}}}_{(1)}\|^{2},$
$\displaystyle=(1-L)^{2}\mathbb{E}\|\widetilde{{\bm{x}}}-\widetilde{{\bm{x}}}_{(1)}\|^{2},$

where in the last inequality, we used the relation that
${\bm{\theta}}_{\perp}=(1-L){\bm{\theta}}_{\star}$. Proposition D.1 suggests
that

$\mathbb{E}\|\widetilde{{\bm{x}}}-\widetilde{{\bm{x}}}_{(1)}\|^{2}\lesssim
d^{2}\left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},$ (41)

which implies

$T_{2}\lesssim
d^{2}(1-L)^{2}\left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.$ (42)

######  Remark E.1 (Comparison with pure nearest neighbor and ERM).

If we rely solely on nearest neighbor method, the prediction error is

$\mathbb{E}[f_{\star}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}}_{(1)})]^{2}=\mathbb{E}\langle\widetilde{{\bm{x}}}-\widetilde{{\bm{x}}}_{(1)},{\bm{\theta}}_{\star}\rangle^{2}\leq\mathbb{E}\|\widetilde{{\bm{x}}}-\widetilde{{\bm{x}}}_{(1)}\|^{2}.$
(43)

On the other hand, if we solely rely on ERM, even with infinite sample, we get

$\mathbb{E}[f_{\star}(\widetilde{{\bm{x}}})-f_{\infty}(\widetilde{{\bm{x}}})]^{2}=\mathbb{E}\langle\widetilde{{\bm{x}}},{\bm{\theta}}_{\star}-{\bm{\theta}}_{\infty}\rangle^{2}\leq(1-L)^{2}\mathbb{E}\|\widetilde{{\bm{x}}}\|^{2}.$
(44)

We can see from the upper bound that ResMem takes advantage of both

  * •

Projecting $f_{\star}$ onto $f_{\infty}$, so that the dependence on the
prediction function is reduced from $1$ to $(1-L)^{2}$.

  * •

Memorizing the residuals using nearest neighbor, so that the variance is
reduced from $\mathbb{E}\|\widetilde{{\bm{x}}}\|^{2}$ to
$\mathbb{E}\|\widetilde{{\bm{x}}}_{(1)}-\widetilde{{\bm{x}}}\|^{2}$.

###  E.4 Test loss for ResMem.

If we combine the previous two parts together, we get

$\displaystyle\mathbb{E}\left[\hat{f}(\widetilde{{\bm{x}}})-f_{\star}(\widetilde{{\bm{x}}})\right]^{2}$
$\displaystyle\lesssim
d^{2}L^{2}n^{-2/3}+d^{2}(1-L)^{2}\left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.$
(45)

This completes the proof of Theorem 5.3.