Title: ResMem: Learn what you can and memorize the rest
Venue: NeurIPS 2023
Accepted: yes
Keywords: deep learning

generalization

memorization

deep learning theory

boosting

nearest neighbor
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 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.
We start by formulating a stylized linear regression problem and rigorously show that ResMem results in a more favorable test risk over a base linear neural network.
Then, we empirically show that ResMem consistently improves the test set generalization of the original prediction model across standard vision and natural language processing benchmarks.

ResMem: Learn what you can and memorize the rest

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, even when the base model has low capacity. We start by formulating a stylized linear regression problem and rigorously show that ResMem results in a more favorable test risk over a base linear neural network. Then, we empirically show that ResMem consistently improves the test set generalization of the original prediction model across standard vision and natural language processing benchmarks.

1 Introduction

Large neural networks achieve remarkable generalization on test samples despite memorization of training samples, in the sense of achieving zero training error [54]. Several recent analyses have established that, under certain settings, memorization is sufficient to achieve generalization [3, 15, 5, 40, 4], and, more surprisingly, can even be necessary [17, 19, 11]. These works suggest 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 (cf. Figure 1). Compared to the implicit memorization performed by large neural models, the key idea behind ResMem is to perform explicit memorization via a separate k -nearest neighbor component. Specifically, ResMem involves first training a standard neural network f DeepNet , and then explicitly memorizing the model's residuals with a k -nearest neighbor based regressor r k NN .

Vaishnavh Nagarajan

f

Response

Features

DeepNet

Features

Residual

r

Residual

Features

Figure 1: Illustration of the residual memorization (ResMem) algorithm. In a nutshell, we first fit a small deep network f DeepNet on the training sample (Figure 1(a)). When this network is nonmemorizing, 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 .

Memorization through k -nearest neighbor can be efficiently computed with various approximation schemes (e.g. [24]). Subsequently, the ResMem prediction on an instance x is given by the sum of the two components, i.e., f DeepNet ( x ) + r k NN ( x ) .

We start by formulating a stylized linear regression problem that captures the essence behind ResMem (cf. Section 3). Our analysis (Theorem 3.3) shows that, without ResMem, the test risk of the base linear neural network decreases to an irreducible constant as the number of samples goes to infinity. In contrast, the test risk of ResMem decreases to zero. The insight of theoretical analysis is that ResMem augments the capacity of the parametric linear network by adding a non-parametric component (i.e., nearest-neighbor).

Empirically, we show that such explicit memorization indeed leads to generalization benefits: ResMem consistently improves the test accuracy of a baseline DeepNet on image classification tasks with CIFAR100 [33], and autoregressive language modeling on C4 [42] (Section 4). Towards understanding this improved performance, we hypothesize that ResMem works by learning in twostages (cf. Section 4.4). Specifically, we posit that the initial DeepNet f DeepNet learns some coarse structure, and ResMem r k NN supplements the DeepNet prediction with fine-grained details (cf. Figure 3). We verify our hypothesis via qualitative analysis on CIFAR100 and C4 (Section 4.4).

To summarize, our contributions are:

We propose residual-memorization ( ResMem ), a two-stage learning algorithm that combines a base prediction model with a nearest neighbor regressor (cf. Figure 1);

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

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

1.1 Applicable scenarios of ResMem

From our theoretical and empirical analysis, we posit that ResMem (Figure 1) yields the largest margin of improvement over a base DeepNet when it is infeasible to perform implicit memorization with the latter. We discuss three such scenarios below. Each of our main empirical or theoretical results roughly corresponds to one of these settings.

Complex dataset. 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 of linear neural networks (cf. Section 3).

Large sample size. Here, the number of training samples is large enough to make training set interpolation (i.e., achieving zero training error) infeasible for a given neural network model. For example, current large language models (LLMs) may be trained for at most a single epoch over trillions of examples [12]. By contrast, ResMem can circumvent this issue by explicitly memorizing the training samples. We emulate this scenario by considering a causal language modeling task on the C4 dataset (cf. Section 4.3).

Small model. 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.2).

2 Related work

We discuss two types of related work: Section 2.1 for literature on memorization and generalization that motivates the ResMem algorithm; Section 2.2 for other related algorithms similar to ResMem.

2.1 Memorization for generalization: prior work

Memorization is compitable for generalization. Overparameterized neural models with many more parameters than training samples have the capacity to perfectly fit (or interpolate ) even random training labels [54]; i.e., they can drive the empirical loss to zero for any training set. At the same time, when trained on real-world datasets, increasing model complexity tends to improve model performance [40, 52]; 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 [2] and theoretically [3, 15, 8, 5, 40, 36, 38, 4, 48, 50, 53]. 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 [17, 11] showed that memorization - either in the sense of interpolation, or in a more general sense of stability [18] may be necessary for generalization. Feldman [16] 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. [11] 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).

2.2 Relation to existing algorithms

Nearest neighbor method. The k -nearest neighbor ( k -NN) [14, 32, 26, 7] 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 [22, 39], this method remains popular. However, the performance of k -NN drops as data becomes high dimensional [10, 39]. Therefore, to apply it to high dimensional data such as image and text [55], one approach is to learn a representation of data using neural networks [44]. Following this approach, [13] finds that applying k -NN directly to memorize the training labels y i yields similar performance with the original softmax based neural network classification. In contrast, ResMem applies k -NN to memorize the residual r i over the predictions of a base network.

