SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\hat {L} _ {\ell} (f) \doteq \frac {1}{n} \sum_ {i \in [ n ]} \ell \left(y _ {i}, f \left(x _ {i}\right)\right),
$$
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))\doteq 1(y\neq$  pred  $(f(x)))$  ; for computational tractability, it is popular to instead use a surrogate loss, such as the softmax cross-entropy.
Given the notation above, ResMem operates as follows:
(a) Test(left)/Training (right) v.s. architectures.
Figure 2: ResMem improvement on CIFAR100 with respect to training sample size and deep network architecture. (a): Using progressively larger CIFAR-ResNet architecture. (b): Using  $10\%$ ,  $20\%$ , ...,  $100\%$  of training data.
(b) Test(left)/Training (right) acc. v.s. sample size.
1. Train the base DeepNet. Train a neural network  $f_{\text{DeepNet}}$  on the training samples  $S$  as usual.
2. Prepare the residual data. Compute the residual  ${}^{2}$  prediction of each training example as
$$
r _ {i} = \operatorname {o n e h o t} \left(y _ {i}\right) - \operatorname {s o f t m a x} \left(f _ {\text {D e e p N e t}} \left(x _ {i}\right) / T\right), \forall i \in [ n ],
$$
where onehot:  $\mathcal{Y}\to \mathbb{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_{\mathrm{DeepNet}}$ , 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.
3. Predict via memorized residuals. To obtain a prediction on a test sample  $\widetilde{x} \in \mathcal{X}$ , first compute its embedding  $\widetilde{z} = \phi(\widetilde{x})$ . Then, use soft  $k$ -nearest neighbor method to build a function  $r_{\mathrm{kNN}}$  defined by weights  $\overline{w}_i(\widetilde{x})$ :
$$
r _ {\mathrm {k N N}} (\widetilde {x}) = \sum_ {i = 1} ^ {n} \bar {w} _ {i} (\widetilde {x}) \cdot r _ {i}. \tag {7}
$$
The weights  $\overline{w}_i(\widetilde{x})$  satisfy  $\sum_{i}\overline{w}_{i}(\widetilde{x}) = 1$ , and are computed from raw weights  $w_{i}$  as follows:
$$
w _ {i} = \exp (- \| \widetilde {z} - z _ {i} \| _ {2} / \sigma), \quad \overline {{w}} _ {i} (\widetilde {x}) \propto \mathbb {1} \left(w _ {i} \geq w _ {(k)}\right) w _ {i},
$$
where  $w_{(k)}$  represents the  $k$ -th largest entry of  $w_{i}$ 's. Note that  $k$  and  $\sigma$  are two hyperparameters that collectively controls the locality of nearest neighbor search.
We make the final prediction based on the following scorer:
$$
f _ {\text {R e s M e m}} (\widetilde {x}) = \operatorname {s o f t m a x} \left(f _ {\text {D e e p N e t}} (\widetilde {x}) / T\right) + r _ {\mathrm {k N N}} (\widetilde {x}). \tag {8}
$$
Remark 4.1 (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_{\mathrm{ResMem}}$  achieves exact memorization of the training sample, i.e.,  $\mathsf{pred}(f_{\mathrm{ResMem}}(x_i)) = y_i$  for all  $i \in [n]$ .
Remark 4.2 (Computation cost). The vision experiments have moderate training sample size, so we perform exact nearest neighbor search and discuss the computation cost in Section 4.2. For language experiments, the training sample size is so large that the exact nearest neighbor computation is inflexible, so we rely on approximate nearest neighbor search discussed in Section 4.2.
In this subsection, we mainly consider image classification task on ResNet [27] with CIFAR100 [33] dataset. We provide additional ImageNet [43] results in Appendix D.
Setup. We use CIFAR-ResNet-\{8, 14, 20, 32, 44, 56\} as the base DeepNet. For all six DeepNet training, we use SGD with batch size 128, trained for 256 epochs. We use a peak learning rate 0.4, and momentum 0.9. We warm up the learning rate linearly for the first 15 epochs, and decay the learning rate by 0.1 after epochs \{96, 192, 224\}. For ResMem, we use the pre-logit layer as the
For an overparameterized network that perfectly fits the training sample, the residuals will all be 0. However, we are interested in either smaller networks or extremely large datasets where implicit memorization is infesible.
vResMem vDeepNet
rose
poppy
language modeling
vResMem vDeepNet
height length
home residence
...allow for plenty of headroom inside whilst still being less than 2.5m in height.
Graphic now consists of all cities with greater than 30,000 locals. Acquiring a home in Spain...
cup
plate
squirrel
rabbit
March June
butterfly
bee
possible feasible
palm
pine
tree
tree
Figure 3: Examples from CIFAR100 and C4 test set with the property that (i)  $y^{\text{ResMem}}$  is correct; (ii)  $y^{\text{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 coarse structure (e.g., predicting poppy for a sample whose true label is rose), which can be refined by ResMem capturing the remaining fine-grained structure.
Filmed around 7:30-8:30 a.m. on Friday, March 9, 2012.
...that will not affect the superior quality of your job. That is possible because we understand how to save...
...answer your questions and schedule the initial meeting. We consistently arrive at the scheduled hour...
consistently always
image embedding, which has dimension 64. For the nearest neighbor search (Step 3, Section 4.1), 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 C (cf. Figure 5).
Results. The results for CIFAR-ResNet-\{8, 14, 20, 32, 44, 56\} are reported in Figure 2(a). 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. To access the statistical reliability of the improvement, we repeat the CIFAR-ResNet-8 experiment 5 times over random initialization of DeepNet etc. We and that the average ResMem accuracy is  $59\%$  with standard deviation  $0.7\%$ , and the average DeepNet accuracy is  $56.5\%$  with standard deviation  $0.8\%$ .
Computationally, we estimate the CPU latency of a CIFAR-ResNet-8 to be  $15.9\mathrm{ms}$  for a single test image. By contrast, the  $k$ -NN step takes  $4.8\mathrm{ms}$  for the same test image. To contextualize the latency cost, the total cost of ResMem with ResNet-8  $(15.9\mathrm{ms} + 4.8\mathrm{ms})$  is lower than the cost of the next-sized model, i.e., ResNet-14  $(26.2\mathrm{ms})$ . Regarding the memory cost, for a batch size of 1 and images of size  $32\times 32$ , a ResNet-8 (68K params) requires 2.5MB, while a ResNet-14 (128K params) requires 4MB. Embeddings from a ResNet-8 and ResNet-14 are both 64 dimensional. To embed the entire CIFAR100 training set (50K examples) requires 15MB of disk space.
Varying sample size. We repeat the above experiment on CIFAR-ResNet-8 with subsets  $(10\%, 20\%, \ldots, 100\%)$  of CIFAR100 training data (subsampled uniformly across different classes). The size of the index set for nearest-neighbor search is the same as the training set for base neural networks (e.g., model with  $10\%$  CIFAR100 data also uses the same  $10\%$  data for nearest-neighbor search). On the left (right) of Figure 2(b), 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 datasets). We can see that ResMem yields 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.
4.3 Language modeling
Setup. For the language experiment, we use a Decoder-Only T5-{small, large} [42] model and C4 [42] 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 kNN 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 approximate nearest neighbor search algorithm ScaNN [24]. 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.
Results. On T5-small, ResMem boosts test accuracy from  $38.01\%$  to  $40.87\%$ , which is around the accuracy  $(40.08\%)$  of a T5-base model without ResMem. On T5-large, ResMem boosts the test accuracy from  $44.8\%$  to  $45.6\%$ . This demonstrates that with explicit memorization, we may leverage smaller base language models while reaping the performance benefits of large language models. Computationally, as the index set is quite large (1.6 billion tokens), exact k-nearest neighbor search is infeasible. So we use the approximate nearest neighbor search algorithm ScaNN [24] to reduce compute time. Please see Appendix E for details on base model training and data processing.
4.4 Where does the improvement come from?
In this section, we identify test samples that contributes to the accuracy improvement of CIFAR100 with CIFAR-ResNet-8 and C4 with T5-small. Let  $\text{Gain}_{\text{ResMem}}$  be the difference between the test accuracy of ResMem and baseline DeepNet:
$$
\operatorname {G a i n} _ {\text {R e s M e m}} = L _ {0 1} \left(f _ {\text {R e s M e m}}\right) - L _ {0 1} \left(f _ {\text {D e e p N e t}}\right),
$$
where  $L_{01}$  is the misclassification error as defined in equation (6). We offer a decomposition of  $\text{Gain}_{\text{ResMem}}$  that sheds light into the mechanism behind ResMem. For a test set  $\{(x_i, y_i)\}_{i=1}^m$ , let  $y_i^{\text{ResMem}}$  be the ResMem prediction on instance  $x_i$  and let  $y_i^{\text{DeepNet}}$  be the baseline neural network prediction on  $x_i$ . When  $y_i^{\text{ResMem}} = y_i^{\text{DeepNet}}$ , sample  $x_i$  does not contribute to  $\text{Gain}_{\text{ResMem}}$ . When  $y_i^{\text{ResMem}} \neq y_i^{\text{DeepNet}}$ , this could arise either from the desirable event where the deep network mis-classifies while ResMem classifies correctly; or from the undesirable event where the ResMem mis-classifies, while the deep network classifies correctly. These can be summarized by the TPR (true positive rate) and FPR (false positive rate) respectively:
$$
\mathrm {T P R} = \frac {1}{m} \sum_ {i = 1} ^ {m} \mathbb {1} \left\{y _ {i} ^ {\text {D e e p N e t}} \neq y _ {i} \text {a n d} y _ {i} ^ {\text {R e s M e m}} = y _ {i} \right\}. \tag {9}
$$
$$
\mathsf {F P R} = \frac {1}{m} \sum_ {i = 1} ^ {m} \mathbb {1} \left\{y _ {i} ^ {\text {D e e p N e t}} = y _ {i} \text {a n d} y _ {i} ^ {\text {R e s M e m}} \neq y _ {i} \right\}. \tag {10}
$$
Note that  $\text{Gain}_{\text{ResMem}} = \text{TPR} - \text{FPR}$ . The decomposition of  $\text{Gain}_{\text{ResMem}}$  says that the gain of ResMem came from the TPR samples, provided they outweigh the FPR samples.
On CIFAR-ResNet-8, we find  $\mathrm{TPR} = 5.89\%$  and  $\mathrm{FPR} = 2.70\%$ , leading to  $\mathrm{Gain}_{\mathrm{ResMem}} = 3.19\%$ . On T5-small with C4 validation split, we find  $\mathrm{TPR} = 5.37\%$  and  $\mathrm{FPR} = 2.44\%$ , leading to  $\mathrm{Gain}_{\mathrm{ResMem}} = 2.93\%$ .
Analysis of TPR samples Focusing on the test samples where ResMem helps  $(y_{i} = y_{i}^{\text{ResMem}} \neq y_{i}^{\text{DeepNet}})$ , we identify a common underlying pattern: while the DeepNet makes an incorrect prediction, it still captures some coarse structure. For example, in CIFAR100, one sample has correct label  $y_{i} = y_{i}^{\text{ResMem}} = \text{rose}$ , but the DeepNet predicts  $y_{i}^{\text{DeepNet}} = \text{poppy}$ , i.e., the label of a different type of flower. (cf. Figure 3). We find similar behavior for the language modeling task (cf. Figure 3).
This 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: ResMem helps memorize information in the training label that the DeepNet cannot learn.
Additional insights from the decomposition. In this paper, we choose the ResMem hyperparameters that minimizes the test error on the validation set or, equivalently, maximize  $\text{Gain}_{\text{ResMem}}$ . Inspired by the decomposition of  $\text{Gain}_{\text{ResMem}}$ , we propose an alternative hyperparameter selection procedure based on the following optimization problem:
$$
\text {m a x i m i z e} _ {\mathsf {F P R} (\text {h y p e r p a r a m .}) <   0. 0 5} \mathsf {T P R} (\text {h y p e r p a r a m .}),
$$
which ensures that ResMem modifies the DeepNet predictions in a more conservative manner. In particular, bounding FPR implies that ResMem has minimal impact on the examples where DeepNet already makes correct predictions. At the same time, a higher value of TPR corresponds maximizing the desirable occurrences where ResMem can correct a wrong prediction by DeepNet.
5 Discussion and future works
Joint training of kNN and DeepNet. The current formulation of ResMem builds the base DeepNet and kNN components sequentially. Consequently, the DeepNet is trained completely oblivious to the fact that there is a subsequent kNN model that will memorize its residuals. A natural direction of future work is to consider the joint training of DeepNet and kNN, so that the models can dynamically interact during training to determine which portion of label is for DeepNet to learn, and the remaining is for kNN to memorize.
To explore the role of training during the first stage, we re-evaluate the CIFAR-ResNet-8 experiment by stopping DeepNet training at different epochs (Table 1). We can see that when the #epoch is small,
Table 1: Comparison of DeepNet and ResMem accuracy over epochs on CIFAR-ResNet-8 experiment.
#epoch	128	160	192	224	256
DeepNet acc.	34.0%	56.2%	55.6%	57.2%	56.6%
ResMem acc.	49.3%	60.2%	58.6%	59.2%	59.5%
ResMem has a dramatic improvement in accuracy. One of the key roles of the first training phase is to learn good representations of the training data so the nearest neighbor retrieval is performed on more meaningful representations. This simple experiments suggests that the proposed direction has the potential to dramatically reduce the training time of DeepNet – while obtaining similar test accuracy with the help of ResMem.
Calibration of ResMem. A potential problem with applying ResMem to classification is scorer mis-calibration. The output of the ResMem prediction vector  $f_{\mathrm{ResMem}}(x)$  (8) is not guaranteed to lie on the probability simplex. This is not an issue when we only care about the predicted class membership, since we take the argmax of  $f_{\mathrm{ResMem}}(x)$ . However, this limitation hinders us to access the confidence of the ResMem prediction. To remedy this, a possible future work is to consider alternative notions of residual. For example, we can do memorization in the logit space instead of the probability space. Then, the one-hot encoding of the true label may be replaced by class mean when defining the residual.
Distribution shift. Finally, ResMem can be a promising approach to tackle test-time covariate shift. The nearest neighbor modifies the prediction of DeepNet based on the training covariate that are closer to the test covariate, making the algorithm more adaptive to the specific test covariate [46].
Part of the work is done while Zitong Yang is at Google Research, New York. We would like to thank Chong You, Yu Sun, Yaodong Yu and anonymous reviewers for their feedback on the final draft. Zitong Yang would like to thank Shuangping Li for discussion regarding the proof of Lemma A.1. Zitong Yang would also like to acknowledge the support of Albion Walter Hewlett Stanford Graduate Fellowship.
A Some concentration results for uniform random variables
In this section, we state some concentration results that are useful for the theoretical analysis in Section 3. Let  $\widetilde{\boldsymbol{x}},\boldsymbol{x}_1,\ldots ,\boldsymbol{x}_n\stackrel {\mathrm{i.i.d.}}{\sim}\mathrm{Unif}(\mathcal{B}_{0,\sqrt{d + 2}})$  be i.i.d. samples from the uniform distribution over the Euclidean norm ball of radius  $\sqrt{d + 2}$  in  $\mathbb{R}^d$ . Let
$$
Z _ {n} = \min  _ {\boldsymbol {x} \in \left\{\boldsymbol {x} _ {1}, \boldsymbol {x} _ {2}, \dots , \boldsymbol {x} _ {n} \right\}} \| \widetilde {\boldsymbol {x}} - \boldsymbol {x} \| ^ {2}. \tag {11}
$$
If  $n = 1$ ,  $\mathbb{E}Z_{1}$  is the sum of the variance of each coordinate of  $\mathrm{Unif}(\mathcal{B}_{0,\sqrt{d + 2}})$ . Therefore,  $\mathbb{E}Z_{n}$  provides a generalized measure of concentration. Intuitively,  $\mathbb{E}Z_{n}\to 0$  as  $n\to \infty$ . The proposition below provides a upper bound on the rate of convergence.
Lemma A.1 (Nearest Neighbor concentration). Given the assumptions above
$$
\mathbb {E} Z _ {n} \lesssim d ^ {2} \left[ \frac {\log (n ^ {1 / d})}{n} \right] ^ {1 / d}, \tag {12}
$$
where  $\lesssim$  means inequality up to an universal constant independent of  $d$  and  $n$ .
$$
\mathcal {E} _ {1} = \left\{Z _ {n} \leq \delta^ {2} \right\}, \tag {13}
$$
$$
\mathcal {E} _ {2} = \{\delta \leq \sqrt {d + 2} - \| \tilde {\boldsymbol {x}} \| \}.
$$
We will compute two probabilities  $\mathbb{P}(\mathcal{E}_1|\mathcal{E}_2)$  and  $\mathbb{P}(\mathcal{E}_2)$  that will be useful latter.
$$
\begin{array}{l} \mathbb {P} \left(\mathcal {E} _ {1} ^ {\mathrm {c}} \mid \mathcal {E} _ {2}\right) = \mathbb {P} \left(Z _ {n} \geq \delta^ {2} \mid \mathcal {E} _ {2}\right) = \mathbb {P} \left(\left\| \widetilde {\boldsymbol {x}} - \boldsymbol {x} _ {i} \right\| \geq \delta , \forall i \mid \mathcal {E} _ {2}\right), \\ = \mathbb {E} _ {\widetilde {\boldsymbol {x}}} \mathbb {P} \left(\| \widetilde {\boldsymbol {x}} - \boldsymbol {x} _ {i} \| \geq \delta \mid \mathcal {E} _ {2}, \widetilde {\boldsymbol {x}}\right) ^ {n} = \mathbb {E} _ {\widetilde {\boldsymbol {x}}} (1 - \mathbb {P} (\| \widetilde {\boldsymbol {x}} - \boldsymbol {x} _ {i} \| \leq \delta \mid \mathcal {E} _ {2}, \widetilde {\boldsymbol {x}})) ^ {n}, \\ = \mathbb {E} _ {\widetilde {x}} \left[ 1 - \frac {\operatorname {V o l} \left(\mathcal {B} _ {\widetilde {x} , \delta}\right)}{\operatorname {V o l} \left(\mathcal {B} _ {\mathbf {0} , \sqrt {d + 2}}\right)} \right] ^ {n} = \left[ 1 - \left(\frac {\delta}{\sqrt {d + 2}}\right) ^ {d} \right] ^ {n}, \tag {14} \\ \leq \exp \left[ - n \left(\frac {\delta}{\sqrt {d + 2}}\right) ^ {d} \right]. \\ \end{array}
$$
Next, we compute  $\mathbb{P}(\mathcal{E}_2)$
$$
\mathbb {P} \left(\mathcal {E} _ {2}\right) = \mathbb {P} \left(\| \widetilde {\boldsymbol {x}} \| \leq \sqrt {d + 2} - \delta\right) = \left(\frac {\sqrt {d + 2} - \delta}{\sqrt {d + 2}}\right) ^ {d} = \left(1 - \frac {\delta}{\sqrt {d + 2}}\right) ^ {d}. \tag {15}
$$
We use  $\mathcal{E}_1$  and  $\mathcal{E}_2$  to compute the following upper bound
$$
\begin{array}{l} \mathbb {E} Z _ {n} = \mathbb {E} \left(Z _ {n} \mid \mathcal {E} _ {1} \cap \mathcal {E} _ {2}\right) \mathbb {P} \left(\mathcal {E} _ {1} \cap \mathcal {E} _ {2}\right) + \mathbb {E} \left(Z _ {n} \mid \left(\mathcal {E} _ {1} \cap \mathcal {E} _ {2}\right) ^ {c}\right) P \left(\left(\mathcal {E} _ {1} \cap \mathcal {E} _ {2}\right) ^ {c}\right), \\ \leq \delta^ {2} + (2 \sqrt {d + 2}) ^ {2} (1 - \mathbb {P} (\mathcal {E} _ {1} \cap \mathcal {E} _ {2})), \tag {16} \\ = \delta^ {2} + 4 (d + 2) \left[ 1 - \mathbb {P} (\mathcal {E} _ {1} | \mathcal {E} _ {2}) \mathbb {P} (\mathcal {E} _ {2}) \right]. \\ \end{array}
$$
To find an upper bound for  $\mathbb{E}Z_{n}$ , we need to find an upper bound for  $1 - \mathbb{P}(\mathcal{E}_1|\mathcal{E}_2)\mathbb{P}(\mathcal{E}_2)$ .
$$
\begin{array}{l} 1 - \mathbb {P} \left(\mathcal {E} _ {1} | \mathcal {E} _ {2}\right) \mathbb {P} \left(\mathcal {E} _ {2}\right) = 1 - \left[ 1 - \mathbb {P} \left(\mathcal {E} _ {1} ^ {\mathrm {c}} \mid \mathcal {E} _ {2}\right) \right] \mathbb {P} \left(\mathcal {E} _ {2}\right), \\ = 1 - \mathbb {P} \left(\mathcal {E} _ {2}\right) + \mathbb {P} \left(\mathcal {E} _ {1} ^ {c} \mid \mathcal {E} _ {2}\right) \mathbb {P} \left(\mathcal {E} _ {2}\right), \tag {17} \\ \leq 1 - \mathbb {P} (\mathcal {E} _ {2}) + \mathbb {P} (\mathcal {E} _ {1} ^ {c} | \mathcal {E} _ {2}). \\ \end{array}
$$
Now choose  $\delta = \sqrt{d + 2} n^{-1 / d}\left[\log \left(n^{1 / d}\right)\right]^{1 / d}$ .
$$
\mathbb {P} \left(\mathcal {E} _ {1} ^ {c} \mid \mathcal {E} _ {2}\right) \leq \exp \left[ - n \left(\frac {\delta}{\sqrt {d + 2}}\right) ^ {d} \right] = \exp \left[ - n n ^ {- 1} \log \left(n ^ {1 / d}\right) \right] = n ^ {- 1 / d}, \tag {18}
$$
$$
\mathbb {P} \left(\mathcal {E} _ {2}\right) = \left(1 - \frac {\delta}{\sqrt {d + 2}}\right) ^ {d} \geq 1 - d \frac {\delta}{\sqrt {d + 2}} = 1 - d n ^ {- 1 / d} \left[ \log \left(n ^ {1 / d}\right) \right] ^ {1 / d}. \tag {19}
$$
$$
1 - \mathbb {P} \left(\mathcal {E} _ {1} \mid \mathcal {E} _ {2}\right) \mathbb {P} \left(\mathcal {E} _ {2}\right) \leq 1 - 1 + d n ^ {- 1 / d} \left[ \log \left(n ^ {1 / d}\right) \right] ^ {1 / d} + n ^ {- 1 / d} \lesssim d n ^ {- 1 / d} \left[ \log \left(n ^ {1 / d}\right) \right] ^ {1 / d}. \tag {20}
$$
$$
\begin{array}{l} \mathbb {E} Z _ {n} \leq (d + 2) n ^ {- 2 / d} \left[ \log \left(n ^ {1 / d}\right) \right] ^ {2 / d} + 4 (d + 2) \times d n ^ {- 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}, \tag {21} \\ = d ^ {2} \left[ \frac {\log (n ^ {1 / d})}{n} \right] ^ {1 / d}. \\ \end{array}
$$
Proposition A.2 ([49] Corollary 6.20). Let  $\boldsymbol{x}_i \stackrel{\mathrm{i.i.d.}}{\sim} \text{Unif}(\mathcal{B}_{\mathbf{0},\sqrt{d + 2}})$  for  $i = 1, \dots, n$  be uniformly distributed over a ball of radius  $B$  in  $\mathbb{R}^d$  centered at  $\mathbf{0}$ . Let
$$
\boldsymbol {\Sigma} _ {n} = \frac {1}{n} \sum_ {i = 1} ^ {n} \boldsymbol {x} _ {i} \boldsymbol {x} _ {i} ^ {\intercal}
$$
be the sample covariance matrix. Then
$$
\mathbb {P} (\| \pmb {\Sigma} _ {n} - \pmb {I} \| _ {\mathrm {o p}} > \varepsilon) \leq 2 d \exp \left[ - \frac {n \varepsilon^ {2}}{2 (d + 2) (1 + \varepsilon)} \right].
$$
In this section, we present the proof of Theorem 3.3. In Section B.1, we provide the detail of the decomposition of the risk into  $T_{1}$  and  $T_{2}$ . Then in Section B.2 we compute an upper bound for  $T_{1}$ , and compute an upper bound for  $T_{2}$  in Section B.3. Finally, we combine everything together in Section B.4 and completes the proof.
B.1 Decomposition of the test risk