

\section{Additional CIFAR100 Results}
\label{sec:add-cifar}
This section includes additional experiment results on applying ResMem to CIFAR100 dataset.
\subsection{Additional robustness results}
In addition to the results already presented in Section \ref{sec:exp-rob}, 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 \ref{sec:exp-setup}.
Generally, we see that ResMem yields larger improvement over the baseline DeepNet when the network is small and dataset is large.
\begin{figure*}[ht]
    \centering
    \subfigure[\label{fig:cifar-sample-resnet8-appendix}CIFAR-ResNet-8]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet8_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet8_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet14}CIFAR-ResNet-14]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet14_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet14_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet20}CIFAR-ResNet-20]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet20_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet20_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet32}CIFAR-ResNet-32]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet32_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet32_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet44}CIFAR-ResNet-44]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet44_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet44_train.pdf}}
    \subfigure[\label{fig:cifar-sample-resnet56}CIFAR-ResNet-56]{
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet56_test.pdf}
    \includegraphics[width=.24\textwidth]{figures/cifar/cifarall/cifar_resnet56_train.pdf}}
    \caption{Test(left)/Training (right) accuracy for different sample sizes.}
    \label{fig:cifar-all}
\end{figure*}


\subsection{Sensitivity analysis for CIFAR100}
\label{sec:sensitivity-cifar}
\begin{figure*}[ht]
    \centering
    \subfigure[\label{fig:cifar-k}\# of neighbours $k$.]{\includegraphics[width=.32\textwidth]{figures/cifar/cifar_k.pdf}}
    \subfigure[\label{fig:cifar-sigma}Radius parameter $\sigma$.]{\includegraphics[width=.32\textwidth]{figures/cifar/cifar_sigma.pdf}}
    \subfigure[\label{fig:cifar-temp}Temperature $T$.]{\includegraphics[width=.32\textwidth]{figures/cifar/cifar_temp.pdf}}
    \caption{
    Sensitivity analysis of ResMem hyperparameters.
    The $y$-axis represents the CIFAR100 test accuracies,
    and the
    $x$-axis represents the sweeping of respective hyperparameters.}
    \label{fig:robustness-hyper}
\end{figure*}

\paragraph{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 \ref{fig:cifar-k}).
The trend of the curve suggests that as $k\rightarrow\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 \ref{fig:cifar-sigma}).

\paragraph{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 \ref{fig:cifar-temp}).
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 \citep{distillation}.
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 \emph{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 \emph{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.

\section{Additional ImageNet results}
\label{sec:add-imagenet}
This section includes additional experiment results on applying ResMem to ImageNet dataset.

\subsection{Sensitivity analysis for ImageNet}
\label{sec:sensitivity-imagenet}
\begin{figure*}[ht]
    \centering
    \subfigure[\label{fig:imagenet-k}\# of neighbours $k$.]{\includegraphics[width=.32\textwidth]{figures/imagenet/imagenet_k.pdf}}
    \subfigure[\label{fig:imagenet-sigma}Radius parameter $\sigma$.]{\includegraphics[width=.32\textwidth]{figures/imagenet/imagenet_sigma.pdf}}
    \subfigure[\label{fig:imagenet-temp}Temperature $T$.]{\includegraphics[width=.32\textwidth]{figures/imagenet/imagenet_temp.pdf}}
    \caption{
    Sensitivity analysis of ResMem hyperparameters.
    The $y$-axis represents the imagenet100 test accuracies,
    and the
    $x$-axis represents the sweeping of respective hyperparameters.}
    \label{fig:imagenet-robustness-hyper}
\end{figure*}

We vary the number of neighbours from $k=15$ to $k=90$.
Similar to CIFAR100 finding,  ResMem test accuracy is relatively stable across the choice of the number of neighbours (cf.~Figure \ref{fig:imagenet-k}).
For $\sigma$, we explored different values of $\sigma \in (0.3, 2.1)$. 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 \ref{fig:imagenet-sigma}).
We tried $T=0.3$ to $T=2.1$, and also identified an unimodal shape for the test accuracy (Figure \ref{fig:imagenet-temp}).

\subsection{ImageNet-LT}
\label{sec:imagenet-lt}
We also evaluated the proposed ResMem on ImageNet-LT~\citep{liu2019large} -- a long-tail version of ImageNet~\citep{ILSVRC15}.
We again utilize ResNet~\citep{resnet} architecture.
In particular, we studied the performance of ResMem with both ResNet-18 and ResNet-50 as the underlying DeepNet.
We trained the DeepNet for 90 epochs by using SGD with Nesterov momentum.
We employ cosine learning rate schedule with the base learning rate of 0.4, and use 0.9 as momentum parameter.
Based on a coarse hyperparameter search, we utilized $\sigma = 2.0$ and temperature $T=.8$ in ResMem for both ResNet18 and ResNet50 as the underlying DeepNet.
As for the number of neighbors, we utlized $k = 70$ and $k = 90$ for ResNet18 and ResNet50, respectively. 

Table~\ref{table:imagenet-lt} showcases the performance of both DeepNet an ResMem on ImageNet-LT. As expected the (balanced) accuracy on ImageNet-LT is worse that that on ImageNet.
However, it is also evident from the table that our proposed ResMem leads to performance improvement even in the a large-scale long-tail setting.

\begin{table}[ht]
\centering
\caption{Test accuracy for ResMem and baseline deep network.}
\label{table:imagenet-lt}
\vskip 0.1in
\scalebox{1}{
\begin{tabular}{@{}cccc@{}} 
\toprule
\multirow{2}{*}{Dataset}  & \multirow{2}{*}{Architecture}  & \multicolumn{2}{c}{Test accuracy} \\
\cmidrule(l){3-4}
& & \nn  & \alg \\
\toprule
ImageNet & ResNet18 & 66.86\%  & \textbf{67.44\%} \\
ImageNet-LT & ResNet18 & 39.56\%  & \textbf{39.91\%} \\
ImageNet-LT & ResNet50 & 46.18\%  & \textbf{46.39\%} \\
\toprule
\end{tabular}
}
\end{table}
\section{Additional details of NLP experiments}
\label{sec:detail-nlp}

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~\citep{2020t5} dataset with auto-regressive language modeling task for 1,000,000 steps, with dropout rate of 0.1 and batch size of 128.
The learning rate for the first 10,000 steps is fixed to 0.01 and the rest steps follow a square root decay schedule.

During the inference for retrieval key, query embeddings and residuals, we ensured every token has at least 64 preceding context by adopting a sliding window strategy, where a window of 256 token slides from the beginning to the end on each of the articles, with a stride of $256 - 64 = 192$.

For residuals, we only stored the top 128 residuals measured by the absolute magnitude, as the residual vector is as large as T5 vocabulary size (i.e., 32128), and storing all 32128 residuals for each token is too demanding for storage. However, when weight-combining the residuals, we zero filled the missing residuals so that all the residual vectors have 32128 elements.

\newpage
\section{Some concentration results for uniform random variables}
In this section, we state some concentration results that are useful for the theoretical analysis in Section \ref{sec:thy}.
Let $\bxt, \bx_1, \dots, \bx_n\overset{\rm i.i.d.}{\sim}$Unif$(\cB_{\bzero, \sqrt{d+2}})$ be i.i.d. samples from the uniform distribution over the Euclidean norm ball of radius $\sqrt{d+2}$ in $\R^d$.
Let
\begin{equation}
  Z_n = \min_{\bx\in \{\bx_1, \bx_2, \dots, \bx_n\}}\|\bxt-\bx\|^2.
\end{equation}
If $n=1$, $\E Z_1$ is the sum of the variance of each coordinate of Unif$(\cB_{\bzero, \sqrt{d+2}})$.
Therefore, $\E Z_n$ provides a generalized measure of concentration.
Intuitively, $\E Z_n \rightarrow 0$ as $n\rightarrow \infty$.
The proposition below provides a upper bound on the rate of convergence.

\begin{lemma}[Nearest Neighbor concentration]
\label{lem:nearestneighbor}
Given the assumptions above
\begin{equation}
\E Z_n \lesssim  d^2 \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},
\end{equation}
where $\lesssim $ means inequality up to an universal constant independent of $d$ and $n$.
\end{lemma}

\begin{proof}
Define
\begin{equation}
\begin{aligned}
\cE_1 &= \{Z_n\leq \delta^2\},\\
\cE_2 &= \{\delta\leq \sqrt{d+2}-\|\bxt\|\}.
\end{aligned}
\end{equation}

We will compute two probabilities $\P(\cE_1|\cE_2)$ and $\P(\cE_2)$ that will be useful latter.
\begin{equation}
\begin{aligned}
\P(\cE_1^c|\cE_2) &= \P(Z_n\geq\delta^2 | \cE_2) = \P(\|\bxt-\bx_i\|\geq\delta,~\forall i|\cE_2),\\
&= \E_{\bxt} \P(\|\bxt-\bx_i\|\geq\delta|\cE_2,\bxt)^n = \E_{\bxt} (1-\P(\|\bxt-\bx_i\|\leq\delta|\cE_2, \bxt))^n,\\
&= \E_{\bxt} \left[1- \frac{\text{Vol}(\cB_{\bxt, \delta})}{\text{Vol}(\cB_{\bzero, \sqrt{d+2}})} \right]^n = \left[1 - \left(\frac{\delta}{\sqrt{d+2}}\right)^d\right]^n,\\
&\leq \exp\left[-n\left(\frac{\delta}{\sqrt{d+2}}\right)^d\right].
\end{aligned}
\end{equation}

Next, we compute $\P(\cE_2)$
\begin{equation}
\begin{aligned}
\P(\cE_2) = \P(\|\bxt\|\leq \sqrt{d+2}-\delta) = \left(\frac{\sqrt{d+2}-\delta}{\sqrt{d+2}}\right)^d = \left(1-\frac{\delta}{\sqrt{d+2}}\right)^d.
\end{aligned}
\end{equation}

We use $\cE_1$ and $\cE_2$ to compute the following upper bound
\begin{equation}
\begin{aligned}
\E Z_n &= \E(Z_n | \cE_1\cap\cE_2) \P(\cE_1\cap\cE_2) + \E (Z_n|(\cE_1\cap\cE_2)^c) P((\cE_1\cap\cE_2)^c), \\
&\leq \delta^2 + (2\sqrt{d+2})^2 \left(1 - \P(\cE_1\cap\cE_2)\right), \\
&= \delta^2 + 4(d+2) \left[1-\P(\cE_1|\cE_2)\P(\cE_2)\right].
\end{aligned} 
\end{equation}

To find an upper bound for $\E Z_n$, we need to find an upper bound for $1-\P(\cE_1|\cE_2)\P(\cE_2)$.
\begin{equation}
\begin{aligned}
1-\P(\cE_1|\cE_2)\P(\cE_2) &= 1-\left[1-\P(\cE_1^c|\cE_2)\right]\P(\cE_2), \\
&= 1 -\P(\cE_2) + \P(\cE_1^c|\cE_2)\P(\cE_2),\\
&\leq 1-\P(\cE_2) + \P(\cE_1^c|\cE_2).
\end{aligned}
\end{equation}
Now choose $\delta=\sqrt{d+2}n^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d}$.
\begin{equation}
\P(\cE_1^c|\cE_2)\leq \exp\left[-n\left(\frac{\delta}{\sqrt{d+2}}\right)^d\right] = \exp\left[-nn^{-1}\log\left(n^{1/d}\right)\right] = n^{-1/d},
\end{equation}

