Title: Training Data Attribution via Approximate Unrolling
Venue: NeurIPS 2024
Accepted: yes
Keywords: training data attribution

influence functions
Abstract: Many training data attribution (TDA) methods aim to estimate how a model's behavior would change if one or more data points were removed from the training set. Methods based on implicit differentiation, such as influence functions, can be made computationally efficient, but fail to account for underspecification, the implicit bias of the optimization algorithm, or multi-stage training pipelines. By contrast, methods based on unrolling address these issues but face scalability challenges. In this work, we connect the implicit-differentiation-based and unrolling-based approaches and combine their benefits by introducing Source, an approximate unrolling-based TDA method that is computed using an influence-function-like formula. While being computationally efficient compared to unrolling-based approaches, Source is suitable in cases where implicit-differentiation-based approaches struggle, such as in non-converged models and multi-stage training pipelines. Empirically, Source outperforms existing TDA techniques in counterfactual prediction, especially in settings where implicit-differentiation-based approaches fall short.

Training Data Attribution via Approximate Unrolling

Abstract

Many training data attribution (TDA) methods aim to estimate how a model's behavior would change if one or more data points were removed from the training set. Methods based on implicit differentiation, such as influence functions, can be made computationally efficient, but fail to account for underspecification, the implicit bias of the optimization algorithm, or multi-stage training pipelines. By contrast, methods based on unrolling address these issues but face scalability challenges. In this work, we connect the implicit-differentiation-based and unrolling-based approaches and combine their benefits by introducing SOURCE, an approximate unrolling-based TDA method that is computed using an influence-function-like formula. While being computationally efficient compared to unrolling-based approaches, SOURCE is suitable in cases where implicit-differentiation-based approaches struggle, such as in non-converged models and multi-stage training pipelines. Empirically, SOURCE outperforms existing TDA techniques in counterfactual prediction, especially in settings where implicit-differentiation-based approaches fall short.

1 Introduction

Training data attribution (TDA) techniques are motivated by understanding the relationship between training data and the properties of trained models [92, 17, 29, 35, 70, 24, 51]. They have diverse applications in machine learning, such as detecting mislabeled data points [72, 50, 41], crafting data poisoning attacks [16, 38, 69], and curating datasets [60, 90, 13]. Many TDA methods aim to perform a counterfactual prediction , which estimates how a trained model's behavior would change if certain data points were removed from (or added to) the training dataset. Unlike sampling-based approaches, which require repeated model retraining with different subsets of the dataset, gradient-based TDA techniques estimate an infinitesimal version of the counterfactual without model retraining. Two main strategies for gradient-based counterfactual TDA are implicit differentiation and unrolling .

Implicit-differentiation-based TDA, most notably influence functions [28, 49], uses the Implicit Function Theorem [52] to estimate the sensitivity of the optimal solution to downweighting a training data point. These methods are well-motivated for models with strongly convex objectives and provide convenient estimation algorithms that depend solely on the optimal model parameters rather than intermediate checkpoints throughout training. However, the classical formulation relies on assumptions such as the uniqueness of and convergence to the optimal solution, which limits their applicability to modern neural networks [6, 3, 77].

By contrast, unrolling-based methods, such as SGD-INFLUENCE [31], approximate the impact of downweighting a data point's gradient update on the final model parameters by backpropagating through the preceding optimization steps. Unrolling is conceptually appealing in modern neural networks because it does not rely on the uniqueness of or convergence to the optimal solution. Furthermore, it can incorporate details of the training process, such as the choice of optimizer, learning rate schedules, or a data point's position during training. For example, unrolling-based

approaches can support TDA for multi-stage training procedures, such as in continual learning or foundation models, where the model undergoes multiple training phases with different objectives or datasets. However, they require storing all intermediate variables generated during the training process ( e . g ., parameter vectors for each optimization step) for backpropagation, which can be prohibitively expensive for large-scale models. Past works have considered applying unrolling to only the last epoch for large-scale models [31, 10], restricting their applicability in analyzing the effect of removing a data point at the beginning of training or in analyzing multi-stage training procedures.

In this work, we connect implicit-differentiation-based and unrolling-based approaches and introduce a novel algorithm that enjoys the advantages of both methods. We start from the unrolled differentiation perspective and, after introducing suitable approximations, arrive at an influence-function-like estimation algorithm. The key idea is to divide the training trajectory into one or more segments and approximate the distributions of gradients and Hessians as stationary within each segment. These segments may represent explicit training stages, such as in continual learning or foundation models, or changes in the Hessian and gradients throughout training. We use these estimated statistical summaries for each segment to approximate unrolling. Hence, we call our method SOURCE ( S egmented stati O nary U n R olling for C ounterfactual E stimation). While our method approximately coincides with influence functions in the simple setting of a deterministic objective optimized to convergence, it applies to more general settings where unrolling is typically required.

SOURCE inherits several key advantages from unrolling. Firstly, it allows the attribution of data points at different stages of training, providing a more comprehensive framework for TDA. Secondly, SOURCE can incorporate algorithmic choices into the analysis, accounting for learning rate schedules and the implicit bias of optimizers such as SGD or Adam [48]. Lastly, it maintains a close connection with the counterfactuals, even in cases where the assumptions made in implicit-differentiation-based methods, such as the optimality of the final parameters, are not met. However, unlike unrolling, SOURCE does not require storing all intermediate optimization variables generated during training; instead, it leverages only a handful of model checkpoints.

We evaluate SOURCE for counterfactual prediction across various tasks, including regression, image classification, and text classification. SOURCE outperforms existing TDA techniques in approximating the effect of retraining the network without groups of data points and identifying training data points that would flip predictions on some test examples when trained without them. SOURCE demonstrates distinct advantages in scenarios where traditional implicit-differentiation-based methods fall short, such as models that have not fully converged or those trained in multiple stages. Our empirical evidence suggests that SOURCE is a valuable TDA tool in various scenarios.

2 Background

Consider a finite training dataset D := { z i } N i =1 . We assume that the model parameters θ ∈ R D are optimized with a gradient-based iterative optimizer to minimize the empirical risk on this dataset:

where L is the (twice-differentiable) loss function. We use the notation θ ⋆ ( S ) to denote the optimal solution obtained when the model is trained on a specific subset of the dataset S ⊆ D , and θ ⋆ to denote the optimal solution on the full dataset. In practice, it is common to employ parameters θ s that approximately minimize the empirical risk ( e . g ., the result of running an optimization algorithm for T iterations), as obtaining the exact optimal solution for neural networks can be challenging and may lead to overfitting [7]. When necessary, we use the notation θ s ( S ; λ , ξ ) to indicate the final parameters obtained by training with the dataset S , along with hyperparameters λ ( e . g ., learning rate and number of epochs) and random choices ξ ( e . g ., parameter initialization and mini-batch order).

2.1 Training Data Attribution

TDA aims to explain model behavior on a query data point z q ( e . g ., test example) by referencing data points used to fit the model. The model behavior is typically quantified using a measurement f ( z q , θ ) , selected based on metrics relevant to the analysis, such as loss, margin, or log probability. Given hyperparameters λ and a training data point z m ∈ D , an attribution method τ ( z q , z m , D ; λ ) assigns

a score to a training data point, indicating its importance in influencing the expected measurable quantity E ξ [ f ( z q , θ s ( D ; λ , ξ ))] , where the expectation is taken over the randomness in the training process. In cases where an optimal solution to Equation (1) exists, is unique, and can be precisely computed, and TDA is performed on this optimal solution, the TDA method is simply written as τ ( z q , z m , D ) .

One idealized TDA method is leave-one-out (LOO) retraining [88], which assesses a data point's importance through counterfactual analysis. Assuming the above optimality condition is satisfied, for a chosen query data point z q and a training data point z m ∈ D , the LOO score can be formulated as:

When the measurement is defined as the loss, a higher absolute LOO score signifies a more substantial change in the query loss when the data point z m is excluded from the training dataset, particularly when the model parameters are optimized for convergence.

2.2 Influence Functions

Influence functions estimate the change in optimal parameters resulting from an infinitesimal perturbation in the weight of a training example z m ∈ D . Assuming that an optimal solution to Equation (1) exists and is unique for various values of the data point's weight ϵ ∈ [ -1 , 1] , the relationship between this weight and the optimal parameters is captured through the response function :

Influence functions approximate Equation (3) using the first-order Taylor expansion around ϵ = 0 :

where the Jacobian of the response function d r / d ϵ | ϵ =0 is obtained using the Implicit Function Theorem [52] and H := ∇ 2 θ J ( θ ⋆ , D ) represents the Hessian of the cost function at the optimal solution. The change in the optimal parameters due to the removal of z m can be approximated by setting ϵ = -1 :

By applying the chain rule of derivatives, influence functions estimate the change in a measurable quantity for a query example z q as:

We refer readers to Koh and Liang [49] for detailed derivations and discussions of influence functions. As observed in Equation (6), influence functions provide algorithms that only depend on the optimal parameters θ ⋆ rather than intermediate checkpoints. However, when applied to neural networks, the connection to the counterfactual prediction is tenuous due to the unrealistic assumptions that the optimal solution exists, is unique, and can be found [6, 3, 77]. In practice, the gradients and Hessian are computed using the final parameters θ s from a single training run instead of the optimal solution.

3 Methods

In this section, we introduce SOURCE, a gradient-based TDA technique combining the advantages of implicit and unrolled differentiation. We motivate our approach from the unrolling perspective and, after introducing suitable approximations, arrive at an influence-function-like algorithm. Finally, we describe a practical instantiation of SOURCE by approximating the Hessian with the Eigenvaluecorrected Kronecker-Factored Approximate Curvature (EK-FAC) [18] parameterization.

3.1 Motivation: Unrolling for Training Data Attribution

Consider optimizing the model parameters using SGD with a fixed batch size B , starting from the initial parameters θ 0 . 1 The update rule at each iteration is expressed as follows:

:    Training

1

2

Gradient

Accumulation

:

Computation Graph

where η k denotes the learning rate for iteration k , z ki is the i -th data point in B k , where B k denotes a mini-batch of examples drawn randomly with replacement from the training dataset D , and T denotes the total number of iterations. We aim to understand the effect of removing a training data point z m ∈ D on the terminal model parameters θ T . To this end, we parameterize the weight of z m as 1 + ϵ , where ϵ = 0 corresponds to the original training run and ϵ = -1 represents the removal of a data point. This parameterization results in the following update rule:

where δ ki := 1 [ z ki = z m ] is the indicator function for having selected z m . The dependence of θ on ϵ will usually be suppressed for brevity. Similarly to other gradient-based TDA methods, such as influence functions, we approximate the change in the terminal parameters due to the data removal θ T ( -1) -θ T (0) with its first-order Taylor approximation d θ T / d ϵ (the notation | ϵ =0 is suppressed as it will always be evaluated at ϵ = 0 ). Let δ k denote the number of times z m is chosen in batch B k . By chain rule, the contribution of iteration k to the total derivative d θ T / d ϵ can be found by multiplying all the Jacobian matrices along the accumulation path, giving the value -η k B δ k J k +1: T g k , where:

Here, H k := 1 B ∑ B i =1 ∇ 2 θ L ( z ki , θ k ) is the mini-batch Hessian for iteration k and we define J k : k := I for any 0 ≤ k < T by convention. A simplified illustration of unrolling is shown in Figure 1.

/epsilon1

=

-

1

/epsilon1

= 0

/epsilon1

= 1

Response Function

Unrolled Approximation

θ

0

tion. Unrolling-based TDA (but not implicit differentiation) accounts for this effect, resulting in a smaller influence along the low curvature direction.

In contrast to influence functions, unrolling does not assume uniqueness or convergence to the optimal solution. An illustrative comparison of the two approaches is shown in Figure 2. Exact influence functions differentiate the response function (Equation (4)), estimating the sensitivity of the optimal solution ( ⋆ ) to downweighting a data point. By contrast, unrolling estimates the sensitivity of the final model parameters (at the end of training) to downweighting a data point; hence, it can account for details of the training process such as learning rate schedules, implicit bias of optimizers, or a data point's position during training. For instance, in our illustrative example, gradient descent optimization is stopped early, such that the optimizer makes much progress in the high curvature direction and little in the low curvature direc-

The effect of removing z m on any single training trajectory may be noisy and idiosyncratic. For stability, we instead consider the expectation over training trajectories, where the selection of training examples in each batch (and all downstream quantities such as the iterates θ k ) are treated as random variables. 2 We are interested in the average treatment effect E [ θ T ( -1) -θ T (0)] , where the expectation is over the batch selection, and approximate this quantity with -E [ d θ T / d ϵ ] . The expected

total derivative can be expanded as a sum over all iterations, applying linearity of expectation:

In principle, we could compute a Monte Carlo estimate of this expectation by averaging many training trajectories. For each trajectory, d θ T / d ϵ can be evaluated using reverse accumulation ( i . e ., backpropagation) on the computation graph. However, this approach is prohibitively expensive as it requires storing all intermediate variables for the backward pass. Furthermore, many Monte Carlo samples may be required to achieve accurate estimates.

3.2 Segmenting the Training Trajectory

To derive a more efficient algorithm for approximating expected total derivative E [ d θ T / d ϵ ] , we now partition the training procedure into L segments and approximate the reverse accumulation computations for each segment with statistical summaries thereof (instead of storing all intermediate variables). Our motivations for segmenting the training procedure are twofold. First, the training procedure may explicitly include multiple stages with distinct objectives or datasets, as in continual learning or foundation models. Second, the Hessians and gradients are likely to evolve significantly over training, and segmenting the training allows us to approximate their distributions as stationary within a segment (rather than over the entire training run).

