SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \mathbb {E} \left[ f ^ {\text {R e s M e m}} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}}) \right] ^ {2} = \mathbb {E} \left[ f _ {n} (\widetilde {\boldsymbol {x}}) + r _ {n} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}}) \right] ^ {2}, \\ = \mathbb {E} \left[ f _ {n} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}}) - f _ {n} (\widetilde {\boldsymbol {x}} _ {(1)}) + f _ {\star} (\widetilde {\boldsymbol {x}} _ {(1)}) \right] ^ {2}, \\ = \mathbb {E} \left[ f _ {n} (\widetilde {\boldsymbol {x}}) - f _ {\infty} (\widetilde {\boldsymbol {x}}) + f _ {\infty} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}}) - f _ {n} (\widetilde {\boldsymbol {x}} _ {(1)}) + f _ {\infty} (\widetilde {\boldsymbol {x}} _ {(1)}) - f _ {\infty} (\widetilde {\boldsymbol {x}} _ {(1)}) + f _ {\star} (\widetilde {\boldsymbol {x}} _ {(1)}) \right] ^ {2}, \\ \leq 3 \times \underbrace {\left[ \mathbb {E} \left(f _ {n} (\widetilde {\boldsymbol {x}}) - f _ {\infty} (\widetilde {\boldsymbol {x}})\right) ^ {2} + \mathbb {E} \left(f _ {n} \left(\widetilde {\boldsymbol {x}} _ {(1)}\right) - f _ {\infty} \left(\widetilde {\boldsymbol {x}} _ {(1)}\right)\right) ^ {2} \right]} _ {T _ {1}} + \underbrace {\mathbb {E} \left(f _ {\infty} (\widetilde {\boldsymbol {x}}) - f _ {\star} (\widetilde {\boldsymbol {x}}) - f _ {\infty} \left(\widetilde {\boldsymbol {x}} _ {(1)}\right) + f _ {\star} \left(\widetilde {\boldsymbol {x}} _ {(1)}\right)\right) ^ {2}} _ {T _ {2}}, \tag {22} \\ \end{array}
$$
where in the last inequality, we used the fact that  $(a + b + c)^2 < 3(a^2 +b^2 +c^2)$  for any  $a,b,c\in \mathbb{R}$
B.2 Upper bound on  $T_{1}$ .
Since  $\mathbb{P}_{\boldsymbol{x}} = \mathrm{Unif}(\mathcal{B}_{0,B})$ , we apply the bound  $\| \widetilde{\boldsymbol{x}} \|, \| \widetilde{\boldsymbol{x}}_{(1)} \| \leq B$  to obtain
$$
\begin{array}{l} T _ {1} = \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}, \\ = \mathbb {E} \langle \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty}, \widetilde {\boldsymbol {x}} \rangle^ {2} + \mathbb {E} \langle \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty}, \widetilde {\boldsymbol {x}} _ {(1)} \rangle^ {2}, \\ \leq \mathbb {E} \| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \| ^ {2} \| \widetilde {\boldsymbol {x}} \| ^ {2} + \mathbb {E} \| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \| ^ {2} \| \widetilde {\boldsymbol {x}} _ {(1)} \| ^ {2}, \tag {23} \\ \leq 2 B ^ {2} \mathbb {E} \| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \| ^ {2}. \\ \end{array}
$$
As  $n$  gets large, the empirical covariance matrix  $\pmb{\Sigma}_n = \pmb{X}^\top \pmb{X} / n$  is concentrated around its mean  $\pmb{I}$ . Let  $\pmb{\Delta}_n = \pmb{I} - \pmb{\Sigma}_n$  denote this deviation. For some  $\varepsilon \in (0,1)$ , define the following "good event" over the randomness in  $\pmb{\Sigma}_n$
$$
\mathcal {A} = \left\{\| \boldsymbol {\Delta} _ {n} \| _ {\mathrm {o p}} <   \varepsilon \right\}, \tag {24}
$$
where  $\| \Delta_n\|_{\mathrm{op}}$  denotes the operator norm of the deviation matrix. The high level idea of the proof is to condition on the event  $\mathcal{A}$  and deduce and upper bound of  $\| \pmb {\theta}_n - \pmb {\theta}_\infty \|$  in terms of  $\varepsilon$ . Then, we use the fact that  $\mathcal{A}$  happens with high probability.
Recall that  $\theta_{\infty} = L\theta_{\star}$ , and
$$
\boldsymbol {\theta} _ {n} = \underset {\| \boldsymbol {\theta} \| \leq L} {\operatorname {a r g m i n}} \frac {1}{n} \| \boldsymbol {X} \boldsymbol {\theta} - \boldsymbol {y} \| ^ {2}. \tag {25}
$$
Since  $\pmb{y} = \pmb{X}\pmb{\theta}_{\star}$  by definition, the Lagrangian of the convex program above is
$$
\mathcal {L} (\boldsymbol {\theta}, \lambda) = \frac {1}{n} \| \boldsymbol {X} \boldsymbol {\theta} - \boldsymbol {X} \boldsymbol {\theta} _ {\star} \| ^ {2} + \lambda \left(\| \boldsymbol {\theta} \| ^ {2} - L\right). \tag {26}
$$
The KKT condition suggests that the primal-dual optimal pair  $(\theta_{n},\lambda_{n})$  is given by
$$
\left\| \boldsymbol {\theta} _ {n} \right\| \leq L,
$$
$$
\lambda_ {n} \geq 0, \tag {27}
$$
$$
\lambda_ {n} (\| \boldsymbol {\theta} _ {n} \| - L) = 0,
$$
and at optimality
$$
\begin{array}{l} \nabla_ {\boldsymbol {\theta}} \mathcal {L} \left(\boldsymbol {\theta} _ {n}, \lambda_ {n}\right) = 0 \Longleftrightarrow \frac {2}{n} \boldsymbol {X} ^ {\top} \boldsymbol {X} \left(\boldsymbol {\theta} - \boldsymbol {\theta} _ {\star}\right) + 2 \lambda_ {n} \boldsymbol {\theta} = 0, \tag {28} \\ \Longleftrightarrow \boldsymbol {\theta} _ {n} = \left(\boldsymbol {\Sigma} _ {n} + \lambda_ {n} \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Sigma} _ {n} \boldsymbol {\theta} _ {\star}. \\ \end{array}
$$
The complementary slackness condition  $\lambda_{n}(\| \pmb{\theta}_{n}\| -L) = 0$  suggests that either  $\lambda_{n} = 0$  or  $\| \pmb{\theta}_n\| = L$  But if  $\lambda_{n} = 0$  , the stationary condition  $\nabla_{\pmb{\theta}}\mathcal{L}(\pmb {\theta},\lambda) = 0$  would suggest that  $\pmb {\theta}_n = \pmb {\Sigma}_n^{-1}\pmb {\Sigma}_n\pmb {\theta}_\star =$ $\pmb{\theta}_{\star}\Rightarrow \| \pmb{\theta}_{n}\| = 1 > L$  , a contradiction. (Note that here  $\pmb{\Sigma}_{n}$  is invertible condition on the event A.) Therefore, we must have  $\| \pmb {\theta}_n\| = L$  . As a result, the primal and dual pair  $(\pmb {\theta}_n,\lambda_n)$  is determined by the system of equations
$$
\left\{ \begin{array}{l l} \boldsymbol {\theta} _ {n} & = \left(\boldsymbol {\Sigma} _ {n} + \lambda_ {n} \boldsymbol {I}\right) ^ {- 1} \boldsymbol {\Sigma} _ {n} \boldsymbol {\theta} _ {\star}, \\ \| \boldsymbol {\theta} _ {n} \| & = L, \\ \lambda_ {n} & > 0. \end{array} \right. \tag {29}
$$
Next, we proceed to compute the deviation  $\| \pmb{\theta}_n - \pmb{\theta}_{\infty}\|$
$$
\begin{array}{l} \pmb {\theta} _ {n} = [ (\lambda_ {n} + 1) \pmb {I} - \pmb {\Delta} _ {n} ] ^ {- 1} \pmb {\Sigma} _ {n} \pmb {\theta} _ {\star}, \\ = \left(\lambda_ {n} + 1\right) ^ {- 1} \left[ I - \frac {\boldsymbol {\Delta} _ {n}}{\lambda_ {n} + 1} \right] ^ {- 1} \boldsymbol {\Sigma} _ {n} \boldsymbol {\theta} _ {\star}, \\ = (\lambda_ {n} + 1) ^ {- 1} \left[ \boldsymbol {I} + \sum_ {k = 1} ^ {\infty} \frac {\boldsymbol {\Delta} _ {n} ^ {k}}{(\lambda_ {n} + 1) ^ {k}} \right] (\boldsymbol {I} - \boldsymbol {\Delta} _ {n}) \boldsymbol {\theta} _ {\star}, \\ = \left(\lambda_ {n} + 1\right) ^ {- 1} \left[ I + \sum_ {k = 1} ^ {\infty} \frac {\boldsymbol {\Delta} _ {n} ^ {k}}{\left(\lambda_ {n} + 1\right) ^ {k}} - \boldsymbol {\Delta} _ {n} - \sum_ {k = 1} ^ {\infty} \frac {\boldsymbol {\Delta} _ {n} ^ {k + 1}}{\left(\lambda_ {n} + 1\right) ^ {k}} \right] \boldsymbol {\theta} _ {\star}, \tag {30} \\ = (\lambda_ {n} + 1) ^ {- 1} \boldsymbol {\theta} _ {\star} + (\lambda_ {n} + 1) ^ {- 1} \boldsymbol {\Delta} _ {n} \left[ \sum_ {k = 1} ^ {\infty} \frac {\boldsymbol {\Delta} _ {n} ^ {k - 1}}{(\lambda_ {n} + 1) ^ {k}} - \boldsymbol {I} - \sum_ {k = 1} ^ {\infty} \frac {\boldsymbol {\Delta} _ {n} ^ {k}}{(\lambda_ {n} + 1) ^ {k}} \right] \boldsymbol {\theta} _ {\star}, \\ = (\lambda_ {n} + 1) ^ {- 1} \pmb {\theta} _ {\star} + (\lambda_ {n} + 1) ^ {- 1} \pmb {\Delta} _ {n} \left[ \sum_ {k = 1} ^ {\infty} \frac {\pmb {\Delta} _ {n} ^ {k - 1} - \pmb {\Delta} _ {n} ^ {k}}{(\lambda_ {n} + 1) ^ {k}} - \pmb {I} \right] \pmb {\theta} _ {\star}. \\ \end{array}
$$
$$
\boldsymbol {D} _ {n} = \boldsymbol {\Delta} _ {n} \left[ \sum_ {k = 1} ^ {\infty} \frac {\boldsymbol {\Delta} _ {n} ^ {k - 1} - \boldsymbol {\Delta} _ {n} ^ {k}}{(\lambda_ {n} + 1) ^ {k}} - \boldsymbol {I} \right]. \tag {31}
$$
Then  $\pmb{\theta}_n = (\lambda_n + 1)^{-1}\pmb{\theta}_{\star} + (\lambda_n + 1)^{-1}\pmb{D}_n\pmb{\theta}_{\star}$ , and
$$
\begin{array}{l} \left\| \boldsymbol {D} _ {n} \right\| \leq \left\| \boldsymbol {\Delta} _ {n} \right\| \left[ 1 + \sum_ {k = 1} ^ {\infty} \frac {\left\| \boldsymbol {\Delta} _ {n} \right\| ^ {k - 1} + \left\| \boldsymbol {\Delta} _ {n} \right\| ^ {k}}{\left(\lambda_ {n} + 1\right) ^ {k}} \right], \\ \leq \varepsilon \left[ 1 + 2 \left(1 + \lambda_ {n}\right) ^ {- 1} \sum_ {k = 1} ^ {\infty} \left(\frac {\varepsilon}{1 + \lambda_ {n}}\right) ^ {k} \right], \tag {32} \\ = \varepsilon \left(1 + \frac {2}{1 + \lambda_ {n}} \frac {1}{1 - \frac {\varepsilon}{1 + \lambda_ {n}}}\right) \leq 3 \varepsilon . \\ \end{array}
$$
$$
\begin{array}{l} L = \left\| \boldsymbol {\theta} _ {n} \right\| ^ {2} = (\lambda_ {n} + 1) ^ {- 2} + (\lambda_ {n} + 1) ^ {- 2} \boldsymbol {\theta} _ {\star} ^ {\mathsf {T}} \boldsymbol {D} _ {n} ^ {\mathsf {T}} \boldsymbol {D} _ {n} \boldsymbol {\theta} _ {\star} + 2 (\lambda_ {n} + 1) ^ {- 2} \boldsymbol {\theta} _ {\star} \boldsymbol {D} _ {n} \boldsymbol {\theta} _ {\star}, \tag {33} \\ \Rightarrow (\lambda_ {n} + 1) ^ {2} L ^ {2} = 1 + \delta_ {n}, \delta_ {n} = \boldsymbol {\theta} _ {\star} ^ {\mathsf {T}} \boldsymbol {D} _ {n} ^ {\mathsf {T}} \boldsymbol {D} _ {n} \boldsymbol {\theta} _ {\star} + 2 \boldsymbol {\theta} _ {\star} ^ {\mathsf {T}} \boldsymbol {D} _ {n} \boldsymbol {\theta} _ {\star}. \\ \end{array}
$$
We can obtain the following bound for  $\delta_{n}$ :
$$
\left| \delta_ {n} \right| \leq \left\| \boldsymbol {\theta} _ {\star} \right\| ^ {2} \left\| \boldsymbol {D} _ {n} \right\| ^ {2} + 2 \left\| \boldsymbol {\theta} _ {\star} \right\| ^ {2} \left\| \boldsymbol {D} _ {n} \right\| \leq 9 \varepsilon^ {2} + 6 \varepsilon \leq 1 5 \varepsilon . \tag {34}
$$
Since  $1 - \delta_n / 2 \leq \sqrt{1 + \delta_n} \leq 1 + \delta_n / 2$ , and  $|\delta_n| \leq 15\varepsilon$ , we obtain
$$
\left| \left(\lambda_ {n} + 1\right) L - 1 \right| \leq \frac {1 5 \varepsilon}{2} \Rightarrow \left| L - \left(\lambda_ {n} + 1\right) ^ {- 1} \right| \leq \frac {1 5 \varepsilon}{2} \left(\lambda_ {n} + 1\right) ^ {- 1} \leq \frac {1 5 \varepsilon}{2}, \tag {35}
$$
where the last inequality follows as we have  $\lambda_{n} > 0$ . Finally,
$$
\begin{array}{l} \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} = (\lambda_ {n} + 1) ^ {- 1} \boldsymbol {\theta} _ {\star} - L \boldsymbol {\theta} _ {\star} + (\lambda_ {n} + 1) ^ {- 1} \boldsymbol {D} _ {n} \boldsymbol {\theta} _ {\star}, \\ \Rightarrow \| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \| ^ {2} = \left[ \left(1 + \lambda_ {n}\right) ^ {- 1} - L \right] ^ {2} + \left(1 + \lambda_ {n}\right) ^ {- 2} \boldsymbol {\theta} _ {\star} \boldsymbol {D} _ {n} ^ {\mathsf {T}} \boldsymbol {D} _ {n} \boldsymbol {\theta} _ {\star} + 2 \left(\lambda_ {n} + 1\right) ^ {- 1} \left[ \left(1 + \lambda_ {n}\right) ^ {- 1} - L \right] \boldsymbol {\theta} _ {\star} \boldsymbol {D} _ {n} \boldsymbol {\theta} _ {\star}, \\ \leq 6 4 \varepsilon^ {2} + 9 \varepsilon^ {2} + 4 5 \varepsilon^ {2} = 1 1 8 \varepsilon^ {2}, \\ \Rightarrow \left\| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \right\| ^ {2} \lesssim \varepsilon^ {2}. \tag {36} \\ \end{array}
$$
Combine the above result with Proposition A.2, we get that
$$
\begin{array}{l} \mathbb {E} \left\| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \right\| ^ {2} = \mathbb {E} \left(\left\| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \right\| ^ {2} | \mathcal {A}\right) \mathbb {P} (\mathcal {A}) + \mathbb {E} \left(\left\| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \right\| ^ {2} | \mathcal {A} ^ {c}\right) \mathbb {P} \left(\mathcal {A} ^ {c}\right), \\ \leq \varepsilon^ {2} + 4 L ^ {2} \times 4 d \exp \left[ - \frac {n \varepsilon^ {2}}{2 (d + 2) (1 + \varepsilon)} \right], \tag {37} \\ \end{array}
$$
If we choose  $\varepsilon = n^{-1 / 3}$ , we get
$$
\mathbb {E} \| \boldsymbol {\theta} _ {n} - \boldsymbol {\theta} _ {\infty} \| ^ {2} \lesssim d L ^ {2} n ^ {- 2 / 3}, \tag {38}
$$
$$
T _ {1} \lesssim d ^ {2} L ^ {2} n ^ {- 2 / 3}. \tag {39}
$$
B.3 Upper bound on  $T_{2}$ .
Plugging in the formula for  $f_{\perp}(\widetilde{\boldsymbol{x}}) = f_{\star}(\widetilde{\boldsymbol{x}}) - f_{\infty}(\widetilde{\boldsymbol{x}}) = \langle \widetilde{\boldsymbol{x}},\boldsymbol{\theta}_{\perp}\rangle$  , we get
$$
\begin{array}{l} T _ {2} = \mathbb {E} \left[ f _ {\perp} \left(\widetilde {\boldsymbol {x}} _ {(1)}\right) - f _ {\perp} (\widetilde {\boldsymbol {x}}) \right] ^ {2}, \\ = \mathbb {E} \left\langle \boldsymbol {\theta} _ {\perp}, \widetilde {\boldsymbol {x}} _ {(1)} - \widetilde {\boldsymbol {x}} \right\rangle^ {2}, (40) \\ \leq (1 - L) ^ {2} \| \boldsymbol {\theta} _ {\star} \| ^ {2} \mathbb {E} \| \widetilde {\boldsymbol {x}} - \widetilde {\boldsymbol {x}} _ {(1)} \| ^ {2}, (40) \\ = (1 - L) ^ {2} \mathbb {E} \| \widetilde {\boldsymbol {x}} - \widetilde {\boldsymbol {x}} _ {(1)} \| ^ {2}, \\ \end{array}
$$
where in the last inequality, we used the relation that  $\pmb{\theta}_{\perp} = (1 - L)\pmb{\theta}_{\star}$ . Proposition A.1 suggests that
$$
\mathbb {E} \| \widetilde {\boldsymbol {x}} - \widetilde {\boldsymbol {x}} _ {(1)} \| ^ {2} \lesssim d ^ {2} \left[ \frac {\log (n ^ {1 / d})}{n} \right] ^ {1 / d}, \tag {41}
$$
$$
T _ {2} \lesssim d ^ {2} (1 - L) ^ {2} \left[ \frac {\log (n ^ {1 / d})}{n} \right] ^ {1 / d}. \tag {42}
$$
Remark B.1 (Comparison with pure nearest neighbor and ERM). If we rely solely on nearest neighbor method, the prediction error is
$$
\mathbb {E} \left[ f _ {\star} (\widetilde {\boldsymbol {x}}) - f _ {\star} \left(\widetilde {\boldsymbol {x}} _ {(1)}\right) \right] ^ {2} = \mathbb {E} \left\langle \widetilde {\boldsymbol {x}} - \widetilde {\boldsymbol {x}} _ {(1)}, \boldsymbol {\theta} _ {\star} \right\rangle^ {2} \leq \mathbb {E} \| \widetilde {\boldsymbol {x}} - \widetilde {\boldsymbol {x}} _ {(1)} \| ^ {2}. \tag {43}
$$
On the other hand, if we solely rely on ERM, even with infinite sample, we get
$$
\mathbb {E} \left[ f _ {\star} (\widetilde {\boldsymbol {x}}) - f _ {\infty} (\widetilde {\boldsymbol {x}}) \right] ^ {2} = \mathbb {E} \langle \widetilde {\boldsymbol {x}}, \boldsymbol {\theta} _ {\star} - \boldsymbol {\theta} _ {\infty} \rangle^ {2} \leq (1 - L) ^ {2} \mathbb {E} \| \widetilde {\boldsymbol {x}} \| ^ {2}. \tag {44}
$$
We can see from the upper bound that ResMem takes advantage of both
- Projecting  $f_{\star}$  onto  $f_{\infty}$ , so that the dependence on the prediction function is reduced from 1 to  $(1 - L)^2$ .
- Memorizing the residuals using nearest neighbor, so that the variance is reduced from  $\mathbb{E}\|\widetilde{\boldsymbol{x}}\|^2$  to  $\mathbb{E}\|\widetilde{\boldsymbol{x}}_{(1)} - \widetilde{\boldsymbol{x}}\|^2$ .
B.4 Test loss for ResMem.
If we combine the previous two parts together, we get
$$
\mathbb {E} \left[ \hat {f} (\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}. \tag {45}
$$
This completes the proof of Theorem 3.3.
C Additional CIFAR100 Results
This section includes additional experiment results on applying ResMem to CIFAR100 dataset.
C.1 Additional robustness results
In addition to the results already presented in Section 4.2, we also evaluate ResMem performance for each architecture in CIFAR-ResNet{8, 14, 20, 32, 44, 56} and each subset (10%, 20%, ..., 100%) of CIFAR100 training data. We use the same training hyperparameter and the ResMem hyperparameter as described in Section 4.2. Generally, we see that ResMem yields larger improvement over the baseline DeepNet when the network is small and dataset is large.
(a) CIFAR-ResNet-8
(b) CIFAR-ResNet-14
(c) CIFAR-ResNet-20
(d) CIFAR-ResNet-32
(e) CIFAR-ResNet-44
Figure 4: Test(left)/Training (right) accuracy for different sample sizes.
(f) CIFAR-ResNet-56
C.2 Sensitivity analysis for CIFAR100
(a) # of neighbours  $k$ .
(b) Radius parameter  $\sigma$
(c) Temperature  $T$
Figure 5: Sensitivity analysis of ResMem hyperparameters. The  $y$ -axis represents the CIFAR100 test accuracies, and the  $x$ -axis represents the sweeping of respective hyperparameters.
Varying locality parameter  $k$  and  $\sigma$ . We vary the number of neighbours from  $k = 27$  to  $k = 500$ . We find that ResMem test accuracy is relatively stable across the choice of the number of neighbours (cf. Figure 5(a)). The trend of the curve suggests that as  $k \to \infty$ , the ResMem test accuracy seems to be converging to a constant level. For  $\sigma$ , we explored different values of  $\sigma \in (0.1, 2.0)$ . We observe that the test accuracy has a unimodal shape as a function of  $\sigma$ , suggesting that there is an optimal choice of  $\sigma$  (cf. Figure 5(b)).
Varying temperature  $T$  and connection to distillation. We tried  $T = 0.1$  to  $T = 5$ , and also identified an unimodal shape for the test accuracy (Figure 5(c)). The fact that we can use different temperatures for (a) training the network and (b) constructing the  $k$ -NN predictor reminds us of the well-established knowledge distillation procedure [28]. In knowledge distillation, we first use one model (the teacher network) to generate targets at a higher temperature, and then train a second model (the student network) using the combination of the true labels and the output of the first network.
ResMem operates in a reversed direction: Here we have a second model (kNN) that learns the difference between true labels and the output of the first model. In both cases, we can tune the temperature of the first model to control how much information is retained. This connection offers an alternative perspective that regards ResMem as a "dual procedure" to knowledge distillation.
D ResMem on ImageNet
This section includes additional experiment results on applying ResMem to ImageNet dataset.
ImageNet. In addition to CIFAR100, we also evaluate the performance of ResMem on ImageNet [43]. We employ a family of pre-trained MobileNet-V2 models [45] from Keras<sup>3</sup>, with varying widths controlled by a multiplier  $a$ . For ResMem, we again use the second last layer of DeepNet as a 1280-dimensional embedding of an image and rely on the  $\ell_2$  distance between the embeddings for nearest neighbor search (Step 3, Section 4.1). We specify the ResMem parameter of  $(k,\sigma,T)$  in the table below. We repeat the experiment over several MobileNet-V2 architectures, with MobileNet-V2-a0.35 being the smallest model and MobileNet-V2-a1.3 being the largest one.
Table 2: Test accuracy for ResMem and baseline deep network for ImageNet data.
Architecture	ResMem param.	Test accuracy
k	σ	T	DeepNet	ResMem
MobileNet-V2-a0.35	10	0.6	0.4	60.2%	61.2%
MobileNet-V2-a0.5	10	0.6	0.4	65.3%	66.1%
MobileNet-V2-a0.75	10	0.8	0.6	69.6%	70.1%
MobileNet-V2-a1.0	20	0.4	0.4	71.3%	71.8%
MobileNet-V2-a1.3	30	0.4	0.4	74.7%	75.1%
We can see that (c.f. Table 2) ResMem boosts the test accuracy by  $1\%$  on the smallest model and by  $0.4\%$  on the largest model.
E Additional details of NLP experiments
The Decoder-Only model used in our experiments is essentially the normal Encoder-Decoder architecture with Encoder and Cross-Attention removed. We pretrained both the T5-small and T5-base model on C4 [42] dataset with auto-regressive language modeling task for 1,000,000 steps, with dropout rate of 0.1 and batch size of 128. The learning rate for the first 10,000 steps is fixed to 0.01 and the rest steps follow a square root decay schedule.
During the inference for retrieval key, query embeddings and residuals, we ensured every token has at least 64 preceding context by adopting a sliding window strategy, where a window of 256 token slides from the beginning to the end on each of the articles, with a stride of  $256 - 64 = 192$ .
<sup>3</sup>https://keras.io/api/applications/mobilenet/
For residuals, we only stored the top 128 residuals measured by the absolute magnitude, as the residual vector is as large as T5 vocabulary size (i.e., 32128), and storing all 32128 residuals for each token is too demanding for storage. However, when weight-combining the residuals, we zero filled the missing residuals so that all the residual vectors have 32128 elements.
F Comparison with other algorithms
We mainly compare ResMem against [31], where the algorithm uses kNN to retrieve labels directly instead of the residual of the label. In their algorithm, a key aparameter is  $\lambda \in [0, 1]$  which specifies how much weight to give to the neural network and how much for the kNN component. In the extreme case of  $\lambda = 1$ , their algorithm reduces to using kNN to memorize data directly.
For the language modeling task, we use the C4 dataset and T5-large architecture. As we change the weight [31, Equation (3)] of the DeepNet component, we find the best performing kNN-LM methods has accuracy  $44.88\%$  which is lower accuracy than the ResMem accuracy  $45.55\%$ . In particular, we obtain the table below
Table 3: Test accuracy for kNN-LM (ResMem accuracy  ${45.55}\%$  )
kNN weight Λ	0	0.2	0.4	0.5	0.6	0.8	1
kNN-LM accuracy	44.76%	44.88%	44.83%	44.66%	44.27%	42.97%	40.95%
ResMem acc. - 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\%$ .