and
\begin{equation}
\P(\cE_2) = \left(1-\frac{\delta}{\sqrt{d+2}}\right)^d\geq 1-d\frac{\delta}{\sqrt{d+2}} = 1-dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d}.
\end{equation}
Thus
\begin{equation}
1-\P(\cE_1|\cE_2)\P(\cE_2) \leq 1 - 1 + dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d} + n^{-1/d} \lesssim dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d}.
\end{equation}
Combining everything together, we get
\begin{equation}
\begin{aligned}
\E Z_n &\leq (d+2) n^{-2/d} \left[\log\left(n^{1/d}\right)\right]^{2/d} + 4(d+2)\times dn^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d},\\
&\lesssim d^2 n^{-1/d}\left[\log\left(n^{1/d}\right)\right]^{1/d},\\
&= d^2 \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.
\end{aligned} 
\end{equation}
This completes the proof.
\end{proof}

\begin{proposition}[\cite{wainwright2019high} Corollary 6.20]
    \label{prop:matcon}
    Let $\bx_i\overset{\rm i.i.d.}{\sim}$Unif$~(\cB_{\bzero, \sqrt{d+2}})$ for $i=1, \dots, n$ be uniformly distributed over a ball of radius $B$ in $\R^d$ centered at $\bzero$.
    Let
    \[
        \bSigma_n = \frac{1}{n} \sum_{i=1}^n \bx_i\bx_i^\sT
    \]
    be the sample covariance matrix. Then
    \[
        \P(\|\bSigma_n - \bI\|_{\rm op} > \eps) \leq 2d\exp\left[ - \frac{n\eps^2}{2(d+2)(1+\eps)}\right].
    \]
\end{proposition}


\newpage
\section{Proof of Theorem \ref{thm:test-loss-ltm}}
\label{sec:lin-proof}
In this section, we present the proof of Theorem \ref{thm:test-loss-ltm}.
In Section \ref{sec:lin-decomp-detail}, we provide the detail of the decomposition of the risk into $T_1$ and $T_2$.
Then in Section \ref{sec:upb-t1} we compute an upper bound for $T_1$, and compute an upper bound for $T_2$ in Section \ref{sec:upd-t2}.
Finally, we combine everything together in Section \ref{sec:final-thm} and completes the proof.

\subsection{Decomposition of the test risk}
\label{sec:lin-decomp-detail}
\begin{equation}
\begin{aligned}
&~~~~~\E  \left[f^{\alg}(\bxt)-f_\star(\bxt)\right]^2 
    = \E  \left[f_n(\bxt)+r_n(\bxt)-f_\star(\bxt)\right]^2,\\
    &= \E  \left[ f_n(\bxt) - f_\star(\bxt)
                      -f_n(\bxnn) + f_\star(\bxnn)\right]^2,\\
    &= \E  \left[ f_n(\bxt)
                      -f_\infty(\bxt) + f_\infty(\bxt)
                      -f_\star(\bxt)
                      -f_n(\bxnn)
                      +f_\infty(\bxnn) - f_\infty(\bxnn)
                      +f_\star(\bxnn)
                 \right]^2,\\
    &\leq 3\times [ \underbrace{\E  (f_n(\bxt)-f_\infty(\bxt))^2
                        +\E  (f_n(\bxnn)-f_\infty(\bxnn))^2}_{T_1}
   
                    +\underbrace{\E  (f_\infty(\bxt)
                       -f_\star(\bxt)- f_\infty(\bxnn)
                       +f_\star(\bxnn))^2}_{T_2}
                  ],