We index the segments as ℓ = 1 , . . . , L , with segment boundaries denoted as T ℓ . By convention, T L := T and T 0 := 0 denote the end of training and beginning of training, respectively, and K ℓ := T ℓ -T ℓ -1 denotes the total number of iterations within a segment. Conceptually, we can compute d θ T / d ϵ using reverse accumulation over a coarse-grained computation graph represented in terms of segments rather than individual iterations. The Jacobian associated with each segment is denoted as S ℓ := J T ℓ -1 : T ℓ . To approximate the expected total derivative E [ d θ T / d ϵ ] , we first rewrite Equation (10) using the segment notation introduced. We then approximate the Jacobians of different segments as statistically independent (see discussion below):

where ≈ uses our independence approximation to push the expectations inward. Note that our product notation ∏ ℓ +1 ℓ ′ = L takes ℓ ′ in decreasing order from L down to ℓ +1 .

To obtain tractable approximations for E [ S ℓ ] and E [ r ℓ ] , we approximate the Hessian and gradients distributions as stationary within each segment. This implies that the Hessians within a segment share a common mean ¯ H ℓ := E [ H k ] for T ℓ -1 ≤ k < T ℓ . Analogously, the gradients within a segment share a common mean ¯ g ℓ := E [ g k ] . Moreover, we approximate the step sizes within each segment with their mean ¯ η ℓ . If these stationarity approximations are too inaccurate ( e . g ., E [ H k ] and/or E [ g k ] change rapidly throughout the segment), one can improve the fidelity by carving the training trajectory into a larger number of segments, at the expense of increased computational and memory requirements. Finally, we approximate the Hessians and gradients in different time steps as statistically independent. 3

Approximation of E [ S ℓ ] . We approximate E [ S ℓ ] in Equation (11) as follows:

where the first ≈ uses the stationary and independence approximations, and the second ≈ uses the definition of matrix exponential. One can gain an intuition for ¯ S ℓ by observing that it is a matrix

function of ¯ H ℓ . 4 Let ¯ H ℓ = QΛQ ⊤ be the eigendecomposition of ¯ H ℓ and let σ j be the j -th eigenvalue of ¯ H ℓ . The expression in Equation (12) can be seen as applying the function F S ( σ ) := exp( -¯ η ℓ K ℓ σ ) to each of the eigenvalues σ of ¯ H ℓ . The value is close to zero in high-curvature directions, so the training procedure 'forgets' the components of θ which lie in these directions. However, information about θ is retained throughout the ℓ -th segment for low-curvature directions.

Approximation of E [ r ℓ ] . We further approximate E [ r ℓ ] as follows:

where Equation (13) uses the stationary and independence approximations and E [ δ k ] = B / N , and Equation (14) uses the finite series 5 and the definition of the matrix exponential. Because all the matrices commute, ¯ r ℓ can also be written in terms of a matrix function, defined as F r ( σ ) := (1 -exp( -¯ η ℓ K ℓ σ )) /σ . In high-curvature directions, this term approaches to 1 / σ , whereas in lowcurvature directions, it approaches to ¯ η ℓ K ℓ . The qualitative behavior of F r can be captured with the function F inv ( σ ) := 1 / ( σ + λ ) , where λ = ¯ η -1 ℓ K -1 ℓ , as shown in Figure 6 (Appendix C). Applying this to ¯ H ℓ results in approximating Equation (14) with the damped inverse-Hessian-vector product ( ¯ H ℓ + λ I ) -1 ¯ g ℓ . This is essentially the formula for influence functions, except that ¯ H ℓ and ¯ g ℓ represent the expected Hessian and gradient rather than the terminal one, and our analysis yields an explicit formula for the damping parameter λ . Hence, influence functions can be regarded approximately as a special case with only a single segment, so our damped unrolling analysis gives an alternative motivation for influence functions.

Full Procedure. We derived a closed-form term to approximate the expected total derivative:

where ¯ S ℓ and ¯ r ℓ are obtained with Equation (12) and Equation (14), respectively. We term our algorithm SOURCE ( S egmented stati O nary U n R olling for C ounterfactual E stimation) and refer readers to Figure 3 for a visual illustration. Similarly to unrolling, SOURCE can incorporate finegrained information about optimization trajectories into the analysis. For instance, SOURCE can support TDA for non-converged models, accounting for the total number of iterations T the model was trained with. It can also support TDA for multi-stage training pipelines: when the model was sequentially trained with two datasets D 1 and D 2 , SOURCE can compute the contribution of a data point z m ∈ D 1 that appeared in the first segment by partitioning the training trajectory into two segments and computing the expected total derivative at the first segment with -1 N 1 ¯ S 2 ¯ r 1 , where N 1 is the size of the first training dataset.

Given terminal parameters θ T from a single training run and a query data point z q , SOURCE approximates the change in the measurable quantity due to the removal of a training data point z m as:

Unlike the single-training-run estimator for unrolling-based approaches, SOURCE does not require access to the exact location where the data point z m was used during training, as it estimates the averaged effect of removing a data point within a given segment. To further account for other sources of randomness, such as model initialization, the multiple-training-run estimator for SOURCE averages the final scores in Equation (16) obtained for each training run with different random choices.

Segment 3

Segment 2

at Segment 1

at Segment 2

at Segment 3

Segment 1

3.3 Practical Algorithm for SOURCE

We now describe an instantiation of SOURCE which is practical to implement. Given the C model checkpoints saved during training, SOURCE begins by organizing them into L distinct segments. These segments may represent explicit stages in training ( e . g ., continual learning) or account for the change in Hessian and gradient throughout training. Within each segment ℓ , SOURCE estimates the stationary Hessian ¯ H ℓ and gradient ¯ g ℓ by averaging the Hessian and gradient across all checkpoints in the segment. SOURCE further estimates ¯ η ℓ by averaging the learning rates used within a segment.

However, computing Equation (15) has two practical bottlenecks for neural networks: computation of the Hessian and its matrix exponential. We fit a parametric approximation to the Hessian using EK-FAC [18]. EK-FAC parameterization is convenient for SOURCE as the approximate Hessian has an explicit eigendecomposition, which enables efficient computation of ¯ S ℓ and ¯ r ℓ by applying appropriate matrix functions to the eigenvalues. Note that EK-FAC approximates the Hessian with the Gauss-Newton Hessian (GNH) [63]. Unlike the Hessian, the GNH is guaranteed to be positive semidefinite, as long as the loss function is convex in the model outputs [62]. The GNH approximation within EK-FAC is also advantageous for SOURCE as it can avoid numerical instability in computing Equation (15), especially when the Hessian has negative eigenvalues. The implementation details are provided in Appendix D.

Compared to influence functions with the same EK-FAC approximation [24], SOURCE requires computing the EK-FAC factors and training gradients for each model checkpoint when performing TDA on all segments. Hence, SOURCE is C times more computationally expensive, where C is the number of checkpoints. In Appendix F.2, we introduce a more computationally efficient version of SOURCE, where we average the parameters within a segment instead of averaging Hessians and gradients. This variant of SOURCE is L times more computationally expensive than influence functions, as the EK-FAC factors and gradients only need to be computed once for each segment.

While we described one instantiation of SOURCE with the EK-FAC approximation, we note that SOURCE can be integrated with other techniques used for approximating implicit-differentiationbased TDA methods, such as TRAK [70], DATAINF [55], and LOGRA [11]. For example, with TRAK, we can use random projection [43] and efficiently compute the averaged Hessian and gradients in the lower-dimensional space. TRAK can be advantageous over the EK-FAC approximation when there are a large number of query data points, as it caches the compressed training gradients in memory, avoiding the need to recompute them for each query.

4 Related Works

Modern TDA techniques for neural networks can be broadly categorized into three main groups: sampling-based, representation-based, and gradient-based. For a comprehensive overview of TDA, including practical applications, we refer the reader to Hammoudeh and Lowd [27] and Mucsányi et al. [66]. Sampling-based (or retraining-based) approaches, such as Shapley-value estimators [78, 21, 39, 54, 87], DOWNSAMPLING [17, 96], DATAMODELS [35], and DATA BANZHAF [5, 86], approximate counterfactuals by repeatedly retraining models on different data subsets. Although effective, these methods are often impractical for modern neural networks due to the significant computational cost of repeated model retraining.

Representation-based techniques evaluate the relevance between a training and query data point by examining the similarity in their representation space ( e . g ., the output of the last hidden layer) [9, 30]. These techniques offer computational advantages compared to other attribution methods, as they only require forward passes through the trained network. Rajani et al. [74] further improves efficiency by caching all hidden representations of the training dataset and using approximate nearest neighbor search [42]. Past works have also proposed model-agnostic TDA approaches, such as computing the similarity between query and training sequences with BM25 [75] for language models [1, 56] or with an embedding vector obtained from a separate pre-trained self-supervised model for image classification tasks [79]. However, representation-based and input-similarity-based techniques lack a connection to the counterfactual and do not provide a notion of negatively influential data points.

Two main strategies for gradient-based TDA are implicit differentiation and unrolling. To the best of our knowledge, the largest model to which exact unrolling has been applied is a 300 thousand parameter model [31]. Our experiments in Section 5 cover TDA for models ranging from 560 thousand parameters (MNIST & MLP) to 120 million parameters (WikiText-2 & GPT-2). SGD-INFLUENCE [31] also considers applying unrolling to only the last epoch for large-scale models. However, this limits its applicability in analyzing the effect of removing a data point at the beginning of training or analyzing multi-stage training processes. In contrast, HYDRA [10] approximates the mini-batch Hessian H k in Equation (10) as zero when computing the total derivatives, avoiding the need to compute Hessian-vector products (HVPs) for each optimization step. However, in Appendix F.1, we empirically observe that an accurate approximation of the Hessian is important to achieve good TDA performance. Both approaches require storing a large number of optimization variables during training. Relatedly, Nickl et al. [68] use local perturbation methods [37] to approximate the data point's sensitivity to the training trajectory.

Apart from implicit-differentiation-based and unrolling-based approaches, TRACIN [72] is another prominent gradient-based TDA technique, which estimates the importance of a training data point by approximating the total change in the query's measurable quantity with the gradient update from this data point throughout training. Similarly to SOURCE, the practical version of TRACIN (TRACINCP) leverages intermediate checkpoints saved during training. While TRACINCP is straightforward to implement as it does not involve approximation of the Hessians, its connection to the counterfactual is unclear [27, 77]. However, past works have shown its strengths in downstream tasks, such as mislabeled data detection [72] and curating fine-tuning data [90].

5 Experiments

Our experiments investigate two key questions: (1) How does SOURCE compare to existing TDA techniques, as measured by the linear datamodeling score (LDS) [70] and through subset removal counterfactual evaluation [33, 93, 35, 97, 70, 8, 79, 19]? (2) Can SOURCE support data attribution in situations where implicit-differentiation-based approaches struggle, particularly with models that have not converged or have been trained in multiple stages with different objectives or datasets?

We compare SOURCE against existing TDA techniques: representation similarity (REPSIM) [9, 30], TRACIN [72], TRAK [70], and influence functions (IF) with the EK-FAC approximation [24]. For consistency with Park et al. [70], the measurement f is defined as the margin for classification tasks and the absolute error for regression tasks. Our evaluations are conducted under two separate settings. First is a single model setup, where TDA techniques use model checkpoints from a single training run. Unless specified otherwise, REPSIM, TRAK, and IF are computed at the final training checkpoint, and TRACIN and SOURCE use at most 6 intermediate checkpoints saved throughout training. In the second setting, TDA techniques use checkpoints from 10 distinct models, each trained with varying sources of randomness. Past works have shown that ensembling attribution scores across models can improve TDA performance [70, 67]. For all TDA techniques, including SOURCE, we simply average the final attribution scores from distinctly trained models with the full dataset, except for TRAK, which uses its custom ensembling procedures with models each trained on sampled 50% of the original dataset.

Our experiments consider diverse machine learning tasks, including: (a) regression using datasets from the UCI Repository [45], (b) image classification with datasets such as MNIST [57], FashionMNIST [91], CIFAR-10 [53], RotatedMNIST [20], and PACS [58], and (c) text classification using the GLUE benchmark [85]. Our tasks can be categorized into three groups:

0

.

0

0

.

5

Concrete (MLP)

0

.

0

0

.

5

FashionMNIST (MLP)

0

.

0

0

.

2

CIFAR-10 (ResNet-9)

0

.

0

0

.

5

RTE (BERT)

0

.

5

Concrete-N (MLP)

0

.

5

FashionMNIST-N (MLP)

0

.

2

RotatedMNIST (MLP)

0

.

1

0

.

2

PACS (ResNet-50)

LDS

REPSIM

TRACIN

TRAK

IF

SOURCE

1 Model

10 Models

0

.

0

0

.

0

0

.

0

0

.

0

Concrete, FashionMNIST, CIFAR-10, & RTE. Models fully trained using a fixed dataset D , where implicit-differentiation-based methods are expected to perform similarly to unrolling-based methods. We use 6 intermediate checkpoints throughout training for TRACIN and SOURCE. SOURCE use 3 segments ( L = 3 ) equally partitioned at the early, middle, and late stages of training to account for the changes in distributions of Hessian and gradients during training.

