SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zitong Yang
zitong@berkeley.edu
Michal Lukasik
New York, NY, 10011
mlukasik@google.com
Zonglin Li
lizonglin@google.com
Ankit Singh Rawat
ankitsrawat@google.com
Manzil Zaheer
manzilzaheer@google.com
The impressive generalization performance of modern neural networks is attributed in part to their ability to implicitly memorize complex training patterns. Inspired by this, we explore a novel mechanism to improve model generalization via explicit memorization. Specifically, we propose the residual-memorization (ResMem) algorithm, a new method that augments an existing prediction model (e.g., a neural network) by fitting the model's residuals with a  $k$ -nearest neighbor based regressor. The final prediction is then the sum of the original model and the fitted residual regressor. By construction, ResMem can explicitly memorize the training labels, even when the base model has low capacity. We start by formulating a stylized linear regression problem and rigorously show that ResMem results in a more favorable test risk over a base linear neural network. Then, we empirically show that ResMem consistently improves the test set generalization of the original prediction model across standard vision and natural language processing benchmarks.
Large neural networks achieve remarkable generalization on test samples despite memorization of training samples, in the sense of achieving zero training error [54]. Several recent analyses have established that, under certain settings, memorization is sufficient to achieve generalization [3, 15, 5, 40, 4], and, more surprisingly, can even be necessary [17, 19, 11]. These works suggest that suitable memorization can be a valuable desiderata for learning. While increasing model size is a conceptually simple strategy to enable memorization, this has the obvious downside of significantly increasing the cost of model training and serving. This raises a natural question: are there alternate mechanisms to improve the memorization (and thus generalization) of a relatively small model?
In this paper, we propose residual memorization (ResMem), a simple yet effective mechanism that achieves this goal (cf. Figure 1). Compared to the implicit memorization performed by large neural models, the key idea behind ResMem is to perform explicit memorization via a separate  $k$ -nearest neighbor component. Specifically, ResMem involves first training a standard neural network  $f_{\mathrm{DeepNet}}$  and then explicitly memorizing the model's residuals with a  $k$ -nearest neighbor based regressor  $r_{\mathrm{kNN}}$ .
(a) Step 1: learn the training set.
Figure 1: Illustration of the residual memorization (ResMem) algorithm. In a nutshell, we first fit a small deep network  $f_{\mathrm{DeepNet}}$  on the training sample (Figure 1(a)). When this network is non-memorizing, it incurs non-zero residual errors in its predictions (Figure 1(b)). We then fit a  $k$ -nearest neighbor based regressor on these residuals (Figure 1(c)). The final prediction is given by the sum of the initial network and  $k$ -NN regressor predictions. In all three figures, the  $x$ -axis represents the features in a supervised learning problem. In Figure 1(a), the  $y$ -axis represents the targets of prediction. In Figure 1(b) and 1(c), the  $y$ -axis represents the residual of the initial fitting from Step 1.
(b) Step 2: compute the residual.
(c) Step 3: memorize the residual.
Memorization through  $k$ -nearest neighbor can be efficiently computed with various approximation schemes (e.g. [24]). Subsequently, the ResMem prediction on an instance  $x$  is given by the sum of the two components, i.e.,  $f_{\mathrm{DeepNet}}(x) + r_{\mathrm{kNN}}(x)$ .
We start by formulating a stylized linear regression problem that captures the essence behind ResMem (cf. Section 3). Our analysis (Theorem 3.3) shows that, without ResMem, the test risk of the base linear neural network decreases to an irreducible constant as the number of samples goes to infinity. In contrast, the test risk of ResMem decreases to zero. The insight of theoretical analysis is that ResMem augments the capacity of the parametric linear network by adding a non-parametric component (i.e., nearest-neighbor).
Empirically, we show that such explicit memorization indeed leads to generalization benefits: ResMem consistently improves the test accuracy of a baseline DeepNet on image classification tasks with CIFAR100 [33], and autoregressive language modeling on C4 [42] (Section 4). Towards understanding this improved performance, we hypothesize that ResMem works by learning in two-stages (cf. Section 4.4). Specifically, we posit that the initial DeepNet  $f_{\text{DeepNet}}$  learns some coarse structure, and ResMem  $r_{\text{kNN}}$  supplements the DeepNet prediction with fine-grained details (cf. Figure 3). We verify our hypothesis via qualitative analysis on CIFAR100 and C4 (Section 4.4).
(1) We propose residual-memorization (ResMem), a two-stage learning algorithm that combines a base prediction model with a nearest neighbor regressor (cf. Figure 1);
(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. Section 4), particularly when the training set is extremely large;
1.1 Applicable scenarios of ResMem
From our theoretical and empirical analysis, we posit that ResMem (Figure 1) yields the largest margin of improvement over a base DeepNet when it is infeasible to perform implicit memorization with the latter. We discuss three such scenarios below. Each of our main empirical or theoretical results roughly corresponds to one of these settings.
- Complex dataset. In this scenario, the Bayes-optimal decision boundary is very complex, and is beyond the capability of the neural network itself. To demonstrate this, we analyze a theoretical linear regression problem where the target regression function is not contained in the hypothesis class of linear neural networks (cf. Section 3).
- Large sample size. Here, the number of training samples is large enough to make training set interpolation (i.e., achieving zero training error) infeasible for a given neural network model. For example, current large language models (LLMs) may be trained for at most a single epoch over trillions of examples [12]. By contrast, ResMem can circumvent this issue by explicitly memorizing the training samples. We emulate this scenario by considering a causal language modeling task on the C4 dataset (cf. Section 4.3).
- Small model. In many practical settings, one may prefer a smaller model over a state-of-the-art model due to the training and deployment cost constraints. We emulate such a setting through an image classification task where it is indeed feasible to memorize the training data perfectly using state-of-the-art neural networks, but instead, we use smaller neural networks for computational efficiency (cf. Section 4.2).
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 computable for generalization. Overparameterized neural models with many more parameters than training samples have the capacity to perfectly fit (or interpolate) even random training labels [54]; i.e., they can drive the empirical loss to zero for any training set. At the same time, when trained on real-world datasets, increasing model complexity tends to improve model performance [40, 52]; that is, the models do not simply memorize the training sample, but rather learn generalizable patterns. Several works have sought to understand the reasons behind this behaviour, both empirically [2] and theoretically [3, 15, 8, 5, 40, 36, 38, 4, 48, 50, 53]. One recurring message from the theory is that memorization (in the form of interpolation) can be sufficient for generalization.
Memorization can be necessary for generalization. Some recent works [17, 11] showed that memorization — either in the sense of interpolation, or in a more general sense of stability [18] — may be necessary for generalization. Feldman [16] considered a setting where the label distribution exhibits a long-tailed distribution, and showed that to prevent incurring a large error on the large number of under-represented classes, it may be necessary to memorize many of their associated training samples. Cheng et al. [11] considered a linear regression setting where it is beneficial to fit the training targets to error lower than the Bayes-error (i.e., the inherent noise in the targets).
2.2 Relation to existing algorithms
Nearest neighbor method. The  $k$ -nearest neighbor ( $k$ -NN) [14, 32, 26, 7] method assigns label to a test sample based on the label of its nearest neighbor(s) in the training set. Owing to its simplicity, flexibility in defining input similarity, and computational efficiency with various approximation schemes [22, 39], this method remains popular. However, the performance of  $k$ -NN drops as data becomes high dimensional [10, 39]. Therefore, to apply it to high dimensional data such as image and text [55], one approach is to learn a representation of data using neural networks [44]. Following this approach, [13] finds that applying  $k$ -NN directly to memorize the training labels  $y_{i}$  yields similar performance with the original softmax based neural network classification. In contrast, ResMem applies  $k$ -NN to memorize the residual  $r_{i}$  over the predictions of a base network.
Boosting and residual fitting. Boosting algorithms such as AdaBoost [20] seek to construct an ensemble of "weak learner" models with good generalization. AdaBoost achieves this in an iterative manner, and can be interpreted as a particular instantiation of forward stage-wise regression [21], a classical procedure from statistics [23, 1, 47]. Intuitively, at each round, one builds a new weaker learner by fitting the residual of the ensemble of weak learners constructed thus far. This fitting is performed iteratively. ResMem can be loosely regarded as a two round boosting algorithm where the first "weak learner" is the base neural network and the second "weak learner" is the nearest-neighbor component. Note that there is no need for the thrid "weak learner", because the nearest-neighbor component already perfectly memorizes the neural network residuals.
Memory-augmented language models. In the language modelling literature, several works explore combining neural models with an external database or memory, which can be queried to retrieve additional context [34, 25, 6, 35]. Closer to our work, Khandelwal et al. [30] employ a linear combination of neural network and  $k$ -NN classifier components. However, a crucial difference is that our  $k$ -NN components memorizes the residual of the DeepNet prediction, whereas Khandelwal et al. [30] memorizes the target label directly; i.e., their approach is akin to an ensemble of  $k$ -NN and a deep network. Various forms of memory have also been considered in generic classification problems [41, 48, 51]. This line of literature again differs from ResMem in that their memory tries to memorize labels directly, whereas ResMem memorizes the residuals, leading to a natural combination of the neural network and the memory component.
Model compression for large neural networks. Since ResMem boosts the test accuracy of a small, non-memorizing neural network, we can also view it as a technique that allows a small network to match the performance of a larger one. This relates to the model compression literature. Distillation [29, 9] is a popular strategy for compressing a large neural model to a smaller one. For a survey of other effective strategies, including pruning, see Menghani [37]. In Appendix C.2, we discuss how ResMem can be regarded as a "dual procedure" of distillation.
As discussed in Section 1.1, ResMem yields the largest improvement when implicit memorization is infeasible. In this section, we formulate (cf. Section 3.1) and analyze (cf. Theorem 3.3) a stylized linear regression problem that concretizes such a setting.
Recall that ResMem (Figure 1) involves first training a base neural network  $f_{\text{DeepNet}}$ , and then fitting the residual of  $f_{\text{DeepNet}}$  on the same training data using a nearest-neighbor regressor  $r_{\text{kNN}}$ . For feasibility of theoretical analysis, we simplify  $f_{\text{DeepNet}}$  with a single layer linear neural network, i.e. linear regression, and we consider 1-nearest neighbor instead of  $k$ -nearest neighbor to memorize the residual of this network. Our results suggest that ResMem improves test-time generalization by augmenting the capacity of the base model with a non-parametric nearest-neighbor component.
3.1 Assumptions and setting
In this section, we present the setup and assumptions for the stylized linear regression problem. We consider a setting where the function class that we minimize over does not include the ground-truth function that relates the covariates to the response. Therefore, even with infinite samples, the test loss will decay to a positive constant. We exactly characterize the rate of decay, and show that it converges to 0 under ResMem. Our analysis rests on the following assumptions.
Assumption 3.1 (Distribution of covariates). The distribution of covariate  $\boldsymbol{x} \in \mathbb{R}^d$ , denoted by  $\mathbb{P}_{\boldsymbol{x}}$  is the uniform distribution over a Euclidean norm ball centered at the origin of radius  $\sqrt{d + 2}$ . The choice of radius ensures that  $\mathbb{E}_{\boldsymbol{x} \sim \mathbb{P}_{\boldsymbol{x}}} \boldsymbol{x} \boldsymbol{x}^{\top} = \boldsymbol{I}$ .
Assumption 3.2 (Linear regression over norm ball). Consider the problem of learning a linear function  $f_{\star}(\boldsymbol{x}) = \langle \boldsymbol{x}, \boldsymbol{\theta}_{\star} \rangle$  with  $\| \boldsymbol{\theta}_{\star} \| = 1$  from training data  $\{(x_i, y_i)\}_{i=1:n}$  where  $x_i \stackrel{\text{i.i.d.}}{\sim} \mathbb{P}_x$  and  $y_i = f_{\star}(x_i)$  using the function class
$$
\mathcal {F} = \left\{\boldsymbol {x} \mapsto \langle \boldsymbol {x}, \boldsymbol {\theta} \rangle , \| \boldsymbol {\theta} \| <   L \right\}. \tag {1}
$$
We assume  $L < 1$  so that the problem belongs to the "hard generalization" scenario discussed in Section 1.1, where the hypothesis space is inadequate to fit the function on its own.
ResMem proceeds by first learning a linear function  $f_{n}(\boldsymbol{x}) = \langle \boldsymbol{\theta}_{n},\boldsymbol{x}\rangle$  from  $\mathcal{F}$  through empirical risk minimization (ERM):
$$
\boldsymbol {\theta} _ {n} = \underset {\| \boldsymbol {\theta} \| \leq L} {\operatorname {a r g m i n}} \frac {1}{n} \sum_ {i = 1} ^ {n} \left[ \langle \boldsymbol {x} _ {i}, \boldsymbol {\theta} \rangle - y _ {i} \right] ^ {2}. \tag {2}
$$
For more general distributions, the theoretical result will depend on quantities like  $\mathbb{P}_{\pmb{x}}(\mathcal{B}(\widetilde{\pmb{x}},h))$ , where  $\mathcal{B}(\widetilde{\pmb{x}},h)$  means a ball of radius  $h$  that is centered at  $\widetilde{\pmb{x}}$ . We took uniform distribution for simplicity and to obtain exact dependence on  $d$ .
The empirical risk minimizer  $f_{n}$  should be thought of as the analog of  $f_{\mathrm{DeepNet}}$  in the deep learning context. It defines a ground-truth residual function  $r_{\star}(\pmb{x}) = f_{\star}(\pmb{x}) - f_{n}(\pmb{x})$ . Now we fix a test covariate  $\widetilde{\pmb{x}} \sim \mathbb{P}_x$ . ResMem "memorizes" the residual function through the 1-nearest neighbor to  $\widetilde{\pmb{x}}$
$$
r _ {n} (\widetilde {\boldsymbol {x}}) = r _ {\star} (\widetilde {\boldsymbol {x}} _ {(1)}) = f _ {\star} (\widetilde {\boldsymbol {x}} _ {(1)}) - f _ {n} (\widetilde {\boldsymbol {x}} _ {(1)}), \tag {3}
$$
where  $\widetilde{\pmb{x}}_{(1)}$  is the nearest neighbor to  $\widetilde{\pmb{x}}$  among the training covariates  $\pmb{x}_1,\dots ,\pmb{x}_n$  ..
$$
\widetilde {\boldsymbol {x}} _ {(1)} = \operatorname * {a r g m i n} _ {\boldsymbol {x} \in \{\boldsymbol {x} _ {1}, \dots , \boldsymbol {x} _ {n} \}} \| \boldsymbol {x} - \widetilde {\boldsymbol {x}} \|.
$$
The final prediction is
$$
f _ {n} ^ {\text {R e s M e m}} (\widetilde {\boldsymbol {x}}) = f _ {n} (\widetilde {\boldsymbol {x}}) + r _ {n} (\widetilde {\boldsymbol {x}}). \tag {4}
$$
Observe that if  $\widetilde{x}$  coincides with any training sample,  $f_{n}^{\mathrm{ResMem}}(\widetilde{x}) = f_{\star}(\widetilde{x})$ , i.e., we have explicit memorization. Note that we worked with 1-nearest neighbor regressor for simplicity instead of the general  $k$ -nearest neighbor algorithm. The effect of choosing different  $k$  is not the main focus of this theoretical analysis.
3.2 A decomposition of the target function
Next, we introduce a decomposition of  $f_{\star}$ , which will help us analyze various components that make up the risk. Define
$$
\begin{array}{l} \boldsymbol{\theta}_{\infty} = \operatorname *{argmin}_{\| \boldsymbol {\theta}\| \leq L}\mathbb{E}_{\boldsymbol {x}\sim \mathbb{P}_{x}}\left[\langle \boldsymbol {\theta},\boldsymbol {x}\rangle -\langle \boldsymbol{\theta}_{\star},\boldsymbol {x}\rangle \right]^{2}, \\ = \operatorname * {a r g m i n} _ {\| \boldsymbol {\theta} \| \leq L} \| \boldsymbol {\theta} - \boldsymbol {\theta} _ {\star} \| = L \boldsymbol {\theta} _ {\star}, \\ \end{array}
$$
which is what ERM learns in the limit of  $n\to \infty$ . We can think of  $\theta_{\infty}$  as the best function that ERM can learn. Then, we can decompose  $\theta_{\star}$  into  $\theta_{\star} = \theta_{\infty} + \theta_{\perp}$ , where  $\theta_{\perp} = \theta_{\star} - \theta_{\infty}$ . This decomposition can be generalized beyond linear regression. Since  $\theta_{\infty}$  defines a function  $f_{\infty}(\boldsymbol {x}) = \langle \boldsymbol {x},\boldsymbol{\theta}_{\infty}\rangle$ , for general non-linear functions, the argument above can be generalized to the decomposition of  $f_{\star}$  to an learnable and non-learnable part
$$
f _ {\star} = f _ {\infty} + f _ {\perp}.
$$
Intuitively,  $f_{\infty}$  is the best function in  $\mathcal{F}$  that ERM can learn, and  $f_{\perp}$  is beyond the capacity of ERM due to the particular choice of function class. ResMem approximates  $f_{\perp}$  using the non-parametric nearest neighbor method, and therefore expanding the capacity of the original hypothesis 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
$$
\mathbb {E} \left[ \left(f _ {n} ^ {\text {R e s M e m}} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}})\right) ^ {2} \right]. \tag {5}
$$
It can be decomposed into two components:  $\mathbb{E}\left[f_n^{\mathrm{ResMem}}(\widetilde{\boldsymbol{x}}) - f_\star (\widetilde{\boldsymbol{x}})\right]^2\leq 3\times$
$$
\left[ \underbrace {\mathbb {E} (f _ {n} (\widetilde {\boldsymbol {x}}) - f _ {\infty} (\widetilde {\boldsymbol {x}})) ^ {2} + \mathbb {E} (f _ {n} (\widetilde {\boldsymbol {x}} _ {(1)}) - f _ {\infty} (\widetilde {\boldsymbol {x}} _ {(1)})) ^ {2}} _ {T _ {1}} + \underbrace {\mathbb {E} (f _ {\infty} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}}) - f _ {\infty} (\widetilde {\boldsymbol {x}} _ {(1)}) + f _ {\star} (\widetilde {\boldsymbol {x}} _ {(1)})) ^ {2}} _ {T _ {2}} \right].
$$
We provide the detail of the decomposition in Section B.1. We can see that  $T_{1}$  arises due to the difference between  $f_{n}$  and  $f_{\infty}$  (i.e., the estimation error), which, as we will show later, goes to 0 as  $n$  goes to infinity:
$$
T _ {1} \rightarrow 0 \text {a s} n \rightarrow \infty .
$$
On the other hand,  $T_{2}$  arises due to the limited capacity of  $\mathcal{F}$ . It captures an irreducible error of the risk, which in general is not asymptotically zero. However, because of the explicit memorization ResMem algorithm introduces  $(\widetilde{\boldsymbol{x}}_{(1)} \to \widetilde{\boldsymbol{x}}$  as  $n \to \infty$ ), we also have
$$
T _ {2} \rightarrow 0 \text {a s} n \rightarrow \infty .
$$
This decomposition provides a statistical perspective on ResMem: it preserves the asymptotic consistency of  $T_{1}$  as in classical learning problems while enforcing the asymptotic consistency of  $T_{2}$  through the nearest-neighbor method.
3.4 Main theoretical result
Given the set up above, we are ready to state the main theoretical result of the paper, which characterizes the rate at which test risk of ResMem approaches 0. The proof is in Appendix B.
Theorem 3.3 (Risk for ResMem algorithm). For the problem defined in Assumption 3.2 with covariates distribution in Assumption 3.1, the ResMem prediction rule  $f_{n}^{ResMem}(\widetilde{\boldsymbol{x}})$  defined in equation (4) achieves risk (5)
$$
\mathbb {E} \left[ f _ {n} ^ {R e s M e m} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}}) \right] ^ {2} \lesssim d ^ {2} L ^ {2} n ^ {- 2 / 3} + d ^ {2} (1 - L) ^ {2} \left[ \frac {\log \left(n ^ {1 / d}\right)}{n} \right] ^ {1 / d},
$$
where  $\lesssim$  denotes inequality up to a universal constant independent of  $d, n$  and  $L$ .
The result includes contribution from two terms introduced in Section 3.3:
-  $T_{1} \lesssim d^{2}L^{2}n^{-2 / 3}$  that arises due to the difference between  $f_{n}$  and  $f_{\infty}$ .
-  $T_{2} \lesssim \left[\log \left(n^{1 / d}\right) / n\right]^{1 / d}$  that vanishes as the nearest neighbor of the test point approaches the test point itself  $\widetilde{\boldsymbol{x}}_{(1)} \rightarrow \widetilde{\boldsymbol{x}}$ .
The two terms  $T_{1}$  and  $T_{2}$  can be viewed as "two stages of learning". Without the ResMem memorization component, we have the usual story of machine learning:  $T_{1} \rightarrow 0$  at the usual parametric rate, and  $T_{2}$  stays as an irreducible error, so the overall test error diminishes to a constant at a very fast rate. With the introduction of nearest neighbor memorization procedure,  $T_{2}$  can also be reduced to 0 at a slower rate, whereas the fast decay of  $T_{1}$  is still preserved.
This result shows why it is not favorable to use the  $k$ -nearest neighbor component to memorize the response directly: as a corollary of setting  $L = 0$  in Theorem 3.3, pure nearest neighbor would result in an overall slow rate of  $\approx n^{-1/d}$ . However, with ResMem, we can enjoy benefit of having the test loss being asymptotically 0, while also enjoying the fast rate of  $n^{-2/3}$  for smaller sample sizes.
In this section, we present empirical results on image classification and language modeling that showcase the efficacy of ResMem. In Section 4.1, we present details of applying the ResMem algorithm to classification problems on real dataset. In Section 4.2 and Section 4.3, we present the setup and the result for vision and language experiments, respectively. In Section 4.4 we conduct an empirical analysis to explain where the improved accuracy of ResMem comes from. Finally, in addition to evaluating the improvement ResMem algorithm over DeepNet itself, we compare ResMem with other reasonable baselines including [31] in Appendix F.
4.1 Details of ResMem algorithm for classification
We consider multi-class classification problems over instances  $\mathcal{X}$  and labels  $\mathcal{Y} \doteq \{1,2,\dots,L\} = [L]$ . Given training examples  $S = \{(x_{i},y_{i})\}_{i\in [n]} \in (\mathcal{X} \times \mathcal{Y})^{n}$ , the goal is to learn a scorer  $f\colon \mathcal{X} \to \mathbb{R}^{L}$  that, given an instance, assigns an affinity score for each label. Such an  $f$  should minimize the misclassification error on test samples:
$$
L _ {0 1} (f) \doteq \mathbb {P} _ {(x, y)} (y \neq \operatorname {p r e d} (f (x))), \tag {6}
$$
where  $\mathbf{pred}(z) \doteq \arg \max_{y' \in [L]} z_{y'}$ , and  $\mathbb{P}$  is the distribution over labelled instances. To achieve this, one typically minimizes the empirical loss