\end{aligned}
\end{equation}
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\R$.

\ActivateWarningFilters[pdftoc]
\subsection{Upper bound on $T_1$.}
\DeactivateWarningFilters[pdftoc]

\label{sec:upb-t1}
\noindent Since $\P_\bx = \text{Unif}(\cB_{\bzero, B}~ )$, we apply the bound $\|\bxt\|, \|\bxnn\|\leq B$ to obtain
\begin{equation}
\begin{aligned}
T_1
&= \E [f_n(\bxt) - f_\infty (\bxt)]^2 + \E [f_n(\bxnn) - f_\infty(\bxnn)]^2, \\
&= \E \<\btheta_n-\btheta_\infty, \bxt\>^2 + \E \<\btheta_n - \btheta_\infty, \bxnn\>^2, \\
&\leq \E\|\btheta_n-\btheta_\infty\|^2 \|\bxt\|^2 + \E \|\btheta_n-\btheta_\infty\|^2 \|\bxnn\|^2,\\
&\leq 2B^2\E\|\btheta_n - \btheta_\infty\|^2.
\end{aligned}
\end{equation}

As $n$ gets large, the empirical covariance matrix $\bSigma_n = \bX^\sT \bX/n$ is concentrated around its mean $\bI$.
Let $\bDelta_n = \bI - \bSigma_n$ denote this deviation.
For some $\eps \in (0, 1)$, define the following ``good event'' over the randomness in $\bSigma_n$
\begin{equation}
\cA = \{\|\bDelta_n\|_{\rm op} <\eps\},
\end{equation}

where $\|\bDelta_n\|_{\rm op}$ denotes the operator norm of the deviation matrix.
The high level idea of the proof is to condition on the event $\cA$ and deduce and upper bound of $\|\btheta_n-\btheta_\infty\|$ in terms of $\eps$.
Then, we use the fact that $\cA$ happens with high probability.

Recall that $\btheta_\infty = L \btheta_\star$, and 
\begin{equation}
\btheta_n = \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \, \frac{1}{n} \|\bX\btheta-\by\|^2.
\end{equation}

Since $\by = \bX\btheta_\star$ by definition,
the Lagrangian of the convex program above is
\begin{equation}
\mathcal{L}(\btheta, \lambda) = \frac{1}{n} \|\bX\btheta-\bX\btheta_\star\|^2 + \lambda (\|\btheta\|^2 - L).
\end{equation}

The KKT condition suggests that the primal-dual optimal pair $(\btheta_n, \lambda_n)$ is given by
\begin{equation}
\begin{aligned}
\|\btheta_n\|&\leq L ,\\
\lambda_n    &\geq 0 ,\\
\lambda_n(\|\btheta_n\|-L)&=0 ,\\
\end{aligned}
\end{equation}

and at optimality
\begin{equation}
\begin{aligned}
\grad_\btheta \mathcal{L}(\btheta_n, \lambda_n) = 0 &\iff \frac{2}{n}\bX^\sT\bX(\btheta-\btheta_\star) + 2\lambda_n \btheta = 0,  \\
&\iff \btheta_n = (\bSigma_n+\lambda_n \bI)^{-1}\bSigma_n\btheta_\star.\\
\end{aligned}
\end{equation}

The complementary slackness condition 
$\lambda_n(\|\btheta_n\|-L)=0$ suggests that either $\lambda_n=0$ or $\|\btheta_n\|=L$.
But if $\lambda_n = 0$, the stationary condition $\grad_\btheta \mathcal{L}(\btheta, \lambda) = 0$ would suggest that 
$\btheta_n = \bSigma_n^{-1}\bSigma_n\btheta_\star = \btheta_\star \Rightarrow \|\btheta_n\|=1>L,$ a contradiction.
(Note that here $\bSigma_n$ is invertible condition on the event $\cA$.)
Therefore, we must have $\|\btheta_n\|=L$. 
As a result, the primal and dual pair $(\btheta_n, \lambda_n)$ is determined by the system of equations
\begin{equation}
\begin{cases}
\btheta_n &= (\bSigma_n + \lambda_n \bI)^{-1} \bSigma_n \btheta_\star, \\
\|\btheta_n\| &= L, \\
\lambda_n &> 0.
\end{cases}
\end{equation}
Next, we proceed to compute the deviation $\|\btheta_n - \btheta_\infty\|$.
\begin{equation}
\begin{aligned}
\btheta_n
&= \left[ (\lambda_n+1)\bI - \bDelta_n\right]^{-1} \bSigma_n \btheta_\star, \\
&= (\lambda_n+1)^{-1} \left[ \bI - \frac{\bDelta_n}{\lambda_n+1}\right]^{-1}\bSigma_n \btheta_\star, \\
&= (\lambda_n+1)^{-1} \left[ \bI + \sum_{k=1}^{\infty} \frac{\bDelta_n^k}{(\lambda_n+1)^k}\right](\bI-\bDelta_n) \btheta_\star,\\
&=(\lambda_n+1)^{-1} \left[ \bI + \sum_{k=1}^\infty \frac{\bDelta_n^k}{(\lambda_n+1)^k} - \bDelta_n - \sum_{k=1}^\infty \frac{\bDelta_n^{k+1}}{(\lambda_n+1)^k} \right]\btheta_\star,\\
&=(\lambda_n+1)^{-1}\btheta_\star + (\lambda_n+1)^{-1} \bDelta_n \left[\sum_{k=1}^\infty \frac{\bDelta_n^{k-1}}{(\lambda_n+1)^k} - \bI - \sum_{k=1}^\infty \frac{\bDelta_n^{k}}{(\lambda_n+1)^k} \right]\btheta_\star,\\
&= (\lambda_n+1)^{-1}\btheta_\star + (\lambda_n+1)^{-1} \bDelta_n\left[\sum_{k=1}^{\infty} \frac{\bDelta_n^{k-1}-\bDelta_n^k}{(\lambda_n+1)^k} - \bI\right]\btheta_\star.
\end{aligned}
\end{equation}

Define
\begin{equation}
\bD_n = \bDelta_n\left[\sum_{k=1}^{\infty} \frac{\bDelta_n^{k-1}-\bDelta_n^k}{(\lambda_n+1)^k} - \bI\right].
\end{equation}

Then $\btheta_n = (\lambda_n+1)^{-1}\btheta_\star + (\lambda_n+1)^{-1}\bD_n\btheta_\star$, and
\begin{equation}
\begin{aligned}
\|\bD_n\| &\leq \|\bDelta_n\| \left[1+\sum_{k=1}^{\infty} \frac{\|\bDelta_n\|^{k-1}+\|\bDelta_n\|^k}{(\lambda_n+1)^k}\right] ,\\
&\leq \eps \left[1+2(1+\lambda_n)^{-1}\sum_{k=1}^\infty \left(\frac{\eps}{1+\lambda_n}\right)^k\right],\\
&= \eps\left(1 + \frac{2}{1+\lambda_n} \frac{1}{1 - \frac{\eps}{1+\lambda_n}}\right) \leq 3\eps.
\end{aligned}
\end{equation}