Concrete-N & FashionMNIST-N. Non-converged models trained with a smaller number of update steps. This is a challenging setup for implicit-differentiation-based methods, such as TRAK and IF, as they inherently assume that TDA is performed on the optimal solution. We use versions of the Concrete and FashionMNIST datasets that have been modified, either by corrupting target values or relabeling 30% of the data points. Then, we train the models for only 3 epochs to avoid overfitting. We use 3 checkpoints (at the end of each epoch) for TRACIN and SOURCE ( L = 3 ).

RotatedMNIST & PACS. Models initially trained with a dataset D 1 , and subsequently trained with another dataset D 2 (a common setup in continual learning). We use test examples from D 2 for query data points and attribute the final model's behavior to the first dataset. Since implicit-differentiationbased methods do not provide any way to separate multiple stages of training, for TRAK and IF, we simply combine the data from both stages into a larger dataset for TDA. We use two segments for SOURCE, partitioned at different stages, and perform TDA only for the first segment. Our experiments use the RotatedMNIST and PACS datasets, both containing multiple data distributions. We select one of these domains for the second training stage, while the remaining ones are used in the first stage.

The detailed description of the experimental setup is provided in Appendix E. Additional results, including comparisons on additional tasks and with additional baselines, further analysis on linear models, and visualizations of the top influential images obtained by each TDA technique, are shown in Appendix F.

5.1 TDA Evaluations with Linear Datamodeling Score (LDS)

We evaluate TDA techniques using the linear datamodeling score (LDS) from Park et al. [70]. To compute LDS, we first generate M random subsets {S j } j M =1 from the training dataset, each containing ⌈ αN ⌉ data points for some α ∈ (0 , 1) . Given a query data point z q and hyperparameters λ used to train the original model, the expected measurable quantity for each data subset E ξ [ f ( z q , θ s ( S j ; λ , ξ ))] is estimated by retraining the model R times under different random choices (which requires MR model retrainings in total). The LDS measures the Spearman correlation [81] between the estimated quantities and the predictions made by the TDA method. Note that, although a TDA method in Section 2.1 assigns a score to each pair of a query and training data point, the inherently additive nature of most TDA techniques allows for the computation of a group prediction score for the data subset S by summing the individual scores attributed to each data point within this subset. The final LDS is obtained by averaging the scores across many (typically up to 2000 ) query data points. We use 100 data subsets ( M = 100 ) and conduct a minimum of 5 retraining iterations ( R ≥ 5 ) for each subset. We refer readers to Appendix A for the detailed formulation and to Appendix E.2 for the practical procedures.

The LDS at α = 0 . 5 for SOURCE and baseline TDA techniques are shown in Figure 4. SOURCE consistently outperforms all baseline methods in a single model setup, achieving high LDS. When aggregating TDA scores from multiple models, we observe a large improvement in the LDS, particu-

0

100

200

300

0

.

00

0

.

25

0

.

50

FashionMNIST (MLP)

0

500

1000

0

.

00

0

.

25

0

.

50

CIFAR-10 (ResNet-9)

0

50

100

0

.

0

0

.

5

RTE (BERT)

0

100

200

300

0

.

0

0

.

5

FashionMNIST-N (MLP)

0

1000

2000

0

.

0

0

.

1

RotatedMNIST (MLP)

Frac. Misclassified Test Examples

Number of Training Examples Removed to Flip Predictions

RANDOM

REPSIM

TRACIN

TRAK

IF

SOURCE

1 Model

10 Models

larly for TRAK, IF, and SOURCE. Our method achieves the highest LDS across all tasks, except for the CIFAR-10 classification task using ResNet-9. SOURCE especially performs strongly against other baseline techniques on settings that pose challenges to implicit-differentiation-based approaches ( e . g ., non-converged models and models trained with multiple stages), and indeed, even the non-ensembled version of SOURCE typically outperforms the ensembled versions of the competing methods.

5.2 TDA Evaluations with Subset Removal

Subset removal counterfactual evaluation examines the change in model behavior before and after removing data points highly ranked by a TDA technique. For classification tasks, we consider 100 test data points that are correctly classified when trained with the full dataset (across all 5 random seeds) and, for each test data point z q , examine if removing and retraining without the topk positively influential data points can cause misclassification on average (over 3 random seeds). By assessing the impact of removing influential training examples on the model's performance, counterfactual evaluation provides a direct measure of the effectiveness of TDA techniques in identifying data points that significantly contribute to the model's behavior. The detailed procedures are described in Appendix E.3. In Figure 5, we show the fraction of test examples (out of the selected 100 test points) that get misclassified on average after removing at most k positively influential training examples identified by each TDA method. We observe that SOURCE better identifies the top influential data points causing misclassification than other baseline TDA techniques. The improvement is more substantial for settings that pose challenges to implicit-differentiation-based methods.

6 Conclusion

We introduced SOURCE ( S egmented stati O nary U n R olling for C ounterfactual E stimation), a novel TDA technique that combines the strengths of implicit-differentiation-based and unrolling-based techniques. SOURCE approximates unrolled differentiation by partitioning the training trajectory into one or more segments and approximating the gradients and Hessians as stationary within each segment, yielding an influence-function-like estimation algorithm. We showed one instantiation of SOURCE by approximating the Hessian with the EK-FAC parameterization. On a diverse task set, we demonstrated SOURCE's effectiveness compared to existing data attribution techniques, especially when the network has not converged or has been trained with multiple stages.

References

E. Akyürek, T. Bolukbasi, F. Liu, B. Xiong, I. Tenney, J. Andreas, and K. Guu. Towards tracing knowledge in language models back to the training data. In Findings of the Association for Computational Linguistics: EMNLP 2022 , pages 2429-2446, 2022.

W. E. Arnoldi. The principle of minimized iterations in the solution of the matrix eigenvalue problem. Quarterly of applied mathematics , 9(1):17-29, 1951.

J. Bae, N. Ng, A. Lo, M. Ghassemi, and R. B. Grosse. If influence functions are the answer, then what is the question? Advances in Neural Information Processing Systems , 35:17953-17967, 2022.

J. Bae, P. Vicol, J. Z. HaoChen, and R. B. Grosse. Amortized proximal optimization. Advances in Neural Information Processing Systems , 35:8982-8997, 2022.

J. F. Banzhaf III. Weighted voting doesn't work: A mathematical analysis. Rutgers L. Rev. , 19: 317, 1964.

S. Basu, P. Pope, and S. Feizi. Influence functions in deep learning are fragile. In International Conference on Learning Representations , 2020.

Y. Bengio. Practical recommendations for gradient-based training of deep architectures. In Neural Networks: Tricks of the Trade: Second Edition , pages 437-478. Springer, 2012.

J. Brophy, Z. Hammoudeh, and D. Lowd. Adapting and evaluating influence-estimation methods for gradient-boosted decision trees. Journal of Machine Learning Research , 24(154):1-48, 2023.

R. Caruana, H. Kangarloo, J. D. Dionisio, U. Sinha, and D. Johnson. Case-based explanation of non-case-based learning methods. In Proceedings of the AMIA Symposium , page 212. American Medical Informatics Association, 1999.

Y. Chen, B. Li, H. Yu, P. Wu, and C. Miao. Hydra: Hypergradient data relevance analysis for interpreting deep neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 35, pages 7081-7089, 2021.

S. K. Choe, H. Ahn, J. Bae, K. Zhao, M. Kang, Y. Chung, A. Pratapa, W. Neiswanger, E. Strubell, T. Mitamura, et al. What is your data worth to gpt? llm-scale data valuation with influence functions. arXiv preprint arXiv:2405.13954 , 2024.

J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding, 2018.

L. Engstrom, A. Feldmann, and A. Madry. DsDm: Model-aware dataset selection with datamodels, 2024.

J. R. Epifano, R. P. Ramachandran, A. J. Masino, and G. Rasool. Revisiting the fragility of influence functions. Neural Networks , 162:581-588, 2023.

R. Eschenhagen, A. Immer, R. Turner, F. Schneider, and P. Hennig. Kronecker-factored approximate curvature for modern neural network architectures. Advances in Neural Information Processing Systems , 36, 2024.

M. Fang, N. Z. Gong, and J. Liu. Influence function based data poisoning attacks to topn recommender systems. In Proceedings of The Web Conference 2020 , pages 3019-3025, 2020.

V. Feldman and C. Zhang. What neural networks memorize and why: Discovering the long tail via influence estimation. Advances in Neural Information Processing Systems , 33:2881-2891, 2020.

T. George, C. Laurent, X. Bouthillier, N. Ballas, and P. Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems , 31, 2018.

K. Georgiev, J. Vendrow, H. Salman, S. M. Park, and A. Madry. The journey, not the destination: How data guides diffusion models, 2023.

M. Ghifary, W. B. Kleijn, M. Zhang, and D. Balduzzi. Domain generalization for object recognition with multi-task autoencoders. In Proceedings of the IEEE international conference on computer vision , pages 2551-2559, 2015.

A. Ghorbani and J. Zou. Data Shapley: Equitable valuation of data for machine learning. In International Conference on Machine Learning , pages 2242-2251. PMLR, 2019.

R. Grosse. University of Toronto CSC2541, Topics in Machine Learning: Neural Net Training Dynamics, Chapter 4: Second-Order Optimization. Lecture Notes, 2021. URL https://www.cs.toronto.edu/~rgrosse/courses/csc2541_2021/ readings/L04_second_order.pdf .

R. Grosse and J. Martens. A kronecker-factored approximate fisher matrix for convolution layers. In International Conference on Machine Learning , pages 573-582. PMLR, 2016.

R. Grosse, J. Bae, C. Anil, N. Elhage, A. Tamkin, A. Tajdini, B. Steiner, D. Li, E. Durmus, E. Perez, E. Hubinger, K. Lukoši¯ ut˙ e, K. Nguyen, N. Joseph, S. McCandlish, J. Kaplan, and S. R. Bowman. Studying large language model generalization with influence functions, 2023.

I. Gulrajani and D. Lopez-Paz. In search of lost domain generalization, 2020.

V. Gupta, T. Koren, and Y. Singer. Shampoo: Preconditioned stochastic tensor optimization. In International Conference on Machine Learning , pages 1842-1850. PMLR, 2018.

Z. Hammoudeh and D. Lowd. Training data influence analysis and estimation: A survey. Machine Learning , pages 1-53, 2024.

F. R. Hampel. The influence curve and its role in robust estimation. Journal of the american statistical association , 69(346):383-393, 1974.

X. Han, B. C. Wallace, and Y. Tsvetkov. Explaining black box predictions and unveiling data artifacts through influence functions. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics , pages 5553-5563, 2020.

K. Hanawa, S. Yokoi, S. Hara, and K. Inui. Evaluation of similarity-based explanations. In International Conference on Learning Representations , 2020.

S. Hara, A. Nitanda, and T. Maehara. Data cleansing for models trained with sgd. Advances in Neural Information Processing Systems , 32, 2019.

K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 770778, 2016.

S. Hooker, D. Erhan, P.-J. Kindermans, and B. Kim. A benchmark for interpretability methods in deep neural networks. Advances in neural information processing systems , 32, 2019.

E. J. Hu, Y. Shen, P. Wallis, Z. Allen-Zhu, Y. Li, S. Wang, L. Wang, and W. Chen. Lora: Low-rank adaptation of large language models, 2021.

A. Ilyas, S. M. Park, L. Engstrom, G. Leclerc, and A. Madry. Datamodels: Predicting predictions from training data. In International Conference on Machine Learning , 2022.

M. R. Izadi, Y. Fang, R. Stevenson, and L. Lin. Optimization of graph neural networks with natural gradient descent. In 2020 IEEE international conference on big data , pages 171-179. IEEE, 2020.

L. A. Jaeckel. The infinitesimal jackknife . Bell Telephone Laboratories, 1972.

M. Jagielski, G. Severi, N. Pousette Harger, and A. Oprea. Subpopulation data poisoning attacks. In Proceedings of the 2021 ACM SIGSAC Conference on Computer and Communications Security , pages 3104-3122, 2021.

R. Jia, D. Dao, B. Wang, F. A. Hubis, N. Hynes, N. M. Gürel, B. Li, C. Zhang, D. Song, and C. J. Spanos. Towards efficient data valuation based on the shapley value. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 1167-1176. PMLR, 2019.

R. Jia, F. Wu, X. Sun, J. Xu, D. Dao, B. Kailkhura, C. Zhang, B. Li, and D. Song. Scalability vs. utility: Do we have to sacrifice one for the other in data importance quantification? In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 8239-8247, 2021.

K. F. Jiang, W. Liang, J. Zou, and Y. Kwon. Opendataval: A unified benchmark for data valuation, 2023.

J. Johnson, M. Douze, and H. Jégou. Billion-scale similarity search with GPUs. IEEE Transactions on Big Data , 7(3):535-547, 2019.

W. B. Johnson, J. Lindenstrauss, and G. Schechtman. Extensions of lipschitz maps into banach spaces. Israel Journal of Mathematics , 54(2):129-138, 1986.

K. K and A. Søgaard. Revisiting methods for finding influential examples, 2021.

M. Kelly, R. Longjohn, and K. Nottingham. The UCI machine learning repository, 2023. URL https://archive.ics.uci.edu .

R. Khanna, B. Kim, J. Ghosh, and S. Koyejo. Interpreting black box predictions using fisher kernels. In The 22nd International Conference on Artificial Intelligence and Statistics , pages 3382-3390. PMLR, 2019.

S. Kim, K. Kim, and E. Yang. GEX: A flexible method for approximating influence via geometric ensemble. Advances in Neural Information Processing Systems , 36, 2024.

D. P. Kingma and J. Ba. Adam: A method for stochastic optimization, 2014.