Boosting and residual fitting. Boosting algorithms such as AdaBoost [20] 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 [21], a classical procedure from statistics [23, 1, 47]. Intuitively, at each round, one builds a new weaker learner by fitting the residual of the ensemble of weak learners constructed thus far. This fitting is performed iteratively. ResMem can be loosely regarded as a two round boosting algorithm where the first 'weak learner' is the base neural network and the second 'weak learner' is the nearest-neighbor component. Note that there is no need for the thrid 'weak learner', because the nearest-neighbor component already perfectly memorizes the neural network residuals.

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 [34, 25, 6, 35]. Closer to our work, Khandelwal et al. [30] employ a linear combination of neural network and k -NN classifier components. However, a crucial difference is that our k -NN components memorizes the residual of the DeepNet prediction, whereas Khandelwal et al. [30] memorizes the target label directly; i.e., their approach is akin to an ensemble of k -NN and a deep network. Various forms of memory have also been considered in generic classification problems [41, 48, 51]. 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 [29, 9] 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 [37]. In Appendix C.2, we discuss how ResMem can be regarded as a 'dual procedure' of distillation.

3 Theoretical results

As discussed in Section 1.1, ResMem yields the largest improvement when implicit memorization is infeasible. In this section, we formulate (cf. Section 3.1) and analyze (cf. Theorem 3.3) a stylized linear regression problem that concretizes such a setting.

Recall that ResMem (Figure 1) involves first training a base neural network f DeepNet , and then fitting the residual of f DeepNet on the same training data using a nearest-neighbor regressor r k NN . For feasibility of theoretical analysis, we simplify f DeepNet with a single layer linear neural network, i.e. linear regression, and we consider 1 -nearest neighbor instead of k -nearest neighbor to memorize the residual of this network. Our results suggests that ResMem improves test-time generalization by augmenting the capacity of the base model with a non-parametric nearest-neighbor component.

3.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 3.1 (Distribution of covariates) . The distribution of covariate x ∈ R d , denoted by P x , is the uniform distribution 1 over a Euclidean norm ball centered at the origin of radius √ d +2 . The choice of radius ensures that E x ∼ P x xx T = I .

Assumption 3.2 (Linear regression over norm ball) . Consider the problem of learning a linear function f ⋆ ( x ) = ⟨ x , θ ⋆ ⟩ with ∥ θ ⋆ ∥ = 1 from training data { ( x i , y i ) } i =1: n where x i i . i . d . ∼ P x and y i = f ⋆ ( x i ) using the function class

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

ResMem proceeds by first learning a linear function f n ( x ) = ⟨ θ n , x ⟩ from F through empirical risk minimization (ERM):

The empirical risk minimizer f n should be thought of as the analog of f DeepNet in the deep learning context. It defines a ground-truth residual function r ⋆ ( x ) = f ⋆ ( x ) -f n ( x ) . Now we fix a test covariate ˜ x ∼ P x . ResMem 'memorizes' the residual function through the 1 -nearest neighbor to ˜ x

where ˜ x (1) is the nearest neighbor to ˜ x among the training covariates x 1 , . . . , x n :