Therefore
\begin{equation}
\begin{aligned}
&L=\|\btheta_n\|^2 = (\lambda_n+1)^{-2} + (\lambda_n+1)^{-2} \btheta_\star^\sT \bD_n^\sT \bD_n \btheta_\star + 2(\lambda_n+1)^{-2} \btheta_\star \bD_n \btheta_\star, \\
\Rightarrow & (\lambda_n+1)^2L^2 = 1 + \delta_n,~ \delta_n = \btheta_\star^\sT\bD_n^\sT\bD_n\btheta_\star + 2\btheta_\star^\sT\bD_n\btheta_\star.
\end{aligned}
\end{equation}

We can obtain the following bound for $\delta_n$:
\begin{equation}
|\delta_n| \leq \|\btheta_\star\|^2 \|\bD_n\|^2 + 2\|\btheta_\star\|^2 \|\bD_n\| \leq 9\eps^2 + 6\eps \leq 15\eps.
\end{equation}

Since $1-\delta_n/2\leq \sqrt{1+\delta_n}\leq 1+\delta_n/2$, and $|\delta_n|\leq 15\eps$, we obtain
\begin{equation}
|(\lambda_n+1)L-1|\leq \frac{15\eps}{2} \Rightarrow \left|L-(\lambda_n+1)^{-1}\right|\leq \frac{15\eps}{2}(\lambda_n+1)^{-1} \leq \frac{15\eps}{2},
\end{equation}

where the last inequality follows as we have $\lambda_n > 0$.
Finally,
\begin{equation}
\begin{aligned}
\btheta_n -\btheta_\infty~~~ &= (\lambda_n+1)^{-1}\btheta_\star - L \btheta_\star + (\lambda_n+1)^{-1}\bD_n\btheta_\star,\\
\Rightarrow \|\btheta_n-\btheta_\infty\|^2 &= [(1+\lambda_n)^{-1}-L]^2 + (1+\lambda_n)^{-2} \btheta_\star \bD_n^\sT\bD_n\btheta_\star + 2(\lambda_n+1)^{-1}[(1+\lambda_n)^{-1}-L] \btheta_\star \bD_n \btheta_\star,\\
&\leq 64\eps^2 + 9\eps^2 + 45\eps^2 = 118\eps^2, \\
\Rightarrow \|\btheta_n - \btheta_\infty\|^2 &\lesssim \eps^2.
\end{aligned}
\end{equation}

Combine the above result with Proposition \ref{prop:matcon}, we get that

\begin{equation}
\begin{aligned}
\E\|\btheta_n - \btheta_\infty\|^2 &= \E (\|\btheta_n - \btheta_\infty\|^2|\cA)\P(\cA) + \E(\|\btheta_n - \btheta_\infty\|^2|\cA^c)\P(\cA^c), \\
&\leq \eps^2 + 4L^2 \times 4d\exp\left[ -\frac{n\eps^2}{2(d+2)(1+\eps)} \right], \\
\end{aligned}
\end{equation}

If we choose $\eps = n^{-1/3}$, we get
\begin{equation}
\E\|\btheta_n - \btheta_\infty\|^2 \lesssim dL^2n^{-2/3},
\end{equation}

which implies that
\begin{equation}
T_1 \lesssim d^2L^2n^{-2/3}.
\end{equation}

\ActivateWarningFilters[pdftoc]
\subsection{Upper bound on $T_2$.}
\DeactivateWarningFilters[pdftoc]

\label{sec:upd-t2}
Plugging in the formula for $f_\perp(\bxt) = f_\star(\bxt)-f_\infty(\bxt) = \<\bxt, \btheta_\perp\>$, we get
\begin{equation}
\begin{aligned}
  T_2 &= \E [f_\perp(\bxnn)-f_\perp(\bxt)]^2,\\
  &= \E\<\btheta_\perp, \bxnn-\bxt\>^2,\\
  &\leq (1-L)^2\|\btheta_\star\|^2 \E \|\bxt-\bxnn\|^2,\\
  &= (1-L)^2 \E \|\bxt-\bxnn\|^2,
\end{aligned}
\end{equation}

where in the last inequality, we used the relation that $\btheta_\perp = (1-L)\btheta_\star$.
Proposition \ref{lem:nearestneighbor} suggests that
\begin{equation}
\E \|\bxt-\bxnn\|^2 \lesssim d^2 \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},
\end{equation}

which implies 
\begin{equation}
T_2 \lesssim d^2(1-L)^2  \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.
\end{equation}

\begin{remark}[Comparison with pure nearest neighbor and ERM]
If we rely solely on nearest neighbor method, the prediction error is 
\begin{equation}
  \E[f_\star(\bxt)-f_\star(\bxnn)]^2 = \E \<\bxt-\bxnn, \btheta_\star\>^2 \leq \E \|\bxt-\bxnn\|^2.
\end{equation}

On the other hand, if we solely rely on ERM, even with infinite sample, we get
\begin{equation}
  \E[f_\star(\bxt)-f_\infty(\bxt)]^2 = \E \<\bxt, \btheta_\star-\btheta_\infty\>^2 \leq (1-L)^2\E \|\bxt\|^2.
\end{equation}

We can see from the upper bound that $\alg$ takes advantage of both 
\begin{itemize}
    \item Projecting $f_\star$ onto $f_\infty$, so that the dependence on the prediction function is reduced from $1$ to $(1-L)^2$.
    \item Memorizing the residuals using nearest neighbor, so that the variance is reduced from $\E\|\bxt\|^2$ to $\E\|\bxnn-\bxt\|^2$.
\end{itemize}
\end{remark}

\ActivateWarningFilters[pdftoc]
\subsection{Test loss for $\alg$.}
\DeactivateWarningFilters[pdftoc]
\label{sec:final-thm}

If we combine the previous two parts together, we get
\begin{equation}
\begin{aligned}
\E \left[\hat{f}(\bxt) - f_\star(\bxt)\right]^2 &\lesssim d^2L^2n^{-2/3} + d^2(1-L)^2  \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d}.
\end{aligned}
\end{equation}

\noindent This completes the proof of Theorem~\ref{thm:test-loss-ltm}.
\section{Introduction}
\input{intro}

\input{method}

\section{Empirical results}
\label{sec:exp}
\input{experiments}

\section{Theoretical results}
\label{sec:thy}
\input{theory}

\section{Discussion and future work}
\label{sec:diss}
\input{discussion}

\section*{Acknowledgements}
\input{ack}

\subsection{Experimental setup}
\label{sec:exp-setup}

\paragraph{CIFAR100.}
For the CIFAR100~\citep{cifar}, we use CIFAR-ResNet8 \citep{resnet} as the base DeepNet.
For the 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 image embedding, which has dimension 64.
For the nearest neighbor search (Step 3, Section \ref{sec:method}), we define the distance between two images to be the $\ell_2$ distance between their embeddings.
We use $\sigma=0.7$, $k=53$, and $T=1.4$ to compute the weights for the nearest neighbor regressor.
We provide the sensitivity analysis of test accuracy against ResMem parameters in Appendix \ref{sec:add-cifar} (cf. Figure \ref{fig:robustness-hyper}).
We can see that ResMem boosts the test accuracy of CIFAR-ResNet8 from 56.46\% to 59.66\%, which is between the base DeepNet test accuracy for CIFAR-ResNet8 and CIFAR-ResNet14 (cf. Figure \ref{fig:cifar-arch}).