P. W. Koh and P. Liang. Understanding black-box predictions via influence functions. In International Conference on Machine Learning , pages 1885-1894. PMLR, 2017.

S. Kong, Y. Shen, and L. Huang. Resolving training biases via influence-based data relabeling. In International Conference on Learning Representations , 2021.

N. Konz, C. Godfrey, M. Shapiro, J. Tu, H. Kvinge, and D. Brown. Attributing learned concepts in neural networks to training data, 2023.

S. G. Krantz and H. R. Parks. The Implicit Function Theorem: History, theory, and applications . Springer Science & Business Media, 2002.

A. Krizhevsky and G. Hinton. Learning multiple layers of features from tiny images. 2009.

Y. Kwon and J. Zou. Beta Shapley: A unified and noise-reduced data valuation framework for machine learning. In International Conference on Artificial Intelligence and Statistics , pages 8780-8802. PMLR, 2022.

Y. Kwon, E. Wu, K. Wu, and J. Zou. DataInf: Efficiently estimating data influence in lora-tuned llms and diffusion models. In International Conference on Learning Representations , 2023.

F. Ladhak, E. Durmus, and T. B. Hashimoto. Contrastive error attribution for finetuned language models. In Proceedings of the 61st Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers) , pages 11482-11498, 2023.

Y. LeCun, C. Cortes, and C. Burges. MNIST handwritten digit database. ATT Labs , 2, 2010.

D. Li, Y. Yang, Y.-Z. Song, and T. M. Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conference on computer vision , pages 5542-5550, 2017.

D. C. Liu and J. Nocedal. On the limited memory BFGS method for large scale optimization. Mathematical programming , 45(1-3):503-528, 1989.

Z. Liu, H. Ding, H. Zhong, W. Li, J. Dai, and C. He. Influence selection for active learning. In Proceedings of the IEEE/CVF International Conference on Computer Vision , pages 9274-9283, 2021.

I. Loshchilov and F. Hutter. Decoupled weight decay regularization. In International Conference on Learning Representations , 2018.

J. Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research , 21(146):1-76, 2020.

J. Martens and R. Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International Conference on Machine Learning , pages 2408-2417. PMLR, 2015.

J. Martens, J. Ba, and M. Johnson. Kronecker-factored curvature approximations for recurrent neural networks. In International Conference on Learning Representations , 2018.

S. Merity, C. Xiong, J. Bradbury, and R. Socher. Pointer sentinel mixture models. In International Conference on Learning Representations , 2016.

B. Mucsányi, M. Kirchhof, E. Nguyen, A. Rubinstein, and S. J. Oh. Trustworthy machine learning, 2023.

E. Nguyen, M. Seo, and S. J. Oh. A bayesian approach to analysing training data attribution in deep learning. Advances in Neural Information Processing Systems , 36, 2024.

P. Nickl, L. Xu, D. Tailor, T. Möllenhoff, and M. E. E. Khan. The memory-perturbation equation: Understanding model's sensitivity to data. Advances in Neural Information Processing Systems , 36, 2024.

S. Oh, B. Ustun, J. McAuley, and S. Kumar. Rank list sensitivity of recommender systems to interaction perturbations. In Proceedings of the 31st ACM International Conference on Information & Knowledge Management , pages 1584-1594, 2022.

S. M. Park, K. Georgiev, A. Ilyas, G. Leclerc, and A. Madry. TRAK: Attributing model behavior at scale. In International Conference on Machine Learning , pages 27074-27113. PMLR, 2023.

A. Paszke, S. Gross, S. Chintala, G. Chanan, E. Yang, Z. DeVito, Z. Lin, A. Desmaison, L. Antiga, and A. Lerer. Automatic differentiation in PyTorch. 2017.

G. Pruthi, F. Liu, S. Kale, and M. Sundararajan. Estimating training data influence by tracing gradient descent. Advances in Neural Information Processing Systems , 33:19920-19930, 2020.

A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever. Language models are unsupervised multitask learners. 2019.

N. F. Rajani, B. Krause, W. Yin, T. Niu, R. Socher, and C. Xiong. Explaining and improving model behavior with k nearest neighbor representations, 2020.

S. E. Robertson, S. Walker, S. Jones, M. M. Hancock-Beaulieu, and M. Gatford. Okapi at TREC-3. Nist Special Publication Sp , 109:109, 1995.

A. Schioppa, P. Zablotskaia, D. Vilar, and A. Sokolov. Scaling up influence functions. In Proceedings of the AAAI Conference on Artificial Intelligence , volume 36, pages 8179-8186, 2022.

A. Schioppa, K. Filippova, I. Titov, and P. Zablotskaia. Theoretical and practical perspectives on what influence functions do. Advances in Neural Information Processing Systems , 36, 2024.

L. Shapley. A value for n -person games. 1953.

V. Singla, P. Sandoval-Segura, M. Goldblum, J. Geiping, and T. Goldstein. A simple and efficient baseline for data attribution on images. arXiv preprint arXiv:2311.03386 , 2023.

J. W. Smith, J. E. Everhart, W. Dickson, W. C. Knowler, and R. S. Johannes. Using the adap learning algorithm to forecast the onset of diabetes mellitus. In Proceedings of the annual symposium on computer application in medical care , page 261. American Medical Informatics Association, 1988.

C. Spearman. The proof and measurement of association between two things. The American journal of psychology , 100(3/4):441-471, 1987.

T. Tieleman and G. Hinton. Lecture 6.5-rmsprop: Divide the gradient by a running average of its recent magnitude. COURSERA: Neural networks for machine learning , 4(2):26-31, 2012.

A. Tsanas and M. Little. Parkinsons Telemonitoring. UCI Machine Learning Repository, 2009. DOI: https://doi.org/10.24432/C5ZS3N.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

A. Wang, A. Singh, J. Michael, F. Hill, O. Levy, and S. R. Bowman. GLUE: A multi-task benchmark and analysis platform for natural language understanding, 2019.

J. T. Wang and R. Jia. Data Banzhaf: A robust data valuation framework for machine learning. In International Conference on Artificial Intelligence and Statistics , pages 6388-6421. PMLR, 2023.

J. T. Wang, Y. Zhu, Y.-X. Wang, R. Jia, and P. Mittal. A privacy-friendly approach to data valuation. Advances in Neural Information Processing Systems , 36, 2024.

S. Weisberg and R. D. Cook. Residuals and influence in regression. 1982.

T. Wolf, L. Debut, V. Sanh, J. Chaumond, C. Delangue, A. Moi, P. Cistac, T. Rault, R. Louf, M. Funtowicz, J. Davison, S. Shleifer, P. von Platen, C. Ma, Y. Jernite, J. Plu, C. Xu, T. L. Scao, S. Gugger, M. Drame, Q. Lhoest, and A. M. Rush. Transformers: State-of-the-art natural language processing. In Proceedings of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations , pages 38-45, Online, Oct. 2020. Association for Computational Linguistics. URL https://www.aclweb.org/anthology/ 2020.emnlp-demos.6 .

M. Xia, S. Malladi, S. Gururangan, S. Arora, and D. Chen. LESS: Selecting influential data for targeted instruction tuning, 2024.

H. Xiao, K. Rasul, and R. Vollgraf. Fashion-MNIST: A novel image dataset for benchmarking machine learning algorithms, 2017.

C.-K. Yeh, J. Kim, I. E.-H. Yen, and P. K. Ravikumar. Representer point selection for explaining deep neural networks. Advances in neural information processing systems , 31, 2018.

C.-K. Yeh, A. Taly, M. Sundararajan, F. Liu, and P. Ravikumar. First is better than last for language data influence. Advances in Neural Information Processing Systems , 35:32285-32298, 2022.

I.-C. Yeh. Concrete Compressive Strength. UCI Machine Learning Repository, 2007. DOI: https://doi.org/10.24432/C5PK67.

S. Zagoruyko and N. Komodakis. Wide residual networks, 2017.

C. Zhang, D. Ippolito, K. Lee, M. Jagielski, F. Tramèr, and N. Carlini. Counterfactual memorization in neural language models. Advances in Neural Information Processing Systems , 36: 39321-39362, 2023.

X. Zheng, T. Pang, C. Du, J. Jiang, and M. Lin. Intriguing properties of data attribution on diffusion models. In International Conference on Learning Representations , 2023.

A Evaluation of TDA Techniques

Given the focus on counterfactual prediction in many TDA methods, LOO estimates, defined in Equation (2), are often considered a ground truth for evaluating these techniques. However, the computation of LOO scores in neural networks encounters several computational and conceptual challenges, as detailed in Appendix B. For a robust and standardized measure for evaluating TDA techniques, we instead use the linear datamodeling score (LDS) from Park et al. [70] as well as subset removal counterfactual evaluation [33, 93, 35, 97, 70, 8, 79, 19].

Linear Datamodeling Score (LDS). A TDA method τ , as detailed in Section 2.1, assigns a score to each pair of a query and training data point. The inherently additive nature of most TDA techniques allows for the computation of a group attribution score for a specific training data subset S ⊂ D . The importance of S on the measurable quantity f is estimated by summing the individual scores attributed to each data point within this subset. The group attribution is expressed as follows:

Consider M random subsets {S j } j M =1 from the training dataset, each containing ⌈ αN ⌉ data points for some α ∈ (0 , 1) . Given a hyperparameter configuration λ to train the model, the LDS for a query point z q is defined as:

where ρ represents the Spearman correlation [81]. This expected measurable quantity is approximated by retraining the network R times under different random choices. The final LDS is obtained by averaging the scores across many (typically up to 2000 ) query data points. In our experiments, we use 100 data subsets ( M = 100 ) and conduct a maximum of 100 retraining iterations ( R ∈ { 5 , 10 , 20 , 100 } ) for each subset to compute the LDS.

Subset Removal Counterfactual Evaluation. Subset removal counterfactual evaluation examines the change in model behavior before and after removing data points that are highly ranked by an attribution technique. For classification tasks, we consider 100 test data points that are correctly classified when trained with the full dataset and, for each test data point, examine if removing and retraining without the topk positively influential data points can cause misclassification on average (trained under different random choices). 6 By assessing the impact of removing influential data points on the model's performance, counterfactual evaluation provides a direct measure of the effectiveness of TDA techniques in identifying data points that significantly contribute to the model's behavior.

Downstream Task Evaluation. TDA techniques have also been evaluated on their performance on downstream tasks, such as mislabeled data detection [46, 72, 47], class detection [30, 55], finding hallucinations in the training dataset [56], and retrieving factual knowledge from the training dataset [1]. These tasks can offer additional insights into the effectiveness and applicability of data attribution methods in practical scenarios. However, the connections between these tasks and counterfactual prediction are often unclear [44, 70], and it is uncertain whether algorithmic improvements in counterfactual prediction will directly result in improved performance on these downstream tasks.

B Limitations of Leave-One-Out Estimates

The computation of leave-one-out (LOO) scores in Equation (2) presents several computational and conceptual challenges for neural networks. Firstly, calculating the LOO score for all training data points requires retraining the model N times, where N is the size of the training dataset. This process can be prohibitively expensive for large datasets and network architectures.

Moreover, the formulation of LOO assumes that an optimal solution to Equation (1) exists, is unique, and can be precisely computed, and that TDA is performed on this optimal solution. However, within the context of neural networks, these assumptions often do not hold, leading to ambiguities in the

computation of LOO estimates. Previous works have investigated various LOO variants as a means to establish counterfactual ground truths [49, 6, 44, 40, 3, 14, 67]. For example, Koh and Liang [49] and Basu et al. [6] formulated the LOO ground truth by training the network for an additional number of steps from the final parameters θ s without a specific training data point. However, as noted by Bae et al. [3], these estimates may reflect the effect of training the network for additional steps instead of retraining without a data point, especially when the network has not converged.

A more standardized extension of LOO for neural networks is the expected leave-one-out (ELOO) retraining [44], formulated as:

where λ denotes the hyperparameters used to train the model, and the expectation is taken over the randomness in the training process (typically estimated by retraining the network R times). Note that the ELOO can also be seen as the ground truth for the linear datamodeling score (LDS) (defined in Appendix A) with α = 1 -1 / N . Past works have demonstrated the unreliability of ELOO estimates due to stochasticity in model training, such as model initialization and batch ordering [44, 14, 67]. Specifically, Nguyen et al. [67] observed that the noise from the stochasticity often overshadows the actual signal of removing a single data point. In Appendix F.3, we also observe that the LDS significantly drops at α = 1 -1 / N , suggesting that the counterfactual ground truth for removing a single data point can be difficult to obtain, as it can be extremely noisy for most training data examples.

C SOURCE with Preconditioning Matrix

In Section 3.1, we motivated our proposed algorithm, SOURCE, for cases where the parameters are optimized using stochastic gradient descent (SGD). In this section, we present the formulation of SOURCE when preconditioned optimizers, such as RMSProp [82], Adam [48], and K-FAC [63], are used to train the model.

To investigate the impact of removing a training data point z m ∈ D , we follow a similar derivation as in Section 3.1, but now considering the preconditioning matrix:

where P k is a (positive definite) preconditioning matrix and δ ki := 1 [ z ki = z m ] is the indicator function for having selected z m .

By applying the chain rule of derivatives, the contribution of iteration k to the total derivative can be found by multiplying all the Jacobian matrices along the backward accumulation path, giving the value -η k B J k +1: T P k g k , where we have:

Hence, by applying the linearity of expectation, the expected total derivative of the terminal parameters θ T with respect to the perturbation ϵ is expressed as:

As discussed in Section 3.2, we group the training trajectories into multiple segments to approximate the expected total derivative for each segment with statistical summaries thereof. In addition to the approximations introduced in Section 3.2, we approximate preconditioning matrices as stationary within a segment and represent it as ˜ P ℓ := P k for T ℓ -1 ≤ k < T ℓ .

Approximation of E [ S ℓ ] . We approximate E [ S ℓ ] in Equation (11) as follows:

Note that the last line uses the properties of the matrix exponential. 7

Approximation of E [ r ℓ ] . We further approximate E [ r ℓ ] as follows:

where we define M ℓ := ¯ P 1 / 2 ℓ ¯ H ℓ ¯ P 1 / 2 ℓ and the last line uses the properties of matrix exponential, as done in Equation (23). Similarly to our analysis presented in Section 3.2, we can represent ¯ r ℓ with the matrix function of M ℓ . Let M ℓ = QΛQ ⊤ be the eigendecomposition of M ℓ and let σ j be the j -th eigenvalue of M ℓ . The expression can be seen as applying the matrix function, defined as:

The qualitative behavior of F r can be captured with the function F inv ( σ ) := 1 / ( σ + λ ) , where λ = ¯ η -1 ℓ K -1 ℓ (see Section 3.2 for details). Hence, one way to understand Equation (28) is by expressing it as the damped inverse Hessian-vector product (iHVP):

In a case where ˜ P ℓ is a diagonal matrix, Equation (32) can be seen as a special case for influence functions with a specific diagonal damping term λ ˜ P -1 ℓ . Using the derived ¯ S ℓ and ¯ r ℓ , we approximate the total expected derivative using Equation (15).

D Implementation Details

This section describes the Eigenvalue-corrected Kronecker-Factored Approximate Curvature (EKFAC) [18] and how we computed SOURCE using this EK-FAC parameterization. The code for implementing SOURCE (as well as baseline techniques) will be provided at https://github.com/ pomonam/kronfluence . For details on the EK-FAC approximation specific to influence functions, we refer readers to Grosse et al. [24].

10

-

4

10

-

2

10

0

10

2

σ

10

-

2

10

-

1

10

0

10

1

F

r

F

inv

D.1 Eigenvalue-corrected Kronecker-Factored Approximate Curvature (EK-FAC)

Kronecker-Factored Approximate Curvature (K-FAC) [63] and EK-FAC [18] introduce a parametric approximation to the Fisher information matrix (FIM) of a neural network, defined as:

where p data is the data distribution and P ˆ y | x ( θ ) is the model's output distribution. For many commonly used loss functions, such as softmax-cross-entropy and squared-error, the FIM is equivalent to the Gauss-Newton Hessian (GNH) [62], denoted as G . The GNH can be seen as an approximation to the Hessian H , where the network is linearized around the current parameters [22]. Different from the Hessian, the GNH is guaranteed to be positive semi-definite (PSD) when the loss function is convex with respect to the model output.

While K-FAC and EK-FAC were originally formulated for multilayer perceptrons (MLPs), they were later extended to other architectures, such as convolutional neural networks [23], recurrent neural networks [64], graph neural networks [36], or to be learnable by gradient-based optimizers [4]. We refer readers to Eschenhagen et al. [15] for a comprehensive overview. This section describes the EK-FAC formulation in the context of MLPs.