Observe that if ˜ x coincides with any training sample, f ResMem n ( ˜ x ) = f ⋆ ( ˜ 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.

3.2 A decomposition of the target function

Next, we introduce a decomposition of f ⋆ , which will help us analyze various components that make up the risk. Define

which is what ERM learns in the limit of n → ∞ . We can think of θ ∞ as the best function that ERM can learn. Then, we can decompose θ ⋆ into θ ⋆ = θ ∞ + θ ⊥ , where θ ⊥ = θ ⋆ -θ ∞ . This decomposition can be generalized beyond linear regression. Since θ ∞ defines a function f ∞ ( x ) = ⟨ x , θ ∞ ⟩ , for general non-linear functions, the argument above can be generalized to the decomposition of f ⋆ to an learnable and non-learnable part

Intuitively, f ∞ is the best function in F that ERM can learn, and f ⊥ is beyond the capacity of ERM due to the particular choice of function class. ResMem approximates f ⊥ using the non-parametric nearest neighbor method, and therefore expanding the capacity of the original hypthesis class.

3.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

It can be decomposed into two components: E [ f ResMem n ( ˜ x ) -f ⋆ ( ˜ x ) ] 2 ≤ 3 ×

We provide the detail of the decomposition in Section B.1. We can see that T 1 arises due to the difference between f n and f ∞ (i.e., the estimation error), which, as we will show later, goes to 0 as n goes to infinity:

On the other hand, T 2 arises due to the limited capacity of 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 ( ˜ x (1) → ˜ x as n →∞ ), we also have

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.

The final prediction is

3.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 is in Appendix B.

Theorem 3.3 (Risk for ResMem algorithm) . For the problem defined in Assumption 3.2 with covariates distribution in Assumption 3.1, the ResMem prediction rule f ResMem n ( ˜ x ) defined in equation (4) achieves risk (5)

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

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

T 1 ≲ d 2 L 2 n -2 / 3 that arises due to the difference between f n and f ∞ .

T 2 ≲ [ log ( n 1 /d ) /n ] 1 /d that vanishes as the nearest neighbor of the test point approaches the test point itself ˜ x (1) → ˜ 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 → 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 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 3.3, pure nearest neighbor would result in an overall slow rate of ≈ n -1 /d . However, with ResMem, we can enjoy benefit of having the test loss being asymptotically 0 , while also enjoying the fast rate of n -2 / 3 for smaller sample sizes.

4 Empirical results

In this section, we present empirical results on image classification and language modeling that showcase the efficacy of ResMem. In Section 4.1, we present details of applying the ResMem algorithm to classification problems on real dataset. In Section 4.2 and Section 4.3, we present the setup and the result for vision and language experiments, respectively. In Section 4.4 we conduct an empirical analysis to explain where the improved accuracy of ResMem comes from. Finally, in addition to evaluating the improvement ResMem algorithm over DeepNet itself, we compare ResMem with other reasonable baselines including [31] in Appendix F.

4.1 Details of ResMem algorithm for classification

We consider multi-class classification problems over instances X and labels Y . = { 1 , 2 , . . . , L } = [ L ] . Given training examples S = { ( x i , y i ) } i ∈ [ n ] ∈ ( X × Y ) n , the goal is to learn a scorer f : X → R L that, given an instance, assigns an affinity score for each label. Such an f should minimize the misclassification error on test samples:

̸

where pred ( z ) . = arg max y ′ ∈ [ L ] z y ′ , and P is the distribution over labelled instances. To achieve this, one typically minimizes the empirical loss

where ℓ : [ L ] × R L → R + is a loss function. Ideally, one would like to use ℓ 01 ( y, f ( x )) . = 1( y = pred ( f ( x ))) ; for computational tractability, it is popular to instead use a surrogate loss , such as the softmax cross-entropy.

Given the notation above, ResMem operates as follows:

̸

resnet8

resnet14

resnet20

resnet32

resnet44

resnet56

DeepNet architecture

0.575

0.600

0.625

0.650

0.675

0.700

0.725

Test accuracy

DeepNet

ResMem

resnet8

resnet14

resnet20

resnet32

resnet44

resnet56

DeepNet architecture

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Training accuracy

DeepNet

ResMem

10%

20%

30%

40%

50%

60%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.40

0.45

0.50

0.55

0.60

Test accuracy

DeepNet

ResMem

10%

20%

30%

40%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Training accuracy

DeepNet

ResMem

(a) Test(left)/Training (right) v.s. architectures.

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

Train the base DeepNet. Train a neural network f DeepNet on the training samples S as usual.

Prepare the residual data. Compute the residual 2 prediction of each training example as

where onehot : Y → 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 DeepNet , denoted by z i = ϕ ( x i ) , as an embedding for the training instance x i . These embeddings are utilized for the nearest neighbor search in the next step.

Predict via memorized residuals. To obtain a prediction on a test sample ˜ x ∈ X , first compute its embedding ˜ z = ϕ ( ˜ x ) . Then, use soft k -nearest neighbor method to build a function r k NN defined by weights w i ( ˜ x ) :

The weights w i ( ˜ x ) satisfy ∑ i w i ( ˜ x ) = 1 , and are computed from raw weights w i as follows:

where w ( k ) represents the k -th largest entry of w i 's. Note that k and σ are two hyperparameters that collectively controls the locality of nearest neighbor search.

We make the final prediction based on the following scorer:

Remark 4.1 (Explicit memorization) . Smaller k or σ corresponds to putting higher weight on residuals of the closest neighboring training examples. For sufficiently small k and σ , f ResMem achieves exact memorization of the training sample, i.e., pred ( f ResMem ( x i )) = y i for all i ∈ [ n ] .

Remark 4.2 (Computation cost) . The vision experiments have moderate training sample size, so we perform exact nearest neighbor search and discuss the computation cost in Section 4.2. For language experiments, the training sample size is so large that the exact nearest neighbor computation is infesible, so we rely on approximate nearest neighbor search discussed in Section 4.2.

4.2 Image classification

In this subsection, we mainly consider image classification task on ResNet [27] with CIFAR100 [33] dataset. We provide additional ImageNet [43] results in Apendix D.

Setup. We use CIFAR-ResNet-{ 8 , 14 , 20 , 32 , 44 , 56 } as the base DeepNet. For all six 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

...allow for plenty of headroom

inside whilst still being less than

2.5m in height.

rose

poppy

cup

plate

squirrel

rabbit

butterfly

bee

palm

tree

pine

tree

Graphic now consists of all cities

with greater than 30,000 locals.

Acquiring a home in Spain…

Filmed around 7:30-8:30 a.m. on

Friday, March 9, 2012.

...that will not affect the superior

quality of your job. That is possible

because we understand how to save...

...answer your questions and schedule

the initial meeting. We consistently

arrive at the scheduled hour...

height

length

home

residence

March

June

possible

feasible

consistently always

Image classification

Language modeling

y

ResMem

y

ResMem

y

DeepNet

y

DeepNet

image embedding, which has dimension 64. For the nearest neighbor search (Step 3, Section 4.1), we define the distance between two images to be the ℓ 2 distance between their embeddings. We use σ = 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 C (cf. Figure 5).

Results. The results for CIFAR-ResNet-{ 8 , 14 , 20 , 32 , 44 , 56 } are reported in Figure 2(a). 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. To access the statistical reliability of the improvement, we repeat the CIFAR-ResNet-8 experiment 5 times over random initialization of DeepNet etc. We and that the average ResMem accuracy is 59% with standard deviation 0 . 7% , and the average DeepNet accuracy is 56 . 5% with standard deviation 0 . 8% .

Computationally, we estimate the CPU latency of a CIFAR-ResNet-8 to be 15.9 ms for a single test image. By contrast, the k -NN step takes 4.8 ms for the same test image. To contextualize the latency cost, the total cost of ResMem with ResNet-8 (15.9 ms + 4.8 ms) is lower than the cost of the next-sized model, i.e., ResNet-14 (26.2 ms). Regarding the memory cost, for a batch size of 1 and images of size 32 x 32, a ResNet-8 ( 68K params) requires 2.5MB, while a ResNet-14 ( 128K params) requires 4MB. Embeddings from a ResNet-8 and ResNet-14 are both 64 dimensional. To embed the entire CIFAR100 training set (50K examples) requires 15MB of disk space.

Varying sample size. We repeat the above experiment on CIFAR-ResNet-8 with subsets ( 10% , 20% , . . . , 100% ) of CIFAR100 training data (subsampled uniformly across different classes). The size of the index set for nearest-neighbor search is the same as the training set for base neural networks (e.g., model with 10% CIFAR100 data also uses the same 10% data for nearest-neighbor search). On the left (right) of Figure 2(b), 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.

4.3 Language modeling

Setup. For the language experiment, we use a Decoder-Only T5-{small, large} [42] model and C4 [42] 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 autoregressive 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 approximate nearest neighbor search algorithm ScaNN [24]. We used the temperature T = 1 for the residual computation and σ = 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.

Results. On T5-small, ResMem boosts test accuracy from 38.01% to 40.87% , which is around the accuracy (40.08%) of a T5-base model without ResMem. On T5-large, ResMem boosts the test accuracy from 44.8% to 45.6% . This demonstrates that with explicit memorization, we may leverage smaller base language models while reaping the performance benefits of large language models. Computationally, as the index set is quite large (1.6 billion tokens), exact k-nearest neighbor search is infeasible. So we use the approximate nearest neighbor search algorithm ScaNN [24] to reduce compute time. Please see Appendix E for details on base model training and data processing.

4.4 Where does the improvement come from?

In this section, we identify test samples that contributes to the accuracy improvement of CIFAR100 with CIFAR-ResNet-8 and C4 with T5-small. Let Gain ResMem be the difference between the test accuracy of ResMem and baseline DeepNet:

where L 01 is the misclassification error as defined in equation (6). We offer a decomposition of Gain ResMem that sheds light into the mechanism behind ResMem. For a test set { ( x i , y i ) } i m =1 , let y ResMem i be the ResMem prediction on instance x i and let y DeepNet i be the baseline neural network prediction on x i . When y ResMem i = y DeepNet i , sample x i does not contribute to Gain ResMem . When y ResMem i = y 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:

̸

̸

̸

Note that Gain ResMem = TPR -FPR . The decomposition of Gain ResMem says that the gain of ResMem came from the TPR samples, provided they outweigh the FPR samples.

On CIFAR-ResNet-8, we find TPR = 5.89% and FPR = 2.70%, leading to Gain ResMem =3.19%. On T5small with C4 validation split, we find TPR = 5.37% and FPR = 2.44%, leading to Gain ResMem =2.93%.

̸

Analysis of TPR samples Focusing on the test samples where ResMem helps ( y i = y ResMem i = y DeepNet i ), we identify a common underlying pattern: while the DeepNet makes an incorrect prediction, it still captures some coarse structure. For example, in CIFAR100, one sample has correct label y i = y ResMem i = rose , but the DeepNet predicts y DeepNet i = poppy , i.e., the label of a different type of flower. (cf. Figure 3). We find similar behavior for the language modeling task (cf. Figure 3).

This 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. In this paper, we choose the ResMem hyperparameters that minimizes the test error on the validation set or, equivalently, maximize Gain ResMem . Inspired by the decomposition of Gain ResMem , we propose an alternative hyperparameter selection procedure based on the following optimization problem:

which ensures that ResMem modifies the DeepNet predictions 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 maximizing the desirable occurrences where ResMem can correct a wrong prediction by DeepNet.

5 Discussion and future works

Joint training of k NNand DeepNet. 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.

To explore the role of training during the first stage, we re-evaluate the CIFAR-ResNet-8 experiment by stopping DeepNet training at different epochs (Table 1). We can see that when the #epoch is small,

ResMem has a dramatic improvement in accuracy. One of the key roles of the first training phase is to learn good representations of the training data so the nearest neighbor retrieval is performed on more meaningful representations. This simple experiments suggests that the proposed direction has the potential to dramatically reduce the training time of DeepNet - while obtaining similar test accuracy with the help of ResMem.

Calibration of ResMem. A potential problem with applying ResMem to classification is scorer mis-calibration . The output of the ResMem prediction vector f ResMem ( x ) (8) 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 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.

Distribution shift. 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 [46].

References

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 .

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.

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.

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 .

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.

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 .

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 .

Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz. SGD learns overparameterized networks that provably generalize on linearly separable data. In International Conference on Learning Representations , 2018. URL https://openreview.net/forum? id=rJ33wwxRb .

Cristian Bucilˇ a, 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.

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 .

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 .

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 .

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 .

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.

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.

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 .

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 .

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 28812891. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/ 2020/file/1e14bfe2714193e7af5abc64ecbd6b46-Paper.pdf .

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, 2020. Curran Associates Inc. ISBN 9781713829546.

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.

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 .

A. Gionis, Piotr Indyk, and Rajeev Motwani. Similarity search in high dimensions via hashing. In Very Large Data Bases Conference , 1999.

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 .

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 .

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.

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.

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.

Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015. URL https://arxiv.org/abs/1503.02531 .

Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in a neural network. CoRR , abs/1503.02531, 2015.

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 .

Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generalization through memorization: Nearest neighbor language models. In ICLR , 2020. URL https://openreview.net/forum?id=HklBjCEKvH .

Donald Knuth. The Art Of Computer Programming, vol. 3: Sorting And Searching . AddisonWesley, 1973.

Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, CIFAR, 2009.

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 .

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 .

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 .

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 .

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 .

Marius Muja and David G. Lowe. Fast approximate nearest neighbors with automatic algorithm configuration. In International Conference on Computer Vision Theory and Applications , 2009.

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.

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 .

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 .

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.

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 .

Mark Sandler, Andrew Howard, Menglong Zhu, Andrey Zhmoginov, and Liang-Chieh Chen. Mobilenetv2: Inverted residuals and linear bottlenecks. In 2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 4510-4520, 2018. doi: 10.1109/CVPR.2018. 00474.

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.

Ryan J. Tibshirani. A general framework for fast stagewise algorithms. J. Mach. Learn. Res. , 16(1):2543-2588, jan 2015. ISSN 1532-4435.

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 .

Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint , volume 48. Cambridge University Press, 2019.

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 .

Zhen Wang and Yuan-Hai Shao. Generalization-memorization machines, 2022. URL https: //arxiv.org/abs/2207.03976 .

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 .

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 .

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.

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 .

A Some concentration results for uniform random variables

In this section, we state some concentration results that are useful for the theoretical analysis in Section 3. Let ˜ x , x 1 , . . . , x n i . i . d . ∼ Unif ( B 0 , √ d +2 ) be i.i.d. samples from the uniform distribution over the Euclidean norm ball of radius √ d +2 in R d . Let

If n = 1 , E Z 1 is the sum of the variance of each coordinate of Unif ( B 0 , √ d +2 ) . Therefore, E Z n provides a generalized measure of concentration. Intuitively, E Z n → 0 as n →∞ . The proposition below provides a upper bound on the rate of convergence.

Lemma A.1 (Nearest Neighbor concentration) . Given the assumptions above

where ≲ means inequality up to an universal constant independent of d and n .

Proof. Define

We will compute two probabilities P ( E 1 |E 2 ) and P ( E 2 ) that will be useful latter.

Next, we compute P ( E 2 )

We use E 1 and E 2 to compute the following upper bound

To find an upper bound for E Z n , we need to find an upper bound for 1 -P ( E 1 |E 2 ) P ( E 2 ) .

Now choose δ = √ d +2 n -1 /d [ log ( n 1 /d )] 1 /d .

and

Thus

Combining everything together, we get

This completes the proof.

Proposition A.2 ([49] Corollary 6.20) . Let x i i . i . d . ∼ Unif ( B 0 , √ d +2 ) for i = 1 , . . . , n be uniformly distributed over a ball of radius B in R d centered at 0 . Let

be the sample covariance matrix. Then

B Proof of Theorem 3.3

In this section, we present the proof of Theorem 3.3. In Section B.1, we provide the detail of the decomposition of the risk into T 1 and T 2 . Then in Section B.2 we compute an upper bound for T 1 , and compute an upper bound for T 2 in Section B.3. Finally, we combine everything together in Section B.4 and completes the proof.

B.1 Decomposition of the test risk

where in the last inequality, we used the fact that ( a + b + c ) < 3( a + b + c ) for any a, b, c ∈ .

B.2 Upper bound on T 1 .

Since P x = Unif ( B 0 ,B ) , we apply the bound ∥ ˜ x ∥ , ∥ ˜ x (1) ∥ ≤ B to obtain

As n gets large, the empirical covariance matrix Σ n = X T X /n is concentrated around its mean I . Let ∆ n = I -Σ n denote this deviation. For some ε ∈ (0 , 1) , define the following 'good event' over the randomness in Σ n

where ∥ ∆ n ∥ op denotes the operator norm of the deviation matrix. The high level idea of the proof is to condition on the event A and deduce and upper bound of ∥ θ n -θ ∞ ∥ in terms of ε . Then, we use the fact that A happens with high probability.

Recall that θ ∞ = L θ ⋆ , and and at optimality

Since y = Xθ ⋆ by definition, the Lagrangian of the convex program above is

The KKT condition suggests that the primal-dual optimal pair ( θ n , λ n ) is given by

The complementary slackness condition λ n ( ∥ θ n ∥-L ) = 0 suggests that either λ n = 0 or ∥ θ n ∥ = L . But if λ n = 0 , the stationary condition ∇ θ L ( θ , λ ) = 0 would suggest that θ n = Σ -1 n Σ n θ ⋆ = θ ⋆ ⇒∥ θ n ∥ = 1 > L, a contradiction. (Note that here Σ n is invertible condition on the event A .) Therefore, we must have ∥ θ n ∥ = L . As a result, the primal and dual pair ( θ n , λ n ) is determined by the system of equations

Next, we proceed to compute the deviation ∥ θ n -θ ∞ ∥ .

Define

Then θ n = ( λ n +1) -1 θ ⋆ +( λ n +1) -1 D n θ ⋆ , and

Therefore

We can obtain the following bound for δ n :

Since 1 -δ n / 2 ≤ √ 1 + δ n ≤ 1 + δ n / 2 , and | δ n | ≤ 15 ε , we obtain

where the last inequality follows as we have λ n > 0 . Finally,

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

If we choose ε = n -1 / 3 , we get which implies that

B.3 Upper bound on T 2 .

Plugging in the formula for f ⊥ ( ˜ x ) = f ⋆ ( ˜ x ) -f ∞ ( ˜ x ) = ⟨ ˜ x , θ ⊥ ⟩ , we get

where in the last inequality, we used the relation that θ ⊥ = (1 -L ) θ ⋆ . Proposition A.1 suggests that

Remark B.1 (Comparison with pure nearest neighbor and ERM) . If we rely solely on nearest neighbor method, the prediction error is

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

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

Projecting f ⋆ onto f ∞ , 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 E ∥ ˜ x ∥ 2 to E ∥ ˜ x (1) -˜ x ∥ 2 .

which implies

B.4 Test loss for ResMem .

If we combine the previous two parts together, we get

This completes the proof of Theorem 3.3.

C Additional CIFAR100 Results

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

C.1 Additional robustness results

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

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.40

0.45

0.50

0.55

0.60

Test accuracy

DeepNet

ResMem

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.65

0.70

0.75

0.80

0.85

0.90

0.95

1.00

Training accuracy

DeepNet

ResMem

(a) CIFAR-ResNet-8

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.40

0.45

0.50

0.55

0.60

0.65

Test accuracy

DeepNet

ResMem

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.80

0.85

0.90

0.95

1.00

Training accuracy

DeepNet

ResMem

(b) CIFAR-ResNet-14

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.45

0.50

0.55

0.60

0.65

0.70

Test accuracy

DeepNet

ResMem

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.825

0.850

0.875

0.900

0.925

0.950

0.975

1.000

Training accuracy

DeepNet

ResMem

(c) CIFAR-ResNet-20

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.45

0.50

0.55

0.60

0.65

0.70

Test accuracy

DeepNet

ResMem

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.90

0.92

0.94

0.96

0.98

1.00

Training accuracy

DeepNet

ResMem

(d) CIFAR-ResNet-32

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.45

0.50

0.55

0.60

0.65

0.70

Test accuracy

DeepNet

ResMem

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.93

0.94

0.95

0.96

0.97

0.98

0.99

1.00

Training accuracy

DeepNet

ResMem

(e) CIFAR-ResNet-44

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.45

0.50

0.55

0.60

0.65

0.70

0.75

Test accuracy

DeepNet

ResMem

10%

20%

30%

44%

60%

50%

70%

90%

80%

100%

Percentage of CIFAR100 used

0.96

0.97

0.98

0.99

1.00

Training accuracy

DeepNet

ResMem

(f) CIFAR-ResNet-56

C.2 Sensitivity analysis for CIFAR100

100

200

300

400

500

#neighbours

k

in ResMem

0.565

0.570

0.575

0.580

0.585

0.590

0.595

Test accuracy

ResMem

DeepNet

(a) # of neighbours

k

.

0.25 0.50 0.75 1.00 1.25 1.50 1.75 2.00

parameter in ResMem

0.46

0.48

0.50

0.52

0.54

0.56

0.58

0.60

Test accuracy

ResMem

DeepNet

(b) Radius parameter

σ

.

0

1

2

3

4

5

Softmax temperature

0.55

0.56

0.57

0.58

0.59

Test accuracy

ResMem

DeepNet

(c) Temperature

T

.

Varying locality parameter k and σ . 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 →∞ , the ResMem test accuracy seems to be converging to a constant level. For σ , we explored different values of σ ∈ (0 . 1 , 2 . 0) . We observe that the test accuracy has a unimodal shape as a function of σ , suggesting that there is an optimal choice of σ (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 [28]. 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 (k NN ) 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.

D ResMem on ImageNet

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

ImageNet. In addition to CIFAR100, we also evaluate the performance of ResMem on ImageNet [43]. We employ a family of pre-trained MobileNet-V2 models [45] from Keras 3 , with varying widths controlled by a multiplier a . For ResMem, we again use the second last layer of DeepNet as a 1280-dimensional embedding of an image and rely on the ℓ 2 distance between the embeddings for nearest neighbor search (Step 3, Section 4.1). We specify the ResMem parameter of ( k, σ, T ) in the table below. We repeat the experiment over several MobileNet-V2 architectures, with MobileNet-V2-a0.35 being the smallest model and MobileNet-V2-a1.3 being the largest one.

We can see that (c.f. Table 2) ResMem boosts the test accuracy by 1% on the smallest model and by 0 . 4% on the largest model.

E 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 [42] 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.

F Comparison with other algorithms

We mainly compare ResMem against [31], where the algorithm uses k NN to retrive labels directly instead of the residual of the label. In their algorithm, a key aparameter is λ ∈ [0 , 1] which specifeis how much weight to give to the neural network and how much for the k NN component. In the extreme case of λ =1, their algorithm reduces to using k NN to memorize data directly.

For the language modeling task, we use the C4 dataset and T5-large architecture. As we change the weight [31, Equation (3)] of the DeepNet component, we find the best performing kNN-LM methods has accuracy 44.88% which is lower accuracy than the ResMem accuracy 45.55%. In particular, we obtain the table below

For image classification with CIFAR-ResNet-8, we run the simple baseline of using k -nearest neighbor to directly memorize the labels . We observe the performance: we observe that pure DeepNet has accuracy 56.46%; pure k NN memorization has accuracy 54.44%; and ResMem has accuracy 59.66%.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nSeveral recent works investigated the role of memorization in Deep Neural Networks as long as its interplay with generalization capabilities. In this context, this work proposes a two-stage procedure where a standard prediction model is augmented  via a  k-NN based mechanism that explicitly memorizes the model residuals. The authors carried on an interesting theoretical analysis as long as an empirical evaluation of the proposed approach.\n\nSTRENGTHS:\nThe proposed two-stage method is a novel and original approach to leverage memorization as a tool for improving neural network performances. The papaer is well written and structured. The theoretical claims are well supported by the empirical findings and their presentation is easy to follow.  I appreciated the clarirty in the theoretical claims exposition and the broad experimental setup.\n\nWEAKNESSES:\nI have few remarks with the proposed approach. While reading the paper, the first doubt that comes to mind is why the authors chose to devise a two-stage process instead of the joint training of DeepNet and kNN, as noticed in the Future Work section. Moreover, from similar consideration begs a question on the role of the first training stage: what is the role of the duration of the first training phase? What happens when the second stage is applied too early or too late? (if the first stage comprehends less/more epochs)  \n\nFinally, I have some concerns regarding the statistical reliability of the proposed solution: from the current text it is not clear whether the reported results are the average over several runs (with different weights initialization) or not\n\nQUESTIONS:\n1) See \"Weaknesses\" section: what is the role of the duration of the first training phase? What happens when the second stage based on kNN is applied too early or too late? \n\n2) I have some concerns regarding the statistical reliability of the proposed solution: from the current text it is not clear whether the reported results are the average over several runs (with different weights initialization) or not -- given also that in bigger datasets such as ImageNet the performance improvement is very limited \n\n3) Line 252: from the current text, the parameters seem cherry picked (the ones achieving the best performance on the Test set, from Appendix C). The authors should clarify this.\n\n4) The memory (mainly, givent that the kNN need to keep track of training data)/computational burden of the proposed method should be better discussed (enriching lines 257-260)\n\n5) From the current text, it is not clear whether Assumption 3.1 is easy to be fullfilled or not.\n\nLIMITATIONS:\nLimitations have been adequately addressed.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes ResMem (residual memorization) to improve neural network predictions by explicitly memorizing training residual patterns. The final prediction is computed by summing the neural network's output and the fitted residual obtained by KNN. The paper provides theories proving that ResMem yields favorable test risk.  Image classification and language modelling experiments demonstrate that ResMem outperforms the original prediction model.\n\nSTRENGTHS:\n- The idea of using KNN to memorize the residual is novel and interesting\n- There is a risk bound proof\n\nWEAKNESSES:\n- The assumption of the proof (linear function, k=1) is not practical and not aligned with the experiment\n- The experiment should compare with approaches that use KNN to store the response/label (e.g. paper [30] cited by the authors should be compared in the NLP task, and please compare with deep KNN paper (Papernot et al., 2018)  in the image classification task )\n- The current DeepNet is weak, especially in CIFAR100. The test accuracy is low (<73%) compared to the current SOTA such as Vision Transformer (Dosovitskiy et al., 2020). The DeepNet cannot fit the training dataset (did you tune the model properly and let it train long enough?). Fig. 2a shows that the performance gain reduces as the DeepNet gets bigger. It raises the concern that ResMem may not scale well with big/SOTA backbones and will be not useful in practice.   \n\nPapernot, Nicolas, and Patrick McDaniel. \"Deep k-nearest neighbors: Towards confident, interpretable and robust deep learning.\" arXiv preprint arXiv:1803.04765 (2018).\n\nDosovitskiy, Alexey, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani et al. \"An image is worth 16x16 words: Transformers for image recognition at scale.\" arXiv preprint arXiv:2010.11929 (2020).\n\nQUESTIONS:\n- The setup in Sec.4.2: Why do you set up this way? \n- Fig. 5: the results seem insensitive to K. Can you explain?\n\nLIMITATIONS:\nNo section on Limitations. One potential limitation is the practicality of ResMem. It is unclear if this research direction can lead to better results with strong deep networks.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes a new learning algorithm, \"ResMem,\" that combines deep neural networks (DNNs) with k-Nearest Neighbors (kNNs). The algorithm is a form of two-stage boosting: we first train a neural network. Then we construct a kNN algorithm that is trained to predict the residuals. The final output combines these two predictions. The nearest-neighbor search is performed in the feature space defined by an intermediate layer of the neural network. The intuition behind this algorithm is the growing body of literature which attributes some of deep learning's success to memorization: why not make it an explicit part of the pipeline?\n\nThe first result in the paper is theoretical, giving a toy problem where (a constrained but natural algorithm for) linear regression performs poorly, vanilla kNN also performs poorly, but ResMem does reasonably well.\n\nThe remaining results are empirical, applying their algorithm to image classification and language modeling. They show that ResMem provides modest but consistent and significant improvements over the deep networks alone.\n\nSTRENGTHS:\nI think this is a nice idea. As they acknowledge, prior work has explicitly explored connections between kNNs and DNNs, but this paper aims to use the kNN to boost where the DNN fails.\n\nThe algorithm is practical and simple, it seems quite reasonable to fit into existing pipelines.\n\nBeyond practice, I think this is also a paper to spark discussion. It provides a complement to the work of Cohen, Sapiro, and Giryes (2018), who used kNNs to probe the memorization capabilities of DNNs. \n\nI felt the introduction and related work sections were well-written.\n\nWEAKNESSES:\nOne might use ResMem to better understand deep networks, perhaps following Sections 3 and 4.4. Instead, the paper is more like a proposal for a practical algorithm. However, I am not sure where ResMem is the best algorithm available. In Section 1.1, \"Applicable Scenarios,\" the paper gives three ideal settings. The first is \"complex dataset,\" where the task is beyond the capacity of the learner. We only see this setting theoretically in the paper. The other two examples are \"large sample size\" and \"[s]mall model,\" but in both of these cases it is not clear to me that ResMem makes practical sense. ResMem requires storing representations of the entire training set, so when the sample size is large or the model is small I expect this would be serious overhead. \n\nI found it hard to find the algorithm in the paper: Section 3 analyzes \"ResMem for regression,\" which is described in text and equations. Section 4.1 gives the algorithm for classification in a enumerated list, plus some additional details in the surrounding text. I strongly prefer putting the algorithm in a separate environment, making it clear to the reader what you propose.\n\nThe paragraph \"Memorization is sufficient for generalization\" (line 76) seems extremely confused. I have not read all the cited papers, but my understanding is that they focus on how memorization is *compatible* with generalization. It seems to me that memorization alone *suffices* for generalization only in trivial learning tasks. Concretely, consider an algorithm that uses a lookup table on it training data and answers randomly on unseen points. This always memorizes but only generalizes when the training data captures most of the probability mass of the data distribution.\n\nQUESTIONS:\nLines 82-83 say \"One recurring message from the theory is that memorization (in the form of interpolation) can be sufficient for generalization.\" Can you connect this statement more directly with an example or two from the preceeding citations? \n\nHow do you pick the intermediate layer to use as the representations for the kNN?\n\nResMem requires storing representations of the training data. Do you think this is an obstacle to its practical application?\n\nCan you present some comparisons about space usage? For instance, in in line 257, you compare the computation of ResNet-14 to that of ResMem (on ResNet-8). What is the footprint of these two models?\n\nLIMITATIONS:\nAside from potential issues around space usage, which are addressed above, I feel the authors adequately address the limitations of their work.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes the residual-memorization (ResMem) algorithm.\nThis algorithm aims to improve the generalization performance by explicitly memorizing the residual between the neural network's outputs and the training labels with a kNN regressor.\nThe paper demonstrates the validity of the algorithm from both the theoretical and the empirical perspectives.\nThe main theoretical result shows the convergence rate of ResMem.\nThe empirical results show the performance increase of ResMem compared with standard neural networks on image and language datasets.\n\nSTRENGTHS:\nThe proposed method is quite novel.\nThe writing is clear.\nThe paper consists of both theoretical and empirical justification.\nThis paper provides an additional perspective on improving generalization in neural networks.\n\nWEAKNESSES:\nWhile the method proposed in this paper is novel and the empirical results look promising, I am concerned about whether the theoretical result justifies the message that the proposed method improves model generalization via memorization.\n\nOn the theoretical result, the paper does show the asymptotic behavior of the prediction error of ResMem. \nHowever, it is not clear if the asymptotic behavior shows that it improves model generalization, since the paper does not include the comparison with the case without ResMem.\nIn particular, the paper claims in Section 3.3 that $T_2$ captures an irreducible error of the risk that is in general not asymptotically zero; \nhowever, it is known that for the linear model with ERM investigated in this paper in the overparameterized regime, the generalization error is asymptotically zero under a range of certain circumstances.\n\nQUESTIONS:\nI am happy to raise the score if the authors address my concern. \n\nQ1: In Section 3.2, why is the equality in the second line of equation under Line 163 correct and consistent with the first line?\n\nQ2: In Section 3.3, what does the $T_2$ part of error represent?\nThe author claims that $T_2$ arises due to the limited capacity of $F$, and captures an irreducible error of the risk.\nCould the authors please justify the argument?\nCan this $T_2$ term represent the generalization error?\n\nQ3: In Section 3.4, could the authors compare the bound with the case without ResMem?\nIn particular, what is the error difference with or without ResMem in the $T_2$ term?\n\nQ4: On the empirical result, what is the additional computation and memory cost of the ResMem algorithm?\nIn particular, how do the additional memory cost, training time and estimation time of the kNN component compare with those of the DNN component?\n\nMinor issues & suggestions:\n- Typo: In Line 67, \"mall model\"\n- The training error comparison in Figure 2 does not seem to be a fair comparison when the proposed ResMem guarantees a near 100% training accuracy. The test error comparison is the more informative one.\n\n---\nThank you for the clarification, especially on clarifying Assumption 3.2.\nI updated the score accordingly.\n\nLIMITATIONS:\nDoes not apply", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}