\paragraph{ImageNet.}
The experimental setting for ImageNet~\citep{ILSVRC15} is similar to CIFAR100. 
We employ a ResNet-18 model \citep{resnet} and employ the following training procedure.
We train the DeepNet for 120 epochs with the batch size of 90. We use an initial learning rate equal $0.8$ and an annealing schedule, decreasing the learning rate 10 times after 30, 60 and 80 epochs.
We also use warm-up for the first 5 epochs of training, linearly increasing the learning rate from 0.1 up to its initial learning rate.
For ResMem, we again use the second last layer of DeepNet as a $512$-dimensional embedding of an image and rely on the $\ell_2$ distance between the embeddings for nearest neighbor search (Step 3, Section \ref{sec:method}).
We use $\sigma=1$, temperature $T=1$ and $k=80$ to compute the weights for the nearest neighbor regressor.
We provide a similar sensitivity analysis of ResMem hyperparameter for ImageNet in Appendix \ref{sec:sensitivity-imagenet} (cf. Figure \ref{fig:imagenet-robustness-hyper}).
The ResMem improvement on ImageNet is relatively small compared with CIFAR100, since the base ResNet18 test accuracy at 66.86\% is already high compared with other base DeepNet accuracies.
In addition to the standard ImageNet dataset, we also evaluate ResMem in ImageNet-LT~\citep{liu2019large}, a long-tailed version of ImageNet, and observed similar improvement (cf. Appendix \ref{sec:imagenet-lt}).

\paragraph{Natural language experiments}
For the language experiment, we use a Decoder-Only T5-small~\citep{2020t5} model and C4~\citep{2020t5} dataset.
C4 is generated from scraping the internet and commonly used as a pretraining dataset or part of the pretraining mix.
We pre-trained the DeepNet on C4 training split with 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 $k$NN search, and we created the query embeddings using the whole validation split and the same representation location. 
For each query, we retrieved 50 neighbors with $L_2$ distance using ScaNN~\citep{avq_2020}.
We used the temperature $T = 1$ for the residual computation and $\sigma = 1$ for computing the neighbor weights. 
The predicted token is the one with highest probability, similar to greedy decoding, and we measured the prediction accuracy to match the vision experiments.
The result in Table~\ref{table:main-table} shows that ResMem improved test accuracy by 2.86\% which is around the accuracy (40.08\%) of a T5-base Decoder-Only model trained in a similar way. Please see Appendix~\ref{sec:detail-nlp} for details on base model training, data processing etc.


\subsection{Where does the improvement come from?}
\label{sec:emp-ana}

\begin{figure*}[t]
    \centering
    \includegraphics[width=\textwidth]{figures/tpr.pdf}
    \caption{Examples from CIFAR100 and C4 test set with the property that 
    \textbf{(i)} $y^\textsf{ResMem}$ is correct;
    \textbf{(ii)} $y^\textsf{DeepNet}$ is wrong but close in meaning.
    We use red to denote the ResMem prediction and blue to denote DeepNet prediction.
    The DeepNet predictions capture \emph{coarse} structure (e.g., predicting \texttt{poppy} for a sample whose true label is \texttt{rose}),
    which can be refined by ResMem capturing the remaining \emph{fine-grained} structure.}
    \label{fig:cifar-nlp-analysis}
\end{figure*}


In this section, we discuss the accuracy improvement in Table \ref{table:main-table}.
Let ${\sf Gain}_{\sf ResMem}$ be the difference between the test accuracy of ResMem and baseline DeepNet:
\[
{\sf Gain}_{\sf ResMem} = L_{01}(f_\alg) - L_{01}(f_\nn),
\]
where $L_{01}$ is the misclassification error as defined in equation \eqref{eqn:L01}.
We offer a decomposition of ${\sf Gain}_{\sf ResMem}$ that sheds light into the mechanism behind ResMem.
For a test set $\{(x_i, y_i)\}_{i=1}^m$, let $y_i^\alg$ be the ResMem prediction on instance $x_i$ and let $y_i^\nn$ be the baseline neural network prediction on $x_i$.
When $y^\alg_i = y^\nn_i$, sample $x_i$ does not contribute to ${\sf Gain}_{\sf ResMem}$.
When $y^\alg_i\neq y^\nn_i$, this could arise 
either from the desirable event where the deep network misclassifies while ResMem classifies correctly;
or from the undesirable event where the ResMem misclassifies, while the deep network classifies correctly. 
These can be summarized by the \textsf{\small TPR} (\textit{true positive rate})
and \textsf{\small FPR} (\textit{false positive rate})
respectively:
\begin{equation}
    \textsf{\small TPR} = \frac{1}{m}\sum_{i=1}^m \mathds{1}\{y^\nn_i\neq y_i ~\text{and}~ y^\alg_i = y_i\}.
\end{equation}
\begin{equation}
    \textsf{\small FPR} = \frac{1}{m}\sum_{i=1}^m \mathds{1}\{y^\nn_i= y_i ~\text{and}~ y^\alg_i \neq y_i\}.
\end{equation}

Note that~${\sf Gain}_{\sf ResMem} = \textsf{\small TPR} - \textsf{\small FPR}.$
Table~\ref{tab:TPR_FPR} summarizes $\textsf{\small TPR}$, $\textsf{\small FPR}$, and ${\sf Gain}_{\sf ResMem}$ for the three main experiments discussed in Table~\ref{table:main-table}.
The decomposition of ${\sf Gain}_{\sf ResMem}$ says that the gain of ResMem came from the \textsf{\small TPR} samples, provided they outweigh the \textsf{\small FPR} samples.


\begin{table}[ht]
\centering
\caption{True positive rate (TPR) and false positive rate (FPR) for vision and language experiments in Table~\ref{table:main-table}.}
\label{tab:TPR_FPR}
\vskip 0.1in
\scalebox{1}{
\begin{tabular}{@{}ccccc@{}} 
\toprule
Dataset & Architecture & \textsf{\small TPR} & \textsf{\small FPR} & ${\sf Gain}_{\sf ResMem}$ \\
\toprule
CIFAR100 & ResNet8  & 5.89\%  & 2.70\%  & \textbf{3.19\%} \\
ImageNet & ResNet18 & 3.15\%  & 2.57\%  & \textbf{0.58\%} \\
C4-subset & T5-small & 5.37\%  & 2.44\%  & \textbf{2.93\%} \\
\toprule
\end{tabular}
}
\vskip -0.1in
\end{table}

\paragraph{Analysis of \textsf{\small TPR} samples}
Focusing on the test samples where ResMem indeed helps ($y^\nn_i\neq y_i$ and $y^\alg_i = y_i$), we identify a common underlying pattern: 
while the DeepNet makes an incorrect prediction, it still captures some coarse structure.
For example, in the CIFAR100, 
one sample has correct label
$
y_i = y^\alg_i = \texttt{rose},
$
but baseline DeepNet predicts
$
y^\nn_i = \texttt{poppy}
$
(cf. Figure~\ref{fig:cifar-nlp-analysis}).
In this case, both \texttt{rose} and \texttt{poppy} are flowers.
We find similar behavior for the language modeling task (cf. Figure~\ref{fig:cifar-nlp-analysis}).