Consider a l -th layer of the network with input activations a l -1 ∈ R I and pre-activation output s l ∈ R O such that s l := W l a l -1 , where W ∈ R O × I is the weight matrix (we drop the layer subscript to avoid clutter and ignore the bias term for simplicity). The pseudo-gradient (where the target is sampled from the model's output distribution; see Equation (33)) is given by D W := D sa ⊤ . K-FAC makes two core approximations: (1) layerwise independence approximation, where GNH is approximated as block-diagonal with each block corresponding to GNH of some specific layer, and (2) input activations a and pseudo-gradient of the pre-activations D s are independent under the model's predictive distribution. The layerwise GNH can be approximated as:

where ⊗ denotes the Kronecker product. The matrices A ∈ R I × I and S ∈ R O × O in Equation (34) represent the uncentered covariance matrices of the activations and the pseudo-gradients with respect to the pre-activations, respectively. These covariance matrices can be estimated by computing the statistics over many data batches and taking the average.

Denoting the eigendecomposition of these covariance matrices as A = Q A Λ A Q ⊤ A and S = Q S Λ S Q ⊤ S , using properties of the Kronecker product, we can express the eigendecomposition of A ⊗ B as:

EK-FAC introduces a more accurate approximation to the GNH by introducing a compact representation of the eigenvalues (instead of representing them as the Kronecker product Λ A ⊗ Λ S ). The layerwise GNH for EK-FAC is represented as follows:

Here, the corrected eigenvalues Λ ∈ R IO × IO are defined as:

The corrected eigenvalues in Equation (37) minimize the approximation error with the GNH measured by the Frobenius norm, where we refer readers to George et al. [18] for the derivations.

D.2 EK-FAC Computations for SOURCE

As detailed in Section 3.3, our practical instantiation of SOURCE requires averaging the Hessians across checkpoints within a segment. We use a common averaging scheme in the optimization literature [63, 18, 26] to compute the averaged EK-FAC factors. We first compute the activation covariance matrices A and pseudo-gradient covariance matrices S for all model checkpoints. These matrices are obtained by computing the statistics over all data points once ( 1 epoch). Then, we take the average over these covariance matrices to obtain ¯ A = 1 C ℓ ∑ C ℓ k =1 A k and ¯ S = 1 C ℓ ∑ C ℓ k =1 S k , where C ℓ is the total number of model checkpoints for the ℓ -th segment and A k and S k are covariance matrices for the k -th checkpoint. Then, we perform eigendecomposition on these averaged covariance matrices to obtain the eigenvectors ¯ Q A and ¯ Q S . Under the eigenbasis ¯ Q A ⊗ ¯ Q S , we compute the corrected eigenvalues Λ k for each model checkpoint (Equation (37)) and then average the eigenvalues to obtain ¯ Λ . In summary, the averaged (Gauss-Newton) Hessian for a particular segment is approximated as:

SOURCE requires computing the covariance matrices and corrected eigenvalues for each model checkpoint. Moreover, calculating the TDA scores for all training data points requires computing the training gradients C times, where C is the total number of checkpoints. Hence, SOURCE is approximately C times more computationally expensive than influence functions evaluated at the final checkpoint. In Section 3.3, we introduced a more efficient variant, which averages the parameters within a segment instead. This variant only needs to compute the EK-FAC factors once for each segment and requires computing the EK-FAC factors and gradients L times. Hence, it is L times more computationally expensive than influence functions.

When the model is trained with SGD with a heavy ball momentum β (SGDm), we scaled the learning rate used in SOURCE as ¯ η ℓ (1 -β ) -1 to account for the effective learning rate (terminal velocity). In cases where AdamW optimizers are used as in Appendix C, computing the matrix exponential for ¯ P 1 / 2 ℓ ¯ H ℓ ¯ P 1 / 2 ℓ is challenging with EK-FAC. We additionally keep track of the diagonal Hessian approximation (which can be easily and efficiently obtained when computing the corrected eigenvalues in Equation (37)) and use the diagonal Hessian approximation for computing the matrix exponential in Equation (23) and Equation (27). Note that we still use the EK-FAC factors to compute ¯ H -1 ℓ ¯ g ℓ in Equation (27).

D.3 Applicability to Other Approximation Techniques

While we described one instantiation of SOURCE with the EK-FAC approximation, SOURCE can be integrated with other techniques used for approximating implicit-differentiation-based TDA methods, such as TRAK [70] and DATAINF [55]. For example, as in TRAK, we can use random projection [43] to efficiently compute the averaged Hessian and gradients in a lower-dimensional space. TRAK is advantageous over the EK-FAC approximation when there are many query data points, as it caches compressed training gradients in memory, avoiding recomputing them for each query.

E Experimental Setup

This section describes the experimental setup used to obtain the results presented in Section 5. This includes a description of each task (Appendix E.1) and the methodology for computing the linear datamodeling score (LDS) (Appendix E.2). Implementation details of the subset removal counterfactual evaluation and baseline techniques are provided in Appendix E.3 and Appendix E.4, respectively. All experiments were conducted using PYTORCH version 2 . 1 . 0 [71]. We used CPUs to conduct UCI regression experiments, A100 (80GB) GPUs to conduct GLUE and WikiText-2 experiments, and A6000 (48GB) GPUs for other experiments. A single GPU was used to run SOURCE and baseline techniques. The internal cluster was used to run the experiments. Since our ground truth requires a lot of model retraining with different random choices ( e . g ., initialization and batch ordering), we considered tasks that train (or fine-tune) with less than 20 minutes using the abovementioned compute resources. For example, the generation of the LDS ground truth for a given task and data sampling ratio α ∈ (0 , 1) takes at most 210 hours of computational resources.

E.1 Datasets and Models

We conducted systematic hyperparameter optimization for all tasks. This process involved conducting grid searches to find hyperparameter configurations that achieve the best average validation performance (accuracy for classification tasks and loss for others). The average validation performance was obtained by retraining the network 5 times using different random seeds. For models trained with SGD with a heavy ball momentum of 0 . 9 (SGDm), our search spaces for learning rate and weight decay were { 3 e1 , 1 e1 , 3 e2 , 1 e2 , 3 e3 , 1 e3 , 3 e4 , 1 e4 , 3 e5 , 1 e5 } and { 3 e2 , 1 e2 , 3 e3 , 1 e3 , 3 e4 , 1 e4 , 3 e5 , 1 e5 , 0 . 0 }, respectively. For models trained with AdamW [61], the search spaces were { 1 e2 , 3 e3 , 1 e3 , 3 e4 , 1 e4 , 3 e5 , 1 e5 } for learning rate and { 3 e2 , 1 e2 , 3 e3 , 1 e3 , 3 e4 , 1 e4 , 3 e5 , 1 e5 , 0.0} for weight decay. In cases where the original experimental setup from which we adapted had a pre-specified learning rate and weight decay, these hyperparameters were incorporated into our search space.

UCI Datasets (Regression). For regression tasks, we used the Concrete [94] and Parkinson [83] datasets from the UCI Machine Learning Repository [45]. Both datasets were pre-processed to have a zero mean and unit variance for input features and targets. We trained a three-layer multilayer perceptron (MLP), where each layer consisted of 128 hidden units and the RELU activation function. The models were optimized using SGDm for 20 epochs with a batch size of 32 and a constant learning rate schedule. A learning rate of 3 e2 and a weight decay of 1 e5 were used for the Concrete dataset. For the Parkinson dataset, the learning rate was set to 1 e2 with a weight decay value of 3 e5 . We saved 6 intermediate checkpoints throughout training. For the noisy Concrete (Concrete-N) dataset, we randomly modified 30% of the targets by sampling from a Normal distribution with zero mean and unit variance. We used the same hyperparameters but trained the models for 3 epochs.

MNIST & FashionMNIST (Image Classification). Following the experimental setup from Koh and Liang [49] and Bae et al. [3], we trained a three-layer multilayer perceptron (MLP) on approximately 10% of MNIST [57] and FashionMNIST [91] datasets. Smaller versions of these datasets were used to compute the counterfactual ground truth more efficiently. The models were trained with SGDm for 20 epochs with a batch size of 64 and a constant learning rate. The learning rate and weight decay were set for both datasets to 3 e2 and 1 e3 , respectively. We saved 6 checkpoints during training and utilized them for TRACIN and SOURCE. For the noisy FashionMNIST (FashionMNST-N) experiment, we randomly relabeled 30% of the training dataset. The network was only trained for 3 epochs with a learning rate 1 e2 and weight decay 3 e5 .

CIFAR-10 (Image Classification). For the CIFAR-10 dataset [53], we trained the ResNet-9 model [32], 8 following the standard data augmentation procedure from Zagoruyko and Komodakis [95]. This included extracting images from a random 32 × 32 crop after applying zero-padding of 4 pixels, with a 50% probability of horizontal flipping. The network was trained for 25 epochs using SGDm with a batch size of 512 and a cyclic learning rate schedule, peaking at 0 . 5 . The initial learning rate was set to 0 . 4 with a weight decay of 1 e3 , and 6 intermediate checkpoints were saved throughout training.

GLUE (Text Classification). We fine-tuned the BERT model [12] on SST-2, RTE, and QNLI datasets from the GLUE benchmark [85] with the training script from the Transformers library [89]. 9 Following the experimental setup from Park et al. [70], we capped the training dataset at a maximum of 51200 examples to compute the LDS efficiently. However, we did not modify the original architecture ( e . g ., removing the last TANH layer) and trained the network with the AdamW optimizer. The weight decay was set to 1 e2 for all tasks, and the learning rates were set as follows: 3 e5 for SST-2, 1 e5 for QNLI, and 2 e5 for RTE. We saved 6 intermediate checkpoints for each training run.

WikiText-2 (Language Modeling). For the language modeling task, we fine-tuned the GPT-2 model [73] using the WikiText-2 dataset [65]. We followed the training script from the Transformer

library but set the maximum sequence length to 512 . 10 During fine-tuning with AdamW, we saved 6 intermediate checkpoints for data attribution. The learning rate, weight decay, and batch size were set to 3 e5 , 1 e2 , and 8 , respectively. We set the measurement f as a loss for the language modeling task.

RotatedMNIST & PACS (Image Classification). We used the RotatedMNIST dataset [20] and the PACS dataset [58], following the data pre-processing procedures from Gulrajani and Lopez-Paz [25]. 11 The training process was divided into two distinct stages for both tasks. During the initial stage of the training, we trained the network with the dataset D 1 , while the second stage used dataset D 2 . For RotatedMNIST, the first dataset D 1 was comprised of images rotated at 0 , 15 , 45 , and 60 degrees, whereas the second dataset D 2 contained images rotated at 30 degrees. We trained a three-layer MLP for 30 ( 20 / 10 ) epochs using SGDm and a batch size of 128 . The learning rate and weight decay were set to 1 e1 and 1 e5 . For PACS, the first dataset D 1 included images from the cartoon, photo, and sketch categories, and the second dataset D 2 had art paintings. We fine-tuned ResNet-50 [32], initialized from the pre-trained parameters, 12 using SGDm for 40 ( 30 / 10 ) epochs with a batch size of 128 , a learning rate of 1 e4 , and a weight decay of 3 e5 .

E.2 Linear Datamodeling Score

We follow a methodology proposed by Park et al. [70] to compute the linear datamodeling score (LDS). Let λ represent the set of hyperparameters used for training the model on a specified task, such as the choice of optimizer and the number of training epochs. Let α ∈ (0 , 1) denote the data sampling ratio. The process for obtaining the LDS involves several steps:

We generate M data subsets, denoted as {S j } j M =1 , each being a uniformly sampled subset of the original training dataset D . Each subset S j ⊂ D contains ⌈ αN ⌉ data points, where N denotes the total number of training data points.

For each data subset S j , the model is trained R times using different random seeds { ξ r } R r =1 ( e . g ., model initialization and batch ordering).

Given an attribution method τ and a query example z q , we measure the Spearman correlations [81] between the prediction and the estimated expected measurable quantity:

where g represents the group attribution prediction, expressed as:

To obtain the final LDS, we average the correlations over a set of query data points (up to 2000 in our experiments) and report the score with 95% bootstrap confidence intervals, which accounts for resampling of the data subset S j (see Park et al. [70] for details).

For a given data sampling ratio α , the networks must be retrained MR times in total to compute the LDS ground truth. In our experiments, we used 100 subsets ( M = 100 ). The repeat R was set to 100 for UCI regression tasks, 10 for MNIST classification tasks, 20 for CIFAR-10 image classification task, 5 for GLUE text classification and WikiText language modeling task, and 20 for RotatedMNIST and PACS image classification tasks. We used the largest feasible R based on our computational budget because we observed improvements in LDS for baseline techniques (especially TRAK, IF, and SOURCE) with larger R .

E.3 Subset Removal Counterfactual Evaluation

For the subset removal counterfactual evaluation, we first train the model with the full dataset D under different random choices (over 5 random seeds) and select 100 test data points correctly classified on

all random choices. Then, for each test data point and attribution technique, we remove the topk data points from the pre-defined interval k 1 , . . . , k I (such that k 1 < · · · < k I ), as indicated as highly positively influential by the data attribution technique, retrain the network with this modified dataset, and examine if the original test data point gets misclassified on average under different random choices (over 3 random seeds). Finally, for each value of k in the pre-defined interval, we report the fraction of test data points that get misclassified after removing at most topk training data points and retraining the network with the modified dataset.

For each TDA technique, this process requires retraining the model 100 × I × 3 times, where I is the pre-defined interval size. We set I = 6 for all experiments, leading to the retraining of the model 1800 times. To reduce the computational cost, we start from the smallest subset removal size k 1 , and if the test data point gets misclassified under the current subset, we do not consider it for the larger subset removal size ( e . g ., k 2 ). Hence, this can be seen as the fraction of test data points that get misclassified by removing at most k training data points (evaluated at a fixed interval). We note that Singla et al. [79] instead use a bisection search to find the smallest subset size in which a test data point can be misclassified, whereas Ilyas et al. [35] use more fine-grained intervals with more number of seeds ( e . g ., 8 intervals and 20 seeds). While it is possible to use a larger number of seeds ( 20 seeds as in Ilyas et al. [35]), because of computational limitations, we use 3 seeds to estimate the averaged misclassification. We used SOURCE and baseline techniques described in Appendix E.4 to identify positively influential training data points. We also included a RANDOM baseline, where we removed the training data points belonging to the same class as the target test example.

E.4 Baselines

This section describes the baseline techniques used in Section 5. Unless specified otherwise, we describe them in the context of a single-training-run estimator, where the TDA techniques use the final parameters θ s obtained with hyperparameters λ and some random choice ξ (the multiple-trainingruns estimators simply average the TDA scores obtained from models trained with different random choices ξ ).

Representation Similarity (REPSIM). Representation similarity technique [9] evaluates the importance of a training data point z m ∈ D to a specific query data point z q by comparing the latent representations of these data point pairs. This can be formulated as follows:

Here, similarity ( v 1 , v 2 ) , where v 1 and v 2 are some vectors, is typically defined through the ℓ 2 metric, dot metric, or cosine metric [30]. In our experiments, the function ϕ θ s ( z ) was designed to map a data point to its last hidden activations (before the final output layer), using a forward pass through the final parameters θ s . We used the cosine metric to compute the attribution score but observed similar performance when using the ℓ 2 metric, aligning with observations in previous studies [35, 70, 79].

TRACIN. We used the TRACINCP estimator from Pruthi et al. [72], defined as:

where C represents the number of checkpoints, ˆ θ k represents the parameters at the k -th checkpoint, and η k is the learning rate applied at the corresponding checkpoint. The last checkpoint is typically set to the final model parameters θ s . While there is an option to compress the gradients using a random projection as suggested by Pruthi et al. [72], our experiments used the full gradients to obtain a stronger baseline. The checkpoint selection details are described in Appendix E.1.

Influence Functions (IF). As detailed in Section 2.2, training data attribution with influence functions is formulated as follows:

where H denotes the Hessian of the cost at the final parameters θ s . To make influence functions scalable to large neural networks, we used the Eigenvalue-corrected Kronecker-Factored Approximate Curvature (EK-FAC) parameterization [18] to approximate the Hessian, as proposed by Grosse et al.

[24]. We refer readers to Grosse et al. [24] and Appendix D for details on the EK-FAC computation. Relatedly, Schioppa et al. [76] use Arnoldi iterations [2], and Kwon et al. [55] utilize the parameterefficient fine-tuning (PEFT) [34] strategy to efficiently approximate influence functions. More recently, Choe et al. [11] further utilized low-dimensional gradient projection to compute influence functions more efficiently.

While Grosse et al. [24] only consider the computation of influence scores to the MLP layers of transformers [84], in our experiments, we extended this computation to include the attention layers as well. We excluded layer normalization, batch normalization, and embedding layers from the influence computation. Influence functions have an additional hyperparameter λ > 0 , which is used to compute the damped inverse Hessian-vector product (IHVP), denoted as ( H + λ I ) -1 v for some vector v . We used a small damping term for consistency with TRAK [70] and set it to 1 e8 to avoid numerical instability (note that TRAK sets the damping term to 0 ).

TRAK. In contrast to the traditional formulation of influence functions, TRAK [70] leverages random projections [43], Generalized Gauss-Newton approximation, and ensembling for data attribution. Specifically, given a random projection matrix P ∼ N (0 , 1) M × K , where K denotes the projection dimension, the final model parameters θ s , and a model output function f ( z , θ ) , TRAK projects all training and query gradients into K -dimensional vectors. The feature map is defined as:

We further define Φ := [ ϕ 1 ; . . . ; ϕ N ] ∈ R N × K as stacked projected gradients for all training data points, where each ϕ i corresponds to ϕ ( z i ) . Subsequently, TRAK's single model estimator is formulated as:

with Q being a N × N diagonal matrix for weightings. Here, τ TRAK represents a vector of dimension N , containing attribution score for each training data point. TRAK uses an ensemble of single model estimators, each derived from models trained with distinct configurations and projection matrices. We refer readers to Park et al. [70] and Engstrom et al. [13] for detailed derivations and discussions of TRAK.

We used the final checkpoints for TRAK in our experimental setup involving a single model. We computed TRAK using the last checkpoint of 10 differently trained models (each trained with 50% of the dataset) for experiments with multiple model setups. TRAK has a hyperparameter that determines the dimension of the random projection K . We set the projection dimension to 20480 for ResNet-9 and RotatedMNIST, 8192 for ResNet-50 on the PACS dataset, 1024 for BERT trained on the RTE dataset and 512 for MLP trained on the Concrete dataset (due to the datasets' smaller size), and 4096 for all other tasks. All experiments were conducted using TRAK's official implementation. 13

Empirical Influence (EI). To compute the empirical influence (DOWNSAMPLING) [17], we first create M data subsets {S j } j M =1 , each being a uniformly sampled subset of the original training dataset. Each subset S i contains ⌈ αN ⌉ data points, where α ∈ (0 , 1) is the data sampling ratio. Given a training data point z m ∈ D , we define M m as the total number of data subsets containing z m . The empirical influence scores are formulated as follows:

where 1 [ · ] is an indicator function to determine if the training data point z m is contained in the j -th data subset S j . Intuitively, Equation (46) computes the averaged query measurement when data point z m is not used in training, whereas Equation (47) computes the averaged measurement when the data point is used in training. Following Zheng et al. [97], we created 512 data subsets ( M = 512 ) with a sampling ratio α = 0 . 5 , which requires retraining the model 512 times with 50% of training data points removed.

0

.

00

0

.

25

LDS

FashionMNIST (MLP)

0

.

0

0

.

1

CIFAR-10 (ResNet-9)

0

.

00

0

.

25

RTE (BERT)

0

.

0

0

.

5

FashionMNIST-N (MLP)

IF

FAST-SOURCE

SOURCE

HYDRA. We used the fast version of HYDRA [10], formulated as:

where T represents the total number of gradient update steps, θ k denotes the parameters at the k -th iteration, and η k is the corresponding learning rate. Here, B k denotes the batch of data points used at the corresponding update, and 1 [ z m ∈ B k ] is the indicator function for having selected z m in the update. Note that HYDRA requires storing all parameter vectors used for training. We refer readers to Hammoudeh and Lowd [27] for derivations and detailed discussions of HYDRA.

F Additional Results

In this section, we present additional experimental results, including a comparison with additional baseline TDA techniques (Appendix F.1), an LDS evaluation of a computationally faster variant of SOURCE (Appendix F.2), an LDS evaluation at various sampling ratios and for more tasks (Appendix F.3), counterfactual evaluation on linear models (Appendix F.4), and visualizations of the top positively and negatively influential training data points for each TDA technique (Appendix F.5).

F.1 Additional Baseline Comparisons

We compare SOURCE with empirical influence (DOWNSAMPLING) [17] and the fast version of HYDRA [10] on the FashionMNIST task. Results for these techniques on other tasks were omitted, since DOWNSAMPLING requires retraining the model over 500 times and HYDRA necessitates saving all intermediate checkpoints throughout training. The implementation details are provided in Appendix E.4, and the results are shown in Table 1. SOURCE achieves the highest LDS on both single and multiple model setups compared to existing baseline TDA techniques we considered.

F.2 SOURCE with Averaged Parameters

In Section 3.3, we introduced a more computationally efficient version of SOURCE, which averages the parameters within a segment instead of Hessians and gradients. Here, we present the LDS results at α = 0 . 5 for the faster version, termed FAST-SOURCE, for FashionMNIST, CIFAR-10, RTE,

0.3

0.5

0.7

0.9

N

-

1

N

0

.

0

0

.

5

LDS

Concrete (MLP)

0.3

0.5

0.7

0.9

N

-

1

N

0

.

0

0

.

1

CIFAR-10 (ResNet-9)

0.3

0.5

0.7

0.9

N

-

1

N

0

.

00

0

.

25

RTE (BERT)

REPSIM

TRACIN

TRAK

IF

SOURCE (

L

= 1

)

SOURCE (

L

= 3

)

α

0

.

0

0

.

5

Concrete (MLP)

0

.

0

0

.

5

Parkinsons (MLP)

0

.

0

0

.

5

MNIST (MLP)

0

.

0

0

.

5

FashionMNIST (MLP)

0

.

0

0

.

2

LDS

CIFAR-10 (ResNet-9)

0

.

0

0

.

5

RTE (BERT)

0

.

0

0

.

2

SST-2 (BERT)

0

.

0

0

.

2

QNLI (BERT)

0

.

25

WikiText-2 (GPT-2)

REPSIM

TRACIN

TRAK

IF

SOURCE

1 Model

10 Models

0

.

00

and FashionMNIST-N tasks. The results are shown in Figure 7. We observe that FAST-SOURCE outperforms influence functions on these tasks, while it generally achieves a lower LDS compared to SOURCE.

F.3 Additional LDS Results

Here, we present additional results, evaluating TDA techniques with LDS at various sampling ratios and considering more tasks, where models are fully trained using a fixed dataset. We first consider computing the LDS across a range of data sampling ratios α . (The procedures to compute the LDS are described in Appendix E.2.) The performance of SOURCE and other baseline attribution methods is shown in Figure 8. SOURCE consistently achieves higher LDS than the baseline methods across diverse α values. However, an exception is noted at α = 1 -1 / N ( e . g ., removing a single training data point), where a significant drop in correlations is observed for all TDA methods. This finding is consistent with previous studies that highlight the limitations of LOO estimates in reliably evaluating attribution techniques [44, 14, 67] (see Appendix B for a detailed discussion). Additionally, our results suggest that while SOURCE with a single segment can be effective, using multiple segments typically improves LDS performance. Lastly, we observe that the relative rankings of TDA techniques typically remain consistent across various α values

Next, we present the LDS results at α = 0 . 5 for additional tasks in Figure 9. SOURCE consistently outperforms baseline methods in a single model setup, achieving higher correlations with the ground truth. When aggregating TDA scores from multiple models, we observe a large improvement in the LDS. Our method obtains the highest LDS across all tasks, except for the CIFAR-10 classification task using ResNet-9. However, we show that our method outperforms baseline methods on the CIFAR-10 task for subset removal counterfactual evaluation in Section 5.2. We note that the tasks in Figure 9

0.5

0.7

0.9

N

-

1

N

α

0

.

8

1

.

0

LDS

Diabetes

0.5

0.7

0.9

N

-

1

N

α

0

.

5

1

.

0

Concrete

0

20

40

Epochs

0

.

8

0

.

9

Concrete

TRACIN

IF

SOURCE

0.5

0.7

0.9

N

-

1

N

0

.

99

1

.

00

LDS

Diabetes

0.5

0.7

0.9

N

-

1

N

0

.

98

0

.

99

1

.

00

Concrete

α

α

consider models sufficiently trained near convergence (with carefully chosen hyperparameters), where implicit-differentiation-based methods are expected to perform similarly to unrolling-based methods.

F.4 Counterfactual Evaluations on Linear Models

In this section, we demonstrate the effectiveness of SOURCE on linear models when the model has not been trained until convergence. We trained linear regression on the Concrete dataset and logistic regression on the Diabetes dataset [80] for 3 epochs with a batch size of 32 . We also constructed the LDS ground truth using SGD with the same hyperparameters. We applied TRACIN, IF, and SOURCE (with L = 1 ) to the trained model and computed the LDS for various data sampling ratios α . The results are shown in Figure 10 (Left & Middle). SOURCE achieves higher LDS on all data sampling ratios for both regression and classification tasks. We further show the LDS at α = 0 . 9 with varying numbers of epochs in Figure 10 (Right). (The LDS ground truth is recomputed at each epoch.) We observe a larger LDS gap between SOURCE and IF when the model was only trained for a small number of epochs, and the gap reduces as we train the model for a larger number of iterations. These results show that our formulation for SOURCE better supports TDA when the network has not fully converged, even in the case of linear models.

For completeness, we show the LDS for influence functions when the TDA is performed on the optimal solution in Figure 11. For each model, we computed the optimal solution (for logistic regression, we used the L-BFGS [59]), computed the influence function estimates, and evaluated their accuracy with LDS (also obtained by computing the optimal solution without some data points). As shown in Figure 11, influence functions obtain high correlations with the ground truth across various values of data sampling ratio α . In contrast to neural network experiments in Appendix F.3, we observe an increase in the LDS as the data sampling ratio α increases (predicting the effect of removing a smaller number of data points), as the group influence predictions introduce more approximation error [3]. Notably, we obtain a high LDS when α = ( N -1) / N (removing a single data point), as the LDS is computed at the precise optimal solution (see Appendix B for the discussion). TRACIN and SOURCE are not applicable in these contexts, as we computed the optimal solution with the direct solution or with L-BFGS, instead of with gradient descent.

F.5 Qualitative Results

We first present the top positively and negatively influential data points obtained by each TDA technique on multiple model settings. Note that for these multiple model settings, REPSIM, TRACIN, TRAK, IF, and SOURCE use an ensemble of 10 models trained with different random choices. The

results for FashionMNIST, CIFAR-10, and RotatedMNIST are shown in Figure 12, Figure 13, and Figure 14, respectively. We also show the top positively and negatively influential data points on the CIFAR-10 dataset for a single model setup in Figure 15. In Table 2, we present the top positively and negatively influential data points obtained by SOURCE on the RTE dataset.

G Broader Impact & Limitations of SOURCE

Broader Impact. Our paper focuses on improving training data attribution, especially in cases where traditional implicit-differentiation-based methods, such as influence functions, struggle. While our work does not have a direct societal impact, as it focuses on algorithmic improvements, there may be societal implications for improving TDA techniques. On the positive side, as shown in Grosse et al. [24], TDA techniques can be used to understand and debug the misbehavior of neural networks ( e . g ., LLMs), which can promote trust in deploying machine learning systems. By identifying the training data points responsible for specific model behaviors, TDA can help improve model interpretability, fairness, and robustness. This, in turn, can lead to more reliable and equitable AI systems that benefit society. However, TDA techniques also allow for the analysis of the impact of training data points on trained models, which can be used for crafting data poisoning attacks [16, 38, 69]. Malicious actors could potentially use TDA to identify and manipulate influential data points, leading to the creation of biased or misleading models. This could have negative consequences, such as the spread of disinformation or unfair treatment of specific groups. To mitigate these risks, it is essential to develop responsible practices for using TDA techniques.

Limitations. Compared to the influence function employing the same EK-FAC parameterization [24], the practical implementation of the SOURCE requires the computation of EK-FAC factors and gradients for all checkpoints (when performing TDA on all segments). Note that, when TDA is performed on one specific segment ℓ , the gradients only need to be computed for checkpoints within a segment (instead of all checkpoints). Denoting the total number of checkpoints as C and the total number of segments as L , SOURCE on all segments exhibits an approximate computational cost of C times higher. Our experiments used configurations with C ∈ { 3 , 6 } and L ∈ { 2 , 3 } . We also introduced a faster version of SOURCE in Appendix F.2, which directly averages the parameters instead of averaging the EK-FAC factors and gradients; the faster version is L times computationally expensive compared to the EK-FAC influence functions.

Compared to implicit-differentiation-based TDA techniques, SOURCE requires access to intermediate checkpoints throughout the training process and corresponding hyperparameters such as learning rate, number of iterations, and preconditioning matrix. In cases where the details of the training process are not available, implicit-differentiation-based TDA techniques, such as TRAK [70] and influence functions [49, 24], may be preferable.

Moreover, SOURCE approximates the distributions of the Hessian and gradient as stationary within each segment of the training trajectory. In certain scenarios, this may not be a reasonable approximation. For instance, when pre-training large transformer models, the Hessian or gradients may undergo drastic changes throughout the training process. If the stationarity approximation is too inaccurate, one can enhance the fidelity of SOURCE by dividing the training trajectory into a larger number of segments, albeit at the cost of increased computational requirements. While we used a fixed number of segments and checkpoints, partitioned equally at the early, middle, and late stages of training, we can extend SOURCE by automatically determining when to segment by examining the changes in the Hessian or gradients, which we leave for future work. Lastly, SOURCE approximate the Hessians and gradients at different time steps as statistically independent to obtain a tractable approximation for the expected total derivative in Section 3.2. As discussed, this independence approximation amounts to neglecting the autocorrelation of optimization iterates.

REPSIM

pullover

TRACIN

trouser

pullover

TRAK

trouser

pullover

IF

trouser

pullover

SOURCE

trouser

trouser

trouser

Top Positively Influential Images

trouser

trouser

shirt

t-shirt

Top Negatively Influential Images

dress

pullover

pullover

pullover

pullover

pullover

pullover

shirt

t-shirt

shirt

t-shirt

Top Positively Influential Images

trouser

trouser

trouser

pullover pullover

pullover

Top Positively Influential Images

trouser

trouser

trouser

pullover pullover

pullover

Top Positively Influential Images

pullover shirt

coat pullover

pullover shirt

Top Positively Influential Images trouser

pullover pullover

pullover

trouser

pullover trouser

pullover shirt

pullover trouser

pullover

Top Negatively Influential Images

shirt

bag bag

bag

Top Negatively Influential Images

shirt

shirt shirt

shirt t-shirt

dress

Top Negatively Influential Images shirt

sandal pullover

shirt pullover

t-shirt

Top Negatively Influential Images shirt

shirt shirt

t-shirt

bag dress

shirt shirt

shirt shirt

t-shirt

REPSIM

truck

truck

truck

Top Positively Influential Images

truck

truck

dog

dog

Top Negatively Influential Images

dog

dog

ship

ship

ship

ship

ship

bird

deer

bird

bird

TRACIN

truck

truck

truck

Top Positively Influential Images

truck

truck

car

car

Top Negatively Influential Images

car

car

ship

ship

ship

ship

ship

airplane

airplane

truck

airplane

TRAK

truck

truck

truck

Top Positively Influential Images

car

cat

car

car

Top Negatively Influential Images

dog

bird

ship

ship

ship

ship

ship

car

ship

truck

car

IF

truck

truck

truck

Top Positively Influential Images

truck

truck

car

car

Top Negatively Influential Images

car

car

ship

ship

ship

ship

ship

truck

car

airplane

airplane

SOURCE

truck

truck

truck

Top Positively Influential Images

truck

truck

car

car

Top Negatively Influential Images

car

car

ship

ship

ship

ship

ship

car

car

truck

car

REPSIM

0

0

0

Top Positively Influential Images

0

0

1

1

Top Negatively Influential Images

8

1

5

5

5

5

5

2

2

2

2

TRACIN

0

0

0

Top Positively Influential Images

0

0

4

9

Top Negatively Influential Images

6

9

5

5

5

5

5

3

3

8

6

TRAK

0

0

0

Top Positively Influential Images

0

0

1

6

Top Negatively Influential Images

7

6

5

6

6

5

6

3

6

3

6

IF

0

0

0

Top Positively Influential Images

0

0

4

6

Top Negatively Influential Images

0

0

5

5

5

5

6

3

6

6

6

SOURCE

0

0

0

Top Positively Influential Images

0

0

5

5

Top Negatively Influential Images

5

3

5

5

5

5

5

3

3

3

3

555

CO

36

5

55

3

5

REPSIM

truck

truck

truck

Top Positively Influential Images

truck

truck

dog

deer

Top Negatively Influential Images

deer

deer

ship

ship

ship

ship

ship

deer

deer

deer

deer

TRACIN

truck

truck

truck

Top Positively Influential Images

truck

truck

car

car

Top Negatively Influential Images

car

car

ship

ship

ship

ship

ship

airplane

airplane

airplane

airplane

TRAK

truck

ship

dog

Top Positively Influential Images

dog

frog

dog

frog

Top Negatively Influential Images

bird

dog

ship

deer

dog

bird

cat

cat

deer

bird

horse

IF

truck

truck

truck

Top Positively Influential Images

truck

truck

car

car

Top Negatively Influential Images

car

car

ship

ship

ship

ship

ship

car

airplane

car

car

SOURCE

truck

truck

truck

Top Positively Influential Images

truck

truck

car

car

Top Negatively Influential Images

car

car

ship

ship

ship

ship

ship

airplane

car

car

car

NeurIPS Paper Checklist

1. Claims

Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?

Answer: [Yes]

Justification: Claims made in the abstract and introduction reflect the paper's contribution and scope.

Guidelines:

The answer NA means that the abstract and introduction do not include the claims made in the paper.

The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.

The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.

It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: We discuss the limitations of SOURCE in Appendix G. These limitations are also mentioned throughout the paper, such as in Section 3.2 (core approximations) and in Section 3.3 (additional computational overheads).

Guidelines:

The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.

The authors are encouraged to create a separate "Limitations" section in their paper.

The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.

The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.

The authors should reflect on the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.

The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.

If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.

3. Theory Assumptions and Proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?

Answer: [Yes]

Justification: We provide a full set of approximations in Section 3.2 and describe the equations in depth throughout the paper.

Guidelines:

The answer NA means that the paper does not include theoretical results.

All the theorems, formulas, and proofs in the paper should be numbered and crossreferenced.

All assumptions should be clearly stated or referenced in the statement of any theorems.

The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.

Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.

Theorems and Lemmas that the proof relies upon should be properly referenced.

4. Experimental Result Reproducibility

Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?

Answer: [Yes]

Justification: We provide detailed information to reproduce the main experimental results of the paper in Appendix E. Implementation details are provided in Appendix D.

Guidelines:

The answer NA means that the paper does not include experiments.

If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.

If the contribution is a dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.

Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.

While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example

(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.

(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).

(b) If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.

(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.

5. Open access to data and code

Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material?

Answer: [No]

Justification: We do not release the code. However, we described the data preparation procedures and the URLs to baseline techniques we used as a comparison in Appendix E. We used publicly available data for our experiments.

Guidelines:

The answer NA means that paper does not include experiments requiring code.

Please see the NeurIPS code and data submission guidelines ( https://nips.cc/ public/guides/CodeSubmissionPolicy ) for more details.

While we encourage the release of code and data, we understand that this might not be possible, so 'No' is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).

The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines ( https: //nips.cc/public/guides/CodeSubmissionPolicy ) for more details.

The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.

The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only a subset of experiments are reproducible, they should state which ones are omitted from the script and why.

At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).

Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted.

6. Experimental Setting/Details

Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: We provide detailed information about the training and test details in Appendix E.

Guidelines:

The answer NA means that the paper does not include experiments.

The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.

The full details can be provided either with the code, in appendix, or as supplemental material.

7. Experiment Statistical Significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: For the LDS evaluation, we show the error bars representing 95% bootstrap confidence intervals, which account for resamplings of data subsets. The subset removal evaluation requires at most 1800 model retrainings for each pair of task and TDA technique. The reported numbers are averages from 3 seeds. Details are provided in Appendix E.2 and Appendix E.3, respectively.

Guidelines:

The answer NA means that the paper does not include experiments.

The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)

The assumptions made should be given (e.g., Normally distributed errors).

It should be clear whether the error bar is the standard deviation or the standard error of the mean.

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have a 96% CI, if the hypothesis of Normality of errors is not verified.

For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).

If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.

8. Experiments Compute Resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: We provide information on the computer resources (CPU or GPU), which were used to conduct our experiments Appendix E.

Guidelines:

The answer NA means that the paper does not include experiments.

The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.

The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.

The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).

9. Code Of Ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: Our research conforms to the NeurIPS Code of Ethics.

Guidelines:

The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.

If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.

The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction).

