SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\hat {\mathbf {S}} _ {t} ^ {\backslash \mathcal {M}} \hat {\mathbf {m}} _ {t} ^ {\backslash \mathcal {M}} - \mathbf {S} _ {t} \mathbf {m} _ {t} = \rho \sum_ {j \in \mathcal {M}} \mathbb {E} _ {q _ {t}} [ \nabla \ell_ {j} (\boldsymbol {\theta}) ] - \mathbb {E} _ {q _ {t}} [ \nabla^ {2} \ell_ {j} (\boldsymbol {\theta}) ] \mathbf {m} _ {t}, \quad \mathbf {S} _ {t} - \hat {\mathbf {S}} _ {t} ^ {\backslash \mathcal {M}} = \rho \sum_ {j \in \mathcal {M}} \mathbb {E} _ {q _ {t}} [ \nabla^ {2} \ell_ {j} (\boldsymbol {\theta}) ] \tag {9}
$$
Algorithm	Update	Sensitivity
Newton's method	θt← θt-1 - Ht-1∇L(θt-1)	Ht-1∇l(i)(θt)
Online Newton (ON) [28]	θt← θt-1 - ρSt-1∇L(θt-1)	St-1∇l(i)(θt)
ON (diagonal+minibatch) [28]	θt← θt-1 - ρst-1·ˆ∇L(θt-1)	st-1·∇l(i)(θt)
iBLR (diagonal+minibatch) [35]	mt← mt-1 - ρst-1·ˆ∇L(θt-1)	st-1·∇l(i)(θt)
RMSprop/Adam [30]	θt← θt-1 - ρst-1·ˆ∇L(θt-1)	st-1/2·∇l(i)(θt)
SGD	θt← θt-1 - ρˆ∇L(θt-1)	∇l(i)(θt)
Table 1: A list of algorithms and their sensitivity measures derived using Eq. 10. The second column gives the update, most of which use pre-conditioners that are either matrices  $(\mathbf{H}_t,\mathbf{S}_t)$  or a vector  $(\mathbf{s}_t)$ ; see the full update equations in Eqs. 31 to 34 in App. C. The third column shows the associated sensitivity measure to perturbation in the  $i$ th example which can be interpreted as a step for the  $i$  example but in the opposite direction. We denote the element-wise multiplication between vectors by “.” and the minibatch gradients by  $\hat{\nabla}$ . For iBLR,  $\theta_t$  is either  $\mathbf{m}_t$  or a sample from  $q_t$ .
Plugging  $\mathbf{S}_t$  from the second equation into the first one, we can recover the following expressions,
$$
\hat {\mathbf {m}} _ {t} ^ {\backslash \mathcal {M}} - \mathbf {m} _ {t} = \rho \left(\hat {\mathbf {S}} _ {t} ^ {\backslash \mathcal {M}}\right) ^ {- 1} \mathbb {E} _ {q _ {t}} \left[ \sum_ {j \in \mathcal {M}} \nabla \ell_ {j} (\boldsymbol {\theta}) \right], \quad \left. \frac {\partial \hat {\mathbf {m}} _ {t} ^ {\epsilon_ {i}}}{\partial \epsilon_ {i}} \right| _ {\epsilon_ {i} = 0} = \rho \mathbf {S} _ {t} ^ {- 1} \mathbb {E} _ {q _ {t}} [ \nabla \ell_ {i} (\boldsymbol {\theta}) ] \tag {10}
$$
For the second equation, we omit the proof but it is similar to App. F, resulting in preconditioning with  $\mathbf{S}_t$ . For computational ease, we will approximate  $\hat{\mathbf{S}}_t^{\backslash \mathcal{M}} \approx \mathbf{S}_t$  even in the first equation. We will also approximate the expectation at a sample  $\theta_t \sim q_t$  or simply at the mean  $\theta_t = \mathbf{m}_t$ . Ultimately, the suggestion is to use  $\mathbf{S}_t^{-1}\nabla \ell_i(\theta_t)$  as the sensitivity measure, or variations of it, for example, by using a Monte-Carlo average over multiple samples.
Based on this, a list of algorithms and their corresponding measures is given in Table 1. All of the algorithms can be derived as special instances of the BLR by making specific approximations (see App. C.3). The measures are obtained by applying the exact same approximations to Eq. 10. For example, Newton's method is obtained when  $\mathbf{m}_t = \boldsymbol{\theta}_t$ ,  $\mathbf{S}_t = \mathbf{H}_{t-1}$ , and expectations are approximated by using the delta method at  $\boldsymbol{\theta}_t$  (similarly to Thm. 4). With these, we get
$$
\mathbf {S} _ {t} ^ {- 1} \mathbb {E} _ {q _ {t}} [ \nabla \ell_ {i} (\boldsymbol {\theta}) ] \approx \mathbf {H} _ {t - 1} ^ {- 1} \nabla \ell_ {i} (\boldsymbol {\theta} _ {t}), \tag {11}
$$
which is the measure shown in the first row of the table. In a similar fashion, we can derive measures for other algorithms that use a slightly different approximations leading to a different preconditioner. The exact strategy to update the preconditioners is given in Eqs. 31 to 34 of App. C.3. For all, the sensitivity measure is simply an update step for the  $i$ th example but in the opposite direction.
Table 1 shows an interplay between the training algorithm and sensitivity measures. For instance, it suggests that the measure  $\mathbf{H}_{t-1}^{-1}\nabla\ell_i(\pmb{\theta}_t)$  is justifiable for Newton's method but might be inappropriate otherwise. In general, it is more appropriate to use the algorithm's own preconditioner (if they use one). The quality of preconditioner (and therefore the measure) is tied to the quality of the posterior approximation. For example, RMSprop's preconditioner is not a good estimator of the posterior covariance when minibatch size is large [27, Thm. 1], therefore we should not expect it to work well for large minibatches. In contrast, the ON method [28] explicitly builds a good estimate of  $\mathbf{S}_t$  during training and we expect it to give better (and more faithful) sensitivity estimates.
For SGD, our approach suggests using the gradient. This goes well with many existing approaches [40, 42, 51, 3] but also gives a straightforward way to modify them when the training algorithm is changed. For instance, the TracIn approach [42] builds sensitivity estimates during SGD training by tracing  $\nabla \ell_{j}(\pmb{\theta}_{t})^{\top}\nabla \ell_{i}(\pmb{\theta}_{t})$  for many examples  $i$  and  $j$ . When the algorithm is switched, say to the ON method, we simply need to trace  $\nabla \ell_{j}(\pmb{\theta}_{t})^{\top}\mathbf{S}_{t}^{-1}\nabla \ell_{i}(\pmb{\theta}_{t})$ . Such a modification is speculated in [42, Sec 3.2] and the MPE provides a way to accomplish exactly that. It is also possible to mix and match algorithms with different measures but caution is required. For example, to use the measure in Eq. 11, say within a first-order method, the algorithm must be modified to build a well-conditioned estimate of the Hessian. This can be tricky and can make the sensitivity measure fragile [5].
Extensions to non-differentiable loss functions and discontinuous parameter spaces is straightforward. For example, when using a Gaussian posterior, the measures in Eq. 10 can be modified to
handle non-differentiable loss function by simply replacing  $\mathbb{E}_{q_t}[\nabla \ell_i(\pmb {\theta})]$  with  $\nabla_{\mathbf{m}}\mathbb{E}_{q_t}[\ell_i(\pmb {\theta})]$ , which is a simple application of the Bonnet theorem [44] (see App. G). The resulting approach is more principled than [32] which uses an ad-hoc smoothing of the non-differentiable loss: the smoothing in our approach is automatically done by using the posterior distribution. Handling of discontinuous parameter spaces follows in a similar fashion. For example, binary variables can be handled by measuring the sensitivity through the parameter of the Bernoulli distribution (see App. D).
3.3 Understanding the causes of high sensitivity estimates for the Gaussian case
The MPE can be used to understand the causes of high sensitivity-estimates. We will demonstrate this for Gaussian  $q$  but similar analysis can be done for other distributions. We find that sensitivity measures derived using Gaussian posteriors generally have two causes of high sensitivity.
To see this, consider a loss  $\ell_i(\pmb {\theta}) = -\log p(y_i|\sigma (f_i(\pmb {\theta})))$  where  $p(y_{i}|\mu)$  is an exponential-family distribution with expectation parameter  $\mu$ ,  $f_{i}(\pmb {\theta})$  is the model output for the  $i$ 'th example, and  $\sigma (\cdot)$  is an activation function, for example, the softmax function. For such loss functions, the gradient takes a simple form:  $\nabla \ell_i(\pmb {\theta}) = \nabla f_i(\pmb {\theta})[\sigma (f_i(\pmb {\theta})) - y_i]$  [6, Eq. 4.124]. Using this, we can approximate the deviations in model outputs by using a first-order Taylor approximation,
$$
\underbrace {f _ {i} \left(\boldsymbol {\theta} _ {t} ^ {\backslash i}\right) - f _ {i} \left(\boldsymbol {\theta} _ {t}\right)} _ {\text {D e v i a t i o n i n t h e o u t p u t}} \approx \nabla f _ {i} \left(\boldsymbol {\theta} _ {t}\right) ^ {\top} \left(\boldsymbol {\theta} _ {t} ^ {\backslash i} - \boldsymbol {\theta} _ {t}\right) \approx \underbrace {\nabla f _ {i} \left(\boldsymbol {\theta} _ {t}\right) ^ {\top} \mathbf {H} _ {t - 1} ^ {- 1} \nabla f _ {i} \left(\boldsymbol {\theta} _ {t}\right)} _ {= v _ {i t}, \text {p r e d i c t i o n v a r i a n c e}} \underbrace {\left[ \sigma \left(f _ {i} \left(\boldsymbol {\theta} _ {t}\right)\right) - y _ {i} \right]} _ {= e _ {i t}, \text {p r e d i c t i o n e r r o r}}. \tag {12}
$$
where we used  $\theta_t^{\backslash i} - \theta_t \approx \mathbf{H}_{t-1}^{-1}\nabla \ell_i(\theta_t)$  which is based on the measure in the first row of Table 1. Similarly to Eq. 2, the deviation in the model output is equal to the product of the prediction error and (linearized) prediction variance of  $f_i(\theta_t)$  [25, 20]. The change in the model output is expected to be high, whenever examples with high prediction error and variance are removed.
We can write many such variants with a similar bi-linear relationship. For example, Eq. 12 can be extended to get deviations in predictions as follows:
$$
\sigma \left(f _ {i} \left(\boldsymbol {\theta} _ {t} ^ {\backslash i}\right)\right) - \sigma \left(f _ {i} \left(\boldsymbol {\theta} _ {t}\right)\right) \approx \sigma^ {\prime} \left(f _ {i} \left(\boldsymbol {\theta} _ {t}\right)\right) \nabla f _ {i} \left(\boldsymbol {\theta} _ {t}\right) ^ {\top} \left(\boldsymbol {\theta} _ {t} ^ {\backslash i} - \boldsymbol {\theta} _ {t}\right) \approx \sigma^ {\prime} \left(f _ {i} \left(\boldsymbol {\theta} _ {t}\right)\right) v _ {i t} e _ {i t}. \tag {13}
$$
Eq. 12 estimates the deviation at one example and at a location  $\pmb{\theta}_t$ , but we could also write them for a group of examples and evaluate them at the mean  $\mathbf{m}_t$  or at any sample  $\pmb{\theta} \sim q_t$ . For example, to remove a group  $\mathcal{M}$  of size  $M$ , we can write the deviation of the model-output vector  $\mathbf{f}(\pmb{\theta}) \in \mathbb{R}^M$ ,
$$
\mathbf {f} \left(\mathbf {m} _ {t} ^ {\backslash \mathcal {M}}\right) - \mathbf {f} \left(\mathbf {m} _ {t}\right) \approx \nabla \mathbf {f} \left(\mathbf {m} _ {t}\right) ^ {\top} \mathbf {S} _ {t} ^ {- 1} \nabla \mathbf {f} \left(\mathbf {m} _ {t}\right) [ \sigma (\mathbf {f} \left(\mathbf {m} _ {t}\right) - \mathbf {y} ], \tag {14}
$$
where  $\mathbf{y}$  is the vector of labels and we used the sensitivity measure in Eq. 10. An example for sparse Gaussian process is in App. H. The measure for SGD in Table 1 can also be used which gives  $f_{i}(\pmb{\theta}_{t}^{\backslash i}) - f_{i}(\pmb{\theta}_{t}) \approx \| \nabla f_{i}(\pmb{\theta})\|^{2}e_{it}$  which is similar to the scores used in [40]. The list in Table 1 suggests that such scores can be improved by using  $\mathbf{H}_t$  or  $\mathbf{S}_t$ , essentially, replacing the gradient norm by an estimate of the prediction variance. Additional benefit can be obtained by further employing samples from  $q_{t}$  instead of using a point estimate  $\pmb{\theta}_{t}$  or  $\mathbf{m}_t$ ; see an example in App. H.
It is also clear that all of the deviations above can be obtained cheaply during training by using already computed quantities. The estimation does not add significant computational overhead and can be used to efficiently predict the generalization performance during training. For example, using Eq. 12, we can approximate the leave-one-out (LOO) cross-validation (CV) error as follows,
$$
\operatorname {L O O} \left(\boldsymbol {\theta} _ {t}\right) = \sum_ {i = 1} ^ {N} \ell_ {i} \left(\boldsymbol {\theta} _ {t} ^ {\backslash i}\right) = - \sum_ {i = 1} ^ {N} \log p \left(y _ {i} \mid \sigma \left(f _ {i} \left(\boldsymbol {\theta} _ {t} ^ {\backslash i}\right)\right)\right) \approx - \sum_ {i = 1} ^ {N} \log p \left(y _ {i} \mid \sigma \left(f _ {i} \left(\boldsymbol {\theta} _ {t}\right) + v _ {i t} e _ {i t}\right)\right). \tag {15}
$$
The approximation eliminates the need to train  $N$  models to perform CV, rather just uses  $e_{it}$  and  $v_{it}$  which are extremely cheap to compute within algorithms such as ON, RMSprop, and SGD. Leave-group-out (LGO) estimates can also be built, for example, by using Eq. 14, which enables us to understand the effect of leaving out a big chunk of training data, for example, an entire class for classification. The LOO and LGO estimates are closely related to marginal likelihood and sharpness, both of which are useful to predict generalization performance [22, 12, 19]. Estimates similar to Eq. 15 have been proposed previously [43, 4] but none of them do so during training.
(a) MLP on MNIST
(b) LeNet on FMNIST
(c) CNN on CIFAR-10
Figure 2: The estimated deviation for an example removal correlates well with the true deviations in predictions. Each marker represents an example. For each panel, the histogram at the bottom shows that the majority of examples have low sensitivity and most of the large sensitivities are attributed to a small fraction of data. We show a few images of high and low sensitivity examples from two randomly chosen classes, where we observe the high-sensitivity examples to be more interesting (possibly mislabeled or just ambiguous), while low-sensitivity examples appear more predictable.
We show experimental results to demonstrate the usefulness of the MPE to understand the sensitivity of deep-learning models. We show the following: (1) we verify that the estimated deviations (sensitivities) for data removal correlate with the truth; (2) we predict the effect of class removal on generalization error; (3) we estimate the cross-validation curve for hyperparameter tuning; (4) we predict generalization during training; and (5) we study evolution of sensitivities during training. All details of the experimental setup are included in App. I and the code is available at https://github.com/team-approx-bayes/memory-perturbation.
Estimated deviations correlate with the truth: Fig. 2 shows a good correlation between the true deviations  $\sigma(f_{i}(\pmb{\theta}_{*}^{\backslash i})) - \sigma(f_{i}(\pmb{\theta}_{*}))$  and their estimates  $\sigma'(f_{i}(\pmb{\theta}_{*}))v_{i*}e_{i*}$ , as shown in Eq. 13. We show results for three datasets, each using a different architecture but all trained using SGD. To estimate the Hessian  $\mathbf{H}_{*}$  and compute  $v_{i*} = \nabla f_{i}(\pmb{\theta}_{*})^{\top}\mathbf{H}_{*}^{-1}\nabla f_{i}(\pmb{\theta}_{*})$ , we use a Kronecker-factored (K-FAC) approximation implemented in the laplace [11] and ASDL [39] packages. Each marker represents a data example. The estimate roughly maintains the ranking of examples according to their sensitivity. Below each panel, a histogram of true deviations is included to show that the majority of examples have extremely low sensitivity and most of the large sensitivities are attributed to a small fraction of data. The high-sensitivity examples often include interesting cases (possibly mislabeled or simply ambiguous), some of which are visualized in each panel along with some low-sensitivity examples to show the contrast. High-sensitivity examples characterize the model's memory because perturbing them leads to a large change in the model. Similar trends are observed for removal of groups of examples in Fig. 6 of App. I.2.
Predicting the effect of class removal on generalization: Fig. 3(a) shows that the leave-group-out estimates can be used to faithfully predict the test performance even when a whole class is removed. The x-axis shows the test negative log-likelihood (NLL) on a held-out test set, while the y-axis shows the following leave-one-class-out (LOCO) loss on the set  $\mathcal{C}$  of a left-out class,
$$
\operatorname {L O C O} _ {\mathcal {C}} \left(\boldsymbol {\theta} _ {*}\right) = \sum_ {i \in \mathcal {C}} \ell_ {i} \left(\boldsymbol {\theta} _ {*} ^ {\backslash \mathcal {C}}\right) \approx - \sum_ {i \in \mathcal {C}} \log p \left(y _ {i} \mid \sigma \left(f _ {i} \left(\boldsymbol {\theta} _ {*}\right) + v _ {i *} e _ {i *}\right)\right).
$$
(a) Predicting the effect of class removal
(b) Evolution of sensitivities during training
Figure 3: Panel (a) shows, in the x-axis, the test NLL of trained models with a class removed. In the y-axis, we show the respective leave-one-class-out (LOCO) estimates. Each marker corresponds to a specific class removed (text indicates class names). Results for two models on FMNIST are shown. Both show good correlation between the test NLL and LOCO estimates; see the dashed lines. Panel (b) shows the evolution of estimated sensitivities during training of LeNet5 on FMNIST. As training progresses, the model becomes more and more sensitive to a small fraction of data.
(a) MLP on MNIST
Figure 4: The test NLL (gray) almost perfectly matches the estimated LOO-CV error of Eq. 15 (black). The x-axis shows different values of  $\delta$  parameter of an  $L_{2}$ -regularization  $\delta \| \pmb{\theta} \|^2 / 2$ .
(b) LeNet5 on FMNIST
(c) CNN on CIFAR-10
The estimate uses an approximation:  $f_{i}(\pmb{\theta}_{*}^{\backslash \mathcal{C}}) - f_{i}(\pmb{\theta}_{*}) \approx \nabla f_{i}(\pmb{\theta}_{*})^{\top}\mathbf{H}_{*}^{-1}\sum_{j\in \mathcal{C}}\nabla \ell_{j}(\pmb{\theta}_{*}) \approx v_{i*}e_{i*}$ , which is similar to Eq. 12, but uses an additional approximation  $\sum_{j\in \mathcal{C}}\nabla \ell_{j}(\pmb{\theta}_{*}) \approx \nabla \ell_{i}(\pmb{\theta}_{*})$  to reduce the computation due to matrix-vector multiplications (we rely on the same K-FAC approximation used in the previous experiment). Results might improve when this approximation is relaxed. We show results for two models: MLP and LeNet. Each marker corresponds to a specific class whose names are indicated with the text. The dashed lines indicate the general trends, showing a good correlation between the truth and estimate. The classes Shirt, Pullover are the most sensitive, while the classes Bag, Trousers are least sensitive. A similar result for MNIST is in Fig. 11(d) of App. I.3.
Predicting generalization for hyperparameter tuning: We consider the tuning of the parameter  $\delta$  for the  $L_{2}$ -regularizer of form  $\delta \| \pmb{\theta} \|^{2} / 2$ . Fig. 4 shows an almost perfect match between the test NLL and the estimated LOO-CV error of Eq. 15. Additional figures with the test errors visualized on top are included in Fig. 7 of App. I.4 where we again see a close match to the LOO-CV curves.
Predicting generalization during training: As discussed earlier, existing influence measures are not designed to analyze sensitivity during training and care needs to be taken when using ad-hoc strategies. We first show results for our proposed measure in Eq. 10 which gives reliable sensitivity estimates during training. We use the improved-BLR method [35] which estimates the mean  $\mathbf{m}_t$  and a vector preconditioner  $\mathbf{s}_t$  during training. We can derive an estimate for the LOO error at the mean  $\mathbf{m}_t$  following a derivation similar to Eqs. 14 and 15,
$$
\operatorname {L O O} \left(\mathbf {m} _ {t}\right) \approx - \sum_ {i = 1} ^ {N} \log p \left(y _ {i} \mid \sigma \left(f _ {i} \left(\mathbf {m} _ {t}\right) + v _ {i t} e _ {i t}\right)\right) \tag {16}
$$
(a) iBLR & LOO of Eq. 16
(b) SGD & diagonal-GGN-LOO
Figure 5: We compare faithfulness of LOO estimates during training to predict the test NLL. The first panel shows results for iBLR where a good match is obtained by using the LOO estimate of Eq. 16 which uses a diagonal preconditioner. The next two panels show results for SGD where we use the LOO estimate of Eq. 15 but with different Hessian approximations. Panel (b) uses a diagonal-GGN which does not work very well. Results are improved when K-FAC is used, but they are still not as good as the iBLR, despite using a non-diagonal Hessian approximation.
(c) SGD & K-FAC-LOO
where  $v_{it} = \nabla f_i(\mathbf{m}_t)^\top \mathrm{diag}(\mathbf{s}_t)^{-1}\nabla f_i(\mathbf{m}_t)$  and  $e_{it} = \sigma (f_i(\mathbf{m}_t)) - y_i$
The first panel in Fig. 5 shows a good match between the above LOO estimate and test NLL. For comparison, in the next two panels, we show results for SGD training by using two ad-hoc measures obtained by plugging different Hessian approximations in Eq. 11. The first panel approximates  $\mathbf{H}_t$  with a diagonal Generalized Gauss-Newton (GGN) matrix, while the second panel uses a K-FAC approximation. We see that diagonal-GGN-LOO does not work well at all and, while K-FAC-LOO improves this, it is still not as good as the iBLR result despite using a non-diagonal Hessian approximation. Not to mention, the two measures require an additional pass through the data to compute the Hessian approximation, and also need a careful setting of a damping parameter.
A similar result for iBLR is shown in Fig. 1(b) where we use the larger ResNet-20 on CIFAR10, and more such results are included in Fig. 8 of App. I.5. We also find that both diagonal-GGN-LOO or K-FAC-LOO further deteriorate when the model overfits; see Fig. 9. Results for the Adam optimizer are included in Fig. 10, where we again see that using ad hoc measures may not always work. Overall, these results show the difficulty of estimating sensitivity during training and suggest to take caution when using measures that are not naturally suited to analyze the training algorithm.
Evolution of sensitivities during training: Fig. 3(b) shows the evolution of sensitivities of examples as the training progresses. We use the iBLR algorithm and approximate the deviation as  $\sigma(f_i(\mathbf{m}_t^{\backslash i})) - \sigma(f_i(\mathbf{m}_t)) \approx \sigma'(f_i(\mathbf{m}_t))v_{it}e_{it}$  where  $v_{it}$  and  $e_{it}$  are obtained similarly to Eq. 16. The x-axis corresponds to examples sorted from least sensitive to most sensitive examples at convergence. The y-axis shows the histogram of sensitivity estimates. We observe that, as the training progresses, the distribution concentrates around a small fraction of the data. At the top, we visualize a few examples with high and low sensitivity estimates, where the high-sensitivity examples included interesting cases (similarly to Fig. 2). The result suggests that the model concentrates more and more on a small fraction of high-sensitivity examples, and therefore such examples can be used to characterize the model's memory. Additional experiments of this kind are included in Fig. 11 of App. I.6, along with other experiment details.
We present the memory-perturbation equation by building upon the BLR framework. The equation suggests to take a step in the direction of the natural gradient of the perturbed examples. Using the MPE framework, we unify existing influence measures, generalize them to a wide variety of problems, and unravel useful properties regarding sensitivity. We also show that sensitivity estimation can be done cheaply and use this to predict generalization performance. An interesting avenue for future research is to apply the method to larger models and real-world problems. We also need to understand how our generalization measure compares to other methods, such as those considered in [22]. We would also like to understand the effect of various posterior approximations. Another interesting direction is to apply the method to non-Gaussian cases, for example, to study ensemble methods in deep learning with mixture models.
This work is supported by the Bayes duality project, JST CREST Grant Number JPMJCR2112.