The empirical analysis suggests the DeepNet in isolation can already learn some large scale structures, but is unable to make fine-grained distinctions.
This is where ResMem helps: \emph{ResMem helps memorize information in the training label that the DeepNet cannot learn.}

\paragraph{Additional insights from the decomposition.}
We discuss an alternative way of choosing hyperparameters inspired by the decomposition of ${\sf Gain}_{\sf ResMem}$ into \textsf{\small TPR} and \textsf{\small FPR}.
In this paper, we choose the ResMem hyperparameters that minimizes the test error on the validation set or, equivalently, maximize ${\sf Gain}_{\sf ResMem}$.
With the decomposition of ${\sf Gain}_{\sf ResMem}$, an alternative hyperparameter selection procedure is to use
the following optimization problem:
\[
{\rm maximize}_{\textsf{\small FPR}(\texttt{hyperparam.})<0.05} \textsf{\small TPR}(\texttt{hyperparam.}),
\]
which ensures that ResMem modifies the predictions of DeepNet in a more conservative manner.
In particular, bounding \textsf{\small FPR} implies that ResMem has minimal impact on the examples where DeepNet already makes correct predictions.
At the same time, a higher value of \textsf{\small TPR} corresponds to maximization of the desirable occurrences where ResMem can correct a wrong prediction by DeepNet.

\subsection{When does ResMem help?}
\label{sec:exp-rob}

\begin{figure*}[t]
    \centering
    \subfigure[\label{fig:cifar-sample}Test(left)/Training (right) accuracy for different sample sizes.]{
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_subset.pdf}
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_subset_train.pdf}}
    \subfigure[\label{fig:cifar-arch}Test(left)/Training (right) accuracy for different architectures.]{
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_arch.pdf}
    \includegraphics[width=.45\textwidth]{figures/cifar/cifar_arch_train.pdf}}
    \caption{
    Robustness of ResMem improvement with respect to training sample size and deep network architecture.
    \textbf{(a):} Using $10\%, 20\%, \dots, 100\%$ of training data.
    \textbf{(b):} Using progressively larger CIFAR-ResNet architecture.
    }
    \label{fig:robustness-sample-arch}
\end{figure*}

In this subsection, we present several experiments on CIFAR100 that demonstrates the robustness of the ResMem improvement.

\paragraph{Varying sample size.}
The main experiment~(cf.~Table~\ref{table:main-table}) uses the entire CIFAR100 training data.
We repeat the experiment with subsets ($10\%, 20\%, \dots, 100\%$) of CIFAR100 training data (subsampled uniformly across different classes).
On the left (right) of Figure \ref{fig:cifar-sample}, we report the test (training) accuracy of ResMem and baseline DeepNet.
As a sanity check, we can see that ResMem always achieves perfect training accuracy, and the DeepNet training accuracy decreases as samples increase (since it's harder to fit larger dataset).
We can see that ResMem yields \emph{progressively larger margin of improvement when more data is used}.
This trend suggests a desirable property of ResMem: in real problems where the dataset is extremely large, ResMem is expected to bring even greater benefit.

\paragraph{Varying the DeepNet size.}
We again consider the main CIFAR100 experiment, which was only for a CIFAR-ResNet8 model.
We repeat the ResMem evaluation for different models in the CIFAR-ResNet family, i.e., CIFAR-ResNet-$\{8, 14, 20, 32, 44, 56\}$.
We use the same choice of training hyperparameter for all six CIFAR-ResNet models, so that only the choice of model architecture changes.
DeepNet has progressively better training accuracy as model gets larger, which is intuitive (cf. Figure \ref{fig:cifar-arch}).
We find that ResMem yields the \emph{largest margin of improvement when the smallest architecture is applied}, i.e., when it is difficult for the DeepNet itself to interpolate/memorize the data.


\section{Background and Notation}

In this section, we fix notation and discuss the relevant literature that motivates us to propose the ResMem algorithm.


\subsection{Notation}
We consider multi-class classification problems over instances $\mathscr{X}$ and labels $\mathscr{Y} \stackrel{.}{=} \{1, 2, \dots, L\} = [L]$.
Given training examples $S = \{ ( x_i, y_i ) \}_{i \in [n]} \in ( \mathscr{X} \times \mathscr{Y} )^n$, 
the goal is to learn a scorer $f \colon \mathscr{X} \to \mathbb{R}^L$ that,
given an instance,
assigns an affinity score for each label.
Such an $f$ should minimize the \emph{misclassification error} on test samples:
\begin{equation}
\label{eqn:L01}
L_{01}( f ) \stackrel{.}{=} \Pr_{(x, y)}( y \neq {\tt pred}( f( x ) ) ), 
\end{equation}
where ${\tt pred}( z ) \stackrel{.}{=} \argmax_{y' \in [L]} z_{y'}$,
and $\Pr$ is the distribution over labelled instances.
To achieve this, one typically minimizes the \emph{empirical loss}
\[
\hat{L}_\ell( f )\stackrel{.}{=} \frac{1}{n} \sum_{i \in [n]} \ell( y_i, f( x_i ) ),
\]
where $\ell \colon [L] \times \mathbb{R}^L \to \mathbb{R}_+$ is a loss function.
Ideally, one would like to use 
$\ell_{01}( y, f( x ) ) \stackrel{.}{=} 1( y \neq {\tt pred}( f( x ) ) )$;
for computational tractability, it is popular to instead use a \emph{surrogate loss}, such as the softmax cross-entropy:
$$ \ell( y, f( x ) ) = -\log \frac{\exp( f_y( x ) )}{\sum_{y' \in [L]} \exp( f_{y'}( x ) )}. $$


\subsection{Memorization for generalization: prior work}

\paragraph{Memorization is sufficient for generalization.}
Overparameterised neural models with many more parameters than training samples have the capacity to perfectly fit (or \emph{interpolate}) even random training labels~\citep{Zhang:2017};
i.e., they can drive the empirical loss $\hat{L}_\ell( f ) \to 0$ for \emph{any} training set.
At the same time, when trained on real-world datasets, 
increasing model complexity tends to \emph{improve} model performance~\citep{Neyshabur:2019, pmlr-v119-yang20j};
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~\citep{Arpit:2017} and theoretically~\citep{Bartlett:2017,Dziugaite:2017,Brutzkus:2018,Belkin:2018,Neyshabur:2019,Liang:2020,Montanari:2020,Bartlett:2020,Vapnik:2021,Wang:2021b, pmlr-v139-yang21a}.
One recurring message from the theory is that memorization (in the form of interpolation) can be sufficient for generalization.

\paragraph{Memorization can be necessary for generalization.}
Some recent works~\citep{Feldman2020,CJK2022Mem} showed that memorization
---
either in the sense of interpolation, or in a more general sense of stability~\citep{Feldman2020_DiscoveringLT}
---
may be necessary for generalization.
\citet{Feldman:2019} considered a setting where the label distribution exhibits a \emph{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.
\citet{CJK2022Mem} 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). 

\section{Our method: ResMem}
\label{sec:method}

