ResMem: Learn what you can and memorize the rest
Authors: Zitong Yang
Michal Lukasik
Vaishnavh Nagarajan
Zonglin Li
Ankit Singh Rawat
Manzil Zaheer
Aditya Krishna Menon
Sanjiv Kumar
Venue: NeurIPS 2023 poster (2023)
Keywords: deep learning
generalization
memorization
deep learning theory
boosting
nearest neighbor
TL;DR: We introduce residual-memorization (ResMem), a simple yet effective algorithm that wraps around a prediction model and improves its test accuracy, verified empirically and theoretically.

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
Zitong Yang Stanford University Stanford, CA 94305 ittongbberkeley..... edu
Michal Lukasik Google Research New York, NY, 10011 mlukasik@google.. com
Vaishnavh Nagarajan Google Research New York, NY, 10011 waishnavh@@ooggee.cm
Zonglin Li Google Research New York, NY, 10011
Ankit Singh Rawat Google Research New York, NY, 10011
Manzil Zaheer
Google Research New York, NY, 10011
1 iizonglin@google.. com
anki takktssraat@@gogle.ccmmm com
manzilzzhherr@oogle..omm
Aditya Krishna Menon Google Research New York, NY, 10011
Sanjiv Kumar
Google Research New York, NY, 10011 sanji nanjiv@@oogle.com com
adi adityyakmenon@goggle.cco
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- memorizatinn (ResMem) algorithm, new method that augments an existing prediction model (e.g., 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 more favorable test risk over base linear neural network. Then, empirically show that ResMem consistently improves the test set generalization of the original prediction model across standard vision and natural language processing benchmarks.
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 sufficient to achieve generalization [ 15 5]
40 4. and, more surprisingly, can even be necessary memorization can be valuable desiderata for learning. While increasing model size 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 relatively small model?
[17
19 II These works suggest that suitable
In this paper, we propose residual memorization (ResMem), simple yet effective mechanism that achieves this goal (cf. Figure []. Compared to the implicit memorization performed by large neural models, the key idea behind ResMem perform explicit memorization via separate k-nearest neighbor component. Specifically, ResMem involves first training a standard neural network fDeepNet. and then explicitly memorizing the model's residuals with k-nearest neighbor based regressor TANN-
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

DeepNeet
Response
Residual
Residual
KNNN
Features
Features
Features
(a) Step 1: learn the training set. (b) Step 2: compute the residual.
(c) Step 3: memorize the residual.
Figure Illustration of the residual memorization (ResMem) algorithm. In a nutshell, we first fit small deep network /DeepNet on the training sample (Figure ((a) When this network is non- memorizing, it incurs non-zero residual errors in its predictions (Figure I(b). We then a k-nearest neighbor based regressor these residuals (Figure [(c)). The final predietion is given by the of the initial network and k-NN regressor predictions. In all three figures, the æ-axis represents the features in a supervised learning problem. In Figure [[(]]] the y-axis represents the targets of prediction. In Figure ((b) and [(c)
the y-axis represents the residual the initial fitting from Step 1.
Memorization through n-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 components, i.e., /DeepNet (ar) TKNN (2:).
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 zero. The insight of thcoretical analysis is that ResMem augments the capacity of the parametric linear network by adding a -on-paraettric component (i.e., neiiest-neighbor).
Empirically, we show that such explicit memorization indeed leads to generalization benefits: ResMem consistently improves the test accuracy of baseline DeepNet on image classification tasks with CIFAR 100 [33], and autoregressive language modeling on C4 [42] (Section 4 Towards understanding this improved performance, we hypothesize that ResMem works by learning in two- stages (ef. Section 4.4). Specifically, posit that the initial DeepNet fDeepNet learns some coarse - fine-grained details (cf. Figure 3). We verify hypothesis via qualitative analysis on CIFAR 100 and C4 (Section 4.4)
structure, and ResMem FKNN supplements the DeepNet prediction with
To summarize, our contributions are:
(1) We propose
(Res.em), two-stage learning algorithm that combines base prediction model with a nearest neighbor regressor (cf. Figure D:
(2) 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).
(3) We empirically demonstrate that ResMem improves test performance of neural networks (cf. Sec-
tion 4 particularly when the training set is extremely large;
1.1 Applicable scenarios of ResMem
From our theoretical and empirical analysis, 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 beyond the capability of the neural network itself. To demonstrate this, we analyze theoretical linear regression problem where the target regression function 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 (L..Ms) may be trained for at most single epoch over trillions of examples [2]. 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 state-of-the-ar 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 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 [2]; 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 B 15 5 40 36, 38 4 |8 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 showed that memorization - either in the sense of interpolation, or in a more general sense of stability [8] may be necessary for generalization. Feldman (16) considered setting where the label distribution exhibits a long-tailed distribution, and showed that to prevent incurring large error on the large number of under-represented classes, it may be necessary memorize many of their associated training samples, Cheng et al. [II] considered a linear regression setting where it is beneficial to fit the training targets 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 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 computatiinnl efficiency with various approximation schemes [22 [39], this method remains popular. However, the performance of k-NN drops 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--N directly to memorize the training labels ya yields similar performance with the original softmax based neural network classification. In contrast, ResMem applies k-NN to memorize the residual r; over the predictions of 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 particular instantiation of forward stage-wise regression (21), classical procedure from statisties 23 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-neighhor component already perfectly memorizes the neural network residuals.
3

---

Memooy-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 Closer to our work, Khandelwal et al. (30) employ a línear combination of neural network and k--N 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.c., their approach is akin to an ensemble of k-NN and deep network. Various forms of memory have also been considered in generic classification problems 41 [1]. This line of literature again differs from ResMem in that their memory tries to memorize labels directly, whereas ResMem memorizes the residuals, leading to natural combination of the neural network and the memory component.
additional context [4 25
35
Model compression for large neural networks. Since ResMem boosts the test accuracy of 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 the model compression literature. Distillation [29] 9 is a popular strategy for compressing a large neural model to smaller one. For survey of other effective strategies, including pruning, see Menghani [37]. In Appendix C.2 we discuss how ResMem can be regarded as "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 (ef. Theorem 3.3) stylized linear regression problem that concretizes such setting.
Recall that ResMem (Figure Th involves first training base neural network fDeepNet- and then fitting the residual of fDeepNet on the same training data using a nearest- regressor TyNN- For feasibility of theoretical analysis, we simplify /DeepNet with 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 non-parametric nearest-neighhor 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 the response. Therefore, even with infinite samples, the test loss will decay to positive constant. We exactly characterize the rate of decay, and show that it converges 0 under ResMem. Our analysis rests on the following assumptions.
Assumption 3.1 (Distribution of covariates). The distribution of covariate Rd, denoted by Par the uniform distribution over Euclidean norm ball centered the origin of radius vd + 2. The choice of radius ensures that Eg xx.
Assumption 3.2 (Linear regression over norm ball). Consider the problem of learning a linear i..d.
function (x) (x, 0.) with ||o.|| 1 from training data (()) where æ; P. and = fa (xx) using the function class
F (0 - (x. 0), |0 < L).
We assume L < so that the problem belongs to the "hard generalization" scenario discussed in Section I.I where the hypothesis space is inadequate fit the function its
ResMem proceeds by first learning linear function fa (x) (On, x) from F through empirical risk minimization (ERM):
O, = argmin
Σ [x,, 0) -0:12
(2)
| SL
For more general distributions, the theoretical result will depend on quantities like P.. (B(@, h)), where ( A) means a ball of radius / that is centered at क. We took uniform distribution for simplicity and to obtain exact dependence on d.

---

The empirical risk minimizer fu should be thought of as the analog of fDDepNet in the deep learning
context. defines ground-truth residual function r, xx) f=(= - fn()). Now we fix a test covariate a ~ P2.. ResMem "memorizes" the residual function through the 1-nearest neighbor to æ
( r. (z)) ( (ā
where (1) is the nearest neighbor æ among the training covariates 21...,, Zni
(1) = argmin
|| - T||.
me (mi....)))))
The final prediction
fn
ResMem ( = fn +
(4)
Observe that if T coincides with any training sample, f FReMMem (z) (3). i.c., 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 l 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
0 00 argmin En-,, [(0, x) (0.,
argmin |00 - 000.111 -011==111 T.O.,
|
which is what ERM learns in the limit of that ERM can learn. Then, we can decompose 0, into 0, = 0x + 01, where 01 = 0. 0.c. This decomposition can be generalized beyond linear regression. Since 0x defines function (xx) = (x, >o), for general non-linear functions, the argument above can be generalized to the decomposition of fa to learnable and non-learnable part
We think of 6x the best function
foc + fL.
Intuitively, fox. the best function F that ERM can learn, and fi beyond the capacity ERM due to the particular choice of function class. ResMem approximates f' using the normmaramtt 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
E
ResMem n
(5)
It can be decomposed into two components: E RRsMMMm (ar) - [T)) 3x
[ - foo (2 ((fx(())) foo (²))))
E()xo(() - f.( foe (z) + fa (z)
We provide the detail of the decomposition in Section B.1 We can see that Ti arises due to the difference between fn and foo (i.e., the estimation error), which, as we will show later, goes to as n goes to infinity:
T + 0 as n + 00.
On the other hand, T2 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 (1)1) > T as n - oo), we also have
T 0 as n co.
This decomposition provides statistical perspective ResMem: it preserves the asymptotic consistency Tj as in classical learning problems while enforcing the asymptotic consistency of T2 through the nearest-neighhor method.
5

---

3.4 Main theoretical result
Given the set up above, we are ready to state the main theoretical result of the paper, which character- izes the rate at which test risk of ResMem approaches O. The proof is in Appendix B
Theorem 3.3 (Risk for ResMem algorithm), For the problem defined in Assumption with covari-
ates distribution in Assumption 3.1 the ResMem prediction rule J PeesMMm (3)) defined in equation
achieves risk 5
1.1. n-11,-22/3-0 II + d2 L)2 log 1/)
1/d
E
[
ResMem ② - f-))]
where S denotes inequality up to a universal constant independent of d,, and
The result includes contribution from two terms introduced in Section 3.3
d2. 2/3 that arises due to the difference between fn and fac.
T2 S [log (n2/a) /n]] 1/d that vanishes as the nearest neighbor of the test point approaches the test point itself (1(1) ac.
The two terms Tj and T.2 can be viewed as "two stages of learning". Without the ResMem memoriza- tion component, we have the usual story of machine learning: 7i - 0 at the usual parametric rate, and T2 stays as an irreducible error, so the overall test error diminishes to constant at very fast rate. With the introduction of nearest neighbor memorization procedure, T2 can also be reduced to O at a slower rate, whereas the fast decay of Ti 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 in Theorem 3.3 pure nearest neighbor would result in overall slow rate of n n-/d However, with ResMem, we can enjoy benefit of having the test loss being asymptotically 0, while also enjoying the fast rate of -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 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
4.1 Details of ResMem algorithm for classification
We consider multi-clas. classification problems over instances X and labels = 2)) [L] Given training examples = {x,, y:) hein) (X x y)a, the goal is to learn a scorer f:X R that, given an instance, assigns affinity score for each label. Such an f should minimize the misclassification error on test samples:
Loi (J) = ((ay)) + (f))))
(6)
where pzed(z) arg /ax,elL] Ey,, and P is the distribution over labelled instances. To achieve this, one typically minimizes the empirical loss
Le(f) = 1 Σ f(y..f(x))) fxx)).
sE(mi
where é: [L] x Rd + R, is loss function. Ideally, one would like to use Eoi f()) = 1(y + pred(f(x)) for computational tractability, it popular instead surrogate loss, such the softmax cross-entropy.
Given the notation above, ResMem operates follows:

---

136
186-
0323
23
20.08
-+- Deeplict
316
ResMern
325
Dooplit
Deepliet
2515
ResMen
640
34
Beepket archtecture
Grepliet Nrtheee
Percertage of CIARR10 uoed
Pescestage CIAAR100
(a) TesaainiiTrraiiin (right) v.S. architectures.
(b) Test(effyTrraning (right) acc. V.s. sample size.
Figure 2: ResMem improvement on CIFARI00 with respect to training sample size and deep network architecture. (a): Using progressively larger CIFAR-ResNee architecture. (b): Using 10%, 20%..........................2
100% of training data.
1. Train the base DeepNet. Train neural network fDeepNet on the training samples S as usual.
2. Prepare the residual data. Compute the residual prediction of each training example as
ri onehot(y)) (oftmaxx(/DepNa((:)))
V E [m].
where onehot: y + R _ is the standard encoding that maps the label to a probability vector. Here, T' is hyperparameter corresponding to the "temperature" scaling of the softmax operation. Then, we employ the output of an intermediate layer of the base network fDeepNer. denoted by 2, = (<z;), as an embedding for the training instance xę- These embeddings are utilized for the nearest neighbor search in the next step.
3. Predict via memorized residuals. To obtain a prediction on a test sample T X, first compute its embedding ž = (()), Then, use soft k-nearest neighbor method to build a function TnNN defined by weights (z):
T&NN T
Σ
W
1=1
The weights (.) satisfy Σ; W; (2 1, and are computed from raw weights as follows:
Wy = exp(p11E -1
2í 12/o),
w.(
x Il (w; 2 w(()) wi,
where W(() represents the k-th largest entry of wy wy's. Note that and a are two hyperparameters that collectively controls the locality of nearest neighbor search.
We make the final prediction based on the following scorer:
fReaMem ( (oftmax((DepNNe( + (T).
(8)
Remark 4.1 (Explicit memorization). Smaller l: or a corresponds to putting higher weight on residuals of the closest neighboring training examples. For sufficiently small h and a, fResMem achieves exact memorization of the training sample, i.e. pred(ffesMer))) (x.)) y4 for all i E [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 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 (2 with CIFAR dataset. We provide additional ImageNet [3 results in Apendix
33
14, 20, 32, 44, 56) 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
Setup. We use
For overparameterized network that perfectly fits the training sample, the residuals will all be 0. However. are interested in either smaller networks extremely large dataset where implicit memorization is infesible.
7

---

Image classification
Language modeling
ResMem
DeepNet
ResMem
DeepNet
plenty inelide shilat beinn leee 2.5m in boight.
rose
poppy
he get
length
Reaplike cossialt of all wiliis with uracies Lian 30, 300 looala. Reguicing a home in Spain.
plate
home
residence
Eiined around T:30-8:30 Fescay, Mazch 9, 2012.
squirrel
rabbit
March
June
..othat will not offect the supeeses quiltey your Jeb, that ie possible zave...
butterfly ly bes
possible
feasible
palm
pine
exvaneve your cueations and achedule scheduled hear...
consistently always
tree
tree
Figure 3: Examples from CIFARIO0 and C4 test set with the property that (i) ResMem is correct; (ii) y DeepNNet 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 sample whose true label is rose), which can be refined by ResMem capturing the remaining fine- fine-grineed structure.
image embedding, which has dimension 64. For the nearest neighbor search (Step 3. Section 4.1). we define the distance between two images to the & distance between their embeddings. We use 0.7, le 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 (cf. Figure
The results for CIFAR-ResNet-ff 14, 22, 44, 56} are reported in Figure 2(a) We from 56.46% to 59.66 %, which is between the base DeepNet test accuracy for Neste8 and CIFAR-ResNeto To access the statistical reliability of the improvement, we repeat the CIFAR-ResNee experiment times random initialization of DeepNet etc. We and that the average ResMem accuracy is 59% with standard deviation 0. 0.7% and the average DeepNet accuracy is 56.5% with standard deviation 0.8%.
Results.
can see that ResMem boosts the test accuracy of
Computationally, we estimate the CPU latency of 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 and images of size 32 32, R ( 68K params) requires 2.5MB, while a ResNet-14 ( 128KK params) requires 4MB. Embeddings from ResNet-8 and ResNet- are both 64 dimensional. To embed the entire CIFAR 100 training set (50K examples) requires 15MB of disk space.
to be 15,9 ms for a single
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 for nearest-neighhor search is the same the training for base neural networks (e.g., model with 10% CIFAR 100 data also uses the same 10% data for nearest- neighot search). On the left (right) of Figure 2(b] we report the test (training) accuracy ResMem and baseline DeepNet. As sanity check, we can see that ResMem always achieves perfect training accuracy, and the DeepNet training aceuracy decreases as samples increase (since it's harder fit larger dataset). We can see that ResMem yields progressively larger margin of improvement when more data used. This trend suggests desirable property ResMem: in real problems where the dataset is extremely large, ResMem is expected to bring even greater benefit.
8

---

4.3 Language modeling
Setup. For the language experiment, we use a Decoder-Only T3-(small, large] [42 model and C4 [42] dataset. C4 is generated from scraping the internet and commonly used as 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 ANN 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 L2 distance using approximate nearest neighbor search algorithm ScaNN [24]. We used the temperature = for the residual computation and 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 T5-base model without ResMem. On TS-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 use the approximate nearest neighbor search algorithm ScaNN [24) reduce compute time. Please see Appendix H 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 -este- and C4 with T5-small. Let GainResMem be the difference between the test accuracy of ResMem and baseline DeepNet:
GainResMem = Loi (fResMem) - Loi (fDeepNet):
where Loi is the misclassification error as defined in equation (6. We offer a decomposition of GainResMem that sheds light into the mechanism behind ResMem. For test {x y:) ) let y ResMem be the ResMem prediction on instance 24 and let y, DeepNet be the baseline neural network prediction on When yf ResMemm y/; DeepNet sample x. does not contribute to GainResMem- When yf ResMem + y/ DeepNet this could arise either from the desirable event where the deep network mis- classifies while ResMem classifies correctly; or from the undesirable event where the ResMem misclassifies, while the deep network classifies correctly. These can summarized by the TPR (Irue positive rate) and FPR (false posilive rate) respectively:
TPR
1 15;
DeepNet # y; and y
ResMem }/
(9)
FPR
DeepNet = y; and y
ResMem + y:)-
(10)
1
Note that GainResMem = came from the TPR samples, provided they outweigh the FPR samples.
The decomposition of Gain says that the gain of ResMem
we find TPR-5.89% and FPR-2.70% leading to GainResMem 19%. On T5- small with C4 validation split, we find 5.3-5.37% and FPR-2.44%. leading to Gain
On
Analysis of TPR samples Focusing the test samples where ResMem helps (y: y, ResMem + V/
DeepNet, we identify 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 V: yr ResMem- rose, but the DeepNet predicts y;
DeepNet poppy, i.e., the label of different type of flower. (cf. Figure 33. We find similar behavior for the language modeling task (cf. Figure B)
This empirical analysis suggests the DeepNet in isolation can already learn some large scale structures. but is unable to make fine- ane-ggained distinctions. This is where ResMem helps: RexMem helps memorize information in the training label that the DeepNet cannot learn.

---

Additional insights from the decomposition. In this paper, we choose the ResMem hyperparame- ters that minimizes the test error on the validation set or. equivalently, maximize Gain Inspired by the decomposition of GeimeesMem: we propose an alternative hyperparameter selection procedure based on the following optimization problem:
maximiikFPPRttperrrraa..m .c0.05
which ensures that ResMem modifies the DeepNet predictions in 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.
Discussion and future works
Joint training of &NN and DeepNet. The current formulation of ResMem builds the base DeepNet and kNN components sequentially. Consequently, the DeepNet is trained completely oblivious the fact that there a subsequent kNN model that will memorize its residuals. A natural direction of future work is to consider the joint training of DeepNet and &NN, that the models can dynamically interact during training to determine which portion label is for Deep-Net to learn, and the remaining for &NN 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 [ We can see that when the #epoch is small,
Table Comparison of DeepNet and ResMem accuracy over epochs on CIeAR-ResNet-& experiment.
#epoch
128
160
192
224
256
DeepNet acc.
34.0%
56.2%
55.6%
57.2%
56.6%
ResMem acc.
49.3%
60.2%
58.6%
59.2%
59.5%
ResMem has a dramatic improvement 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'ReMem (8) 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 ('ResMem(a)) However, this limitation hinders us to access the confidence of the ResMem prediction. To remedy this, possible future work is to consider alternative notions of residual. For example, we can do memorization in the logit space instead 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 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).
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, Yaodong Yu and anonymous reviewers for their feedback on the final draft. Zitong Yang would like to thank Shuangping Li for discussion regarding the proof of Lemma A.T Zitong Yang would also like to acknowledge the support of Albion Walter Hewlett Stanford Graduate Fellowship.
10

---

References
[1] William M. Alley. A note on stagewise regression. The American Statistician, 41(2):132-134,
1987. doi: 10. 00.10800000310055 1355. 198. URL URLhttps: httph: /tttps://www... tandfonline cwt.coicoooooooooooo abs/10. 1080/00031305
[2] Devansh Arpit, Stanistaw Jastrzebski, Nicolas Ballas, David Krueger, Emmanuel Bengio, Maxinder S. Kanwal, Tegan Maharaj, Asja Fischer, Aaron Courville, Yoshua Bengio, and Simon Lacoste-Julien. closer look at memorization in deep networks. In Proceedings of the 34th International Conference Machine Learning- Volume 70, ICML' 17, pages 233-242. Jogg 2017.
[3] Peter L. Bartlett, Dylan J. Foster, and Matus Telgarsky. Spectrally-normalize margin bounds for neural networks. 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
2017.
[4] Peter I.. Bartlett, Philip M. Long, Gábor Lugosi, and Alexander Tiglerr. Benign overfitting in linear regression. Proceedings of the National Academy of Sciences, 117(48))30063-300) 30070. 2020. doi: 10. 10.1073/ppas. 190737811 pnas. 1907378177
URL httph:ttt//wttt www.wwppaa..om wwwppppppp.....
[5] 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.
[6] Sehastian 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/21 12.04426, 2021.
URL URLhttps: //arxivaa....////a wwwpsiii.rrrggeeee abwgg/b////ua 2112 04426
[7] Léon Bottou and Vladimir Vapnik. Local Learning Algorithms. Neural Computation, (6):888 900, 11 1992. ISSN 0899-76677 //doi.org//i 10. 1162/ www////co........... neco 1992. 4.6. 888
doi: 10.1 10.1162/neco. http///cco.192...66.. 6.8.888. URL https
[8] Alon Brutzkus, Amir Globerson, Eran Malach, and Shai Shalev-Shwartz. SGD) learns over- parameterized networks that provably generalize linearly separable data. In International Conference on Learning Representations, 2018. URL httptst///peeeeeiiiieeee wtwpeeeereieieeeeeefereeeeee net/forum?
id=r
[9] Cristian Bucilă, Rich Caruana, and Alexandru Niculescu-Mizit 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.
[10] Kamalika Chaudhuri and Sanjoy Dasgupta. Rates of convergence for nearest neighbor classification. Z. Ghahramani, M. Welling, C. Cortes, N. Lawrence, and K.Q. Wein- berger, editors, Advances in Neural Information Processing Systems, volume 27. Curran Associates, Inc., 2014. URL (1b577cf6688d11222233aaddbbbb
neurins /wwerr/ps..c/pppppp////////// Paper. pdf
[11] Chen Cheng, John Duchi, and Rohith Kuditipudi. Memorize to generalize; on the necessity of interpolation in high dimensional linear regression, 2022. URL https://arxxii org/ abs/ 2202.09889
[12] 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 Tsyyashchenko, Joshua Maynez, Abhishek Rao, Parker Barnes, Yi Tay, Noam Shazeer, Vinodkumar Prabhakaran, Emily Reif, Nan Du, Ben Hutchinson.
11

---

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 Lewkowyez, Erica Moreira, Rewon Child, Oleksandr Polozov, Katherine Lee, Zongwei Zhou, Xuezhi Wang. Brennan Saeta, Mark Diaz, Orhan Firat, Michele Catasta, Jason Wei, Kathy Meier-Hellstem,, Douglas Eck, Jeff Dean, Slav Petrov, and Noah Fiedel. Palm: Scaling language modeling with pathways, 2022.
//axivv org abs 2204.. 02311
[13] Gilad Cohen, Guillermo Sapiro, and Raja Giryes. Dnn or k-nn: That is the generalize vs. memorize question, 2018.
org. www.or/gga///.......o/ 06822
[14] T. Cover and P. Hart. Nearest neighbor pattern classification. IEEE Transactions on Information Theory, 13(1):21-27, 1967. doi: 10.1 109//09TTTTII 19T.1//T..96..111.511..... 1973.69
[15] 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.
[16] Vitaly Feldman. Does learning require memorization? A short tale about long tail. CoRR, abs/1906.05271, 2019. htt.httt://arrwv.rr///aaa wwwp://gg.a.....oo..ooo
[17] Vitaly Feldman. Does Learning Require Memorization? A Short Tale about Long Tail, page 954-959, Association for Computing Machinery, New York, NY, USA, 2020. ISBN 9781450369794, URL/attss://aoi org/10. 10. 1 11141 338499
[18] 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., 2020. URL https: pttttppp://proceedinng. 2020/file/le14bffe22
www.isseurrpp.... cttprippccccppprr htteeee/fffaabbb6eeebbebeee Papers pdf
[19] 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 9781713829549
[20] Yoav Freund and Robert E. Schapire. A desicion-theorettc generalization of on-line learning and an application to boosting. In Paul Vitányi, editor, Computattional Learning Theory, pages 23-37, Berlin, Heidelberg, 1995. Springer Berlin Heidelberg. ISBN 978-3-540-49195-88
[21] 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: (0.1224//os/101el 1016218223
URL https: :/dp::://di.... oroeii..rg11... 1214/ /
[22] A. Gionis, Piotr Indyk, and Rajeev Motwani. Similarity search in high dimensions via hashing. In Very Large Data Bases Conference, 1999.
[23] Arthur S. Goldberger and D. B. Jochems. Note stepwise least squares. Journal of the Ameri- can Statistical Association, 50529) 1961. doi: 10. 1001000011111551.991.. 11w.4999996611..182009900 www.:/wwwwwnnfoolll.nammmmm wttp://ccooooccoccceoaco http://ww..aaaaaiooooooo//////a/ http:/10011000000221100100 19111449..99999...1 10482095
[24] Ruiqi Guo, Philip Sun, Erik Lindgren, Quan Geng, David Simcha, Felix Chern, and Sanjiy Kumar. Accelerating large-scale inference with anisotropic vector quantization. In International Conference on Machine Learning, 2020. URL httph:ttttss/tiiii http:/aaxxvvvv..reeee// wttpe/iww.wr///b/////9////// 10396
[25] Kelvin Guu, Kenton Lee, Zora Tung, Panupong Pasupat, and Ming- Wei Chang. Realm: Retrieval-augmeentd language model pre-training. Proceedings of the 37th International Conference Machine Learning, ICML'20. JMLR.org, 2020.
12

---

[26] Trevor Hastie. Robert Tibshirani, and Jerome Friedman. The Elements of Statistical Learning.
Springer Series in Statisties. Springer New York Inc., New York, NY, USA, 2001.
[27] 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.1 1...1009/PP..206...00
[28] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean. Distilling the knowledge in a neural network, 2015. URL httllhttt:::araiiti org/ aaw.a 02531
[29] Geoffrey E. Hinton, Oriol Vinyals, and Jeffrey Dean. Distilling the knowledge in neural network. CoRR, abs//1503.0251,, 2015.
[30] Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Generaliza-
tion through memorization: Nearest neighbor language models. In International Conference Learning Representations, 2020. URL httpp:://openreeieer net/ fort/formmm foriummmi--iilk B
[31] Urvashi Khandelwal, Omer Levy, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Gener- alization through memorization: Nearest neighbor language models. In ICLR, 2020. URL
hnttps://openrevies: Fitee///fooiiiddH1BBBB
[32] Donald Knuth. The Art Of Computer Programming, vol. 3: Sorting And Searching. Addison- Wesley, 1973.
[33] Alex Krizhevsky. Learning multiple layers of features from tiny images. Technical report, CIFAR, 2009.
[34] 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, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pages 8546-8557, 2019, URL https://prooceednngss wwwpinggg.neurpppe newn cwc.iissccccceeeeeee ctt.cc/cpppao (tpppp///19//haaeaeeeeepap 9d8df 7ccdd7738cccb3336667766666111113 htt.://99....1111...aaaaaaaa.... Abstract. html
[35] 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 httpp:// opwnreview.net. wtt.://www...eettoooormddddddddd
[36] Tengyuan Liang and Alexander Rakhlin. Just interpolate: Kernel "Ridgeless" regression can generalize. The Annals of Statisties, 48(3):1329 1347, 2020. doi: 10.121441--AOS1888 URL https://doi.
[37] Gaurav Menghani. Efficient deep learning: A survey on making deep learning models smaller, faster, and better. CoRR, abs/2106.08862, 2021. URL. https://arxiv.org//bbbl wwt.://rrr/ab/////////// 08962
[38] Andrea Montanari and Yiqiao Zhong. The interpolation phase transition in neural networks: Memorization and generalization under lazy training. 2020. URL httpp::/a// org/ abs/
2007. 12826
[39] Marius Muja and David G. Lowe, Fast approximate nearest neighbors with automatic algorithm configuration. In International Conference on Computer Vision Theory and Applications, 2009,
[40] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun, and Nathan Srebro. The role of over- preepaaamttriitooonn in generalization of neural networks. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 2019. OpenReview.net, 2019.
[41] 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. BtR.hhttp::/ppoceediin
mlr. wtw.peees//v113///iigaaaaaaa/a//
13

---

[42] 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-61) 2020. URL
pttp:/paparre/20--peppe http:://aal..orgpppeer///////
htmml
[43] 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-225, 2015. doi: 10. 100011001113-0011-1111
[44] Ruslan Salakhutdinov and Geoff Hinton. Learning 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. /ttLstttt///pooceeeiii
mlr wwwpsrweeeeepepp http/v//lalllaaaa wwwtaaaoootdaaaaaaaaaoaooooo html
[45] 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 Puttern Recognition, pages 4510-4520, 2018. doi: 10.1 101/ 00474.
[46] 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.
[47] Ryan J. Tibshirani. A general framework for fast stagewise algorithms. J. Mach. Learn, Res., 16(1):2543-2588 jan 2015. ISSN 1532-4435.
[48] Vladimir Vapnik and Rauf Izmailov. Reinforced SVM method and memorization mechanisms. Pattern Recognition, 119:108018. 2021. ISSN 0031-3203. doi: https///doi.org/10010111 patcog. 20801.. URL https:://////t http://www..ceencdreeettttt httpr//tteeeomeeeeeeeeeeeerrtttt
50031820522002053
[49] Martin J Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cambridge University Press, 2019.
[50] Ke Wang, Vidya Muthukumar, and Christos Thrampoulidis. Benign overfitting in multiclass classification: All roads lead to interpolation, 2021. URL UR..Ittess//rxii.oor///////000
10865
[51] Zhen Wang and Yuan- Shao. Gemeralizatioon-mmoriiiaa machines, 2022. URL https
//arxiv. org/ abs 2207 03976
[52] 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 Pro- ceedings of Machine Learning Research, pages 10767-10777. PMLR, 13-18 Jul 2020. URL mlr ppesss/119/yanggea .hmml html
[53] 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.no mlr press wwt.:/prrssseeeeeeeeeee v139/yang21a...... html
[54] 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.
[55] Hao Zhang, Alexander C. Berg, Michael Maire, and Jitendra Malik. Svm-knn: Discriminative nearest neighbor classification for visual category recognition. In CVPR (2), pages 122--1166 2006. URLLhttpst //do:i///oi// wtt/////............. 1109/C11R. 2006. 301
14

---

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 Let a -1,-- 2n 11.d On. be i.i.d. samples from the uniform distribution the Euclidean norm ball of radius vd + in Rdd Let
Zn
min
E ---2- ==|
(11)
If n = 1, EZ is the sum of the variance of each coordinate of Unif(B. /. Therefore, EZZ, provides a generalized measure of concentration. Intuitively, EZn + 0 as n + co. The proposition below provides upper bound on the rate of convergence.
Lemma A.1 (Nearest Neighbor concentration). Given the assumptions above
1/d
log (n
EZ S d?
(12)
where E means inequality up to an universal constant independent of and n.
Proof. Define
& 12, s 8
(13)
& {6 - Vd + 2 1-
We will compute two probabilitiess (2) and P(22) that will be useful latter.
P( P(Z, 2 (2)) P( j x; 8, Vil
P xx , = Eg - P(I -xill S 81E2, 2))",
Vol(Bzz66
E2
(14)
Vol(Bo,
v +
< exp
Vd + 2
Next, we compute P(E2)
vd+2 -Va+22-0 -(1-
- -6)
8
((2) P( S vd+2 6)
(15)
Vd + 2
We use E, and &g to compute the following upper bound
EZ,, B(Z,|&, (&))(( ( + (((2)((()))
s 82
(16)
82 4(d + 2) (()
To find an upper bound for EZn, we need to find an upper bound for 1 - P(E)))))(()))
1 - (()
= 1 - - F()
= 1 - (P())))))))))) P(ET))(
(17)
S 1 1 PP
Now choose 8 Vd + 2n
1/d -1/d [log (n 1/4
d
E2 & exp
= cxp
mm
log (n' [1)]
n -1/d
(18)
Vdl +
15

---

and
PP(2)) (1
d
dn 1/d log (d)) 1/a
1
2
(19)
Vd+ 2
2
+ vd+2
Thus
1/d [og (n2 1// +
) -1/d [og (n2//a] 1/d (20)
1 1 P(11(()))(()) 11-1+ -1/d
1/d 11/d
2 dn 1/d
Combining everything together, we get
EZ - (d + 2)m 2/d 1og
(n//a)
2/d
+ 4(d + 2) dn 1/d log (n)//a) 1//
S n 1/d [Log (nl/a)
1/d
(21)
log (n2/d) de
1/d
n
This completes the proof.
Proposition A.2 [49] Corollary 6.20). Let H; 1id Unif (242) for 1,.... be uniformly
distributed over ball of radius Rdd centered O.
En
1=1
be the sample covariance matrix. Then
ne 2(d + 2) E)
P( - Iloo 6) S 2dexp
16

---

B Proof of Theorem 3.3
In this section, we present the proof of Theorem 3.3
Section B.1 we provide the detail of the
decomposition of the risk into Ii and 72. Then in Section B.2 we compute upper bound for Ti, and compute an upper bound for T2 in Section B.3 Section B 3.4 and completes the proof.
Finally, we combine everything together in
B.1 Decomposition of the test risk
+ E ResMem ( - (m) [=()))
E [f(()) fx( Jn (Z)))) +.. [4
EE [f((0) - + - (.(() ((()) + foo (1) fx(()) (())
3 x f + ((f((()))
+ f. (2)1)) foo (z))) ((foo fx()) - (
T2
(22)
where in the last inequality, we used the fact that (a + +b+ < 3( + 6 + c²) for any a, b, c E R.
B.2 Upper bound on Ti.
Since Pa Unif(Bo, B ). we apply the bound || |لا _ B to obtain
Ti - fzo
-
+ (1) for (3(1))
E(O, 0,x, xz) E x1)) )
(23)
< El 10 - 00 10 (1)
$ 2
- 100
As gets large, the empirical covariance matrix E, =XTX/ T7X/n n is concentrated around its mean I. Let A, = I - E, denote this deviation. For some & E (0, 1). define the following "good event" the randomness in E,,
A = Allpp <e).
(24)
where || |Axllop denotes the operator norm of the deviation matrix. The high level idea of the proof is condition on the event A and deduce and upper bound of O 0. terms of e. Then, use the fact that A happens with high probability.
Recall that 0x LO., and
O, argmin - Ixo y1
(25)
| n
Since X0, by definition, the Lagrangian of the convex program above is
C(A, A) = - xx 1000x001--0 + 2(110)2 L).
(26)
The KKT condition suggests that the primal-dual optimal pair (0,, An) is given by
||0,|| S L,
An 2 0,
(27)
A, e - L) 0,
and at optimality
2 T 0.) 2X, 0,
VoL(On, Aa) =0 +=
n
(28)
O, (En + An O..
17

---

The complementary slackness condition O1 - L) = 0 suggests that either A, = 0 or 0 = L. But if A, 0, the stationary condition (o 0, = 1101111111
0 would suggest that 0,, Σ, 0. L, a contradiction. (Note that here E, is invertible condition on the event A.) Therefore, we must have ||On|] = I.. As a result, the primal and dual pair (0n; An) is determined by the system of equations
=
n
= =EE + 1)
|10m11
L,
(29)
0.
Next, we proceed to compute the deviation 110, 0zll
0n [(nn 1)I
An n
-1
0,,
An
(An + 1)
I
An +1
k
+1
I+
A.),.,
(An 1)&e
k=
k+1
(30)
= =mm
I+
0.,
(An
(An 1)bh
k=1
k=1
k-1
k
(An, + 10 + (An + 1)
0.,
1)*
(An
1)*
k=
k-1
A
(An 10 + (An A
I
O..
(An + 1)*
Define
8 k- A
D2 An
(31)
(An + 1)t
Then On (An 10- + +(A, + +,, and
O,,
k- + A,
|Da| S Anll |l
1+
(nn + +1)k
k=1
E
<
+ 2(1 + Am
(32)
1 1 An
k=1
2
| (1+
3E.
1 An
1+Am
Therefore
= ||0 (An 11 (An +1) -0 2(A, -0
(33)
+ (An, 1) + 8n, 8n
We can obtain the following bound for on
|6 < ||0.1| D. 2 + 0. 1. D.. < 9 + 65 < 15e.
(34)
Since 1 - -../25 v1+ On < 1 + 8,/2, and | 15s, we obtain
15s
+ L (2n + 1) s 15a (An + -1-
15s
1( + +) 11 -
(35)
2
2
18

---

where the last inequality follows as we have Xn > O. Finally,
0n 000 =A,, 1) -0.- I0. (An + 1 0,
+ + 0x2
[1 + 11)-1111 (1+ Am) --- D. 2(An, - 1)1 A,)) Lle.D,0..
< 64s² + -9 + + =118==
+ | - 0x11 E².
(36)
Combine the above result with Proposition A.2 we get that
0r 00l =(( 0x 11 (AA))(A) + O( 0x (x113119))P((99))
NE² 2(d 2) 8)
(37)
5 AL2 4dexp
If we choose n 1/3, we get
E On 0x01 11 1 - 2/3
(38)
which implies that
Ti T 12.
(39)
B.3
Upper bound on T2.
Plugging in the formula for fz) = f.(z) - fā) = (T, 01). we get
T- = ((f(( fl (z)l²
E[01,a01) 3)2
(40)
$ - =)1-)) | BB@ 2n))
= (1 2) {1))
where in the last inequality, we used the relation that 0_ = (1 - L)O. Proposition A.1
suggests that
log (n)/a) 1/d
Eli (1 S de
(41)
which implies
log (n=
1/d
T2 S d= -,))²
(42)
n
Remark B. (Comparison with pure nearest neighbor and ERM). If we rely solely on nearest neighbor method, the prediction error is
B ( (z E (()) 0.)² < E क (1)
(43)
On the other hand, if we solely rely on ERM, even with infinite sample, we get
(ā) fax (=))² = E(a 6, 0x)² < (1 - =)
(44)
We can see from the upper bound that ResMem takes advantage of both
Projecting f. onto fac. 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 reduced from [R] x||2 to [1) I
19

---

B.4
Test loss for ResMem.
we combine the previous two parts together, we get
log (n)/a)
1/d
E - S d 2/3 d2 L)²
(45)
n
This completes the proof of Theorem 3.3
20

---

C Additional CIFAR100 Results
This section includes additional experiment results on applying ResMem to CIFARI00 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-ResNer(8 8, 14, 20, 32, 44, 56) and each subset (10%, 20%, ... 100%) of
CIFARIO0 100 training data. We use the same training hyperparameter and the ResMem hyperparameter described in Section 4.2 Generally, we that ResMem yields larger improvement the baseline DeepNet when the network is small and dataset is large.
3.55
Semeee 655
ssu ResMom
Despher
Deepleet RasMer
0.40
645-
Paroentage of CINRRDO
Percentage CHARIDO usod
Pecentage CHAR200 used
Pacertage
(a) CeeAR-ResNet-8
(b) CIFAR-ResNet-14
3.00
3925
Dosshic
23300
655-
Deeplit
Deeplitt
Deep Vet Testeer
ReeMee
3858
845
Raidlem
4 $ o'
3325
of Peccertage copARtnduund
CIARTOD
(c) 20
(d) 3eFAR-ResNet-32
130
1.0
3..5
6.00-
655
3.5
236
P350
Dospfeci
Docplitt 4
345
Puroeragg: JICIHAK10 used
Parcetog:: JICIAR100 u500
Percertage of CHRALIOO used
Paratlage of CHAAPL0O usel
(e) CIeAR-ResNet-44
(f) CIFAR-ResNet-56
Figure 4: Test(eett/Training (right) accuracy for different sample sizes.
C.2 Sensitivity analysis for CIFAR100
0.595
6.60
ResMem
accuracy
0.56
0.58
DecpNet
0.585
- ResMem
6.54
Text arrureyy
0.58
0575
DeepNet
[ 0.50
0.52
0.57
0.48
ResMem
0.56
0.65
0.46
DeepNet
6.55
0.25 0.50 0.75 1.00 1,55 1550 1.75 2.00 o parameter ResMam
#neighbours ResMem
Softaax temperature
(a) # neighbours
(b) Radius parameter a.
(c) Temperature
Figure 5: Sensitivity analysis of ResMem hyperparameters. The y-axis represents the CIFARI00 test accuracies, and the z-axis represents the sweeping of respective hyperparameters.
21

---

Varying locality parameter k and a. We vary the number of neighbours from h: = 27 to k = 500. We find that ResMem test accuracy is relatively stable across the choice of the number of neighbours (ef. Figure 5(a)). The trend of the curve suggests that as k -> co, the ResMem test accuracy seems to be converging a constant level. For a, we explored different values of (0.1, 2.0). We observe that the test accuracy has a unimodal shape function of a, suggesting that there is an optimal choice of a (cf. Figure 5()).
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(()). The fact that we can use different temperatures for (a) training the network and (b) constructing the k-NN predictor reminds of the well-cstablished knowledge distillation procedure [28]. In knowledge distillation, we first use one model (the teacher network) to generate targets a higher temperature, and then train 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 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 "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 CIFARIO0, we also evaluate the performance of ResMem on Ima- geNet 43 We employ a family of pre-trained MobileNet- V2 models 45| from Keras with varying widths controlled by multiplier a. For ResMem, we again use the second last layer of DeepNet as 2200dimensioonll embedding of an image and rely on the & distance between the embeddings for nearest neighbor search (Step 3. Section 4.1). We specify the ResMem parameter of (k: or 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-V222 33 being the largest one.
Table 2: Test accuracy for ResMem and baseline deep network for ImageNet data.
ResMem param.
Test accuracy
Architecture
T
DeepNet
ResMem
MobileNet- V2-a0.35
MobileNet-V2-a0.355
10
0.6
0.4
60.2%
61.2%
MobileNet- V2-a0.5
10
0.6
0.4
65.3%
66.1 %
MobileNet-
10
0.8
0.6
69.6%
70.1 %
MobileNet-V2-al V2-a1.0
0.4
0.4
71.3%
71.8 %
MobileNet-V22a1.33
30
0.4
0.4
74.7%
75.1 %
We can see that (c.f. Table 2 0.4% on the largest model.
ResMem boosts the test accuracy by 1% on the smallest model and by
E
Additional details of NLP experiments
The Decoder-Only model used in our experiments is essentially the normal Encoder-Decoder archi- tecture with Encoder and Cross-Attention removed. We pretrained both the T5-small and b5-base model on C4 42 dataset 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.
regressive language modeling task for 1,000,000 steps, with
During the inference for retrieval key, query embeddings and residuals, we ensured every token has at least 64 preceding context by adopting sliding window strategy, where window of 256 token slides from the beginning to the end on each of the articles, with stride of 256 64 = 192.
22

---

For residuals, we only stored the top 128 residuals measured by the absolute magnitude, as the residual vector is as large 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.
Comparison with other algorithms
We mainly compare ResMem against [3I. where the algorithm uses kNN to retrive labels directly instead of the residual of the label. In their algorithm, key aparameter E [0, 1] which specifeis how much weight give the neural network and how much for the ANN component. In the extreme case of x=1, their algorithm reduces to using &NN to memorize data directly.
For the language modeling task, use the C4 dataset and T5- 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
Table 3: Test accuracy for KNN-LM (ResMem accuracy 45.55%)
kNN weight A
0.2
0.4
0.5
0.6
kNN-LM accuracy
44.76%
44.88%
44.83%
44.66%
44.27%
42.97%
40.95%
ResMem KNN-LM acc.
0.79%
0.67%
0.72%
0.89%
1.28%
2.58%
4.60%
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 kNN memorization has accuracy 54.44%; and ResMem has accuracy 59.66%.
23


Reviews:
[{"review_id": "tGG3c7szCL", "reviewer": "Reviewer_1sN3", "summary": "The 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.", "questions": "I 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.", "limitations": "Does not apply", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The 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.", "weaknesses": "While 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for seeing the novelty of ResMem algorithm. We will fix the minor issues and adopt the suggestions.\nWe address Q1-Q4 below.\n\n> ReQ1: equation under line 163.\n\nWe use the linearity of expectation and use the fact that \n$\\mathbb{E}\\_{\\mathbf{x}\\sim\\mathbb{P}\\_{\\mathbf{x}}} [\\mathbf{x}\\mathbf{x}^\\mathsf{T}]=\\mathbf{I}$ \nfrom Assumption 3.1.\n\n$$\n\\arg\\min\\_{\\Vert\\mathbf{\\theta}\\Vert\\leq L} \\mathbb{E}\\_{\\mathbf{x}\\sim\\mathbb{P}\\_\\mathbf{x}} [\\langle\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star, \\mathbf{x}\\rangle]^2,\n$$\n\n$$\n=\\arg\\min\\_{\\Vert\\mathbf{\\theta}\\Vert\\leq L} \\mathbb{E}\\_{\\mathbf{x}\\sim\\mathbb{P}\\_\\mathbf{x}} [(\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star)^\\mathsf{T} \\mathbf{x}\\mathbf{x}^\\mathsf{T}(\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star)^\\mathsf{T}\\rangle],\n$$\n\n$$\n=\\arg\\min\\_{\\Vert\\mathbf{\\theta}\\Vert\\leq L} \\mathbb{E}\\_{\\mathbf{x}\\sim\\mathbb{P}\\_\\mathbf{x}} [(\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star)^\\mathsf{T} \\mathbf{x}\\mathbf{x}^\\mathsf{T}(\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star)^\\mathsf{T}],\n$$\n\n$$\n=\\arg\\min\\_{\\Vert\\mathbf{\\theta}\\Vert\\leq L} (\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star)^\\mathsf{T} \\mathbb{E}\\_{\\mathbf{x}\\sim\\mathbb{P}\\_\\mathbf{x}} [\\mathbf{x}\\mathbf{x}^\\mathsf{T}] (\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star)^\\mathsf{T},\n$$\n\n$$\n=\\arg\\min_{\\Vert\\mathbf{\\theta}\\Vert\\leq L} \\Vert\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star\\Vert^2 \n$$\n\n$$\n=\\arg\\min_{\\Vert\\mathbf{\\theta}\\Vert\\leq L} \\Vert\\mathbf{\\theta}-\\mathbf{\\theta}\\_\\star\\Vert = L\\mathbb{\\theta}\\_\\star\n$$\n\nThe last equality comes from the fact that the projection of the vector $\\mathbf{\\theta}\\_\\star$ with norm $\\Vert\\mathbf{\\theta}\\_\\star\\Vert=1$  onto a norm ball of radius $L<1$ (Assumption 3.2) is simply the rescaled $L\\cdot\\mathbf{\\theta}\\_\\star$.\n\n> ReQ2: $T_2$ at line 175.\n\n$T\\_2$ alone doesn't represent generalization error. The sum of $T\\_1$ and $T\\_2$ (c.f. line 175) constitutes an upper bound for the generalization error.  Looking at the the expression for $T\\_2$ under line 175, the term stems from the difference between $f\\_\\infty$ and $f\\_\\star$. As discussed in line 162-168, $f\\_\\infty(\\mathbf{x})=\\langle\\mathbf{\\theta}\\_\\infty, \\mathbf{x}\\rangle$ minimizes the population error  (ERM with infinite samples. This is why we describe it as \"irreducible error of risk\"). In this case, we would have $f\\_\\infty=f\\_\\star$ if the function class is large enough that $f\\_\\star\\in\\mathcal{F}$. However, because of the Assumption 3.2 that we imposed, we don't have $f\\_\\star\\in\\mathcal{F}$, so $T_2$ is not asymptotically zero. This justifies the claim that $T_2$ appears due to limited capacity of $\\mathcal{F}$. Assumption 3.2 is precisely what distinguishes our conclusion from most papers that analyzes linear regression risk.\n\n\n\n> ReQ3: bound without ResMem.\n\nWe first acknowledge that this is a good comparison and we will add to the camera-ready version. Looking at the generalization error under line 190, without ResMem, $T_1\\lesssim d^2 L^2 n^{-2/3}$ would stay the same modulo minor change in constants. In contrast, $T_2$ would be independent of $n$: $T_2\\lesssim d^2(1-L)^2$. Therefore, when $n$ is large, the risk with ResMem would be smaller.\n\n> ReQ4: computation and memory cost in empirical experiments.\n\nWe discuss the computation cost for image classification experiment between line 257 and 260.\n\n\"\"\"\n\nComputationally, we estimate the CPU latency of a CIFAR-ResNet-8 to be 15.9 ms for a single\ntest image. By contrast, the k-NN step takes 4.8 ms for the same test image. To contextualize the\nlatency cost, the total cost of ResMem with ResNet-8 (15.9 ms + 4.8 ms) is lower than the cost of the\nnext-sized model, i.e., ResNet-14 (26.2 ms).\n\n\"\"\"\n\nRegarding 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. We will add this discussion on top of the current line 257-260.\n\n\nFor the language experiment, the index set is quite large (1.6 billion tokens) and exact k-nearest neighbor search is infeasible. So we use the approximate nearest neighbor search algorithm ScaNN [1] (discussed in line 29, 280) to cut the kNN compute time. \n\n[1] 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.\n\n\n"}, {"review_id": "97xaDXKJZ8", "reviewer": "Reviewer_M8UD", "summary": "This 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.", "questions": "Lines 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?", "limitations": "Aside from potential issues around space usage, which are addressed above, I feel the authors adequately address the limitations of their work.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "I 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. ", "weaknesses": "One 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks for seeing the strength of the proposal and providing insightful comments! We will answer your questions below:\n \n> A separate algorithm section.\n\nWe agree this is a good idea. The dilemma was that ResMem is a very general algorithm and can be applied to any supervised learning problem (regression/classification). As a result, nailing it to a specific classification or regression setting can potentially hide the generality of the algorithm. We will add a separate algorithm section/box that highlights the three steps in Figure 1 as the main \"ResMem Algorithm\". We really appreciate this suggestion!\n\n> Discussion on related work regarding \"memorization is sufficient for generalization\". \n\nWe acknowledge that the title of this paragraph is not mathematically precise. Note that here, by \"memorization\" we mean \"the training data is fit to zero error\". Indeed, it's not true that _for any_ learning algorithm and data distribution, fitting the training data to zero error is sufficient for generalizing to test data. Rather, we intended to convey the idea that training _neural networks_ to fit _real-world_ training data as perfectly as possible is known to lead to good generalization in many situations. We will clarify the usage of the word \"memorization\" in the camera-ready version. For concrete examples from the cited papers,  let's consider the following two:\n\n[A] Behnam Neyshabur, Zhiyuan Li, Srinadh Bhojanapalli, Yann LeCun,Nathan Srebro. Towards understanding the role of over-parametrization in generalization of neural networks, 2018. URL: https://arxiv.org/pdf/1805.12076.pdf\n\n[B] Zitong Yang, Yu Bai, Song Mei. Exact gap between generalization error and uniform convergence in random feature models, 2020. URL: https://arxiv.org/pdf/2103.04554.pdf\n\nIn the left of Figure 1 in A, we see that the test error is smallest when the training error is near zero (meaning that the entire training data has been memorized.) In the left of Figure 1 in B, the red curve is the maximum test error among models that interpolates the training data and the yellow curve is the test error found by the minimum norm interpolator. We see that the red curve (max test err. among memorizing models) captures the general trend of the yellow curve (test err. learned by ERM). These two figures provide complementary views supporting the claim that \"memorization is sufficient for generalization [for neural networks, on real-world data]\".\n\n> Intermediate layer representation\n\nFor the ResNet-CIFAR, we simply use the pre-logit layer as the embedding (line 250). For language modeling, we use the last layers pre-MLP, post LayerNorm embedding (line 277).\n\n> Can 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\nFor 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.\n\n\n> Practicality of ResMem hindered by additional memory cost\n\nWe would like to propose ResMem both as \n   - (a) a conceptual idea that fits into the memorization/generalization puzzle that's been discussed in the literature, and at the same time \n   - (b) as a potentially practical algorithm; this can be engineered with approximate [1] or distributed [2, 3] nearest neighbor search algorithms as an efficient add-on to the original prediction model. \n\nWe acknowledge that the representations of the training data could consume a large amount of disk space. However, we would like to point out that accelerators are in general several orders of magnitude more expensive than storage. We believe it'll be worthwhile to trade more space for less time spent on inferencing the accelerators.\n\n[1] 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.\n\n[2] Jeff Johnson, Matthijs Douze, Herve Jegou. Billion-scale similarity search with GPUs, 2017. URL https://arxiv.org/pdf/2305.18466.pdf\n\n[3] Moritz Hardt, Yu Sun. Test-Time Training on nearest neighbors for large language models, 2023. URL https://arxiv.org/pdf/2305.18466.pdf"}, {"review_id": "MYuegY2qIE", "reviewer": "Reviewer_diiC", "summary": "The 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. ", "questions": "- 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?", "limitations": "No 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. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The idea of using KNN to memorize the residual is novel and interesting\n- There is a risk bound proof\n", "weaknesses": "- 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).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "> The assumption of the proof (linear function, k=1) is not practical and not aligned with the experiment\n\nWe agree that the toy experiment with linear function and 1-nearest neighbor is very different from the actual empirical. The main purpose of the theory is to design an abstract setting that captures how the residual memorization component can \"expand\" the capacity of the underlying ERM function class while preserving the fast rate of ERM (lines 201 - 204).\n\nExtension to more complex models are possible.\n* When the exact minimizer $\\mathbf{\\theta}_n$ is infeasible in closed form (e.g. non-convex objective, model via iterative procedure such as SGD),  one may augment the analysis to additionally consider the effect of optimization error [1] in such cases.\n* For the number of neighbors $k>1$, we can follow the procedure outlined in the section 3.1.1 [2]. The high level logic is that given a test data $\\tilde{x}$, there are two bad events that happens with low probability. First, the variance component of kNN regressor is too large. This can be controlled by setting $k$ sufficient large. Second, the neighbors used to compute the average is too far away from $\\tilde{x}$. This can be controlled by setting the number of samples $n$ sufficient large relative to $k$. As a result, when $n$ is sufficient large, and $k$ is neither too big or too small, we are able to obtain an upper bound in the regression error (Theorem 3.3.1 of [2]). \n\nIntegrating this line of attack to the ResMem setting (where we are memorizing residuals but not labels) would pose additional technical challenges. We believe that the current simplified theory already captures the intuition we wish to convey. We will add a remark to the theory section to discuss further extensions to more complex models in the camera ready version.\n\n[1] Leon Bottou and Olivier Bousquet. The tradeoffs of large scale learning. Advances in Neural Information Processing Systems, volume 20, 2007. URL https://proceedings.neurips.cc/paper/2007/file/0d3180d672e08b4c5312dcdafdf6ef36-Paper.pdf.\n\n[2] George H. Chen; Devavrat Shah, Explaining the Success of Nearest Neighbor Methods in Prediction, 2018. URL: https://ieeexplore.ieee.org/document/8384208\n\n\n> 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 in the image classification task.\n\nWe appreciate the comment regarding the comparative methods. For the NLP task we compared ResMem against the kNN-LM methods from paper [30]. We use the C4 dataset and T5-large architecture. As we change the weight ([equation (3), 30]) of the DeepNet component, we find the best performing kNN-LM methods has accuracy 44.89\\% which is lower accuracy than the ResMem accuracy 45.55\\%.\n\n\n| kNN weight $\\lambda$                    | 0 (pure DeepNet) | 0.1     | 0.2     | 0.3     | 0.4     | 0.5     | 0.6     | 0.7     | 0.8     | 0.9     | 1 (pure kNN) |\n|---|----|---|---|---|---|---|---------|---------|---------|---------|--------------|\n| kNN-LM accuracy| 44.76\\%| 44.84\\% | 44.88\\% | 44.89\\% | 44.83\\% | 44.66\\% | 44.27\\% | 43.73\\% | 42.97\\% | 41.98\\% |      40.95%        |\n| ResMem acc. - kNN-LM acc. |0.79\\% |  0.71\\% |  0.67\\% |  0.66\\% |  0.72\\% |  0.89\\% |  1.28\\% |  1.82\\% |  2.58\\% |  3.57\\% |      4.60%        |\n\nFor image classification, we run the simple baseline of using k-nearest neighbor to directly memorize the labels $y$. We observe the performance: [pure DeepNet: 56.46%; pure kNN memorization: 54.44%; ResMem: 59.66%]. We will add those result to the camera-ready version of the paper.\n\n> The current DeepNet is weak and raises the concern that ResMem may not scale well with big/SOTA backbones and will be not useful in practice.\n\nWe believe that ResMem is most effective *when the dataset size is much larger than the model size*, so there are two ends of the spectrum we can go:\n* (a) Fix the dataset size, and evaluate the smaller models;\n* (b) Study problems which have (essentially) unlimited data, and evaluate the larger models. \n\n| Regime | Task                 | Data size      | Model size      | Test accuracy improvement |\n|--------|----------------------|----------------|-----------------|---------------------------|\n| (a)    | Image classification | 50K (CIFAR100) | 67K (ResNet-8)  | 3.2%                      |\n| (b)    | Language modeling    | 33B (C4)       | 30M (T5-small)  | 2.9%                      |\n\nFor both (a) and (b), the data size is much larger relative to the model size. The difference is that in (b), the language modeling task, the data size can be very large, so we can also evaluate large models (with 30M parameters). However, for image classification, the size of the dataset is limited by the effort of human laborers. As a result, we *intentionally* use small models to conform with the “data>>model” story. \n\n> The setup in Sec.4.2: Why do you set up this way? Did you tune the model properly and let it train long enough?\n\nThe CIFAR-ResNet8 (#param=67K) is a small model for CIFAR100. We follow a standard procedure from [3] with setup details outlined in Section 4.2. Our approach was to reuse a standard approach to optimization, instead of developing one ourselves\n\n[3] Kaiming He, Xiangyu Zhang, Shaoqing Ren and Jian Sun. Deep residual learning for image recognition,\" CVPR 2016\n\n> Fig. 5: the results seem insensitive to K. Can you explain?\n\nThe number of neighbors used in the figure are k = [1,   27,  54,  80, ... , 421, 447, 474, 500]. So there is a big jump in test accuracy from k=0 (no ResMem) to k=1, and the test error becomes insensitive to $k$ for large $k$. One possible explanation is by looking at Theorem 3.3.1 from [2]. They find that the in vanilla kNN (no ResMem), the test error $\\varepsilon$ scales as $\\varepsilon\\sim\\frac{1}{k^2}$  w.r.t. to $k$. So from $k=0$ to $k=1$, there is a big jump, and for large $k$, the sensitivity declines.\n"}, {"review_id": "d7HReIEOnZ", "reviewer": "Reviewer_Kq4J", "summary": "Several 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.", "questions": "1) 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.", "limitations": "Limitations have been adequately addressed.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The 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.", "weaknesses": "I 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", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for seeing the originality of the ResMem algorithm. We address your questions below.\n\n \n> Joint training of kNN and DeepNet. Role of the duration of the first training phase? What happens when the second stage is applied too early or too late?\n\nThis is a very insightful question. We run the suggested experiment on CIFAR100 and CIFAR-ResNet-8 with different training epochs.\n\n| #epoch| 128| 160| 192| 224| 256|\n|---|---|---|---|---|---|\n| DeepNet acc. | 34.0\\% | 56.2\\% | 55.6\\% | 57.2\\% |56.6\\% |\n| ResMem acc.  | 49.3\\% | 60.2\\% | 58.6\\%  | 59.2\\% |59.5\\% |\n\nWe can see that we 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. From the experiment, it seems that just from the first half of the training stage, the ResNet already picks up some good representation of the data, and it's enough to make the nearest neighbor retrieval meaningful.\n\nAs for the joint training of kNN and DeepNet, as pointed out by the reviewer, we have left as future work. The intuition behind two-stage design is from the classical boosting algorithm, where the residual of the previous model is used to fit the current model. ResMem is like one step boosting since there is no residual left after kNN memorization is applied. We believe that both designs of ResMem, joint training or two stage learning, are worth studying. The focus of this paper is on two-stage learning and we left joint training as future work.\n\n> Statistical reliability and poor ImageNet improvement.\n\nWe acknowledge that the reported results are not averaged over several runs of random seeds. For CIFAR100 and NLP C4 this is less of an issue, as the margin of improvement is relatively large. We agree for that ImageNet results reported in the Appendix of main body pdf has a small margin of improvement, so statistical reliability comes into play. As for the ImageNet result, we'd like to bring your attention to our ImageNet results in the separate Supplementary PDF we uploaded, where we evaluate ResMem with a different architecture (MobileNet), showing a larger margin of improvement. We also cite the them below:\n\n| Architecture|$k$ |$\\sigma$|$T$ |DeepNet test acc.|ResMem test acc.|\n|---|---|---|---|--|---|\n| MobileNet-V2-a0.35|10 | 0.6 | 0.4 | 60.2\\% | 61.2\\%   |\n| MobileNet-V2-a0.5|10 | 0.6 | 0.4 | 65.3\\% | 66.1\\% |\n| MobileNet-V2-a0.75|10 | 0.8 | 0.6 | 69.6\\% |70.1\\% |\n| MobileNet-V2-a1.0|20 | 0.4 | 0.4 | 71.3\\% | 71.8\\% |\n| MobileNet-V2-a1.3|30 | 0.4 | 0.4 | 74.7\\% | 75.1\\% |\n\nWe can see that ResMem boosts the test accuracy by $1\\%$ on the smallest model and by $0.4\\%$ on the largest model. We will update the CIFAR100, ImageNet result with confidence interval in the camera ready version. The C4 experiment is more expensive to run, if time permitting we will also report an average improvement. Otherwise we will add a remark in the text clarifying that the result comes from a single run of experiment. \n\n> 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\nFor 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. We will add this discussion on top of the current line 257-260.\n\n>  Line 252: from the current text, the parameters seem cherry picked \n\nWe acknowledge that the ResMem parameters are not selected strictly following the train/validation split. We use CIFAR-ResNet-8 on the test set to select one combination of hyperparameters. That choice is then fixed throughout the CIFAR experiment across different architecture and dataset size. In that sense, our intention is just to use the test set of get a rough range of numerical values hyperparameters should take. To access the effect of test set tuning, we conduct the following experiment on CIFAR-ResNet-8\n\nWe split the 60K (10K+50K) CIFAR100 images into three folds:\n\n* Training Set\n* Validation Set\n* Test Set\n\nThen, we train (run SGD, building the retrieval set) the model $M$ on the Training Set, and make two comparisons. We fix a grid of hyper-param $\\mathcal{H}$. We find the optimal hyper-parameter $h^\\star\\in\\mathcal{H}$ that maximizes the accuracy on  *(a): Test Set* and *(b): Validation Set*, and report the accuracy of $M$ with $h^\\star$ on the Test Set. Under setting (a) we get 59.65\\% accuracy and 59.51\\% under (b). This difference suggests that the accuracy reported in the paper is stable under hyper-parameter selection based on test or validation set.\n\n> From the current text, it is not clear whether Assumption 3.1 is easy to be fullìlled or not.\n\nThe uniform distribution is chosen to simplify the nearest neighbor analysis for a cleaner result. The radius is chosen to ensure that the covariance matrix is identity. For a more general distribution, we can have an upper bound on the test error involving the term $\\mathbb{P}\\_\\mathbf{X}(\\mathcal{B}(x, h))$, where $\\mathbb{P}\\_X$ is the measure for the covariate distribution and $\\mathcal{B}(x, h)$ is a ball of radius $h$ centered at $x$, as in the Theorem 3.1.1 of [1]. The advantage of choosing a specific distribution is that we can see the precise scaling with respect to $d$ and $n$. We will add a sentence explaining the logic behind the choice of covariate distribution in the camera ready version.\n\n[1] George H. Chen; Devavrat Shah, Explaining the success of nearest neighbor methods in prediction, 2018. URL: https://ieeexplore.ieee.org/document/8384208"}]

Author rebuttal:

Meta-review: Four knowledgeable reviewers recommended accepting the paper: Borderline Accept, Weak Accept, Weak Accept, and Weak Accept. Based on the clear positive consensus, it is our recommendation to accept the paper. Authors should attend to the main points in the reviews. when preparing a final version. No basis to overturn the reviews.

The method proposed in this paper is novel, the empirical results are indeed good and supported by theory. This paper is valuable and should be shared within the community to advance research on ANN.
Decision: Accept (poster)