10. Broader Impacts

Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?

Answer: [Yes]

Justification: We provide discussions on the broader impact of our work in Appendix G.

Guidelines:

The answer NA means that there is no societal impact of the work performed.

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.

Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.

The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.

If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).

11. Safeguards

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: Our paper does not pose such risks.

Guidelines:

The answer NA means that the paper poses no such risks.

Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.

Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.

We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.

12. Licenses for existing assets

Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?

Answer: [Yes]

Justification: We properly cite the original owners of the code, data, and models used in the paper (they are all publicly available).

Guidelines:

The answer NA means that the paper does not use existing assets.

The authors should cite the original paper that produced the code package or dataset.

The authors should state which version of the asset is used and, if possible, include a URL.

The name of the license (e.g., CC-BY 4.0) should be included for each asset.

For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

If this information is not available online, the authors are encouraged to reach out to the asset's creators.

13. New Assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [NA]

Justification: We do not release any new assets.

Guidelines:

The answer NA means that the paper does not release new assets.

Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.

The paper should discuss whether and how consent was obtained from people whose asset is used.

At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.

14. Crowdsourcing and Research with Human Subjects

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?

Answer: [NA]

Justification: Our paper does not involve any crowdsourcing, and we do not perform research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Including this information in the supplemental material is fine, but if the main contribution of the paper involves human subjects, then as much detail as possible should be included in the main paper.