Prior works suggests that \emph{implicit} memorization through increase in model size leads to better generalization.
When such implicit memorization is infeasible
---
e.g., because of the increased computational cost of using a large model
---
is it possible to perform \emph{explicit} memorization and get generalization benefits?\\

\noindent This motivates us to the propose the ResMem algorithm (c.f. Figure \ref{fig:cartoon}), which operates as follows:

\begin{tcolorbox}[colback=black!10]
\begin{center}
    {\large \textbf{ResMem} --- Residual Memorization}
\end{center}
\begin{enumerate}
    \item \textbf{Train the base DeepNet.}
    Train a neural network $f_\nn$ on the training samples $S$ as usual.
    
    %
    \item \textbf{Prepare the residual data.} Compute the \emph{residual} \footnote{For an overparameterized network that perfectly fits the training sample, the residuals will all be $0$. However, 
    we are interested in employing
    smaller networks, which will likely have non-zero residuals.} prediction of each training example as
    \[
    r_i = \textsf{onehot}(y_i) - \textsf{softmax}(f_\nn(x_i)/T),~\forall~i\in [n],
    \]
    where $\textsf{onehot} \colon \mathscr{Y} \to \R^L$ is the standard encoding that maps the label to a probability vector.
    Here, $T$ is a hyperparameter corresponding to the ``temperature'' scaling of the softmax operation.
    Then, we employ the output of an intermediate layer of the base network $f_\nn$, denoted by $z_i = \phi(x_i)$, as an embedding for the training instance $x_i$.
    These embeddings are utilized for the nearest neighbor search in the next step.
    %
    \item \textbf{Predict via memorized residuals.}
    To make a prediction on a test sample $\tilde{x}\in\mathscr{X}$, first compute its embedding $\tilde{z} = \phi(\tilde{x})$. 
    Then, use soft $k$-nearest neighbor method to build a function $r_{\knn}$ defined by weights $\overline{w}_i(\tilde{x})$:
    \begin{equation}
        r_{\knn}(\tilde{x}) = \sum_{i=1}^n \overline{w}_i(\tilde{x}) \cdot r_i.
    \end{equation}
    The weights $\overline{w}_i(\tilde{x})$ satisfy $\sum_{i}\overline{w}_i(\tilde{x}) = 1$, and
    are computed from raw weights $w_i$ as follows:
    \[
    w_i = \exp(-\|\tilde{z}-z_i\|_2/\sigma),
    \quad
    \overline{w}_i(\tilde{x}) \propto \mathds{1}\left(w_i\geq w_{(k)}\right) w_i. 
    \]
    Here $w_{(k)}$ represents the $k$-th largest entry of $w_i$'s, and $k, \sigma$ are two hyperparameters that collectively controls the locality of nearest neighbor search.
\end{enumerate}

\noindent We make the final prediction based on the following scorer: 
\begin{equation}
\label{eqn:resmem-pred-classification}
f_\alg(\tilde{x}) = \textsf{softmax}(f_\nn(\tilde{x})/T) + r_{\knn}(\tilde{x}).   
\end{equation}
\end{tcolorbox}

\begin{remark}[Classification versus regression]
The above introduces ResMem in a classification setting,
which serves as our primary focus in the experiments in Section \ref{sec:exp}.
That said, we emphasize that ResMem can easily be extended to other supervised learning problems beyond classification with neural networks.
For example, in Section \ref{sec:thy}, we analyze ResMem in a regression setting. 
\end{remark}

\begin{remark}[Explicit memorization]
Smaller $k$ or $\sigma$ corresponds to putting higher weight on residuals of the closest neighboring training examples.
For sufficiently small $k$ and $\sigma$, $f_\alg$ achieves exact memorization of the training sample, i.e., ${\tt pred}(f_\alg(x_i))=y_i$ for all $i \in [n]$.
\end{remark}


\subsection{Applicable scenarios of ResMem}
We posit that $f_\alg$ yields the largest margin of improvement over $f_\nn$ when the learning task is ``hard'' (and thus \emph{implicit} memorization is infeasible).
We discuss three such scenarios below.
Each of our main empirical or theoretical results roughly corresponds to one of these settings.

\begin{itemize}[leftmargin=5mm, itemsep=1mm, partopsep=0pt,parsep=0pt]
    %
    \item \textbf{Hard generalization.} 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 (cf.~Section \ref{sec:thy}).

    %
    \item \textbf{Hard interpolation.} Here, the number of samples is large enough to make training set interpolation (achieving zero training error) infeasible for a given neural network model.
    In contrast, ResMem can circumvent the large sample size by explicitly memorizing the training samples. 
    For example, current large language models (LLMs) are typically trained for at most a single epoch over trillions of examples~\citep{Chowdhery:2022}.
    %
    We emulate this scenario by considering a causal language modeling task (cf.~Section \ref{sec:exp}).   

    %
    \item \textbf{Efficiency constraint.} 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 \ref{sec:exp}).
\end{itemize}

\subsection{Relation to existing algorithms}

\paragraph{Nearest neighbor method.} The nearest neighbor~\citep{NNCover, KnuthAOP, Hastie:2001, locallearning} method assigns label to a test sample based on the label of its nearest neighbor(s) in the training set.
Owing to its \emph{simplicity}, \emph{flexibility} in defining input similarity, and \emph{computational efficiency} with various approximation schemes~\citep{Gionis1999SimilaritySI, Muja2009FastAN}, this method remains popular.
However, the performance of nearest neighbor drops as data becomes high dimensional~\citep{nnrates, Muja2009FastAN}.
Therefore, to apply it to high dimensional data such as image and text~\citep{knnsvm}, one approach is to learn a representation of data using neural networks~\citep{pmlr-v2-salakhutdinov07a}.
Following this approach, \citep{arXivDnnKnn} finds that applying nearest neighbor methods directly to memorize the training labels $y_i$ yields similar performance with the original softmax based neural network classification.
In contrast, ResMem applies nearest neighbor to memorize the \emph{residual} $r_i$ of the training samples.

\paragraph{Boosting and residual fitting.}
Boosting algorithms such as AdaBoost~\citep{Freund:1997} 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 \emph{forward stage-wise regression}~\citep{Friedman:2000}, a classical procedure from statistics~\citep{Goldberger:1961,Alley:1987,Tibshirani:2015}.
Intuitively, at each round, one fits the residual of the ensemble of weak learners constructed thus far so as to minimize the exponential loss $\ell_{\rm exp}( y, f( x ) ) = e^{-y \cdot f( x ) }$.
Concretely, for binary classification problems,
AdaBoost fits a scorer 
$f( x ) = \sum_{t = 1}^{T} \alpha_t \cdot h_t( x )$,
where $\alpha_t \in \mathbb{R}$,
$h_t \colon \mathscr{X} \to \{ \pm 1 \}$ is a weak classifier (e.g., a decision stump),
and
$T$ is the number of rounds of fitting.
This fitting is performed iteratively, with $(\alpha_t, h_t)$ learned during the $t$-th iteration.


