SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Elisa Nguyen
Tübingen AI Center University of Tübingen
Minjoon Seo
KAIST AI
Seong Joon Oh
Training data attribution (TDA) techniques find influential training data for the model's prediction on the test data of interest. They approximate the impact of down- or up-weighting a particular training sample. While conceptually useful, they are hardly applicable to deep models in practice, particularly because of their sensitivity to different model initialisation. In this paper, we introduce a Bayesian perspective on the TDA task, where the learned model is treated as a Bayesian posterior and the TDA estimates as random variables. From this novel viewpoint, we observe that the influence of an individual training sample is often overshadowed by the noise stemming from model initialisation and SGD batch composition. Based on this observation, we argue that TDA can only be reliably used for explaining deep model predictions that are consistently influenced by certain training data, independent of other noise factors. Our experiments demonstrate the rarity of such noise-independent training-test data pairs but confirm their existence. We recommend that future researchers and practitioners trust TDA estimates only in such cases. Further, we find a disagreement between ground truth and estimated TDA distributions and encourage future work to study this gap. Code is provided at https://github.com/ElisaNguyen/bayesian-tda.
Understanding how machine learning models arrive at decisions is desirable for social, legal and ethical reasons, particularly for opaque deep learning models [1]. One approach to explanations is the data-centric approach of training data attribution (TDA). As the name suggests, TDA finds attributing training samples for a model decision, uncovering which part of the training data is relevant. The attribution  $\tau$  of a training sample  $z_{j}$  on another sample  $z$  is usually defined as the change of model loss  $\mathcal{L}$  on  $z$  when the model is retrained without  $z_{j}$  [2]:
$$
\tau \left(z _ {j}, z\right) := \mathcal {L} \left(z; \theta_ {\backslash j}\right) - \mathcal {L} \left(z; \theta\right) \tag {1}
$$
where  $\theta$  is a model trained on the entire dataset  $\mathcal{D}$  and  $\theta_{\backslash j}$  is a model trained on the same set without  $z_{j}$ . Since the direct computation of Equation 1 is expensive, various TDA techniques for approximating the quantity have been proposed, such as influence functions [3] or TracIn [4]. Their approximations are often based on some form of inner product between the parameter gradients  $\nabla_{\theta}\mathcal{L}(z;\theta)$  and  $\nabla_{\theta}\mathcal{L}(z_j;\theta)$ .
Knowing how training samples attribute to a model decision provides an actionable understanding of the training data distribution, especially in cases of model error. TDA methods can identify the training samples that are most relevant to an error and therefore enable users to understand why the error occurred (e.g. due to domain mismatch of test and training data or wrongly labelled training data) [3]. Additionally, TDA gives them the tool to address the errors by e.g. changing the model directly through the training data. Even in non-erroneous cases, understanding the attributing training
data may enable users affected by model decisions to contest the decisions if the attributing training data is noisy or of low quality [5].
At the same time, TDA methods, especially influence functions [3], have been criticised for their fragility when applied to deep models [6, 7, 8]. The main reasons are model complexity and the stochasticity of deep model training. While the former poses a challenge specifically for influence functions as they rely on strong convexity assumptions, the latter is a more general challenge [6, 9]. The randomness inherent to the training process does not only lead to variation in the learned model parameters but also in TDA scores, which makes them untrustworthy. Hence, K & Søgaard [9] recommend using expected TDA scores for increased stability.
We argue that solely considering the expectation is not sufficient to ensure the reliability of TDA but requires inspecting the variance, too. We introduce a Bayesian perspective on the TDA task, noticing that there is no deterministic mapping from a dataset  $\mathcal{D}$  to the corresponding model  $\theta$  for deep neural networks. The learned model depends on the initialisation and batch composition in the stochastic gradient descent (SGD) optimiser. We capture the resulting randomness via Bayesian model posterior  $p(\theta|\mathcal{D})$  over the parameter space [10, 11, 12]. In turn, the TDA estimate (Equation 1) is a random variable that depends on two posteriors,  $p(\theta|\mathcal{D})$  and  $p(\theta_{\backslash j}|\mathcal{D}_{\backslash j})$ .
This viewpoint leads to a few insights into the practical usage and evaluation of TDA techniques. We confirm quantitatively that the ground-truth influence  $\tau(z_j, z)$  is often dominated by the noise:  $\sqrt{\operatorname{Var}(\tau)} > \mathbb{E}|\tau|$ . We argue that it is practically difficult to apply any TDA technique on pairs  $(z_j, z)$  whose ground-truth attributions  $\tau(z_j, z)$  are noisy in the first place. Likewise, any evaluation of TDA methods on such high-variance pairs would not be reliable.
Nonetheless, we are optimistic that TDA techniques are useful in practice, particularly for train-test pairs with high signal-to-noise ratios:  $\sqrt{\mathrm{Var}(\tau)}\ll \mathbb{E}|\tau |$  . We observe that such pairs are rare but consistently present in multiple experiments. We recommend that researchers and practitioners confine their usage to scenarios where the signal-to-noise ratios are expected to be large enough.
Our contributions are as follows: (1) Bayesian formulation of the training data attribution (TDA) task. (2) Observation that the ground-truth TDA values are often unreliable and highly variable. (3) Recommendation for the community to use the TDA tools only when the expected noise level is low. (4) Experimental analysis of the contributing factors to the variance of ground-truth TDA values. (5) Observation that the TDA estimation methods capture local changes in the model with regard to the counterfactual question of "retraining without training sample  $z_{j}$ ", while LOO retraining itself results in a more global change through the training procedure.
2 Background
We cover the background materials for the paper, including the concept, method, and evaluation of training data attribution (TDA) methods and Bayesian deep learning.
2.1 Training data attribution (TDA)
We introduce the TDA task, a few representative TDA methods, and existing evaluation strategies.
TDA task. Given a deep model  $f_{\theta}$  parametrised by  $\theta$ , a training set  $\mathcal{D} \coloneqq \{z_1, \dots, z_N\}$ , and a test sample  $z$ , one is interested in the impact of a training sample  $z_j$  on the model's behaviour on the test sample  $z$ . In the TDA context, one is often interested in the counterfactual change in the loss value for  $z$  after leave-one-out (LOO) training, when  $z_j$  is excluded from the training set (Equation 1). TDA has been considered in different use cases, such as understanding the bias in word embeddings [13], fact tracing in language model outputs [14] and measuring the robustness of model predictions [5].
TDA methods. The conceptually most straightforward way to compute the difference due to LOO training (Equation 1) is to compute it directly. However, this is computationally expensive, as it involves the learning algorithm for obtaining  $\theta_{\backslash j}$  for every  $j$ . This gives rise to various TDA techniques that find approximate estimates  $\tau'(z_j, z)$  of LOO. A prominent example of such approximation is the influence function (IF) method [3] based on [15]. Under strong smoothness assumptions, they
approximate Equation 1 by:
$$
\tau^ {\prime} \left(z _ {j}, z\right) := - \nabla_ {\theta} \mathcal {L} (z; \theta) ^ {\top} H _ {\theta} ^ {- 1} \nabla_ {\theta} \mathcal {L} \left(z _ {j}; \theta\right) \tag {2}
$$
where  $\nabla_{\theta}\mathcal{L}(z;\theta)$  and  $\nabla_{\theta}\mathcal{L}(z_j;\theta)$  refer to the parameter gradients of  $f_{\theta}$  for  $z$  and  $z_{j}$  respectively. Recognising the difficulty of scaling up the inverse Hessian computation  $H_{\theta}^{-1}$  and the high dimensionality of operations in Equation 2, subsequent papers have proposed further approximations to speed up the computation [16, 17]. Charpiat et al. [18] have analysed the influence of  $z_{j}$  on  $z$  by dropping the need to compute the Hessian and formulating influence as the loss change when an additional training step (ATS) on  $z_{j}$  is taken:
$$
\tau \left(z _ {j}, z\right) := \mathcal {L} \left(z; \theta_ {+ j}\right) - \mathcal {L} \left(z; \theta\right) \tag {3}
$$
where  $\theta_{+j}$  is a learned model parameter with  $\mathcal{D}$  and an additional step on  $z_{j}$ . They propose two approximations:
$$
\text {G r a d - D o t} (\mathbf {G D}): \quad \tau^ {\prime} \left(z _ {j}, z\right) := \nabla_ {\theta} \mathcal {L} \left(z _ {j}; \theta\right) ^ {\top} \nabla_ {\theta} \mathcal {L} (z; \theta) \tag {4}
$$
$$
\text {G r a d - C o s} (\mathbf {G C}): \quad \tau^ {\prime} (z _ {j}, z) := \frac {\nabla_ {\theta} \mathcal {L} (z _ {j} ; \theta)}{\| \nabla_ {\theta} \mathcal {L} (z _ {j} ; \theta) \|} ^ {\top} \frac {\nabla_ {\theta} \mathcal {L} (z ; \theta)}{\| \nabla_ {\theta} \mathcal {L} (z ; \theta) \|} \tag {5}
$$
This method is closely linked to TracIn [4] which computes the Grad-Dot not just at the end of the training, but averages the regular Grad-Dot similarities throughout the model training iterations. We note later in our analysis that within our Bayesian treatment of TDA, the TracIn method coincides conceptually with the Grad-Dot method. In our analysis, we study the sensitivity of LOO and the above TDA methods against noise.
TDA evaluation. The primal aim of TDA methods is to measure how well they approximate the ground-truth LOO values. This is often done by measuring the correlation between the estimates from each TDA method and the ground-truth LOO values (Equation 1) [3, 6, 7, 9]. They use either a linear (Pearson) correlation or a rank (Spearman) correlation over a small number of train-test sample pairs  $(z_{j}, z)$  due to the computational burden of computing the actual LOO values, especially for larger models. Usually, a few samples  $z$  are chosen for a comparison against LOO, e.g. Koh & Liang [3] report results for one  $z$  and Guo et al. [16] for 10 samples  $z$ . In some cases, the ground-truth LOO is obtained by computing the change in loss after training further from the learned model parameters, e.g. [3]. Some works have adopted indirect evaluation metrics such as the retrieval performance of mislabelled or poisoned training data based on the TDA estimates [3, 19, 16, 9, 17]. In this work, we adopt the Pearson and Spearman correlation metrics and discuss ways to extend them when the target (LOO from same initialisation) and estimates (TDA) are both random variables.
2.2 Bayesian deep learning.
Bayesian machine learning treats the learned model as a posterior distribution over the parameter space, rather than a single point:
$$
p (\theta | \mathcal {D}) = p (\mathcal {D} | \theta) p (\theta) / p (\mathcal {D}). \tag {6}
$$
Bayesian ML nicely captures the intuition that the mapping from a training set  $\mathcal{D}$  to the learned model  $p(\theta|\mathcal{D})$  is not a deterministic mapping, especially for non-convex models like deep neural networks (DNNs). Depending on the initialisation, among other factors, DNN training almost always learns vastly different parameters.
The estimation of the true posterior is indeed difficult for complex models like DNNs. The field of Bayesian deep learning is dedicated to the interpretation of certain random elements in DNN training as sources of randomness for the approximated Bayesian posteriors. For example, if Dropout [20] is used for training a model, it may be used at test time to let users sample  $\theta$  from the posterior distribution  $p(\theta | \mathcal{D})$  [21]. More generally used components like stochastic gradient descent (SGD) have also been interpreted as sources of randomness. The random walk induced by SGD iterations in the parameter space can be viewed as a Markov Chain Monte-Carlo sampler from the posterior distribution, after a slight modification of the optimisation algorithm (Stochastic Gradient Langevin Dynamics [10]). Similarly, the last few iterations of the vanilla SGD iterations may also be treated as samples from the posterior, resulting in more widely applicable Bayesian methods like Stochastic Weight Averaging (SWA) [12, 22]. Finally, the random initialisation of DNNs has also been exploited for modelling posterior randomness; training multiple versions of the same model with different initial parameters may be interpreted as samples from the posterior [11]. We show in the next section how the Bayesian viewpoint will help us model the sources of stochasticity for TDA estimates.
Figure 1: A Bayesian interpretation of training data attribution (TDA).
Difference in loss (LOO) is a 1-D random variable.
3 A Bayesian perspective on training data attribution
Training data attribution (TDA)  $\tau(z_j, z)$  is defined as the attribution of one training sample  $z_j$  to another sample  $z$  in terms of how a target metric like the loss of a sample  $\mathcal{L}(z, \theta)$  changes when the model is trained without  $z_j$  (Equation 1). We note here that according to the definition, we are interested in the impact of the change in the dataset from  $\mathcal{D}$  to  $\mathcal{D}_{\backslash j}$ , rather than the change in the model parameter. From a Bayesian perspective, a change in the training dataset leads to a shift in the posterior distribution,  $p(\theta | \mathcal{D}) \to p(\theta_{\backslash j} | \mathcal{D}_{\backslash j})$ , leading to the definition of TDA as a random variable:
$$
\tau \left(z _ {j}, z | \mathcal {D}\right) := \mathcal {L} \left(z; \mathcal {D} _ {\backslash j}\right) - \mathcal {L} (z; \mathcal {D}) = \mathcal {L} \left(z; \theta_ {\backslash j} \mid \mathcal {D} _ {\backslash j}\right) - \mathcal {L} (z; \theta | \mathcal {D}) \tag {7}
$$
where  $\theta \sim p(\theta|\mathcal{D})$  and  $\theta_{\backslash j} \sim p(\theta_{\backslash j}|\mathcal{D}_{\backslash j})$ . This interpretation is more natural, given the non-uniqueness of the mapping from a training dataset  $\mathcal{D}$  to the optimal model parameter  $\theta$  for general, non-convex models like DNNs. Alternatively, one could treat the model built from  $\mathcal{D}$  as a fixed variable rather than a posterior as TDA is applied to a specific model in practice. The change of dataset from  $\mathcal{D}$  to  $\mathcal{D}_{\backslash j}$  however still introduces ambiguity in  $\theta_{\backslash j}$ , which is captured in the Bayesian posterior  $p(\theta_{\backslash j}|\mathcal{D}_{\backslash j})$ . In this study, we use the probabilistic formulation of TDA in Equation 7.
Sampling TDA values. One could plug in various Bayesian DL techniques (§2.2) to compute samples of  $p(\theta | \mathcal{D})$ , which can be used to get the samples of  $\tau(z_j, z)$ . In our work, we use the Stochastic Weight Averaging (SWA) [12, 22] and Deep Ensemble (DE) [11] which are applicable to a wide class of deep models. More specifically, we obtain  $T$  samples  $\theta^{(1)}, \dots, \theta^{(T)} \sim p(\theta | \mathcal{D})$  either by taking the last  $T$  model checkpoints of the SGD iterations (SWA) or by taking the last model checkpoints from  $T$  different model initialisations (DE). The same is done for the counterfactual posterior  $\theta_{\backslash j}^{(1)}, \dots, \theta_{\backslash j}^{(T)} \sim p(\theta_{\backslash j} | \mathcal{D}_{\backslash j})$ . This results in a mixture-of-Gaussian posterior, where DE samples correspond to centroids of the distribution. Our sampling approach is thus a version of stratified sampling, where the number of samples  $T$  from a centroid is fixed and sampled IID.
Statistical analysis on TDA. The simplest statistics for the TDA  $\tau(z_j, z)$  are the mean and variance:
$$
\mathbb {E} [ \tau (z _ {j}, z) ] = \frac {1}{T} \sum_ {t} \mathcal {L} (z; \theta_ {\backslash j} ^ {(t)}) - \mathcal {L} (z; \theta^ {(t)}) \tag {8}
$$
$$
\operatorname {V a r} [ \tau (z _ {j}, z) ] = \frac {1}{T ^ {2}} \sum_ {t, t ^ {\prime}} \left(\mathcal {L} (z; \theta_ {j} ^ {(t)}) - \mathcal {L} (z; \theta^ {(t ^ {\prime})}) - \mathbb {E} [ \tau (z _ {j}, z) ]\right) ^ {2} \tag {9}
$$
Our main interest lies in whether the influence of the training data  $z_{j}$  on the test data  $z$  is statistically significant and not dominated by the inherent noise of deep model training. For this purpose, we design a Student t-test [23] for quantifying the statistical significance. Our null and alternative hypotheses are:
$$
H _ {0}: \mu = 0 \quad H _ {1}: \mu \neq 0. \tag {10}
$$
We consider the test statistic based on sample mean and variance:
$$
t = \frac {\mu - \mathbb {E} [ \tau (z _ {j} , z) ]}{\sqrt {\operatorname {V a r s} [ \tau (z _ {j} , z) ] / T ^ {2}}}. \tag {11}
$$
Vars refers to the sample variance where the denominator in Equation 9 is  $T^2 - 1$  instead. We report the significance of the absolute TDA  $|\tau(z_j, z)|$  for every train-test pair  $(z_j, z)$  by computing the p-value corresponding to the t-test statistic. The greater the p-value is, the greater the dominance of noise is for the TDA estimate.
Figure 2: Sources of randomness for Bayesian posteriors. In each case, the training starts from initialisation  $\theta_0$ . Depending on whether  $z_{j}$  is included in the training data, one has either samples from the original posterior  $p(\theta | \mathcal{D})$  or from the counterfactual posterior  $p(\theta_{\setminus j} | \mathcal{D}_{\setminus j})$ . For deep ensemble [11], the randomness stems either from random initialisation (DE-Init) or from SGD batch composition (DE-Batch). For stochastic weight averaging (SWA) [12, 22], last few checkpoints of the training are treated as posterior samples.
TDA methods likewise estimate random quantities. Approximate TDA methods like influence functions (IF), Grad-Dot, and Grad-Cos (§2.1) also predict random quantities  $\tau'(z_j, z)$ . For example, IF predicts  $\tau(z_j, z) \approx \tau'(z_j, z) \coloneqq -\nabla_\theta \mathcal{L}(z_j; \theta)^\top H_\theta^{-1} \nabla_\theta \mathcal{L}(z; \theta)$ , where one may sample  $\theta$  from the posterior  $p(\theta | \mathcal{D})$ . We note that IF suffers theoretical issues in its application to deep models, as convexity assumptions are not met. In practice, estimation algorithms make use of a damping term to ensure the positive definiteness of the inverse Hessian. Through a Bayesian lens, the damping term could be seen as an isotropic Gaussian prior centred at the origin. Similar statistical analyses on the TDA estimations can be performed as above, including the mean and variance computations and statistical testing for the significance of influence.
Evaluating TDA as a random variable. Previously, the LOO-based TDA values  $\tau(z_j, z)$  and the estimates from various approximate TDA methods  $\tau'(z_j, z)$  are compared via correlation measures like Pearson or Spearman. Our treatment of those quantities as 1-D random variables poses a novel challenge for evaluation because there exists no inborn notion of ordering among 1-D random variables. We address the challenge by examining the approximation ability of TDA methods for both the first and second moments of the true TDA values  $\tau(z_j, z)$ . More specifically, we compute the Pearson and Spearman correlation for both the mean (Equation 8) and variance (Equation 9) between the ground-truth  $\tau(z_j, z)$  and estimated TDA  $\tau'(z_j, z)$  values across multiple train-test pairs  $(z_j, z)$ .
We introduce our experimental settings, present analyses on factors contributing to the reliability of TDA values, compare TDA methods, and draw suggestions on the evaluation practice of TDA.
4.1 Implementation details
We illustrate the specific details of our implementation. See the Appendix for further information.
TDA methods. We study different TDA methods from a Bayesian perspective. We test the methods introduced in §2.1 for estimating TDA: influence functions (IF) [3], Grad-Dot (GD) and Grad-Cos (GC) [18]. We use the PyTorch implementation of IF from Guo et al. [16] and modify it for our models. As the ground-truth target, we consider Leave-one-out training (LOO) [3]. For LOO, we remove of  $z_{j}$  from the training set  $\mathcal{D}$  by zeroing out the weight for sample  $z_{j}$  towards the loss. Additionally, we include Charpiat et al.'s [18] notion of TDA that a training data point  $z_{j}$  attributes more if an additional training step (ATS) on it changes the test loss more significantly.
Inducing randomness in posterior  $p(\theta | \mathcal{D})$ . In §2.2, we have introduced the interpretation of various elements around model training as sources of randomness for Bayesian posterior. We summarise our methods for inducing randomness in Figure 2. We use the notion of the Deep Ensemble (DE) [11] to sample from the posterior. In a variant of DE with the initialisation as the source of randomness (DE-Init), we train each of  $T_{\mathrm{DE}}$  randomly initialised parameters  $\theta_0^{(t)}$  on either  $\mathcal{D}$  or  $\mathcal{D}_{\backslash j}$ . The resulting parameter sets,  $\theta^{(t)}$  and  $\theta_{\backslash j}^{(t)}$ , are treated as samples from respective
Table 1: Stability of TDA estimates. We report p-values for the ground-truth TDA  $\tau (z_j,z)$  (LOO) and the estimated TDA values  $\tau^{\prime}(z_j,z)$  (rest 4 columns). The p-values are averaged across all train-test pairs  $(z_{j},z)$ . We use the CNN model throughout.
Data	Randomness	LOO	ATS	IF	GD	GC
MNIST3	SWA+DE-Init	0.331	0.254	0.352	0.363	0.003
SWA+DE-Batch	0.025	0.039	0.000	0.000	0.000
CIFAR10	SWA+DE-Init	0.692	0.437	0.575	0.587	0.356
SWA+DE-Batch	0.487	0.296	0.484	0.517	0.236
Figure 3: Stability of TDA estimates per train-test pair. Distribution of p-values for ground-truth TDA (LOO) for different experiments.
posteriors. We also consider the batch composition in stochastic gradient descent (SGD) as the source of randomness (DE-Batch). In this case, we train from one initial parameter  $\theta_0$  with  $T_{\mathrm{DE}}$  different random shuffles  $\pi^{(t)}$  of the training sets  $\mathcal{D}$  and  $\mathcal{D}_{\backslash j}$ . This results in two sets of samples from the original and counterfactual posteriors. We increase the number of samples by taking the last  $T_{\mathrm{SWA}}$  checkpoints as the Stochastic Weight Averaging (SWA) samples [12, 22]. For Grad-Dot, this coincides with the definition of TracIn [4] as we average the dot products across checkpoints. In total, we take  $T = T_{\mathrm{DE}} \times T_{\mathrm{SWA}} = 10 \times 5$  samples from  $p(\theta|\mathcal{D})$  and  $p(\theta_{\backslash j}|\mathcal{D}_{\backslash j})$  to estimate  $\tau(z_j, z)$ .
Datasets  $\mathcal{D}$ . To enable an exhaustive analysis of every train-test pair  $(z_{j},z)$ , we define smaller datasets. We use variants of MNIST [24] limited to three classes (MNIST3), and CIFAR10 [25]. For MNIST3, we sample a training set of size 150 and a test set of size 900, i.e. 135,000 train-test pairs. For CIFAR10, we define the training and test set at size 500, i.e. 250,000 train-test pairs.
Models. We consider two types of image classifiers, visual transformers (ViT, [26]) and convolutional neural networks [27], where we primarily study a two-layer (CNN2-L). We also include a three-layer version (CNN3-L) to study the factor of model complexity. For ViT variants, instead of full finetuning, we use LoRA adapter layers [28] to minimise the number of parameters being tuned. The number of trainable parameters of ViT+LoRA (597,514) is comparable to CNN3-L (620,362).
4.2 Reliability of TDA evaluation