According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.

15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects

Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?

Answer: [NA]

Justification: Our paper does not involve any crowdsourcing, and we do not perform research with human subjects.

Guidelines:

The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.

For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a training data attribution (TDA) method based on implicit differentiation and unrolling. The goal is to estimate the effect of removing (or changing the weight) of a training example on the final (not necessarily optimal) parameters, accounting for training details influencing the trajectory of the parameters during training.\n\nThe method approximates the change in the final parameters $\\boldsymbol\\theta_T$ ($T$ is the number of optimization iterations) based on approximating the effect of re-weighting a training example with weight $1+\\epsilon$. This requires computation of $\\partial \\boldsymbol\\theta_T/\\partial\\epsilon$, which is a product of Jacobians $\\partial\\boldsymbol\\theta_{k}/\\partial\\boldsymbol\\theta_{k-1}$ from all training steps $k$., each computed from a Hessian of the loss. To reduce computational requirements, the authors apply different approximations while accounting for all parts of the training trajectory. The resulting algorithm, SOURCE, requires 2-18 more computation (Hessian estimations) than influence functions.\n\nExperiments show that SOURCE outperforms other TDA methods in counterfactual evaluations on linear datamodeling score (LDS) (Park et al., 2023) and on standard subset removal counterfactual evaluation.\n\nSTRENGTHS:\n- The paper is very well written and structured.\n- The paper is technically sound (as far as I have checked) and presents a valuable analysis and novel insights.\n- The proposed training data attribution method, SOURCE, outperforms other TDA methods on counterfactual evaluations on different tasks across different data modalities.\n- The SOURCE code will be published (if I interpret L692 correctly).\n\nWEAKNESSES:\nThere are a few easy-to-fix errors:\n- L6: missing comma before but.\n- $\\mathcal B_k$ is not mentioned before L124.\n- L206-207: \"approaches to\".\n- Appendix G is not referenced in the main text.\n- L1046: \"the SOURCE\".\n- L1051: $C = \\{3, 6\\}$ -> $C \\in \\{3, 6\\}$ (likewise for $L$).\n- (The curly brackets in Fig. 3 have a slight distortion.)\n\nQUESTIONS:\n**Questions:**\n1. Could you point more precisely to the Implicit Function Theorem and its proof?\n2. Eq. (11): Could something more be said about the properties of the approximation?\n\n**Suggestions:**\n- L105: \"optimal solution\" -> \"optimal solution when $\\epsilon=0$\".\n- Fig. 2: \"Each contour\" -> \"Each set of concentric contours\"?\n- FastSOURCE could be explained more clearly, with references to equations.\n- An overview of computation time for different TDA and evaluation methods would be interesting.\n\nLIMITATIONS:\nThe limitations seem to be addressed adequately.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2024-07-09"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper propose a new TDA method (SOURCE) that combines implicit-differentiation-based methods and unrolling-based approaches. The new method is an extension to the SGD-influence and inherits the advantage to support multi-stage training pipelines while reduce the computation complexity to store and calculate the hessian matrix for each optimization step. The paper also carried out several experiments to show that SOURCE outperforms existing TDA methods under several training settings.\n\nSTRENGTHS:\n- The paper is well presented.\n\n\t○ The math derivation is self-contained, section 3.1 is a restate of SGD-influence and the extension to all iterations, 3.2 for segmentation and approximation, which is easy to follow.\n\n\t○ Some figures (especially figure 3) help me to understand the algorithm in a more intuitive way.\n\n- The contextual information of this paper in the whole TDA community is stated clearly.\n\n\t○ It's stated clearly about the relationship between IF/SGD-Influence/SOURCE.\n\n- The experiments for different training settings are convincing.\n\n\t○ This including 4 normal settings, 2 non-converging settings, 2 multi-stage training settings.\n\n- The segmentation trick is quite intuitive and reasonable.\n\nWEAKNESSES:\n- First of all, I am not quite convinced and identify difference between the segmentation proposed in this paper and the one in TracIN (and also using different checkpoints in one training process for IF/TRAK, as stated in the paper as implicit-differentiation-based methods). \n\n\t○ Furthermore on this point, it will be good if the experiment for IF/TRAK could also include a version which ensembles (natively for IF and term-wisely for TRAK) on multiple checkpoints for each independently trained models.\n\n\t○ TRAK also made some experiment to show that it can use non-converge checkpoints to get a comparable LDS, could be mentioned in the paper as well.\n\n- Not quite sure how much performance (accuracy) degradation is involved by the segmentation since there is no comparation directly with sgd-influence.\n\nQUESTIONS:\n1. **I am willing to make the rating higher** if the first weakness is resolved (intuitively or mathematically will be enough) since I understand add new experiment might be time-consuming.\n\nLIMITATIONS:\nI believe there is no negative societal impact of this work.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe article discusses the limitations of existing training data attribution (TDA) methods, which aim to estimate how a model's behavior would change if specific data points were removed from the training set. The authors propose a new method called Source, which combines the benefits of implicit-differentiation-based and unrolling-based approaches. Source is computationally efficient and suitable for cases where implicit-differentiation-based methods struggle, such as non-converged models and multi-stage training pipelines. Empirical results show that Source outperforms existing TDA techniques, particularly in counterfactual prediction, where implicit-differentiation-based approaches fall short.\n\nSTRENGTHS:\n1. The writing in this article is very clear and excellent, making it easy to read and follow.\n\n2. The design of this article avoids the assumption made by Koh et al.'s estimator that the loss is a graph function.\n\nWEAKNESSES:\nSee the question part.\n\nQUESTIONS:\nRegarding this article, as a practitioner in the field, I would like the authors to answer a series of questions.\n\n1. We all hope that Data-attribution/Influence Analysis can find suitable application areas, such as the work done by Gross et al. in explaining the output of LLM. However, currently, we have not seen Data-attribution/Influence Analysis truly solve a series of real problems, especially for large models such as Stable-Diffusion/LLM/LVM.\n\n2. If the time and computational cost required for attribution is much higher than that for training (as it should compute the gradient one-sample-by-one-sample, and even for many checkpoints), is there any practical application significance worth for such high complexities?\n\nLIMITATIONS:\n1. Error Bound: This article lacks an analysis of the error bound for the estimator.\n\n2. Regarding the contribution: The claim \"it allows the attribution of data points at different stages of training\" is not a unique contribution of this article, as some previous works [1,2,3] have achieved this without assuming that a checkpoint with testing is necessarily at gradient=0. Therefore, I do not believe that this claim can be considered a significant contribution to this article.\n\n3. Experiments: This article lacks validation on significant datasets (even without experiments on relatively old datasets like ImageNet) and instead tests on very small toy datasets.\n\n[1]. S Hara, et al. Data Cleansing for Models Trained with SGD. NeurIPS. \n\n[2]. G Pruthi, et al.  Estimating Training Data Influence by Tracing Gradient Descent. NeurIPS. \n\n[3]. H Tan, et al. Data pruning via moving-one-sample-out. NeurIPS.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-12"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces SOURCE, a new data attribution method designed primarily for deep neural networks; more generally, SOURCE is suited for any model class optimized with gradiend-based methods. The authors motivate SOURCE as an approximation to gradient unrolling, i.e., differentiating through the entire training process. They introduce a number of assumptions and simplifications that allow them to arrive at a formulation resembling existing influence-function-based data attribution methods.\n\nThe main assumption the authors employ is that the model Hessian (of the loss w.r.t. the model parameters) remains constant throughout segments of the training process (they set the number of segments $L$ as a hyperparameter). Additionally, the authors use the EK-FAC heuristic approximation for the Hessian [1] to reduce the computational cost of SOURCE. The same approximation has been employed in a previous data attribution method [2], which the authors use as a baseline.\n\nIn the experimental section of the paper, the authors evaluate SOURCE, alongside with an extensive list of baselines, on numerous (albeit rather small) datasets. SOURCE outperforms existing methods on the vast majority of benchmarks on standard data attribution evaluation metrics.\n\n[1] T. George, C. Laurent, X. Bouthillier, N. Ballas, and P. Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018.\n\n[2] R. Grosse, J. Bae, C. Anil, N. Elhage, A. Tamkin, A. Tajdini, B. Steiner, D. Li, E. Durmus, E. Perez, E. Hubinger, K. Lukošiu¯te˙, K. Nguyen, N. Joseph, S. McCandlish, J. Kaplan, and S. R. Bowman. Studying large language model generalization with influence functions, 2023.\n\nSTRENGTHS:\non the selected datasets (Concrete, FashionMNIST, CIFAR-10, RTE, PACS), SOURCE performs consistently well, largely outperforming existing methods\n- the paper is clearly written: the authors contextualize their method well within the existing literature and succinctly describe the novel ideas resulting in their method SOURCE\n- the direction of adapting data attribution methods towards realistic (for deep learning) scenarios like lack of convergence and multi-stage training procedures is very timely and important; the paper does a great job at balancing between addressing these challenges and maintaining computational efficiency\n- the resulting method, SOURCE, is easy to understand and is rather \"obvious\" (in hindsight, of course)\n\nWEAKNESSES:\n- compared to [1], SOURCE requires significantly larger amount of compute (as the number $L$ of training trajectory segment increases). At the same time, I would label the increase in performance from [1] to SOURCE as only \"moderate\". Given, in addition to the previous point, the conceptual similarity between the two methods, the technical contribution of this paper seems to be on the smaller side. Yet, I still appreciate the connection with unrolling that the authors present.\n- if possible, it would be good to see (at least a qualitative) evaluation of SOURCE on a larger scale setting, especially since the mechanically very similar method [1] is employed on larger-scale language models; this said, I understand that the rebuttal period is short, and would not base my score change on this result.\n\n[1] R. Grosse, J. Bae, C. Anil, N. Elhage, A. Tamkin, A. Tajdini, B. Steiner, D. Li, E. Durmus, E. Perez, E. Hubinger, K. Lukošiu¯te˙, K. Nguyen, N. Joseph, S. McCandlish, J. Kaplan, and S. R. Bowman. Studying large language model generalization with influence functions, 2023.\n\nQUESTIONS:\n- can the authors provide an estimate of the computational cost of each method (e.g. wall-clock time) used in Figure 4? My understanding is that SOURCE uses up to 6x the amount of compute as the reported baselines (line 283: \"TRACIN and SOURCE use at most 6 intermediate checkpoints saved throughout training.\"). More generally, can the authors provide an equivalent to Figure 4 where compute, rather than the number of checkpoints, is equalized?\n- how does the performance (e.g., as measured by LDS) of SOURCE scale as 1) the number of segments $L$ increases, and 2) the number of fully-retrained checkpoints increase? In particular, I am curious where LDS asymptotes, and whether the asymptote value for SOURCE is higher than the one for [1, 2]; if that is indeed the case, this would, in my opinion, significantly strengthen the argument that SOURCE has a smaller \"modeling\" error compared to previous approaches.\n- my understanding is that the authors of [1] also propose using intermediate checkpoints along the training trajectory (without motivating this choice as an approximation to gradient unrolling); could the authors comment on the differences between the two estimators (excluding, obviously, the choice of Hessian approximation---EK-FAC and random projections, respectively)? Additionally, can the authors compare the performance of TRAK when intermediate checkpoints are used?\n\n[1] S. M. Park, K. Georgiev, A. Ilyas, G. Leclerc, and A. Madry. TRAK: Attributing model behavior at scale. In International Conference on Machine Learning, pages 27074–27113. PMLR, 2023.\n\n[2] R. Grosse, J. Bae, C. Anil, N. Elhage, A. Tamkin, A. Tajdini, B. Steiner, D. Li, E. Durmus, E. Perez, E. Hubinger, K. Lukošiu¯te˙, K. Nguyen, N. Joseph, S. McCandlish, J. Kaplan, and S. R. Bowman. Studying large language model generalization with influence functions, 2023.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2024-07-15"}