\paragraph{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~\citep{Lample:2019,Guu:2020,Borgeaud:2021,li2022decoupled}.
Closer to our work,~\citet{Khandelwal:2020} employ a linear combination of neural network and $k$-nearest neighbour classifier components.
However, our work is quite different in that our $k$-nearest neighbor components memorizes the \emph{residual} of the DeepNet prediction, whereas \citet{Khandelwal:2020} memorizes the target label directly.
Various forms of memory have also been considered in generic classification problems~\citep{Panigrahy:2021,Vapnik:2021,Wang:2022}.
This line of literature again differs from ResMem in that their memory tries to memorize labels directly, whereas ResMem memorizes the \emph{residuals}, leading to a natural combination of the neural network and the memory component.


\paragraph{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~\citep{Hinton:2015, Bucilua:2006} is a popular strategy for compressing a large neural model to a smaller one.
For a survey of other effective strategies, including pruning, see~\citet{Menghani:2021}.
In Appendix \ref{sec:sensitivity-cifar} we discuss that ResMem can be regarded as a ``dual procedure'' of distillation.
\subsection{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 \emph{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.

\begin{assumption}[Distribution of covariates]
\label{ass:lin-dist}
The distribution of covariate $\bx\in\R^d$, denoted by $\P_\bx$, is the uniform distribution over a Euclidean norm ball centered at the origin of radius $\sqrt{d+2}$. The choice of radius ensures that $\E_{\bx\sim\P_\bx} \bx\bx^\sT = \bI$.
\end{assumption}

\begin{assumption}[Linear regression over norm ball]
\label{ass:lin-prob-def}
Consider the problem of learning a linear function $f_\star(\bx) = \<\bx, \btheta_\star\>$ with $\|\btheta_\star\|=1$ from training data $\{(\bx_i, y_i)\}_{i=1:n}$ where $\bx_i\overset{\rm i.i.d.}{\sim} \P_\bx$ and $y_i=f_\star(\bx_i)$ using the function class
\begin{equation}\label{eqn:linfunc}
  \cF = \{\bx\mapsto\<\bx,\btheta\>, \|\btheta\|<L\}.
\end{equation}
We assume $L<1$ so that the problem belongs to the ``hard generalization'' scenario discussed in Section \ref{sec:method}, where the hypothesis space is inadequate to fit the function on its own.
\end{assumption}

ResMem proceeds by first learning a linear function $f_n(\bx) = \<\btheta_n, \bx\>$ from $\cF$ through empirical risk minimization (ERM):
\begin{equation}
    \btheta_n = \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \, \frac{1}{n} \sum_{i=1}^n \left[\<\bx_i, \btheta\>-y_i\right]^2.
\end{equation}
The empirical risk minimizer $f_n$ should be thought of as the analog of $f_\nn$ in the deep learning context.
It defines a ground-truth residual function $r_\star(\bx)=f_\star(\bx)-f_n(\bx)$.
Now we fix a test covariate $\bxt\sim\P_x$.
ResMem ``memorizes'' the residual function through the $1$-nearest neighbor to $\bxt$
\begin{equation}
    r_n(\bxt) = r_\star(\bxnn) = f_\star(\bxnn) - f_n(\bxnn),
\end{equation}
where $\bxnn$ is the nearest neighbor to $\bxt$ among the training covariates $\bx_1,\dots,\bx_n$:
\[
 \bxnn = \underset{\bx\in\{\bx_1, \dots, \bx_n\}}{\operatorname{argmin}} \, \|\bx - \bxt\|.
\]
The final prediction is
\begin{equation}
\label{eqn:lin-pred-rule}
f_{n}^{\alg}(\bxt) = f_n(\bxt) + r_n(\bxt).
\end{equation}
Observe that if $\bxt$ coincides with any training sample, 
$f_{n}^{\alg}(\bxt) = f_\star(\bxt)$, 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.


\subsection{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{aligned}
\btheta_\infty
    &= \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \,
       \E_{\bx\sim\P_x}\left[\<\btheta, \bx\> - \<\btheta_\star, \bx\>\right]^2,\\
    &= \underset{\|\btheta\|\leq L}{\operatorname{argmin}} \,
       \|\btheta-\btheta_\star\| = L \btheta_\star,
\end{aligned}
\]
which is what ERM learns in the limit of $n\rightarrow\infty$.
We can think of  $\btheta_\infty$ as the best function that ERM can learn.
Then, we can decompose $\btheta_\star$ into $\btheta_\star = \btheta_\infty + \btheta_\perp$, where $\btheta_\perp=\btheta_\star - \btheta_\infty$.
This decomposition can be generalized beyond linear regression.
Since $\btheta_\infty$ defines a function $f_\infty(\bx) = \<\bx, \btheta_\infty\>$, 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.


\subsection{A decomposition of the prediction error}
\label{sec:lin-risk-decomp}

We now introduce a decomposition of the prediction risk that reveals how $\alg$ algorithm boosts generalization.
Note that the prediction error of $\alg$ is
\begin{equation}
\label{eqn:lin-test-loss}
\E  \left[\left(f_{n}^{\alg}(\bxt)-f_\star(\bxt)\right)^2\right].
\end{equation}
It can be decomposed into two components:
\[
\begin{aligned}
&\E  \left[f_{n}^{\alg}(\bxt)-f_\star(\bxt)\right]^2  \leq\\
&3\times[~ \underbrace{\E  (f_n(\bxt)-f_\infty(\bxt))^2
                    +\E  (f_n(\bxnn)-f_\infty(\bxnn))^2}_{T_1}
+\underbrace{\E  (f_\infty(\bxt)
                   -f_\star(\bxt)- f_\infty(\bxnn)
                   +f_\star(\bxnn))^2}_{T_2}
              ~].
\end{aligned}
\]
We provide the detail of the decomposition in Section \ref{sec:lin-decomp-detail}.
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{as}~ n\rightarrow\infty.\]

On the other hand, $T_2$ arises due to the limited capacity of $\cF$.
It captures an irreducible error of the risk, which in general is \textbf{not} asymptotically zero.
However, because of the explicit memorization $\alg$ algorithm introduces ($\bxnn\rightarrow\bxt$ as $n\rightarrow\infty$), we also have
\[T_2\rightarrow 0 ~\text{as}~ 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.

\subsection{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 of the result below can be found in Appendix \ref{sec:lin-proof}.

\begin{theorem}[Risk for $\alg$ algorithm]
\label{thm:test-loss-ltm}
For the problem defined in Assumption \ref{ass:lin-prob-def} with covariates distribution in Assumption \ref{ass:lin-dist}, the $\alg$ prediction rule $f_{n}^{\alg}(\bxt)$ defined in equation \eqref{eqn:lin-pred-rule} achieves risk \eqref{eqn:lin-test-loss}
\[
\begin{aligned}
\E \left[f_{n}^{\alg}(\bxt) - f_\star(\bxt)\right]^2
\lesssim d^2L^2n^{-2/3} + d^2(1-L)^2  \left[\frac{\log\left(n^{1/d}\right)}{n}\right]^{1/d},
\end{aligned}
\]
where $\lesssim$ denotes inequality up to a universal constant independent of $d, n$ and $L$.
\end{theorem}

The result includes contribution from two terms introduced in Section \ref{sec:lin-risk-decomp}:
\begin{itemize}
    \item $T_1\lesssim d^2L^2n^{-2/3}$ that arises due to the difference between $f_n$ and $f_\infty$.
    \item $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 $\bxnn\rightarrow\bxt$.
\end{itemize}

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 theoretical 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 \ref{thm:test-loss-ltm}, 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 enjoy the fast rate of $n^{-2/3}$ for smaller sample sizes.
