SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In Fig. 6(a) we consider a binary USPS dataset consisting of the classes for the digits 3 and 5. Using  $|\mathcal{M}| = 16$ , we show the first and second approximations in Eq. 43 both correlate well with the truth obtained by removing a group and retraining the model. In Fig. 6(b) we do the same on MNIST with  $|\mathcal{M}| = 64$ , where we see similar trends. For the experiment on binary USPS in Fig. 6(a), we train a MLP with three hidden layers with 30 neurons each. The original model  $\theta_{*}$  is trained for 500 epochs with a learning-rate of  $10^{-3}$ , a batch size of 32 and a  $L_{2}$ -regularization parameter  $\delta = 5$ . It has  $100\%$  training accuracy and  $94.8\%$  test accuracy. For the leave-group-out retraining to obtain  $\theta_{*}^{\backslash M}$ , we initialize the model at  $\theta_{*}$ , use a learning-rate of  $10^{-3}$  and train for 1000 epochs. For the MNIST result in Fig. 6(b) we use the MLP (500, 300) model with the same hyperparameters as for
(a) MLP on USPS-3vs5,  $|\mathcal{M}| = 16$
(b) MLP on MNIST,  $|\mathcal{M}| = 64$
Figure 6: Panel (a) and Panel (b) show that the estimated deviation for removal of groups of examples correlates well with the true deviations obtained by retraining. Each marker corresponds to a removed group of examples. The red circles show the second approximation in Eq. 43. In Panel (a), we additionally show (with blue squares) the first approximation of Eq. 43. We see that the second approximation is quite accurate in this case.
$\pmb{\theta}_{*}$  in Table 2. For  $\pmb{\theta}_{*}^{\backslash M}$ , we initialize the model at  $\pmb{\theta}_{*}$  and use a cosine schedule of the learning-rate from  $10^{-4}$  to  $10^{-5}$  over 500 epochs. We do not use data augmentation. Similarly to the experiments on per-example removal, we use a K-FAC approximation.
Dataset	Model	E	LR*	LR*min	LR\C	LR\Cmin
MNIST	MLP (500, 300)	500	10-2	10-3	10-4	10-5
MNIST	LeNet5	300	10-1	10-3	10-5	10-6
FMNIST	MLP (32, 16)	300	10-2	10-3	10-5	10-6
FMNIST	LeNet5	300	10-1	10-3	10-4	10-5
Table 3: Hyperparameters for the class removal experiments in Fig. 3(a) and Fig. 11(d).  $B$ ,  $E$  and  $LR$  denote batch size, training epochs and learning-rates. The superscripts * and  $^{C}$  indicate hyperparameters for training on all data and warmstarted leave-one-class-out retraining, respectively.  $LR_{\mathrm{min}}$  is the minimum learning-rate of the cosine scheduler.
I.3 Details of "Predicting the effect of class removal on generalization"
For the FMNIST experiment in Fig. 3(a), we use the MLP (32, 16) and LeNet5 models. For the MNIST experiment in Fig. 11(d), we use the MLP (500, 300) and LeNet5 models. The hyperparameters are given in Table 3. The MLP on MNIST has a training accuracy of  $99.9\%$  and a test accuracy of  $98.4\%$ . When using LeNet5, the training and test accuracies are  $99.2\%$  and  $99.1\%$ . On FMNIST, the LeNet5 has an accuracy of  $95.0\%$  on the training set, and an accuracy of  $91.2\%$  on the test set. On the same dataset, the MLP has a training accuracy of  $89.9\%$  and a test accuracy of  $86.2\%$ . For all models, we use a regularization parameter of 100 and a batch size of 256. The leave-one-class-out training is run for 1000 epochs and the rest of the training setup is same as the previous experiment.
I.4 Details of "Estimating the leave-one-out cross-validation curves for hyperparameter tuning"
The details of the training setup are in Table 2. Fig. 7 is the same as Fig. 4 but additionally shows the test errors. For visualization purposes, each plot uses a moving average of the plotted lines with a smoothing window. Other training details are similar to previous experiments. All models are
(a) MNIST, MLP
(b) FMNIST, LeNet5
Figure 7: Leave-one-out estimation with sensitivities obtained from MPE (Train-LOO-MPE) can accurately estimate the LOO-CV curve for predicting generalization and tuning of the  $L_{2}$ -regularization parameter on MNIST, FMNIST and CIFAR-10.
(c) CIFAR10, CNN
trained from scratch where we use Adam for FMNIST, AdamW [36] for CIFAR10, and SGD with a momentum parameter of 0.9 for MNIST. We use a cosine learning-rate scheduler to anneal the learning-rate. The other hyperparameters are similar to the settings of the models trained on all data from the leave-one-out experiments in Table 2, except for the number of epochs for CIFAR10 where we train for 150 epochs. Similarly to App. I.2, we use a Kronecker-factored Laplace approximation for variance computation and do not employ data augmentation during training.
Dataset	Model	Number of δs	Range	Smoothing window
MNIST	MLP (500, 300)	96	10^0 – 10^3	3
FMNIST	LeNet5	96	10^1 – 10^3	5
CIFAR10	CNN	30	10^1 – 10^3	3
Table 4: Experimental settings for Fig. 4.
I.5 Details of "Predicting generalization during the training"
Details of the training setup: The experimental details, including test accuracies at the end of training, are listed in Table 5. We use a grid search to determine the regularization parameter  $\delta$ . The learning-rate is decayed according to a cosine schedule. For diagonal-GGN-LOO and K-FAC-LOO, we use the SGD optimizer with an exception on the FMNIST dataset where we use the AdamW optimizer [36]. In that experiment, we use a weight decay factor of  $\delta / N$  replacing the explicit  $L_{2}$ -regularization term in the loss in Eq. 1. The regularizer  $\mathcal{R}(\theta)$  is set to zero. We do not use training data augmentation. For all plots, the LOO-estimate is evaluated periodically during the training, which is indicated with markers.
Additional details on hyperparameters of iBLR are as follows, where  $h_0$  is the initialization of the Hessian:
- MNIST, MLP (32, 16):  $h_0 = 0.1$
- MNIST, LeNet5:  $h_0 = 0.1$
FMNIST, LeNet5:  $h_0 = 0.1$
CIFAR10,CNN:  $h_0 = 0.05$
CIFAR10, ResNet20:  $h_0 = 0.01$
We set  $\beta_{1} = 0.9$  and  $\beta_{2} = 0.99999$  in all of those experiments. The magnitude of the prediction variance can depend on  $h_0$ , which therefore can influence the magnitude of the sensitivities that are perturbing the function outputs in the LOO estimate of Eq. 16. We choose  $h_0$  on a grid of four values [0.01, 0.05, 0.1, 0.5] to obtain sensitivities that result in a good prediction of generalization performance.
Additional Results: In Fig. 8, we show additional results for MNIST and CIFAR10 that are not included in the main text. For MNIST, we evaluate both on a the MLP (32, 16) model and a LeNet5 architecture. For the additional CIFAR10 results, we use the CNN. In Fig. 9 we include an additional
Dataset	Model	Method	LR	\( LR_{\text{min}} \)	B	δ	Test acc.
MNIST	MLP (32, 16)	iBLR	\( 10^{-2} \)	\( 10^{-4} \)	256	80	95.6%
diag.-GGN-LOO	\( 10^{-3} \)	\( 10^{-4} \)	256	80	95.8%
K-FAC-LOO	\( 10^{-3} \)	\( 10^{-4} \)	256	80	95.8%
MNIST	LeNet5	iBLR	\( 10^{-2} \)	\( 10^{-4} \)	256	60	97.5%
diag.-GGN-LOO	\( 10^{-3} \)	\( 10^{-4} \)	256	60	97.4%
K-FAC-LOO	\( 10^{-3} \)	\( 10^{-4} \)	256	60	97.4%
FMNIST	LeNet5	iBLR	\( 10^{-1} \)	0	256	60	90.7%
diag.-GGN-LOO	\( 10^{-2} \)	\( 10^{-4} \)	256	60	91.0%
K-FAC-LOO	\( 10^{-2} \)	\( 10^{-4} \)	256	60	91.0%
CIFAR10	CNN	iBLR	\( 10^{-1} \)	\( 10^{-4} \)	512	250	81.0%
diag.-GGN-LOO	\( 10^{-1} \)	0	512	250	75.4%
K-FAC-LOO	\( 10^{-1} \)	0	512	250	73.6%
CIFAR10	ResNet-20	iBLR	\( 2*10^{-1} \)	0	50	10	83.4%
Table 5: Experimental settings for predicting generalization during the training in Fig. 1(b), Fig. 5 and Fig. 8.  $B$  and  $E$  denote the batch-size and training epochs, respectively.  $LR$  and  $LR_{\mathrm{min}}$  are the start and end learning-rates of the cosine scheduler.  $\delta$  is the regularization parameter. The specification in brackets in the third column indicates the method for computing sensitivities. We use either iBLR or SGD with diagonal GGN (diag.GGN) or K-FAC.
Figure 8: These plots are similar to Fig. 5 but for different model-data pairs. The three rows correspond to MLP on MNIST, LeNet5 on MNIST, and CNN on CIFAR10, respectively. The trends are almost same as those discussed in the main text.
experiment where the model overfits. The K-FAC-LOO estimate deteriorates in this case, but we can still use the LOO as a diagnostic for detecting overfitting and as a stopping criterion. We train
(a) Diagonal-GGN-LOO
(b) K-FAC-LOO
Figure 9: Additional results for training with AdamW where we observe overfitting. We see that K-FAC-LOO deteriorates when the model starts to overfit. Both the LOO measures can still be useful tools for diagnosing overfitting. Details of training setup are given in Table 6
Figure 10: LOO-CV estimates with Adam using the measure suggested in Table 1. The two rows correspond to a batch size of 8 and a batch size of 32, respectively. A smaller batchsize generally decreases the gap between the test NLL and the estimate. Details of the training setup are given in Table 7.
a LeNet5 on FMNIST with AdamW and predict generalization. The trend of the estimated NLL matches the trend of the test NLL in the course of training.
In Fig. 10, we include further results for sensitivity estimation with the Adam optimizer. We use the following update
$\mathbf{r}_t \gets \beta_1 \mathbf{r}_{t-1} + (1 - \beta_1) \mathbf{g}_t, \quad \mathbf{s}_t \gets \beta_2 \mathbf{s}_{t-1} + (1 - \beta_2) (\mathbf{g}_t \cdot \mathbf{g}_t), \quad \boldsymbol{\theta}_t \gets \boldsymbol{\theta}_{t-1} - \rho \mathbf{r}_t / (\sqrt{\mathbf{s}_t} + \epsilon)$ , where  $\mathbf{g}_t$  is the minibatch gradient,  $\beta_1$  and  $\beta_2$  are coefficients for the running averages,  $\rho$  is a learning-rate, and  $\epsilon$  a small damping to stabilize. We construct a diagonal matrix  $\mathbf{S}_t = \mathrm{diag}(N \sqrt{\mathbf{s}_t})$  to estimate sensitivity with MPE as suggested in Table 1 ( $N$  is the number of training examples). Better results are expected by building better estimates of  $\mathbf{S}_t$  as discussed in [27]. As described in section 3.4 of [27], a smaller batch size should improve the estimate, which we also observe in the experiment.
I.6 Details of "evolution of sensitivities during training"
We use the MPE with iBLR for neural network classification on MNIST, FMNIST and CIFAR10, as well as MPE for logistic regression on MNIST. Experiment details are in Table 8.
For the experiment in Fig. 11(a), we consider Bayesian logistic regression. We set  $\delta = 0.1$ . The Hessian is always positive-definite due to the convex loss function therefore we use the VON algo-
Dataset	Model	Method	LR	\( LR_{\text{min}} \)	B	δ	Test acc.
FMNIST	LeNet5	diag., AdamW	10-3	10-3	256	60	88.1%
K-FAC, AdamW	10-3	10-3	256	60	87.6%
Table 6: Experimental settings for predicting generalization during the training in Fig. 9.
Dataset	Model	LR	\( LR_{\text{min}} \)	δ	Test acc. (B = 8)	Test acc. (B = 32)
MNIST	MLP (32, 16)	\( 10^{-3} \)	0	80	97.3%	97.4%
MNIST	LeNet5	\( 10^{-3} \)	0	60	99.2%	99.2%
FMNIST	LeNet5	\( 10^{-3} \)	0	60	91.4%	91.2%
CIFAR10	CNN	\( 10^{-3} \)	0	50	75.2%	78.4%
Table 7: Experimental settings for Fig. 10.
(a) MNIST, Bayesian logistic repr.
(b) MNIST, iBLR with MLP
(c) CIFAR-10, iBLR & ResNet-20
(d) Class removal result for MNIST
Figure 11: Additional experiments similar to Fig. 3(b). In Panel (a), we show the evolution of sensitivities for Bayesian logistic regression on MNIST trained with the VON algorithm. In Panel (b) we use a MLP trained with the iBLR optimizer. In Panel (c), we use a ResNet-20 trained with iBLR on CIFAR-10. Panel (d) shows the class removal result similar to Fig. 3(a), but on MNIST
Dataset	Model	B	δ	E
MNIST	MLP (500, 300)	256	30	100
FMNIST	LeNet5	256	60	100
CIFAR10	ResNet-20	512	35	300
Table 8: Experimental settings for evolution of sensitivities during training in Fig. 3(b), and Fig. 3.
rithm given in Eq. 29. We use 125 updates with batch-size 200, reaching a test accuracy of around  $91\%$  using the mean  $\mathbf{m}_t$ . We use linear learning-rate decay from 0.005 to 0.001 for the mean  $\mathbf{m}$  and a learning-rate of  $10^{-5}$  for the precision  $\mathbf{S}$ . The expectations are approximated using 3 samples drawn from the posterior. We plot sensitivities at iteration  $t = 5, 10, 25, 125$ . For this example, we use samples from  $q_t$  to compute the prediction variance and error (150 samples are used). We sort examples according to their sensitivity at iteration  $t = 125$  and then plot their average sensitivities in 60 groups with 100 examples in each group.
For the experiments in Fig. 3(b), Fig. 11(b) and Fig. 11(c), we consider neural network models  $\mathbf{f}(\boldsymbol{\theta}_t)$  on FMNIST, MNIST and CIFAR10. We do not use training data augmentation. For CIFAR10 we use a ResNet-20. The expectations in the iBLR are approximated using a single sample drawn from the posterior. For prediction, we use the mean  $\mathbf{m}_t$ . The test accuracies are  $91.3\%$  for FMNIST,  $98.5\%$  for MNIST and  $80.9\%$  for CIFAR10. We use a cosine learning-rate scheduler with an initial learning-rate of 0.1 and anneal to zero over the course of training. Other experimental details are stated in Table 8. Similar to before, we use sampling to evaluate sensitivity (150 samples are used).
J Author Contributions Statement
Authors list: Peter Nickl (PN), Lu Xu (LX), Dharmesh Tailor (DT), Thomas Moellenhoff (TM), Mohammad Emtiyaz Khan (MEK)
All co-authors contributed to developing the main idea. MEK and DT first discussed the idea deriving sensitivity measure based on the BLR. MEK derived the MPE and the results in Sec 3 and DT helped in connecting them to influence functions. PN derived the results in 3.3 and came up with the idea to predict generalization error with LOO-CV. LX adapted it to class-removal. PN wrote the code with help from LX. PN and LX did most of the experiments with some help from TM and regular feedback from everybody. TM did the experiment on the sensitivity evolution during training with some help from PN. All authors were involved in writing and proof-reading of the paper.
K Differences Between Camera-Ready Version and Submitted Version
We made several changes to take the feedback of reviewers into account and improve the paper.
1. The writing and organization of the paper were modified to emphasize the generalization to a wide variety of models and algorithms and the applicability of MPE during training.
2. The presentation was changed in Section 3 to emphasize the focus on the conjugate model. Detailed derivations were pushed to the appendices and more focus was put on big picture ideas. Arbitrary perturbations parts were made explicit. Table 1 was added and more focus was put on training algorithms.
3. We added experiments using leave-one-out estimation to predict generalization on unseen test data during training. We also added results to study the evolution of sensitivities during training using MPE with iBLR.