ResMem: Learn what you can and memorize the rest

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.

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) Finally, 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

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? 2) 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 3) 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. 4) 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) 5) From the current text, it is not clear whether Assumption 3.1 is easy to be fullfilled or not.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Limitations have been adequately addressed.

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.

Strengths

- The idea of using KNN to memorize the residual is novel and interesting - There is a risk bound proof

Weaknesses

- The assumption of the proof (linear function, k=1) is not practical and not aligned with the experiment - 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 ) - 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. Papernot, Nicolas, and Patrick McDaniel. "Deep k-nearest neighbors: Towards confident, interpretable and robust deep learning." arXiv preprint arXiv:1803.04765 (2018). Dosovitskiy, 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).

Questions

- The setup in Sec.4.2: Why do you set up this way? - Fig. 5: the results seem insensitive to K. Can you explain?

Rating

5: Borderline accept: Technically solid paper where reasons to accept outweigh reasons to reject, e.g., limited evaluation. Please use sparingly.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

2 fair

Contribution

2 fair

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.

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? The 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. The 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.

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. The algorithm is practical and simple, it seems quite reasonable to fit into existing pipelines. Beyond 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. I 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. I 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. The 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.

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? How do you pick the intermediate layer to use as the representations for the kNN? ResMem requires storing representations of the training data. Do you think this is an obstacle to its practical application? 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?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Aside from potential issues around space usage, which are addressed above, I feel the authors adequately address the limitations of their work.

Summary

The paper proposes the residual-memorization (ResMem) algorithm. This 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. The paper demonstrates the validity of the algorithm from both the theoretical and the empirical perspectives. The main theoretical result shows the convergence rate of ResMem. The empirical results show the performance increase of ResMem compared with standard neural networks on image and language datasets.

Strengths

The proposed method is quite novel. The writing is clear. The paper consists of both theoretical and empirical justification. This 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. On the theoretical result, the paper does show the asymptotic behavior of the prediction error of ResMem. However, 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. In 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; however, 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.

Questions

I am happy to raise the score if the authors address my concern. Q1: In Section 3.2, why is the equality in the second line of equation under Line 163 correct and consistent with the first line? Q2: In Section 3.3, what does the $T_2$ part of error represent? The author claims that $T_2$ arises due to the limited capacity of $F$, and captures an irreducible error of the risk. Could the authors please justify the argument? Can this $T_2$ term represent the generalization error? Q3: In Section 3.4, could the authors compare the bound with the case without ResMem? In particular, what is the error difference with or without ResMem in the $T_2$ term? Q4: On the empirical result, what is the additional computation and memory cost of the ResMem algorithm? In particular, how do the additional memory cost, training time and estimation time of the kNN component compare with those of the DNN component? Minor issues & suggestions: - Typo: In Line 67, "mall model" - 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. --- Thank you for the clarification, especially on clarifying Assumption 3.2. I updated the score accordingly.

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

Does not apply