Better Training Data Attribution via Better Inverse Hessian-Vector Products
Authors: Andrew Wang
Elisa Nguyen
Runshi Yang
Juhan Bae
Sheila A. McIlraith
Roger Baker Grosse
Venue: NeurIPS 2025 poster ()
Keywords: Influence Functions
Interpretability
Primary area: deep_learning
TL;DR: We apply the EKFAC-preconditioner on Neumann series iterations to arrive at an unbiased iHVP approximation for TDA that improves influence function and unrolled differentiation performance.

Abstract: Training data attribution (TDA) provides insights into which training data is responsible for a learned model behavior. Gradient-based TDA methods such as influence functions and unrolled differentiation both involve a computation that resembles an inverse Hessian-vector product (iHVP), which is difficult to approximate efficiently. We introduce an algorithm (ASTRA) which uses the EKFAC-preconditioner on Neumann series iterations to arrive at an accurate iHVP approximation for TDA. ASTRA is easy to tune, requires fewer iterations than Neumann series iterations, and is more accurate than EKFAC-based approximations. Using ASTRA, we show that improving the accuracy of the iHVP approximation can significantly improve TDA performance.

Better Training Data Attribution via Better Inverse Hessian-Vector Products
Andrew Wang Elisa Nguyen'
Runshi Yang 1.2
Juhan Bae! Sheila A. Mcllraith 124 Roger Grosse
University of Toronto Vector Institute for Artificial Intelligence
TTübingen AI Center, University of Tübingen
*Schwartz Reisman Institute for Technology and Society
Abstract
Training data attribution (TDA) provides insights into which training data is respon- sible for a learned model behavior. Gradient-based TDA methods such as influence functions and unrolled differentiation both involve a computation that resembles an inverse Hessian-vector product (iHVP), which is difficult to approximate efficiently. We introduce an algorithm (ASTRA) which the AKFFAC-reccoditioon Neumann series iterations to arrive at an accurate iHVP approximation for TDA. ASTRA easy to tune. requires fewer iterations than Neumann series iterations, and is more accurate than EKFAC-bass approximations, Using ASTRA, we show that improving the accuracy of the iHVP approximation can significantly improve TDA performance.
on
Introduction
Machine learning systems derive their behavior from the data they trained on. Training data attribution (TDA) is a family of techniques that help uncover how individual training examples influence model predictions. As such, TDA is valuable tool with applications in data valuation and curation [1-4], interpreting model behavior [5-1]] 1]. building more equitable and transparent machine learning systems [12, 13] and investigating questions of intellectual property and copyright by tracing outputs back to specific data sources [11, 14, 15], among other applications.
Influence functions (IF) [16, 17] and unrolled differentiation [4. 18-21] are two gradient-based TDA methods that involve, or can be approximated as computing inverse Hessian-vector products (iHVPss. Inverting the Hessian is infeasible but for the smallest of neural networks, so the iHVP is typically computed without explicitly constructing the Hcssian. There are a number of choices available: Koh and Liang [17]. who first introduced influence functions to deep learning, use the iterative algorithm LiSSA [23], which is a method based on stochastic Neumann series iterations (SNI) [24, 25] that can unbiased solution [7, 17). Alternatively, Grosse et al. [7] adopt a tractable parametric approximation to the Hessian [1, 11, 19] using Eigenvalue-corress Kronecker Factorization (EKFAC) [26, 27J. EKFAC makes several simplifying assumptions that hold only approximately in practice, but they dramatically lower both computational and memory cost, making it feasible scale parlamererrrmtter language models [7]. However, EKFAC influence
take thousands of iterations converge
functions (EKFAC-IF) computed on converged models only correlate modestly with ground truth
"Correspondence to andrewwang @s toronto.comm edu. 'Other gradient- based TDA methods such TRAK [22] or LOGRA [1] also involve iHVPs but additional techniques such random/PCA gradient projection.
PLiSSA was introduced in the context of optimization and contains other components, but we refer to the component that computes the iHVP, found in the Koh and Liang [17] implementation. Henceforth, we will the terms LiSSA and SNI interchangeably, with their slight difference descrihed in Appendix B.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

over variety of datasets and model architectures [19] and tend to struggle for architectures involving convolution, suggesting further room for improvement. We aim to improve EKFAC-based TDA methods in this paper by improving the iHVP approximation.
Computing iHVPs in the context of TDA can be seen as finding the minimizer to high- dimensional quadratic optimization problems in parameter space [25, 28, 29]. Curvature ma- trices for large neural networks are known to be ill-conditioned [30, 31], causing slow conver- genee for iterative methods such as SNI. While costly and rather difficult to tune [17, 32, 33], the upside of SNI that it produces a consistent estimator the algorithm converges to the iHVP in the limit as more compute used [23]. In contrast, while EKFAC often provides better cost vs. accuracy tradeoff for TDA [7], there is no simple way to improve its accuracy by apply- ing more compute. Our central insight is that can combine the best of both worlds: we can repurpose the EKFAC decomposition - which needs to be computed for EKFAC- based influ- ence functions and unrolled differentiation anyways - as a preconditioner for SNI, yielding a cost effective procedure for improving the iHVP accuracy (Figure 1). Our contributions are as follows:
MXX
IHVP astimation loss contour without proconditoring
Addtionahasii goer EXFAC
IHVP estimetion losa contour with precanditioing
Truncated Noumarn: (Hay + Al) "v
(Hrave + Aly
(H AI) 0
Figure I: The objective is to compute the damped iHVP (H + AI) 'v. Preconditioning Stochastic Ncumann Iterations (SNI) with EKFAC (ASTRA) improves the convergence speed of the iHVP approximation. Initial- ized at 0, it results in the same approximation as using EKFAC after one iteration. SNI may require thousands of iterations to converge, and truncating early results in undesirable implicit damping.
First, we present an algorithm called ASTRA which uses EKFAC as a preconditioner for SNI with the aim of computing cost-effective and accurate iHVPs. ASTRA can be applied to both influence functions (ASTRA-IF) and unrolled differentiation via an approximation called SOURCE (ASTRA-SOURCE) [19] among other applications. To the best of our knowledge, prior work has applied the EKFAC preconditioner to SNI for the TDA setting. In our experiments, the incremental cost of ASTRA-IF and ASTRA- and EKFAC-SOURCE, respectively. In contrast to other settings in which we would encounter iHVPs, are often willing to pay this extra computational cost to obtain more accurate iHVP for TDA.
was only hundreds of iterations, compared to EKFAC-IF
Second, while past papers have questioned the reliability of influence functions [34, 35], we show that influence functions computed with accurate iHVPss have strong predictive power, even in settings in which the assumptions
its derivations may not hold [35, 36). In our experiments, ASTRA-IF was able to achieve a Spearman correlation score [37] with ground-truth retraining of around 0.5 across many settings, 3 and ensembling these predictions often raised performance to 0.6. ASTRA provides an accurate approximation of the iHVP, which significantly increases the efficacy of ensembling [19, 22] compared to EKFAC and also significantly improves TDA performance for architectures involving convolution layers. For the experiment involving convolution architecture, performance of ensembled ASTRA-IF was 0.6, a large increase from EKFAC-IF's 0.25.
Finally, leveraging EKFAC's eigendecomposition, we show that low curvature directions are essential for high quality influence estimates. Truncating Neumann series has an implicit damping effect [29, 38), which has disproportional impact on low curvature directions. To quantify the downstream impact on influence functions performance, we perform an ablation study by using the EKFAC eigendecompossiion to project the iHVP onto subspaces containing different levels of curva- ture during Neumann series iterations. We show that these low curvature components are essential for high quality influence estimates. This suggests that good influence function performance demands careful hyperparameter tuning, which can be costly, especially for off-the-shelf iterative solvers. In contrast, ASTRA requires much less hyperparameter tuning: we use set of hyperparameters determined by simple heuristics which worked well for all of our experiments. We also note that preconditioner in second-order optimization [26]. its compact
while EKFAC was introduced representation of the eigendecompooition remains relatively underappreciated - the method in this ablation study in which we use the EKFAC eigendecomposition to analyze the quadratic objective could therefore be of independent interest.
"More precisely, we use a widely- used evaluation metric called Linear Datamodeling Score (LDS) [22]. which measures the rank correlation between ground retraining outcomes and TDA algorithm's 's predictions. LDS is defined and performance comparisons are provided in Section 5.
2

---

2
Preliminaries
This section briefly introduces preliminaries relating to 1) computing the iHVP and 2) influence functions. To help the reader navigate the mathematical objects, we have provided a table of notations and acronyms in Appendix A. Further background and expanded discussion is provided in Appendix B.
Given a training dataset D l= where 2; (xx() (0) is an taput-target pair, and model parameterized by 0 E RP, let 000, x() denote the model output on x()) let C(y, t) be convex loss function. We define a training objective J(0, = &V Σ ((g(0,x()) x()), t()) as the average loss over D. Given a query data point 2g and measurement function fz, (0), such as correct-class margin [22], an idealized objective of TDA is to approximate the impact of removing a training example 2m from the training dataset D on /z,- A pointwise TDA method assigns score r =( 2q, D) that measures the impact of removing 2m from 'D on Jz Since modern neural networks often exhibit multiple optima, the stochasticity in the optimization process from sources such parameter initialization [39], sampled dropout masks [40], and mini-batch ordering [41] can result in different learned optima, which we denote 0 Let be random variable which captures this stochasticity in the training procedure [42, 43].
2.1 Computing Inverse Dessian-Vector Products for TDA
Inverse Hessian vector products (iHVPs) are ubiquitous in many machine learning settings beyond TDA [25, 26, 44, 45], such as second-order optimization [26, 46, 47], but different settings have different considerations when trading off cost vs. accuracy. The canonical second-order optimization method - Newton's 's method utilizes an inverse dlessinngrrdden product to compute the Newton step, and may achieve faster local convergence than first-order methods [28]. However, computing an iHVP is substantially more expensive than a standard gradient computation. In optimization, there exists tradeoff between devoting extra compute to obtain better iHVP approximation versus taking more steps in the optimization procedure [48]. Because most deep learning optimizers rely stochastic updates, the extra cost of highly precise curvature estimate often is not justified and popular optimizers such as Adam [47] default to the much cheaper diagonal preconditioner. In contrast, we typically willing pay higher cost for accurate iHVP approximations in the TDA context, since - as we will show - an accurate iHVP may significantly improve TDA performance. Two popular ways of computing the iHVP for TDA are EKFAC [7, 26, 27] and LiSSA [17, 23]. We now briefly describe each method.
Computing the iHVP with EKFAC KFAC [26] and EKFAC [27] make block-diagonal paramet- ric approximation of the Fisher Information Matrix (FIM) so that its eigendecomposition can be done for each layer I independently, which is significantly cheaper than an often infeasible frute-force eigendeconposition. The FIM defined as:
1 N
F
Σ Eg-re(y)xc)) XDDD
Vo log po (y |x Vo log (8(x0))
(1)
where (y() is the model's own distribution over targets. For multi-layer perceptrons, KFAC approximates the /th block of the FIM with:
F
a
a
(Ds:Ds
(2)
A
= (QA,
2 Qs,)
(DA @ (QA, 8 Qs,)
AI,ESC
where a the activations of the 7 1th layers and Ds; are the pseudogradients of the loss with respect to the preactivations S, of the Ith layer, and A = QA, DA. Qh and S, =
"For standard regression and classification tasks whose outputs be as the natural parameters of an exponential family, the Fisher Information Matrix and the Generalized newss-Newton Hessian G coincide, which we will use as a substitute for the Hessian H. For details, see Appendix B.
'The har indicates the use the homogeneous vector notation - = [a]_ 1]
"y pseudogradients, mean the gradient of the loss using sampled target with respect to the parameters.
3

---

Qs, Ds. Qs, are eigendecompositions of A;-1 and S₁ respectively. We denote the block doakk-iiaoonal matrix approximated this way F KFAC- We can then approximate the inverse FIM (F + AI) so (QA-,, Qs.) (AI, + AI) (QA. & Qs.) where damping hyperparameter. The EKFAC approximation F FEKFaC is an improvement over KFAC using the diagonal matrix ALEKFAC instead, whose ith entry along the diagonal is:
A,FAACC = E[((QA,,, @ Qs,) Do:))
(4)
where DO, are the pseudogradients with respect to the parameters in layer (.
To illustrate the compute and memory savings, for P-layer multi-layer perceptron with Ò input dimensions and D output dimensions for all layers, due to the block-diagonal approximation, the EKFAC eigendecomposition only costs O() = O(D time and storing its statisties requires O(P = ( memory [49].² 7 Although EKFAC significantly reduces the time and space complexity of the eigendecomposition, it results in biased iHVP. See Appendix B for a more detailed discussion of EKFAC.
Computing the iHVP with LiSSA In contrast to EKFAC, iterative methods such as LiSSA are based on Neumann series iterations (NI) [24]. NI do not require EKFAC"s assumptions, and thus in principle can produce exact iHVPs as more compute is applied. For an invertible matrix A E RRDDD
with |I 2 1, the Neumann series is approximated as A Σ)= A)', which is generalization of the geometric series a Ei=0( (1 a)² for - al < 1. By substituting A with a scaled positive-definite damped Generalized multiplying both sides by any v E RP, we obtain: (G + vv ンド=
Hessian (GGN) a G AI), and
(I- aG - aAI) v. Here À is positive scalar known as the damping term and > 0 is the learning rate hyperparameter. The learning rate must satisfy < +G) where max (G) is the largest eigenvalue of G so that
||I Alla 1. We can then approximate the iHVP (G + A1)- via the iterative update:
Vk+1 V2 aGG G + AI)v= +v,
(5)
which satisfies the property VA + ( as k oo. Computing G requires two backward
passes over the whole dataset, 8 so an unbiased estimate G2 of G is usually used instead by sampling mini-batch with replacement, which refer to stochastic Neumann series iterations (SNI). Compared to other iterative methods like conjugate gradient (CG) [51], SNI typically preferred to compute the iHVP for TDA since CG tends to struggle with stochastic gradients [17, 26, 52]. LiSSA reduces the variance in SNI by taking an average over multiple trials.
2.2 Training Data Attribution with Influence Functions
Influence functions [16, 53] are derived under the assumption that J is strictly convex in 0 and twice differentiable. Let 0+ == arg ming J (0, D) be the optimal parameters over D. We can define the objective after downweighting a training example 2m by c as: Q(0. c) [(e. D) x L(O, zm) where e is a scalar that specifies the amount of downweighting. When c = O, this corresponds to the original objective J. We can then define the optimal parameters after downweighting function of e: r(e) = arg mingerR e) When Zm E DD and = 1, this corresponds to the downweighted objective in which 2m removed from D. Typically,
assumed be small, and dr
can approximate the leave-one-out (LOO) parameter change as 0* (D) \ { 0 (D)
where dr 1 objective the optimal parameters. To approximate the effect on the measurement function fx,- invoke the chain rule: /z, (0* (0* {=m))) - (0) & Volz, - zm)) which also gives the first-order Taylor approximation of fx, (0*(D) {=m})) after rearranging terms. When applying this approximation to neural networks in which the convexity assumption does not
2m) and H VOJ D) denotes the Hessian of the training
de c=
'The time and space complexity of ordinary eigendecomposition is O(D*) and O(D*) respectively, and PD typically is much smaller than D This analysis treats as constant. An efficient implementation with O(D) time and space complexity requires Jacobian- vector product and Jector-Jacobian product [50].

---

hold, H may not be invertible, so H is typically approximated with the damped GGN G + AI [54],
which is always positive definite for À > 0. and tends to work well in practice [7. 11, 19, 55].² With this substitution, the influence functions attribution score is:
(6) When applying influence functions to neural networks, the model may not have fully converged, so typically compute the gradients and the GGN in Equation 6 with the final parameters 0c instead of 0", We can also ensemble influence functions for better TDA performance, by training models with various seeds and averaging over OF for each seed get an ensembled score [19, 22] (details in Appendix B). Ensembling for other TDA methods, such as unrolled differentiation, can be done analogously.
TiE 4y D) = Vefz, (0
The iHVP in Equation 6 was originally computed with LiSSA by Koh and Liang [17] and has the drawback that its iterative procedure must be carried out once for each vector in the iHVP. Fortunately, it is frequently the case that (Dqueryl < (D) [7, 36, 56), so by choosing v == Volz, (0)) ((*))((G (G + AI)- in Equation 6, 6,10 we can reduce the number of iterative
and first computing Vofz, (0*) procedures to |Pqueryl for the influence function computation. Plugging these values into Equation 5, we can compute Volz, (0*) (G + AI)-1 via the iterative update:"!
Ox+1 f O2 aGG + ΛI)O, + aVelz, (0").
(7)
While [Dqueryl sounds like a large number of iterative procedures, in practice no D Dquery exists, and instead queries are run when the user wants understand specific behavior pertaining to Zq. Nevertheless, get good approximation of Volz, (0%) (0*)T(G + AI)-1 for any particular 2g may require thousands of iterations [7. 17J. limiting its scalability. Furthermore, tuning the hyperparameter [17, 32, 33] due to the ill-conditioning and stochasticity of the gradients. Once the iHVP is computed, its dot product with VoC(O", 2m) ) for every 2m in consideration is taken. If the goal is to simply compute the influence of nm on 2g for given 2. and Zq. then the cost of the dot product is minimal. However, if the goal is to search for the most influential points in 2m E D on Zy. then we must take the dot product with every training example gradient in D, which amounts to backward pass over the entire training dataset and can be a substantial component of the cost of computing influence functions." 12 Arguably, the latter goal is more prevalent in settings such as data-centric model debugging and interpretability, where insight into model behavior is given by retrieving highly influential training samples [57-61].
for SNI is difficult
3 Introducing EKFAC-Accelerated Neumann Series Iterations for TDA
We now have laid sufficient groundwork to introduce a novel algorithm called ASTRA, which preconditions SNI updates with EKFAC. In this section, we present the ASTRA update rule, discuss its computational costs, and discuss extensions to umrolled-diffeeeentattoo based TDA.
Preconditioning Stochastic Neumann Series Iterations with EKFAC The SNI update rule in Equation 7 can be viewed as performing mini-batch gradient descent on the quadratic objective [23, 62]:
1 T 2
f,,, (0)
(G + AII - 0 Volz,(0').
(8)
It is well-known that for a converged neural network, the curvature matrix in the objective in Equation 8 is typically ill-conditioned [30, 31J, which presents challenges for iterative methods. To improve the conditioning of h; (0), we introduce an algorithm which computes accurate iHVPs for use in TDA called (ASTRA) by applying preconditioning. The resulting update rule is:
Neumann Series Iterations for TRaining Data Attribution
θx+1 <- 0₂ aP + Xr)-1(G, + ATJO, + (P + AI) (0",
We will use the approximation G R H throughout, and our use of the term iHVP will generally refer to both the inverse Hessian- vector product and the inverse Heuus--Nwwoo--Hessiiee vector product.
WThis uses the fact that G is symmetric. (Velz, (0*) T (G + AI) = (G + - (O"), We emphasize that G. and V Velz, (0*) are computed using the final parameters 0², use 02 to denote the iterates for SNI since it has the same dimensions as the model's 's parameters.
Procedures such TF-IDF filtering exist prune the potentially vast training dataset [7].
5

---

While number of choices of preconditioners exist, we choose the FIM computed with EKFAC (i.e., P F FEKFAC) for the following reasons: First, the computation cost of the EKFAC eigendecompo- sition is usually much cheaper than full matrix inversion, and its eigendecomposition statistics can be stored compactly and shared across all [Pquery] optimization problems, since the value of G is the same across all objectives. Second, this choice has a close connection with EKFAC influence functions: Equation 5 suggests initializing Og as Vefz, (0*). which is frequently done in public implementations [17]. Observe that if we initialize On- <- 0, choose À = X and a learning rate a 1, we arrive at the iHVP which would be approximated by EKFAC after one step of ASTRA, resulting in the same TDA prediction as EKFAC-IF. We hypothesize that further training using the update rule in Equation 9 will improve the iHVP approximation, a claim we validate empirically in Section 5, This update rule assumes using mini-batch gradient descent, but our formulation is compatible with other optimization algorithms as well.
Time Complexity of ASTRA Computing the update in Equation 9 involves explicitly constructing neither (P + AI)-1 nor (G + AT). Instead, we use Hessian-vector products [50] and first compute (G& + xI)O2. We can then compute (P + ÀI)-I(G, R + AI)0; using the EKFAC preconditioner P F F'KFAC, so the incremental time complexity of each iteration in our algorithm is O(BD) where B the mini-batch size used to sample G2. For both EKFAC-IF and ASTRA-IF, computing FEKFAC is necessary to compute the iHVP, which only needs to be done once per model and can be shared among all queries. When we need to search the entire dataset for highly influential training examples, both methods also need to compute the dot product of the resulting iHVP with the training example gradients over all Zm. in consideration, as discussed in Section 2.2. In our experiments, ASTRA-IF required only few hundred incremental iterations, so its iterative component is relatively small additional cost per query compared to the total cost to compute EKFAC-IF.
Application to Unrolled Differentiation In addition to influence functions, we also study the use of ASTRA in the context of an unrolling-based TDA method. Influence functions make assumptions such as model convergence and unique optimal parameters in its derivation, which may not be satisfied in practice [17, 36). In contrast, unrolled differentiation methods such as SOURCE sidestep this limitation by differentiating through the training trajectory. Here, we only sketch our approach to apply ASTRA to SOURCE, deferring the full discussion of the SOURCE derivation to Appendix B to Appendix C. SOURCE approximates diffcrentiating through
and the details of ASTRA-SOURGE the training trajectory by partitioning it into I. segments, assuming stationary and independent GGNs and gradients within each. Its approximation of the first order effect of downweighting training example contains L. different finite series involving the GGN, each of which can be approximated with an iHVP. Similarly to ASTRA-IF, ASTRA-SOURE improves this approximation by repurposing the EKFAC decompositions as preconditioners, which would have been needed to be computed to implement SOURCE anyways.
4 Relationship to Existing Works
TDA Methods using iHVPs Both influence functions and unrolled differentiation can be viewed as belonging to family of gradient-based TDA methods (for a survey see [58]), which both have connection with iHVPs. Many gradient-basee based TDA methods are variants of the influence functions method proposed by Koh and Liang [17] with aims improve its computational cost [7, 22, 33],
by using techniques such as EKFAC for iHVP approximation [7], Arnoldi iterations [33], gradient projection [1. 22], and rank-one updates [63], instead of iterative algorithms such as LiSSA [23] or CG [51, 64], since the former is expensive and hard to tune [17, 32, 33], and the latter struggles with stochastic gradients [65]. Unrolled differentiation addresses the key derivation assumptions underlying influence functions namely, the convexity of the training objective and the convergence of the final model parameters [36]. Methods include SGID-influence [4], HYDRA [18], SOURCE [19], DVEmb [20] and MAGIC [21]. which all differentiate through the training trajectory, and only differ in their approximations. Rather than comparing the TDA algorithms themselves, our goal in this paper is to show that substantial TDA performance improvements can be achieved by using better iHVP approximations. In some public comparisons, LiSSA based influence functions perform poorly
Going forward, we therefore directly initialize ASTRA with do (P + 'V)='vosz,(0))

---

MNIITT (MILE)
FAsHionMNIST (MLE))
CIFAR- 10 (ResNet-9)
UCI Concrete (MLP))
0.8
118
0.4
0.0
LCI Parkinsons (MIP))
Wild texta (GP1-))
MashioMMMIST-- (MLP)
0.8
DRAK
EXPAC-ET
ASTRA-LE
EXTNG-SOURGT
0.4
10 Modall
Figure 2: TDA Performance. Single model ASTRA-IF and ASTRA-SOURCE beat EKFAC-based counterparts in most settings, as well as other TDA methods such TracIn [83] and TRAK [22] when measured by average LDS the query Duery. ASTRA also enjoys larger performance boost from ensembling. Improvement particularly large convolution architectures such as ResNet-9. standard error. We omit TRAK for GPT-2 due to lack of public implementations.
bars (where available) indicate
[22, 63, 66], sometimes use EKFAC-IF as their method or baseline choice [1, 11, 19]. principle, EKFAC-IF is only an approximation of what LiSSA-base influence functions attemptss to compute, since the latter is only constrained by solver error while the former makes assumptions on the structure of the curvature matrix. We show in this paper that indeed, accurately solving the iHVP typically produces better TDA performance than the EKFAC solution, which what our algorithm ASTRA addresses.
worse than dot products [63, 67] and many have opted to instead
iHVPs beyond TDA iHVPs can also be found in higher-order optimization algorithms such as Newton's 's method [28], quasi-Newton methods [68-71], natural gradient descent [65, 72, 73), KFAC [26]. and Hessian-free optimization [65, 74], which computes Hessian-vector products iteratively with CG [51, 64]. Influence functions also be cast bilevel optimization problem [25, 29, 44, 75], which can be solved via implicit differentiation or unrolled differentiation. Since iHVPs show up frequently in machine learning, there is motivation adapt and develop iHVP computation techniques such as SNI and EKFAC. While these two methods are well-established and preconditioning is a well- established technique in optimization [26. 72], the best of knowledge. no prior TDA method has combined these methods to compute the iHVP. For extended related works, see Appendix D.
Performance Comparisons
This section aims answer the following questions: Do ASTRA-IF and ASTRA-SOURCE outperform their EKFAC-based counterparts? 2) Is ASTRA substantially faster than vanilla SNI? To answer these questions, we run experiments in a number of settings. For regression tasks, we use the UCI datasets Concrete and Parkinsons [76] trained with multi-layer perceptron (MLP). P). For classifi- cation tasks, we use CIFAR-10 [77] trained with ResNet-9 [78], MNIST [79] and FashionMNIST((8)) trained with MI.ss, and GPT-2 [81] fine-tuned with WikiText-2 [82]. We also include non-converged setting, FashionMNIST-NN introduced by Bac et al. [19]. for which SOURCE was specifically de- signed; in this setting, 30% of the training examples were randomly labeled, and the model was trained for only three epochs to avoid overfitting [19]. addition to comparing against EKFAC-IF and EKFAC-SOURCE,
also compare against two popular TDA methods TracIn [83] and TRAK [22]. Details for all experiments can found in Appendix F.
Evaluating Training Data Attribution Performance We evaluate the performance of our TDA algorithms popular evaluation metric called Linear Datamodeling Score (LDS) [22], and use mean absolute error as the measurement function for regression tasks and correct-class margin for classification tasks in line with past works [19, 22]. LDS measures TDA algorithm's 's ability to predict the outcome of counterfactual retraining on a subset of data. Given a collection of uniformly
"This TDA method, sometimes called Hessian free [63, 67]. simply takes the dot product between the training gradient and the query gradient and equivalent to influence functions if the damped GGN set to the identity.

---

1.0
165
165
165
165
0.5
L,88
0.0
0.4
105
0.3
0.4
0.2
0.2
0.1
100
200
160
200
160 Ierationn
200
160
200
Iterations
Ierations
Ierations
MNIST (MLP)
PashionMNIST (MLP)
CIFAR-10 (ResNet-9)
UCI Concrete (MLP)
ASTRA-1F
SNI
Figure 3: Training Curves. Loss and LDS curves for SNI and ASTRA measured over seeds (shaded region standard error) on an arbitrary query point using influence functions as the TDA method, SNI makes slower progress compared to ASTRA as measured by LDS.
randomly sampled subsets {Si,.... SM S. C. D) of fixed size, typically fraction 3 of D and measurement function Z the LDS scores a TDA method 7 as follows:
LDS(T, 2ng (0 = E [M], (S,
[M]),
(10)
expected ground truth predictions on Zy using models trained various S;
TDA method's prodictions on 24 for various S,
where P denotes the Spearman correlation [37]. and the group influence (S,, Zy: D) is defined lin- early as: (S Zy; D) = =E T(Z;, Zg, D). To compute the ground-truth to which we compare our TDA method, we need to retrain a model many times both over various subsets S,, and also over random seeds for every subset to obtain good estimate of the expectation in Equation 10, which can be quite noisy [19, 43]. For example, for the experiment involving GPT-2, computing ground- truth involved fine-tuning 1000 models. We report the average LDS Pauery Σx,eD LDS(T, 2q)
over a test set Dquery containing 100 query points. We randomly sample M = 100 subsets with a subsampling fraction of B 0.5 in line with previous works [1. 19]. We discuss other evaluation methods Appendix E.
LDS evaluation of ASTRA Figure 2 compares LDS across various TDA methods. We com- pare ASTRA-IF and ASTRA-SOURCE with their respective EKFAC-versions. For each setting, EKFAC-IF and ASTRA-IF use the same damping value implied by SOURCE for comparability (details in Appendix F). In almost all settings, ASTRA improves TDA performance by LDS, strongly suggesting that better iHVP approximations are responsible. We also observe an especially large improvement over EKFAC for CIFAR-10 10 trained on ResNet-9 [78]. When applied to convolution layers, EKFAC makes additional simplifying assumptions, 15 which can cause EKFAC-IF and EKFAC-SOURCE to underperform that ASTRA effectively addresses. Figure also reveals increased benefits from ensembling when applying ASTRA, which computes an unbiased estimator of the iHVP. In some cases, such Fash- ionMNIST, the advantages using precise iHVPs become much more pronounced in conjunction with ensembling. We hypothesize that this is due to the various bottlenecks in TDA methods: in some cases, the primary performance bottleneck lies in computing the iHVP accurately: in others, stems from other factors such as the method's underlying assumptions (e.g., unique optimal parameters), which can be mitigated by ensembling.
measured
architectures involving convolution layers
issue
ASTRA speeds up iHVP approximation The top row of Figure shows the loss curves for SNI and ASTRA as each iHVP solver progresses. The bottom row corresponds to the LDS that influence functions achieves based on the current progress of each iHVP solver. For SNI, we follow public implementations [17], which typically initialize SNI using the query gradient Volz, (0*), and results
In addition to wise independence and independence of activations and pseudogradients [26]. assumes spatially uncorrelated derivatives and spatial homogeneity [84].

---

in an initial influence functions prediction that approximates the damped-GGa with the identity in Equation 6. For both methods, we conduct a hyperparameter sweep for the learning rate over 10", .... 10-6 in steps of one order of magnitude, and use the best hyperparameter based on the average training loss performance on the same query point over the last 10 iterations. We find that EKFAC preconditioning reduces the notorious challenge of tuning learning rates [17, 32, 33] in all of the settings reported in Figure 3 the learning rate for ASTRA used was 10 10-2 In comparison, the reported (and best) SNI learning rates were 1, 0.1, 0.01, 0.1 for MNIST, FashionMNIST CIFAR- 10, and UCI Concrete respectively and LDS performance was very sensitive to the learning rate. Figure shows that SNI makes slow progress while ASTRA usually converges in fewer than 200 iterations.
6 Investigating the Role of Low Curvature Directions in Influence Functions Performance
Our results in the previous section lead us hypothesize that preconditioning accelerates convergence in directions of low curvature, which is important for influence function performance. We can analyze how directions of low curvature affected by NI (without preconditioning) when truncated early, something that is tempting to do as it usually takes long to converge. We derive the following expression for the truncated Neumann series with J iterations:
aG aATY (G AD) 'I oG (I-oG-aAD
(11)
2 (G (A+a
(12)
where the equality utilizes the definition for finite series, and the approximation is identical to that used in [19] (see Appendix G for the derivation). From the last equation, we can see that truncating Neumann series effectively adds an implicil damping term 1/aJ, which disproportionately affects directions of low curvature.
This insight prompts an investigation into the role of low curvature directions influence func- tions performance since, in addition to the implicit damping effect mentioned above, some meth- ods may discard them when projecting gradients into lover-dimensional subspaces [1, 33]. Let FEKFAC = QDQ be the EKFAC eigendeconposition of the GGN at the final parameters. We study the importance of directions of varying curvature by doing the following: We bin the D eigenval- ues given by FEKFAC into 5 bins labeled S1, S2,S ... S5, where each bin S, holds eigenvalues larger than 10- Let Qs: E Dx(5a be the projection matrix whose columns are the associated orthonormal eigenvectors of the eigenvalues in each bin. We investigate the values of h (0x) and
the corresponding LDS for each S: where ff, is defined as:
1. S. (0) = ( 0) Q3, (G (G+ AT (Q, (Q3, T VoIz, (0").
(13)
We use the EKFAC eigendecomposition since the true cigendecomposition intractable for the settings we report. We conduct the experiment in the MNIST and FashionMNIST settings and use small damping hyperparameter of X = 10-4 to be able to observe the impact of directions of low curvature on influence functions performance (details in Appendix F).
The top of Figure 4 shows the outcome when we ASTRA and SNI on the objective in Equation 8 obtain sequence of 0r. and plot the value of each h; (0x) for each S,. The bottom row shows the LDS, which are computed by projecting the iterates 0₁ into each curvature subspace defined by S; and computing influence functions using these projected vectors, which we can write as:
TPROI-IF A (Zm, Zy, D) (0))=(Q].00 ON) TQ,, VeL(O" 2m))
(14)
Figure 4 reveals that low curvature directions play large role in the performance of influence functions: projecting to high-curvature eigenspaces degrades LDS performance, as evidenced by the large vertical gaps between lines representing different level of curvature. When the iHVP is computed via SNI, large eigenvalue directions converge quickly during training, but it takes longer for small eigenvalue directions to converge, evidenced by the earlier plateau of the loss curves in

---

MNIST (MLP)
PashionMNIST (MLP)
165
167
1.0
0.0
0.0
1888
0.5
0.5
0.0
-1.0
0.0
-1.0
0.6
LDS
0.2
0.2
0.2
0.2
-0.2
0.2
0
500
1000
500
1000
500
1000
500 1900
0> 10
d=10-2
0> 10-3
o> 10
<> 10-5
The values S objective in Equation 8 for an arbitrary zy. The subspaces represented by Si,... S5 are spanned by eigenvectors with cigenvalues a > 10......, o > 10 respectively. The loss of subspaces with large cigenvalue directions tend to platcau first, followed by subspaces with smaller eigenvaluc directions, especially for SNI, which does not use preconditioning. Bottom: LDS of influence functions after projecting corresponding subspace. The objective for high curvature directions plateaus first; continued training further decreases the objective in progressively lower curvature directions, yielding LDS gains even after high curvature directions have converged. Shaded region standard error.
Figure 4: Training Curves in Various Eigenspaces. Top:
/2g as ASTRA and SNI train on the
the high-curvature subspaces in the top row of Figure 4. Nevertheless, as the solution progresses in low-curvature directions, LDS rises substantially, evidenced by the growing gap between lines representing large and small levels of curvature the bottom row of Figure 4. While the slower convergence in low-curvature directions is present for ASTRA, it is substantially diminished due to the preconditioning. Our results highlight that the behavior of estimators in low-curvature subspaces may be substantial factor in the performance of TDA methods."
Conclusion
We presented an algorithm ASTRA that combines the EKFAC preconditioner with SNI for TDA. We compared ASTRA-II and ASTRA-SOURCE with their EKFAC-based counterparts in variety of settings. In many settings, TDA performance measured by LDS improved substantially, especially for convolution architectures. We find that in general. the efficacy of ensembling. ASTRA is casier to tune and converges faster than SNI. Compared with EKFAC, it only incrementally costs hundreds of iterations in our experiments it leverages the same eigendecomposition. We conclude this paper by providing insights into how various curvature directions affect influence functions performance. We show that low curvature directions are important for good influence functions performance by using the EKFAC decomposition to analyze the quadratic objective, a technique that may be of independent interest outside of TDA. Overall, the technical contributions of this work should lead to improved performance in real-world problems such as data curation and interpreting model behavior, among other applications. We discuss limitations and broader implications of our work in Appendix H.
accurate iHVP approximation increases
Acknowledgements
We gratefully acknowledge funding from the Natural Sciences and Engincering Research Coun- cil of Canada (NSERC) and the Canada CIFAR AI Chairs Program. Resources used in prepar- ing this research were provided, in part, by the Province of Ontario, the Government of Canada through CIFAR, and companies sponsoring the Vector Institute for Artificial Intelligence (https: ai/partnershipl
We thank the Schwartz Reisman Institute for Tech- nology and Society for providing rich multi-disciplinary research environment. RG and SM acknowledge support from the Canada CIFAR AT Chairs program and from the Natural Sciences
VeWe note that the difference in performance between ASTRA and SNI depends on the damping hyperparam- eter, and since the damping in this experiment is smaller, the performance boost from ASTRA compared to SNI is notably larger than in Figure 3.
10

---

and Engineering Research Council of Canada (NSERC). RG acknowledges support from Open Philanthropy and the Schmidt Sciences A12050 Fellows Program.
References
[1] Sang Keun Choe, Hwijeen Ahn, Juhan Bae, Kewen Zhao, Minsoo Kang, Youngseog Chung, Adithya Pratapa, Willie Neiswanger, Emma Strubell, Teruko Mitamura, al. What is your data worth to gpt? llm-scale data valuation with influence functions. arXiv preprint arXiv:2405. 1004400..9444 13954, 2024.
[2] Saachi Jain, Kimia Hamidich, Kristian Georgiev, Andrew Ilyas, Marzyeh Ghassemi, and Aleksander Madry. Improving subgroup robustness via data selection. Advances in Neural Information Processing Systems, 11994490-945111 2024.
[3] Stefano Teso, Andrea Bontempelli, Fausto Giunchiglia, and Andrea Passerini. Interactive label cleaning with example- based explanations. Advances in Neural Information Processing Systems, 34: 1996-122977 2021.
[4] Satoshi Hara, Atsushi Nitanda, and Takanori Machara. Data cleansing for models trained with sgd, Advances in Neural Information Processing Systems, 32, 2019.
[5] Pang Wei W Koh, Kai-Siang Ang, Hubert Teo, and Percy S Liang. On the accuracy of influence functions for measuring group effects. Advances in neural information processing systems, 32, 2019.
[6] Chih-Kuan Yeh, Joon Kim, Ian En-Hsy Yen, and Pradeep K Ravikumar. Representer point selection for explaining deep neural networks. Advances in neural information processing systems, 31, 2018.
[7] Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamilé Lukošiūte, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. Studying large language model generalization with influence functions, 2023. URL https:
// //arxiv.er org/ abs/2308 03296.
[8] Ekin Akyürek, Tolga Bolukbasi, Frederick Liu, Binbin Xiong, Ian Tenney, Jacob Andreas, and Kelvin Guu. Tracing knowledge in language models back to the training data. arXiv preprint arXiv:2205.1 /482, 2022.
[9] Fulton Wang. Julius Adebayo, Sarah Tan, Diego Garcia-Olano and Narine Kokhlikyan. Error discovery by clustering influence embeddings. Advances in Neural Information Processing Systems, 36:41 16641765-417777 2023.
[10] Andrew Ilyas, Sung Min Park, Logan Engstrom, Guillaume Leclere, and Aleksander Madry. Datamodels: Predicting predictions from training data. arXiv preprint arXiv:2202.0062 2022.
[11] Bruno Mlodozeniec, Runa Eschenhagen, Juhan Bae, Alexander Immer, David Krueger, and Richard Turner. Influence functions for scalable data attribution in diffusion models. arXiv preprint arXiv:2410. 13850, 2024.
[12] Marc-Etienne Brunet, Colleen Alkalay-Houlihan Ashton Anderson, and Richard Zemel. Understanding the origins of bias in word embeddings. In International conference on machine learning, pages 803-811. PMLR, 2019.
[13] Haonan Wang, Ziwei Wu, and Jingrui He. Fairif: Boosting fairesss in deep learning via influence functions with validation set sensitive attributes. Proceedings of the 17th ACM International Conference on Web Search and Data Mining, pages 721-730, 2024.
[14] Gerrit van den Burg and Chris Williams. On memorization probabilistic deep generative models. Advances in Neural Information Processing Systems, 34:27916-27922 2021.
11

---

[15] Emanuele Mezzi, Asimina Mertzani, Michael P Manis, Siyanna Lilova, Nicholas Vadivoulis, Stamatis Gatirdakis, Styliani Roussou, and Rodayna Hmede. Who owns the output? bridging law and technology in Ilms attribution. arXiv preprint urXiv:2504.01033: 2025.
[16] Frank R Hampel. The influence curve and its role in robust estimation. Journal of the american statistical association, 69(346):383-393, 1974.
[17] Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pages 1885-1894, PMLR, 2017.
[18] Yuanyuan Chen, Boyang Li, Han Yu, Pengcheng Wu, and Chunyan Miao. Hydra: Hypergradi- ent data relevance analysis for interpreting deep neural networks. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 7081-7089, 2021.
[19] Juhan Bae, Wu Lin, Jonathan Lorraine, and Roger Grosse Training data attribution via approximate unrolled differentation. arXiv preprint wrXi:2245.122866 2024.
[20] Jiachen T Wang, Dawn Song, James Zou, Prateek Mittal, and Ruoxi Jia. Capturing the temporal dependence of training data influence. arXiv preprint arXiv:2412.09538, 2024.
[21] Andrew Ilyas and Logan Engstrom. Magic: Near-optimal data attribution for deep learning.
arXiv preprint arXiv:2504. 16430, 2025.
[22] Sung Min Park, Kristian Georgiev, Andrew Hlyas, Guillaume Leclerc, and Aleksander Madry. Trak: Attributing model behavior at scale. arXiv preprint arXiv:2303. 14166 2023.
[23] Naman Agarwal, Brian Bullins, and Elad Hazan. Second-order stochastic optimization for machine learning in linear time. Journal of Machine Learning Research, 18(16)):1-40 2017.
[24] Roger A Horn and Charles R Johnson. Matrix analysis. Cambridge university press, 2012.
[25] Jonathan Lorraine, Paul Vicol, and David Duvenaud. Optimizing millions of hyperparameters by implicit differentiation. In International conference on artificial intelligence and statistics, pages 1540-1552. PMLR, 2020.
[26] James Martens and Roger Grosse. Optimizing neural networks with kronecker-factored approximate curvature. In International conference machine learning, pages 1417. PMLR, 2015.
[27] Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances Neural Information Processing Systems, 31, 2018
[28] Jorge Nocedal and Stephen Wright. Numerical optimization. Springer Series in Operations Research and Financial Engineering. Springer Nature, 2006.
[29] Paul Vicol, Jonathan P Lorraine, Fabian Pedregosa, David Duvenaud, and Roger B Grosse. On implicit bias in overparameterized bilevel optimization. In International Conference on Machine Learning, pages 22234 22259. PMLR, 2022.
[30] Levent Sagun, Leon Bottou, and Yann LeCun. Eigenvalues of the hessian in deep learning: Singularity and beyond. arXiv preprint arXiv:1611.07477 2016.
[31] Behrooz Ghorbani, Shankar Krishnan, and Ying Xiao. An investigation into neural net optimization via hessian eigenvalue density. In International Conference Machine Learning, pages 2232-2241. PMLR, 2019.
[32] Yegor Klochkov and Yang Liu. Revisiting inverse hessian vector products for calculating influence functions. arXiv preprint arXiv:2409.17357, 2024.
[33] Andrea Schioppa, Polina Zablotskaia, David Vilar, and Artem Sokolov. Scaling up influence functions. In Proceedings of the AAAI Conference on Artificial Intelligence, pages 8179-8186 2022.
12

---

[34] Samyadeep Basu, Philip Pope, and Soheil Feizi. Influence functions in deep learning are fragile. arXiv preprint arXiv:2006. 14651, 2020.
[35] Andrea Schioppa, Katja Filippova, Ivan Titov, and Polina Zablotskaia. Theoretical and practical perspectives what influence functions do. Advances in Neural Information Processing Systems, 36, 2024.
[36] Juhan Bae, Nathan Ng, Alston Lo, Marzyeh Ghassemi, and Roger B Grosse. If influence functions are the answer, then what is the question? Advances in Neural Information Processing Systems, 15:17953-17990 2022.
[37] C Spearman. The proof and measurement of association between two things. The American Journal of Psychology, 15(1):72-101, 1904.
[38] L Lo Gerfo, Lorenzo Rosasco, Francesca Odone, E De Vito, and Alessandro Verri. Spectral algorithms for supervised learning. Neural Computation, 20(7):1873-1897 2008.
[39] Xavier Glorot and Yoshua Bengio. Understanding the difficulty of training deep feedforward neural networks. In Proceedings of the thirteenth international conference on artificial intelligence and statistics, pages 249-256. JMLR Workshop and Conference Proceedings, 2010.
[40] Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. Dropout: simple way to prevent neural networks from overfitting. The journal of machine learning research, 15(1):1929-1)) 1958, 2014.
[41] Mu Li, Tong Zhang, Yuqiang Chen, and Alexander J Smola. Efficient mini-batch training for stochastic optimization. In Proceedings of the 20th ACM SIGKDD international conference on Knowledge discovery and data mining, pages 661-670, 2014.
[42] Jacob R Epifano, Ravi P Ramachandran, Aaron J Masino, and Ghulam Rasool. Revisiting the fragility of influence functions. Neural Networks, 162:581-588, 2023.
[43] Elisa Nguyen, Minjoon Seo, and Seong Joon Oh. A bayesian approach to analysing training data attribution in deep learning. Advances in Neural Information Processing Systems, 36, 2024.
[44] Dougal Maclaurin, David Duvenaud, and Ryan Adams. Gradient-based hyperparameter optimization through reversible learning. In International conference on machine learning, pages 21 PMLR, 2015.
[45] Hippolyt Ritter, Aleksandar Botev, and David Barber. scalable laplace approximation for neural networks. In 6th international conference on learning representations, ICLR 2018- conference track proceedings, volume 6. International Conference on Representation Learning, 2018.
[46] Tijmen Tieleman. Lecture 6.5-rmsprop: Divide the gradient by running average of its recent magnitude. COURSERA: Neural networks for machine learning, 4(2):26, 2012.
[47] Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv: 2014.
[48] Richard H Byrd, Samantha L. Hansen, Jorge Nocedal, and Yoram Singer. A stochastic quasi- newton method for large-scale optimization. SIAM Journal on Optimization, 1602);1008-1031, 2016.
[49] Roger Grosse. Neural network training dynamics, 2021.
[50] Barak A Pearlmutter. Fast exact multiplication by the hessian. Neural computation, 6(1): 147-160, 1994.
[51] Magnus Hestenes, Eduard Stiefel, et Methods of conjugate gradients for solving linear systems. Journal of research of the National Bureau of Standards, 49(6):409-43 436, 1952.
13

---

[52] Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. Advances in neural information processing systems, 25, 2012.
[53] R Dennis Cook. Influential observations in linear regression. Journal of the American
Statistical Association, 74(365):169-174, 1979.
[54] Nicol N Schraudolph. Fast curvature matrix-vector products for second-order gradient descent. Neural computation, 14(7):1723-1738, 2002.
[55] James Martens. New insights and perspectives on the natural gradient method. Journal of Machine Learning Research, 21(146):1-76, 2020.
[56] Myeongseob Ko, Feiyang Kang, Weiyan Shi, Ming Jin, Zhou Yu, and Ruoxi Jia. The mirrored influence hypothesis: Efficient data influence estimation by harnessing forward passes. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages (2286-26295, 2024.
[57] Andrea Brennen. What do people really want when they say they want" explainable ai?" we asked 60 stakeholders. In Extended abstracts of the 2020 CHI conference human factors in computing systems, pages 1-7, 2020.
[58] Zayd Hammoudeh and Daniel Lowd. Training data influence analysis and estimation: A survey. Machine Learning, 1 11355):3311-240) 2024.
[59] Pouya Pezeshkpour and Sarthak Jain. Combining feature and instance attribution to detect artifacts. Proceedings of the Association for Computational Linguistics (ACL), 2022.
[60] Elisa Nguyen, Johannes Bertram, Evgenii Kortukov, Jean Y Song, and Seong Joon Oh. Towards user-focused research in training data attribution for human-centered explainable ai. arXiv preprint arXiv:2409.16976 2024.
[61] Shreya Shankar, Rolando Garcia, Joseph M Hellerstein, and Aditya Parameswaran. Opera- tionalizing machine learning: An interview study, arXiv preprint arXiv:2209.09/25 2022.
[62] Yousef Saad. Iterative methods for sparse linear systems. SIAM, 2003.
[63] Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou. Datainf: Efficiently estimating data influence in lora-tuned iims and diffusion models. arXiv preprint arXiv:2310.00902, 2023.
[64] Reeves Fletcher and Colin M Reeves. Function minimization by conjugate gradients. The computer journal, 7()): 149-154, 1964.
[65] James Martens al. Deep learning via hessian-free optimization. In leml, volume 27, pages 735-742, 2010.
[66] Junwei Deng, Ting- Wei Li, Shiyuan Zhang, Shixuan Liu, Yijun Pan, Hao Huang, Xinhe Wang, Pingbang Hu, Xingjian Zhang, and Jiaqi Ma. dattri: A library for efficient data attribution. Advances in Neural Information Processing Systems, 37:136763-13678 2024.
[67] Zhe Li, Wei Zhao, Yige Li, and Jun Sun. Do influence functions work on large language models? arXiv preprint arXiv:2409. 19998, 2024.
[68] R. Fletcher and M. J. D. Powell. A rapidly convergent descent method for minimization. The Computer Journal, 6(2):163-168) 1963.
[69] Charles G Broyden. class of methods for solving nonlinear simultaneous equations. Mathe- maties of computation, 19 19(92):577-593, 1965.
[70] Dong C Liu and Jorge Nocedal, On the limited memory bigs method for large scale optimiza- tion. Mathematical programmingg 45(1):503-528, 1989.
[71] Jorge Nocedal. Updating quasi-newton matrices with limited storage. Mathematics of compu- tation, 35(115))777777782 1980.
14

---

[72] Shun-Ichi Amari. Natural gradient works efficiently in learning. Neural computation, 10(2):
251-276, 1998.
[73] Guodong Zhang, Shengyang Sun, David Duvenaud, and Roger Grosse. Noisy natural gradient variational inference. In International conference machine learning, pages 5852-5882 3852-5861. PMLR, 2018.
[74] James Martens and Ilya Sutskever. Learning recurrent neural networks with hessian-free optimization. In Proceedings of the 28th international conference on machine learning (ICML-11), pages 1033-1040, 2011.
[75] Juhan Bac. Beyond Gradients: Using Curvature Information for Deep Learning. PhD thesis, University of Toronto (Canada), 2025.
[76] Dheeru Dua and C Graff. Uci machine learning repository. university of california, school of information and computer science, irvine, ca (2019), 2019.
[77] Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images, 2009.
[78] Kaiming He, Xiangyu Zhang. Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 770-778, 2016.
[79] Yann LeCun, Corinna Cortes, and CJ Burges, MNIST handwritten digit database. ATT Labs, 2, 2010.
[80] Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion-MNIST: A novel image dataset for benchmarking machine learning algorithms. arXiv preprint 17Xiv:1708.07744 2017.
[81] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[82] Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mixture 2016.
models. arXiv preprint
[83] Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. Estimating training data influence by tracing gradient descent. Advances Neural Information Processing Systems, 13:10091119930, 2020.
[84] Roger Grosse and James Martens. A Cronecker-factoree approximate fisher matrix for con- volution layers. In International Conference on Machine Learning, pages 573-582, PMLR, 2016.
[85] Anders Søgaard et al. Revisiting methods for finding influential examples. arXiv preprint arXiv:2111.04683 2021.
[86] Yuzheng Hu, Pingbang Hu, Han Zhao, and Jiaqi W Ma. Most influential subset selection: Challenges, promises, and beyond. arXiv preprint arXiv:2409. 2024.
[87] Nathan Ng, Roger Grosse, and Marzyeh Ghassemi. Mcasuring stochastic data complexity with boltzmann influence functions. arXiv preprint arXiv:2406.02745. 2024.
[88] Arthur Jacot, Franck Gabricl, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018.
[89] Jaehoon Lee, Lechao Xiao, Samuel Schoenholz, Yasaman Bahri, Roman Novak, Jascha Sohl- Dickstein, and Jeffrey Pennington. Wide neural networks of any depth evolve as linear models under gradient descent. Advances in neural information processing systems, 32, 2019.
[90] Alexander Wei, Wei Hu, and Jacob Steinhardt. More than a toy: Random matrix models predict how real-world neural representations generalize. In International Conference on Machine Learning, pages 23549-23333 23588. PMLR, 2022.
15

---

[91] Frederik Kunstner, Philipp Hennig, and Lukas Balles. Limitations of the empirical fisher approximation for natural gradient descent. Advances in neural information processing systems, 32, 2019.
[92] Nikunj Saunshi, Arushi Gupta, Mark Braverman, and Sanjeev Arora. Understanding influence functions and datamodels via harmonic analysis. In The Eleventh International Conference on Learning Representations, 2022.
[93] Samyadeep Basu, Xuchen You, and Soheil Feizi. On second-order group influence functions for black-box predictions. In International Conference on Machine Learning, pages 715-724. PMLR, 2020.
[94] Balaji Lakshminarayanan, Alexander Pritzel, and Charles Blundell. Simple and scalable predictive uncertainty estimation using deep ensembles. Advances in neural information processing systems, 30, 2017.
[95] Pavel Izmailov, Dmitrii Podoprikhin, Timur Garipov, Dmitry Vetrov, and Andrew Gordon Wilson. Averaging weights leads to wider optima and better generalization. arXiv preprint #rXiv:1803.05407 2018.
[96] Wesley J Maddox, Pavel Izmailov, Timur Garipov, Dmitry P Vetrov, and Andrew Gordon Wilson. A simple baseline for bayesian uncertainty in deep learning. Advances in neural information processing systems, 32, 2019.
[97] Andrew G Wilson and Pavel Izmailov. Bayesian deep learning and a probabilistic perspective of generalization. Advances in neural information processing systems, 33:4697-4708,, 2020.
[98] Yarin Gal and Zoubin Ghahramani. Dropout as a bayesian approximation: Representing model uncertainty in deep learning. In international conference on machine learning, pages 1050-1059, PMIR, 2016.
[99] Kelvin Guu, Albert Webson, Ellie Pavlick, Lucas Dixon, Ian Tenney, and Tolga Bolukbasi. Simfluence: Modeling the influence of individual training examples by simulating training runs. arXiv preprint arXiv:2303.081 14, 2023.
[100] Paul W Holland. Statistics and causal inference. Journal of the American statistical Association,
81 (396)9945-96 1986.
[101] Xiaochuang Han, Byron C Wallace, and Yulia Tsvetkov. Explaining black box predictions and unveiling data artifacts through influence functions. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguisties, pages 5553-5563, 2020.
[102] Jinghan Yang, Sarthak Jain, and Byron C Wallace. How many and which training points would need to be removed to flip this prediction? arXiv preprint 22Xx:23002.02666 2023.
16

---

NeurIPS Paper Checklist
The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No], or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (1-2 sentence) justificcton right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation.
While "[Yes] " is generally preferable to "[No]", ", it is perfectly acceptable to answer "[No] provided proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased in binary way, we acknowledge that the true answer often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Checklist",
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: Yes.
The main claims made in the abstract and introduction accurately reflect the paper's contributions and scope.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made 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]
17

---

We discuss the limitations in Appendix H and throughout the body of the paper. Specifically, Section describes the additional computational cost from running ASTRA compared to EKFAC.
Justification:
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 encouraged to create 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails 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, worse outcome might be that reviewers discover limitations that aren't 't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant 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 proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer:
[Yes]
Justification:
We describe the derivation of ASTRA in Section 2 and Section 3.
Guidelines:
The answer NA means that the paper does not include theoretical results. the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix 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 ex- perimental results the paper to the extent that 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 disclose the information to reproduce the main experimental results in Section 5, Section 6 and Appendix F.
18

---

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 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 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 specific model and empirical evaluation, it may be necessary to either make 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 submis- sions 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. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-sonnee open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe 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 verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [No]
Justification: We do not provide open access the data and code. We use publicly available data for our experiments, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cm cc/
Sawiiddeee/Cdeeuummseeeeeee sassssssssiiiiooo for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https:
/ //nips wttpccccpppcccceeeeeeceec.cceclt htwbbiiiggiddessodduuutiii Swto//oeoeeesiiiieeeeeeee sswonnolicciiooeoooo 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. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
19

---

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) 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, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide all the training and test details necessary to understand the results in Section 5, Section and Appendix F.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to 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, supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: We show eтгоr bars representing standard error for one-model LDS perfor- mance comparisons and all other experiments. Details in Appendix
Guidelines:
The answer NA means that the paper does include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, 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, 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 the standard error of the mean.
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 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 range (e.g. negative error rates).
If error bars 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
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: We provide information on the computer resources needed to reproduce the experiments in Appendix F.
20

---

Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU 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 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 NeurlPS Code of Ethics https: /ttt://eeeiss..eeeeeee
Answer: [Yes]
Justification: The research conducted in this paper conform, in every respect, with 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 deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
Broader impacts
Question: Does the paper discuss both potential positive socictal impacts and negative socictal impacts of the work performed?
Answer: [Yes]
Justification: We discuss both potential positive societal impacts and negative societal impacts of the work performed in Appendix H.
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 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 being used as intended but gives incorrect results, and harms following from (intentional 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI.).
11. Safeguards
21

---

Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow 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 best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: We properly credit the creators original owners of assets used in the paper.
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 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, paperswithcoe thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is available online, the authors are encouraged to reach out to the asset's 's creators.
13. New assets
Question: Are assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl 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 used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
22

---

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: The paper does not involve crowdsoureing nor 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 contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, 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:
The paper does not involve crowdsourcing nor 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.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: addition to using L.MSs as a writing, editing, formatting, and a programming aid, we conduct training data attribution experiment using GPT-2.
Guidelines:
The answer NA means that the method development in this research does not involve L.LMs as any important, original, or non-standard components.
Please refer to LLM policy (https: http://erripss.ccoo/eeeeeeeee for what should or should not be described.
23

---

A Notation & Acronyms
A.1 Notation
Notation
Description
Batch size mini- batih gradient descent
B
Batch size for stochastic Neumann series iterations
D
Number of parameters in neural network
J
Number of iterations for SNI
Number of subsets (masks) to sample LDS computation
N
Number of training data points, N = [D]
P
Number of layers in neural network.
R
Number of trials (repeats) for LiSSA.
D {
Training dataset
Deuery
Query dataset, used to henchmark TDA algorithms small practice Data subset of the training dataset An arbitrary i-th training example training example from the dataset D query data point
S
D
2q
x(1
The inputs (feature vector) of the i-th training example
The neural network output for the i-th training example
to
The ground truth target for the i-th example
500
Sampled target for the i-th example using model probabilities
Source of training procedure randomness Randomness from batch ordering set containing various seeds { The neural network parameters Optimal parameters
0"
0' ('(
Optimal parameters trained data subset S C
6"
Final model parameters (not necessarily optimum)
(E)
Final model parameters (not necessarily optimum) which depends on randomness
02
The parameters of a network after iterations of an algorithm, which depends on context
g0,xx)
The output (logits) of a neural network with parameters and input x Loss function as a function of the neural network output and target (c.g., cross-entropy)
t)
2(0,z)
Loss function function of the parameters and training example
D(O
Cost function, J D) = E'= L(O, z.)
fz, (0)
Measurement function query point Zq. typically correct-class margin or absolute error
l1fan
Neumann series iteration objective
iHVP Val=, (0*) AI)
Neumann series iteration objective projected onto the subspace corresponding to S;
fzn
The Fisher information matrix
Fekfac
The Fisher Information Matrix approximated with EKFAC
The Generalized Gauss-Newton Newton Hessian (GGN) matrix
G
An unbiased sample of the GGN computed with data itcration A
H
Hessian matrix H V2J (0" (
P
Preconditioning matrix used in ASTRA, chosen as the Gerfac
8
Fraction of D used for subsampling when computing LDS Learning rate when training the neural network Learning rate to find the iHVP with SNI ASTRA
Downweighting amount used influence functions and SOURCE formulation Damping parameter to compute iHVPs: a small positive scalar.
A
Ā
Damping parameter added the preconditioner: small positive scalar Training data attribution method
Eigenvalues in the decomposition of the curvature matrix.
Group influence of training data attribution method The Spearman correlation coefficient [37]
P
The Kronecker product
24

---

A.2 SOURCE specific notation
Notation
Description
Ögs
Indicator variable used in SOURCE formulation An index variable indicating the current segment in SOURCE
Kg
The number of iterations within segment in SOURCE
L
The number of segments in SOURCE
The number of optimization steps for the underlying model
The average gradient in segment {
T]
The average learning rate segment
Defined in Equation 20
S
Defined in Equation 20 An approximation of Fş introduced by Bae al. [19]
A.3 Acronyms
Acronym
Description
CG
Conjugate Gradient [51]
EKFAC
Eigenvalue-corrected Cronecker-factored Approximate Curvature [27]
FIM
Fisher Information Matrix
iНVP
Inverse Hessian-vecton product, or inverse Gauss-Newtoon-esssiann product
IF
Influence functions [17]
KFAC
Kronecker-factored Approximate Curvature [26]
LiSSA
Linear time Stochastic Second-Order Algorithm [23]
LOO
Leave-one-out
LDS
Linear Datamodeling Score [22]
MLP
Multi-layer perceptron
NI
Neumann series iterations
PBRF
mookimal-Bregman Response Function [36]
SGD
Stochastic gradient descent
SOURCE
Segmented statiOnary UnRolling for Counterfactual Estimation [19]
SNI
Stochastic Neumann series iterations
TDA
Training data attribution
25

---

Algorithm iHVP approximation with LiSSA
Require: v E RRP, a > 0 (learning rate), J > 0 (number of iterations), R (repeat size),
(damping term), B (batch size), D (training dataset) Initialize the accumulator for final estimation
x +-+-0
for 1 to R do
Vo f- v
D Initialize Vo as per the initial condition
for j 0 to J do
B + SumpleeWttheplleemen(((. B)
Sample a mini-batch of size B from D
f GgV,
> Compute HVP using mini-batch B. SNI update rule
fJ+1 V, ap Av,) +av end for
f- + V.
Accumulate the result of this repetition
end for
+ x/R
D Average the accumulated results over R repetitions i Return final iHVP estimation H H
return
B Extended Preliminaries
B.1 LISSA and SNI
LiSSA [23] is second-order optimization algorithm which involves the computation of iHVP. Koh and Liang [17] choose LiSSA as their iHVP solver, but LiSSA contains other components as well. In this paper, "LiSSA" refers to the iHVP component. Algorithm shows that the primary difference between LiSSA and SNI is that the former repeats the SNI procedure multiple times to reduce variance (highlighted in red). We use R throughout this paper, LiSSA's iHVP component equivalent to SNI and thus we use the two terms interchangeably.
B.2 Influence Functions
Previous papers have shown that influence function estimates often fragile due to the strong assumptions in the influence function derivation [34-36, 42, 85, 86]. In Table 1, we outline the main assumptions.
Table Summary of assumptions of Influence Functions vs. Unrolled Differentiation. Assumption
Influence Functions
Unrolled
Differentiation
First order approximation Objective convex with respect parameters Model trained to optimal parameters
/
X
X
Despite the strong assumptions the derivation of influence functions, a poor iHVP approximation make influence functions estimates appear less reliable than they are. To appreciate these assumptions, we refer the reader to Appendix B.1 [36] for a well-presented derivation of influence functions.
Ensembling Influence Functions TDA scores can be typically ensembled over multiple training trajectories [19, 22] to mitigate the problem of noise in the training procedure [42, 43). This is typically done by training models with various seeds & E E, and approximating the expected first- order downweighting effect with the empirical average of attribution scores
TIF Ensemblee (Zm, Zg, D) =
Σ CEE
Volz, (0 (E) TGE + AT) -0-1v (VL(("
(15)
where G3 is the GGN computed at (E). We perform ensembling in this paper using the procedure in Equation 15, and apply ensembling for SOURCE analogously. Note that Equation 15 ensembles the attribution scores [19, 22], while some other works ensemble the weights [19, 87].
26

---

B.3 Curvature Matrices
We explain the relationships between curvature matrices below for completeness, which is heavily based on Grosse [49] and Martens [55].
Approximating H with G Throughout this paper, we use the approximation G H. Letting z = g x) denote the neural network output", the GGN is equal to the Hessian if we drop the second term from the following decomposition of H [49]]
Oc E Vo ; dzz (6)
H
Σ
[ J2000 Hgg 2000
J200 +
(16)
N
(x(())t())
where J Jg is the Jacobian matrix of the neural network's ' outputs with respect to the parameters for the ith training example, Hg: = (z(( t(²) refers to the Hessian of the loss function with respect to the neural network outputs for the ith training example, ac refers to the derivative Zx,"
of the loss function with respect to the jthh neural network output for the ith training example and (ā ( refers the Hessian of the jth neural network output for the ith training example with respect to the parameters. For a linear neural network, (0g x(x)) = O, so the GGN is equal to the Hessian if we linearize the neural networks with respect to the parameters and only capture the curvature in the loss function. Linearization of the neural network is an approximation documented in previous works [55, 88-90] and used frequently for influence functions if the damped GGN G + AI used [7, 11, 22, 36].
Equivalence of F and G For the machine learning tasks that we consider, such regression and classification tasks, the outputs of the neural network can be seen as the natural parameters of an exponential family. For these cases, the Fisher Information Matrix and the GGN coincide (i.e. F G). G). We will illustrate this for softmax classification, but the case for regression can be derived similarly. The cross-entropy loss for a training example (x, t) is L(0, z) -t' log pe (y/x) (x) whose gradient is:
Vo (( z)= = Ino (po(y x) t). Then the following equalities hold:
F = 1
Σ
Ex-po(vix)))
Vo log Vo (oggo(())x(11)
N
(x(t),t(t))e
1
Σ
E Ex-ro(y)x(0)
J (PO x()) (Po x()) - 3) TJ Jam0
N
(x(0),(0))
1
Σ
0 diag(po(S xx()) ro(s (x()))o(s |x())T)J J2
N
(x(0)
EAD
1
Σ
200 (Z(()
(t)) = G
= N
B.4 Training Data Attribution with Unrolled Differentiation
Influence functions may struggle with models that have not converged [19, 35, 36, 83] (Table I). Fortunately, unrolled differentiation methods [4, 18-21] do not rely on model convergence. Instead, they capture the effect of downweighting training example by differentiating through the entire training trajectory. They can also capture the effect of other sources of randomness, such as batch ordering [41]. Assume
optimization algorithm is mini-batch gradient descent, which uses learning rate 77 and batch size B, and let ōμ be an indicator variable that equals if and only if Zm Zki, where Zgz the i-th training example in batch k: Then the mini-batch gradient descent
I Note the difference between the bold font z, which refers to neural network outputs, with italicized font z, which refers to a data point x (x,t)
We note that O2 in this setting refers the parameters time step l when training the network and distinguish from SNI or ASTRA iterations presented in Section and Section
27

---

update rule is:
B B
Ox+1(c) < (c(()=
116
(1-O0stVVC(((ACC)) zks); 1=1
(17)
where 5x : Exn1 ŌRs. Let Es denote the randomness from batch ordering. Then the expected first-order effect on the parameters Or after T steps of training with Zm downweighted by is:
E
d0r
Es
T-
Σ B
The
OR
II [I ( 1=7-1
k+1
VoL(Ok. zm)
(18)
Bae et al. [19] introduce an algorithm called SOURCE, which approximates Equation 18 much more cheaply by segmenting the trajectory into L. segments, assuming stationary and independent GGNs and gradients within each segment. For the 9th segment which starts at iteration T and ends at iteration Tr, and k satisfying Tı-1 S k < Tr, let G. Eç. [.+]. B. == Eg. V (e zm )}, and î refer to the average GGN, average gradient, and average learning rate in segment e respectively and TT-Tr-- TB-1 refer to the total number of iterations in segment t. Then SOURCE approximates
let K2 == T2 - the first-order effect of downweighting parameters as:
8+1
Te- Σ A=TT-1
d
1
EE
23
(I
Tee Gr))
TeeI G.) Te-1))--1100 N Be
(19)
de
P=L
exp( Ke.
ex-exp(-Thhh -TKK G2
g
(20)
L.
5
F
The stationary and independent assumptions allow factor shared products resulting in the ap- proximation in Equation 19, which can be then approximated with matrix exponentials in Equation 20 and computed with the EKFAC eigendecomposition. Bae al. [19] provide interpretation of the ry term, noticing that ry ~ (G, Te K, II) == ř,, which is an iHVP that we can use to apply ASTRA. We discuss this term in more detail Appendix C. We refer the reader to Bac et al. [19] for a full explanation of SOURCE.
B.5 Kronecker-Factored Approximate Curvature
The KFAC approximation was introduced in [26] in the context of second-order optimization and explained in [7] in the context of influence functions. To understand the cost EKFAC and EKFAC- IF compared to ASTRA-IF,
well as the assumptions EKFAC make which results in biased iHVP approximation, we present the derivation below which is heavily based on Grosse [49] and Grosse et al. [7] and refer readers to Martens and Grosse [26] and George al. [27] for further reading.
Our goal is to compute the iHVP with the Fisher Information Matrix (FIM) F an approximation for the Hessian H. The FIM is defined as:
1
Σ x
Ey-po(rix(t)
Ve ( Vo (og
(21)
where pe (y(x) is the model's own distribution targets. We omit the random variables in the expectation's 's subscripts going forward to reduce clutter. Using the model's own distribution over targets (as opposed to actual targets) is rather important since using the actual targets rather than the model's is distribution results in a matrix called the Empirical Fisher, which does not have the same properties the FIM [-1].
We now describe KFAC for multilayer perceptrons. We refer readers to Grosse and Martens [84] for the derivation for convolution layers. Consider the 1-th layer of neural network,"" 20 which has input
Unless otherwise stated, derivatives taken with respect + 1.
evaluated
0. The notation
-
means taking products decreasing order from Note the difference between f, which refers to segment in SOURCE, and l, which is an index denoting layer of neural network.
28

---

activations a;-1 E R', weights W E Rox!, bias by E RO, and outputs RD. For convenience, we use the notation à = [a[1 i] and W, = : bi] to handle weights and biases together, and we write 0 vec( Wi) to denote the reshaped vector of layer parameters. Then each layer computes:
St W1 au (i(ss,,
(22)
where Φi is the activation function. We will define the
operator as:
De V. log x)
(23)
for notational convenience. Notice that De is a random variable whose randomness arises from sampling y. Using the properties of the Kronecker product, we can write the pseudo-gradient of 0 as:
D0; vec(DW.1) vec(Ds)a__1) a Dsi,
(24)
where & is the Kronecker product. Then the ith block of the FIM can be approximated as:
F, = [D00,D0,]
(25)
[ @ Dsi) a1-1 @ Dsi)
(26)
E [ā-1 a 1 3 Ds:Ds
(27)
~ E | [|]] & E DD: Ds|
(28)
A 2 S.
(29)
where have applied Kronecker product identities the third equality. Our final approximation F
to the FIM is the block-diagonal matrix in which each block is F,. Here, A E(a اااة and
S F[Ds are the uncentered covariance matrices of the activations and the pre-activation pseudo-gradient with dimensions (T + 1) x (T + 1) and O x O, respectively. Practically, we can estimate the expectations via empirical estimate and store the resulting statisties A
1 Σ a a/_, and S, Σ Ds. and we define F, A St. N
To approximate (F + AI) 'v for vector needed to compute influence functions, vv, separately for each layer 1. Let V. be the slice of reshaped match W, 1.
compute (F, XI) - and define v, vee V i). Applying the eigendecompoositions A,-1 QA.,, DA.,, QA and
S = Qs, Dg, O and using the Kronecker identity U& = (Qu & (v) (Du & Dv) QQU @ )
for two symmetric matrices U & V, we write:
AD) = (A X -
(30)
(QA,, & (DA,) Ds, +AIA,- & Is,) Is,)
QQ..1 Qs,) Tv, (31)
Scaling matrix Aseac
Orthonormal cigenbasis
where IA... and Is. represent identity matrices of the same shape A1-1 and S respectively.
Eigenvalue Corrected Kronecker-Factoree Approximate Curvature One simple adjustment to the KFAC approximation yield material improvements to the curvature approximation well as the influence approximation. The KFAC formulation in Equation 31 suggests that after expressing Vi Ds, + AL & Is.)
in the eigenbasis (QA.., & Qs.) we scale each element with (D Observing that for any matrix W E[uu ] et al. [27] propose that a better scaling matrix is the diagonal matrix A with entries:
USUT, it is true that S E[(UT v) George
A E(((Q... Qs.) DO:)
(32)
In practice, this eigenvalue correction results in better influence estimates compared with KFAC.
Assumptions in EKFAC From the equations above, we can see that KFAC makes two critical approximations which EKFAC inherits: First, it assumes that correlations between D0; and DO, are zero if they belong to different layers, yielding block-diagonal approximation of Second, it treats the activations independent from the pre-activation pscudo-gradients the basis of Equation 28.
29

---

Furthermore, the matrix S; depends on the sampled labels y. which for efficiency reasons is usually only sampled once per input x which also may introduce some approximation error. In total, these assumptions cause FEKEAC to differ from the true FIM F, which is an error that ASTRA A corrects. Grosse and Martens [84] introduce the KFAC approximation for convolution layers which introduces two further assumptions- spatially uncorrelated derivatives and spatial homogeneity. Consistent with past works [1. 191, we find that EKFAC-IF struggles for convolution architectures in comparison with MLPs. ASTRA-IF's TDA performance on ResNet-9 achieves particularly large improvement compared to EKFAC-IF.
Cost of Computing EKFAC-IF The three main components of computing EKFAC-IF are: I) collecting the statistics A and S,, which requires backward pass over the whole dataset. 2) computing the eigendecompositions A QA. DA, QAi,, and S Qs, Ds, Qs whose total precise cost depends on the architecture [7]. 3) if we want to search the whole dataset D for influential examples, once we approximate (F + - (0*) via the procedure above, we need to take dot product with every training example gradient under consideration in D. ASTRA adds an incremental iterative procedure the cost of EKFAC-IF which results in stronger TDA performance.
C Introducing ASTRA-SOURCE
Understanding the iHVP We have discussed how to compute iHVPs and therefore influence functions with ASTRA. We can also apply ASTRA to SOURCE by making the substitution rg R (G, + 77 ",7KK'1' 'T) ge == Fy described in Appendix B.4, which replaces the matrix exponential in Fe. To understand the approximation, observe that the following term (with some rearranging) found in Equation 19 v/. - - (I - 77 G. ) gr resembles applying the Neumann series iterations on the vector = -gr and the matrix A G, with learning rate of 7/ for total of K;- iterations. For large enough Kg, the truncation can be seen approximately running Neumann series iterations until convergence, which results in the iHVP G2 ge. However, each segment actually only involves Kr iterations at an average learning rate of 7/g and thus is more closely related to truncated Neumann series iterations (Appendix G), in which the parameters make less progress the low eigenvalue directions. Bae et al. [19] provide interpretation that this can be approximated with damping as follows: T/e Σ (I TT -r------------- gr (G + 7 K, - "D)-'g>> gg, and that over a wide range of eigenvalues, the qualitative behavior matches well. The transformation of this term into an iHVP allows us to apply ASTRA to SOURCE.
Practical Implementation Recall that the final goal in SOURCE to approximate Volz, (0*) SieTTE,, Ee, d'de] To do this, we follow Equation 20 from left to right: for all é == l,..., segments, we first compute 1 Vofz, (0*) I Sp in the same manner as SOURCE. This will be the vector in our iHVP. ASTRA differs from SOURCE in that instead of multiplying this vector
N
by another matrix exponential, we multiply it by the inverse damped GGN (G, + mE kz'r) which we can do using ASTRA in the same manner described previously. Finally, we multiply by the average gradient, g, and accumulate the result over I. segments, which is done in the same manner as SOURCE Compared to ASTRA-IF, there is an additional detail introduced: how to obtain the average GGN Ge. There are a number of options available, but we find that using the average weights in the segment works well, which we discuss below. The full procedure of applying ASTRA to SOURCE requires T. iHVPs per query. In many cases, the number of segments T. is likely to be small. 21 Since the preconditioners used by ASTRA would need to be computed if we were to run EKFAC-SOURCE anyways, the incremental cost of ASTRA-SOURCE, compared to EXFAC-SOURCE, is no more than L times the number of iterations for each iHVP, which is hundreds of iterations in our experiments.
Approximation of Other Terms We have discussed how to improve the approximation of Fr with ASTRA, but have not discussed S.. We found evidence that improving the quality of the term Fg improved TDA performance, but could not find the same evidence for S Therefore we spent
Part of motivation for SOURCE is to devise a scalable TDA algorithm multi-stage training procedures. in which the number of segments is likely to be modest. Bae et al. [19] use L = 3.
30

---

additional compute on improving the iHVP approximation. As a result, we will leave Se as the approximation involving EKFAC.
Computing the Average Gauss-Newton Hessian There are number of options in computing the average GGN G.. Option 1): since Ge = E [&1] for T:-1 S h: < Tr, one can compute the matrix-vector product involving G. and v E RD simply by taking an empirical average over samples within the segment: G.v = से Erz ISkcT Es, [x]vv. This might be costly since one would have to load multiple checkpoints into memory just to compute a forward pass. Option 2): Instead of sampling every checkpoint in the segment, we could reduce the number of samples and take the empirical average of that instead. This is consistent with SOURCE as it uses only a subset of checkpoints in each segment anyways. If the segments chosen in SOURCE are indeed stationary as the derivation approximates, then in the extreme, we could take one checkpoint in each segment as the representative. Option 3): Bae al. [19] provides an alternative averaging scheme called FAST-SOURCE, in which one averages the parameters rather than the gradients, and shows that it works approximately as well as SOURCE. We tested Option 2 and Option 3 for a few settings and could not find meaningful differences, so opted to present the results for Option 3, using the average weights.
D Extended Related Works
Understanding Inffuence Functions A number of previous works [34-36, 42, 85, 86, 92] have studied influence functions accuracy in various modern neural network settings. [34] show that influence functions do not approximate LOO retraining well. [36] discover that the derivation of influence functions actually approximate the Proximal Bregman Response Function (PBRF) rather than LOO retraining, which can be seen an objective which tries to maximize loss on the removed training example, subject to constraints in function space and weight space measured from the final parameters. Two large contributions to influence functions error are the warm-start retraining assumption, which assumes that the counterfactual model is initialized at the final parameters, and the non-convergence gap, which relates to the fact that the derivation assumes the model has converged to an optimal solution. Ensembling can help address the warm-start retraining assumption, while the non-convergence gap addressed by [19]. Others have focused on whether influence functions can accurately approximate group influence [5, 10, 92, 93] as it makes linearity assumption due to the fact that it is first-order approximation. While LOO influence very noisy [19, 43], Ilyas et al. [10] discover that model predictions are approximately linear with respect to training example inclusion, which provides the justification for LDS [22]. Overall, these works typically aim to address the fundamental assumptions surrounding influence functions (Table 1). In contrast, our work shows that a poorly approximated iHVP can cause substantial performance degradation.
Ensembling in Influence Functions Ensembling combines multiple models for improved gener- alization, uncertainty estimation, and calibration. is a common approach to estimate the model posterior 2) in Bayesian deep learning. Different strategies for sampling models as members of an ensemble exist: For example, deep ensembles sample models from varying random initial- izations [94] to represent variations stemming from possible training trajectories. while stochastic weight averaging (SWA) approaches sample model parameters from the final iterations of model training [95, 96], which has the advantage of reduced training cost. Other methods may combine these ideas [97], or use different approaches (e.g.. Dropout [98]). Taking the average across several runs with varying sources of training process randomness is common approach to account for the variability of model training in TDA estimation [10, 19, 22]. This can be seen as sampling from the distribution of true TDA to estimate the average treatment effect of excluding a training subset [43] and has been effective in stabilizing estimations as well as evaluations (c.g. the LDS [22]). The size of the ensemble is generally connected to improved estimation quality of TDA scores, as shown in [22]. Our results show that ASTRA enjoys a larger performance boost from ensembling.
31

---

E
Evaluating TDA
In this paper, we evaluate the performance of TDA methods with the LDS [22], widely used metric which we shortly described in Section 5. Besides LDS, other evaluations for measuring TDA performance also exist. In this section, we present alternate methods of TDA evaluation.
Expected leave-one-out retraining. TDA methods usually define the influence of training sample Zm on the model as the change in the model's 's predictions if 2m were not part of the training set. Hence, a straightforward way to compute the ground-truth to compare against is leave-one-out (IOO) retraining, as done in [6. 17, 99]. However, since the stochasticity inherent model training makes LOO noisy measure [42, 43], the LOO score should be considered in expectation over the training process stochasticity ç. We can then define the expected leave-one-out (ELOO) score as:
(z.: Fyg D) Egl/z, (0 (zm):()) [= (O;
(33)
This score can be viewed the ground-truth average treatment effect (ATE) [100] of the removal of Zm from D. While principled, in practice, the effect of removing a single point highly noisy [19, 43] so that stable estimate of ELOO may only be achieved with an extremely large number of samples to compute the empirical expectation. In contrast, LDS considers the ATE of excluding a group of training samples from training, which has been shown to be more stable in expectation [19, 22].
Top-k removal and retraining. The sign of the TDA scores indicates whether the excluded training samples are positively negatively influential [19], also referred to as proponents and opponents [83],
helpful and harmful samples [17], or excitatory and inhibitory [6], respectively [19]. The idea behind this evaluation is that the removal of positively influential samples {zm) removes support for the
query sample 2, and consequently should lead to a change in prediction confidence on [101]. [101] conduct this evaluation by removing the top and bottom 10% of training samples ranked by influence functions and compare against the removal of the least influential (i.e., smallest influence scores by magnitude) and random samples to see if the resulting models change their predictions in the expected way. Similar to top-k removal and retraining, previous work has tested how many highly influential samples need be removed to flip a prediction [19, 102], called subset removal counterfactual evaluation. Similar to LDS, removing top-k samples is based on counterfactual retraining with excluded groups. However, one core difference is that since LDS involves sum over all the attribution scores for every 2m in each subset (Equation 10), poorly calibrated attribution scores resulting one outlier score T(Zm, Zy. D) may result in poor LDS, demanding that the TDA method assigns calibrated scores across all points nm in consideration.
F
Experiment Details
F.1 Choice of Measurement F Function
While in principle the derivation of our influence scores does not restrict what measurement function fz, is used, in practice some choices measurement functions work better than others. In our experiments, for regression problems,
the measurement function:
f=, (0) = (0 xg) tyy
(34)
where g(, x,) denotes the last layer output of the neural network when X, is the input and ty is scalar output. For all classification problems with the exception GPT-2, we use the measurement function:
0(g(0,x.)) xy))t,
(0) log - (g x.y )),
(35)
= - x,),, + log
Σ g exp g( xy)e,
(36)
where a denotes the softmax function and the subscript ty refers to the taking the entry corresponding to the correct class. This measurement function is identical to the one found in [22]. For GPT-2, use the training loss as the measurement function.
32

---

F.2 Experiment Details
Table 2 shows the architecture and hyperparameters used compute the final parameters 0, which we use to run EKFAC-IF, EKFAC-SOURCE, ASTRA-IF and ASTRA-SOURCE. The first column shows the size of the training dataset and the number of query examples in each setting. The third column shows the hyperparameters used to train the models in each setting; we use the same hyperparameters as reported in Bae al. [19]. We estimate the expected ground-truth in the left- hand-side of Equation 10 with an empirical average the last column in Table shows the number of repeats per mask S, (i.e., number of sampled) estimate this value. All settings use a constant learning rate, with the exception of CIFAR-10, which uses cyclic learning rate schedule.
Table 2: Summary of training details.
Dataset
Architecture
Hyperparameters
Ground-truth Retraining
SGD w/ momentum
Learning rate: 10
UCI Concrete 896 training examples
MLP Layers (128, 128, 128) Hidden Units activation
Weight docay: 10-5
Repeats: 100
Musks; 100
Momentum:
103 query examples
Total: 10,000
Hatch size: 32
Epochs:
w/ momentum
Learning rate: 10-2
UCI Parkinsons
ML.P Layers (128, 128. Hidden Units
Weight decuy: 10-3
Repeats: 100
5,280 training examples 100 query examples
Masks: 100
ReLU activation
Momentum: 0.9
Total: 10,000
Batch 32
Epochs: 20
SGD w/ momentum
Learning rate: 10
MNIST (Suhset)
MLP- Layers (512, 256, 128) Hidden Units
Weight docay: 10-3
Repeats: 50
6, 6,144 training examples
Musks: 100
Momentum:
100 query examples
activation
Total: 5,000
Batch size: 64
Epochs:
SGD w/ momentum
Leurning
10
FashionMNIST (Subset) 6,144 training examples 100 query examples
MLP Layers
Weight decay: 10
Repeats: 50
256. Hidden Units
Masks: 100
Momentum: 0.9
ReLU activation
Total: 5,000
Batch 64
Epochs: 20
SGD w/ momentum
Peak learning rate: 0.4
CIFAR-10 (Subset) 3,072 training examples 100 query examples
Weight decay: 10 a
Repeats: 50
ResNer-9 [78)
Musks: 100
Momentum: 0.9
Total: 5,000
Batch size: 512
Epochs:
Adam
WikiText-2
4,656 truining sequences 512 sequence length 100 query sequences
Leurning rate: 10
Repeuts: 10
GPT22 [81]
Weight decay: 10-2
Masks: 100
Batch size: 8
Total: 1,000
Epochs:
SGD w/
FashionMNIST--N
Learning rute: 10
6, 144 training examples
MI.P- Layers (512. 256, 128) Hidden Units
Weight decay: x 10
Repeats:
Masks: 100
30% of the dataset randomly relabeled
100 query examples
ReLU activation
Momentum:
Batch size: 64
Total: 5,000
Epochs:
Table 3 shows the details for the various TDA algorithms we use in our experiments.
For EKFAC-IF, use the same damping as the corresponding ASTRA-IF for comparability. SOURCE provides a natural value for damping from its derivation: d = 1/ [19], allow- ing to sidestep tuning the damping term. For both EKFAC-IF and ASTRA-IF, we use the damping value implied by SOURCE for comparability between influence functions and SOURCE: we take the average de implied by SOURCE for each segment and weigh them by the total iterations K2 in segment as our influence functions damping value.
EKFAC-IF, EKFAC-SOURCE, ASTRA-IF, and ASTRA-SOURCE all compute influence on the same set of layers. For MLP architectures, we compute influence on all layers. For ResNet-9, we
33

---

compute influence on MLP and convolution layers. For GPT-2, we compute influence only on MLP layers in line with Grosse al. [7].
Table 3: Summary of TDA Algorithm Details
Dataset
ASTRA-IF Details
ASTRA-SOURCE Details
We use 3 segments. For each segment, we run ASTRA for 200 iterations, and use preconditioner damping factor equal the GGN damping factor A = A 1/ learning rate ar 0. batch-size 256, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
We run ASTRA for 200 iterations, and use GGN damping factor 0.0017.
preconditioner damping factor À = 0.0017, lcarning rate 0.1 batch-size 256, decay the learning rate by 0.5 every 50 iterations.
UCI Concrete
SGD 0.9 momentum,
use segments. For each segment, we run ASTRA for 200 iterations, and use preconditioner damping factor equal to the GGN damping factor A A 1/ learning rate az = 0.1 batch-size 256, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
We run ASTRA for 200 iterations, and use GGN damping factor 0.00091, preconditioner damping factor À = 0.00091, lcarning rate 0.1À, batch-size 256, SGD w/ 0.9 momentum, lcarning rate by 0.5 every 50 iterations.
UCI Parkinsons
decay the
We use 3 segments. For each segment, run ASTRA for 200 iterations, and preconditioner damping factor equal to the GGN damping factor de Ae 1/4, Ke. learning rate ar 0.1 batch-size 256, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
We run ASTRA for 200 iterations, and use GGN damping factor À = 0.0052, preconditioner damping factor Ā = 0.0052, learning rate a 0, hatch-size 256,, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
MNIST
We use 3 segments. For each segment, we run ASTRA for 200 iterations, and use preconditioner damping factor equal the GGN damping factor X = x 1/1, Ke, learning rate az 0. Ae, batch-size 256, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
We run ASTRA for 200 iterations, and use GGN damping factor = 0.0052, preconditioner damping factor Ā 0.0052, learning rate a 0.1A, batch-size 256. SGD w/ 0.9 momentum, and decay the lcamning rate by 0.5 every 50 iterations.
MaSSI
We use segments. For each segment, run ASTRA for 200 itcrations, and use preconditioner damping factor equal to the GGN damping factor A ^ Ke, learning rate ax = 0.01 batch-size 128, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
We run ASTRA for 200 iterations, and use GGN damping factor À 0.014, preconditioner damping factor À = 0.014, learning rate 0.01 batch-size 128, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
CIFAR-10
We use 3 segments. For each segment, run ASTRA for 300 iterations, and preconditioner damping factor equal the GGN damping factor de Ae Ke, learning rate Oę Ae, batch-size SGD w/ 0.9 momentum, and decay the learning rate by 0.9 every 100 iterations.
We run ASTRA for 300 iterations, and use GGN damping factor À = 0.011, preconditioner damping factor A = 0.011. learning rate a = X. batch-size 16, SGD w/ 0.9 momentum, and decay the learning rate by 0.9 every 100 iterations.
WikiText-2
We use 3 segments. For each segment, we run ASTRA for 200 iterations, and use preconditioner damping factor equal the GGN damping factor X = A 1/ Ke, learning rate az 0.1 batch-size 256, SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
We run ASTRA for 200 iterations, and use GGN damping factor = 0.10, preconditioner damping factor À = 0.10, learning rate a 0. batch-size 256. SGD w/ 0.9 momentum, and decay the learning rate by 0.5 every 50 iterations.
MNsSTon
34

---

Other Baselines In Figure 2, in addition to the EKFAC-based baselines, we also compare ASTRA against TracIn [83] and TRAK [22]. We use the same checkpoints for TracIn as SOURCE for comparability. TRAK applies random projections to reduce the computation and memory footprint - for MNIST, FashionMNNIST(-N) and UCT Parkinsons, we use projection dimensions of 4,096 and for UCI Concrete, we use projection dimension of 512, consistent with [19]. For CIFAR-10, 10, we do a hyperparameter sweep among [128, 256, 512, 1024, 2048, 4096, 8192, 20480] for the best hyperparameter (512). We omit TRAK's results on GPT-2 due to lack of publicly available implementations.
Training Curve Details In Figure 3. compare ASTRA-IF and SNI on arbitrary Zy. using 10 seeds and report mean values along with shaded regions representing standard error. To compute both ASTRA-IF and the baseline SNI in Figure 3. conduct hyperparameter sweep for the learning rate from 10" to 10 5 in increments of one order magnitude with momentum set to zero, and the best learning rate based on the average training objective Both SNI and ASTRA-IF use the same damping values and batch sizes for Figure 3, as listed in Table 3 for comparability. The learning rates used were 10 10-2 for all settings for ASTRA-IF, and the best learning rates for SNI were 1, 0.1, 0.01, 0.1 for MNIST, FashionMNIST, CIFAR-10, and UCI Concrete respectively.
the last ten iterations.
Eigendecomposition Details In Figure 4, we compare the performance of influence functions computed with ASTRA and SNI after projecting various subspaces. This experiment uses smaller damping of 10-4 for both MNIST and FashionMNIST to be able to discern the impact of directions of low curvature
influence functions performance. We use the same batch sizes disclosed in Table 3 and no momentum. For ASTRA-IF, we used a learning rate of 10 1 for both settings. For the baseline SNI, the learning rates were 0.1 and 0.01 respectively for MNIST and FaShionMNIST. which we found to give the strongest LDS for the subspace corresponding to Sg.
Compute Resources A shared cluster was used (both internal and external), consisting of mix of A6000 (48GB), A100 (80GB) and H100 (80GB) GPUs, which were used to conduct all experiments. Computing the ground truth of the LDS experiments in Figure is an expensive component of the overall compute to replicate the experiments. For the most expensive setting, fine-tuning GPT-2 with WikiText-2, computing ground-truth costs at most 500 hours of compute with the resources listed above. Similarly, running EKFAC-IF, EKFAC-SOURCE, ASTRA-IF and ASTRA-SOURCE for the GPT-2 setting the most expensive of all the settings. Running ASTRA-SOURCE for the GPT-2 setting costs at most 40 hours with the compute resources listed above.
Statistical Significance For LDS experiments, we provide error bars indicating standard error for all one-model predictions estimated with 10 seeds. For the training curve (Figure 3) and the eigendecomposition experiments (Figure 4), we show standard error estimated with seeds. The standard error is computed by taking the standard deviation of the computed metric, divided by V10.
Assets We use pytorch 2.5.0 and the publicly available kronfluence package for our experi- ments, which can be found at https: // thub. wtt.coooo/oo/cc woom/pomonam//roofffloo fuence....
G Implications of Truncated Neumann Series
In Section 2.1, we introduced the connection between NI and the iHVP approximation. Iterative iHVP approximation algorithms like LiSSA [23] usually requires a large number of iterations to achieve good iHVP approximations [7. 17]. In practice, the number of iterations in the algorithm is usually set with the assumption that the approximation converges within the iterations (e.g., 5000 in [17]). While a large number of iterations makes convergence likely, it is not a guarantee, and raising the number of iterations implies additional computational cost. We demonstrate that using fewer iterations and stopping before convergence can be interpreted as adding an additional damping term the iHVP approximation. There is existing work noting that the truncation of Neumann series can be viewed as increased damping [29], but here we present it with the derivation technique found in
35

---

[19]. We derive the following expression for the truncated Neumann series with J iterations:
J-1
- aG - -a1)² = (G XI) - (I (I- oG aad))))
(37)
==
(G + - (1 - exp( exp(-a - aJJAl) ~ (G An)
(38)
(39)
where = a Equation 37 and Equation 38 utilize the finite series and the matrix exponential definition, respectively. Given that the matrices in Equation 38 commute, we express it as the following matrix function:
F(o) 1 expp-aaaeee aJA)
(40)
0 0+
Let G = QDQ QDQ be the eigendecompossttinn where oj denotes the jth eigenvalue of G. The expression in Equation can be interpreted as applying the function F(a) to each eigenvalue a of G.
In high-curvature directions, this term asymptotically approaches 1/0 + x, whereas in low-curvature directions with small values of X, it tends toward a...
The qualitative behavior of the function F can be captured by Fin,, defined as:
Fin (o) a -
(41)
Applying Fi. to the Hessian results in approximating Equation 39 with a modified damping term Å :== Å + 1/J.. Hence, the truncated version can interpreted as incorporating a larger damping term, by 1/a.. The implicitly larger damping term affects the iterative updates and adds additional difficulty to tuning the hyperparameters of iterative iHVP solvers. In ASTRA, we leverage the EKFAC approximation a preconditioner for the iHVP approximation to improve the conditioning of the problem, which mitigates this issue.
H Limitations & Broader Impact
Limitations We have addressed the problem of computing accurate iHVPs for TDA. As outlined Section 2.2, in addition to the computing the iHVP, a large component of the cost in computing influence functions for both EKFAC-IF and ASTRA-IF when scanning the dataset for highly influential examples is taking the dot product with all the training example gradients, which an orthogonal but important issue that we do not address. We also noted our experiments that in some cases, the bottleneck for influence function performance in TDA is not an inaccurate iHVP, but other factors such violating the fundamental assumptions involved in the influence functions derivation. In these cases, the benefits from an improved iHVP approximation may not materialize until other bottlenecks are resolved. For example, the FashionMNIST experiments show a large increase in performance after ensembling relative to the 1-model scores obtained by an improved iHVP approximation, which suggests that stochasticity the training procedure may the dominant factor hindering TDA performance. Finally, for ESTRA-SOURCE we present a way in which we can apply the iHVP computation, but we leave to future work to explore various ways to compute the average GGN Gg which may further improve performance.
Broader impact Our work improves the accuracy of the iHVP approximation for use in TDA. The algorithmic improvements we present do not have direct societal impact. However, improved TDA can provide insights into the relation between training data and model behavior. From this perspective, our work has similar broader impact to other work in TDA, sharing similar potential benefits namely, enhanced interpretability, transparency, and fairness in AI systems and share similar risks. Specifically, advancing TDA can help us understand models through the lens of training data, which can be applied to many domains such as building more equitable and transparent machine learning systems [12, 13], and investigating questions intellectual property and copyright [11, 14, 15]. On the flip side, TDA can also be used to maliciously to craft data poisoning attacks and could result in models with undesirable behavior. It important that TDA algorithms are improved with mitigation of the risks.
36


Reviews:
[{"review_id": "Ty2Oa8qgEp", "reviewer": "Reviewer_Kz7b", "summary": "Influence functions and unrolled differentiation involve a computation that resembles an inverse Hessian-vector product (iHVP), which is difficult to numerically compute efficiently. In this work the authors propos an algorithm which uses an EKFAC-preconditioner with the standard Neumann series iterations to arrive at an accurate iHVP approximation. They numerical show that using their algorithm enhances the  performance of training data attribution methods.", "questions": "Please see my questions above,", "limitations": "N/A", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 4, "strengths_and_weaknesses": "- The main problem that I have with the paper is that I am not convinced that the main bottleneck in influence functions is the accuracy of computing the inverse Hessian. Although the authors show some empirical benefits of their proposed algorithm that improvs the accuracy of  Hessian vector products, I'm not sure how much benefit this method has in real problems. I think the the paper studies a rather limited problem.\n\n- The writing of this paper is not in an acceptable form. It is very hard to read the paper. The paper does not mathematically define many objects, even the actual problem definition is buried in the text. Acronyms are used without proper definition, etc. The writing needs to be significantly improved for the paper to be publishable. Even the proposed algorithm is not well presented. How is G_{EKFAC} calculated (precisely).\n\n- The idea of the proposed algorithm is very simple. The authors observe that  Stochastic Neumann Series Iterations is equivalent to GD applied to the quadratic objective (4). Then they propose to replace GD with E-KFAC that takes into account the curvature of loss function. It is not very surprising that this methods performs better than the baseline.\n\n- Does the modified NSI provably converge? \n\n- How does the performance change as the models are scaled? A \"scaling law\" experiment would be beneficial.\n\n- I am not very familiar with the prior work in this exact area, so I can not comment fully on the novelty of the proposed method for this problem.\n\n- The actual choice of the preconditioner P seems to be arbitrary. How does this choice affect the performance? Is the preconditioner that you choose exactly the E-KFAC preconditioner?\n\n- What happens if diagonal preconditioners are used instead of E-KFAC?\n\n- How sensitive is the results to the choice of learning rate? can the authors provide a plot for one of their experiments that shows the loss at a fixed # of iterations, for different values of learning rate?", "was_revised": true, "final_justification": "I thank the authors for their detailed response. The response, in particular the scaling experiment, resolves most of my concerns. I will increase my score to 4. I need to add two points:\n\nI still think the writing of the paper can be significantly improved. I encourage the authors to enhance the writing of the paper, in particular regarding the points I raised in the review.\n\nAs I mentioned in the review, I am not fully familiar with the literature of this subject, thus I could not comment confidently on the novelty of the paper. After reading other reviews, it seems that the paper is indeed novel, but I cannot confirm indepednently.", "rebuttal": "Thank you for your detailed review! We can generally categorize your comments into two buckets: 1) concerns about the writing style of the paper and clarification questions about certain definitions and 2) doubts whether our contribution helps advance TDA. If there are any places that remain unclear after our discussion, please let us know - we think any ambiguities can be cleared up with very modest modifications to the manuscript.\n\n> The main problem that I have with the paper is that I am not convinced that the main bottleneck in influence functions is the accuracy of computing the inverse Hessian.\n\nIn our experiments in Figure 2, we directly show that the iHVP is one of the bottlenecks -- across a range of settings, we find that improving the iHVP approximation in general leads to better Linear Datamodeling Score (LDS), holding everything else the same. This is especially true for convolution architectures.\n\nThe fact that a better iHVP in general improves TDA outcomes is not a unique observation by us. In the extreme case, consider approximating the Hessian with the identity: Bae et. al. 2024 [1] and Choe et. al. 2024 [2] show that methods involving the Hessian (e.g. EKFAC influence) outperforms methods which do not use the Hessian (e.g. TracIn [7], Graddot). \n\n> I'm not sure how much benefit this method has in real problems. I think the the paper studies a rather limited problem.\n\nThe core goal of TDA algorithms is to estimate how changes to the training dataset would affect the model’s predictions, without actually carrying out expensive retraining. Our evaluation metric LDS measures how well a TDA method correlates with ground-truth retraining, and is agnostic to the downstream task. This method of evaluation is common to many well-known TDA papers [1,2,3,4]. In this paper, we address the efficacy of TDA algorithm itself, rather than how this tool is applied to downstream tasks such as interpretability, data curation etc..\n\n> The writing of this paper is not in an acceptable form. It is very hard to read the paper.\n> The paper does not mathematically define many objects\n\nThank you for this feedback. At the outset, when writing this paper, we recognized it would be difficult to fully convey all the necessary mathematical background in the main body since this paper is at the intersection of many topics: influence functions, unrolled differentiation, EKFAC preconditioning and stochastic Neumann series iterations (SNI). Given the page limit, we had to tradeoff whether to state in the main paper certain preliminaries that have been repeated before in related works (e.g. [8]), or whether to defer some of the preliminaries to Appendix B, and only reference them in the main paper. We do believe all the mathematical objects we discuss have been defined, if one reads the main paper in conjunction with Appendix B. We are happy to refer to Appendix B in a more specific way in the main paper and to pull forward more details about EKFAC.\n\n> even the actual problem definition is buried in the text.\n\nWe are not quite sure what you mean by “actual problem definition”. If you mean that “the actual problem definition” is to compute an iHVP, lines 116-125 are dedicated to describing the LiSSA/SNI algorithm which computes the iHVP (line 125). If you mean “the actual problem definition” is to compute influence scores that correlate with ground-truth retraining, then you can find the evaluation metric LDS defined in Equation 6. Given the extra page for the camera ready, we are happy to improve the paragraph spacing and give certain equations their own lines.\n\n> Acronyms are used without proper definition, etc. \n\nNon-standard acronyms we use (e.g. SNI) are written in parentheses after the full term in the paper, with a full list of acronyms listed in Appendix A. Acronyms corresponding to algorithms used in TDA such as EKFAC and LiSSA are cited, and also its full name is listed in Appendix A. Could you let us know exactly where they are undefined? \n\n> Even the proposed algorithm is not well presented. How is $\\\\mathbf{G}\\_{\\\\text{EKFAC}}$ calculated (precisely).\n\nThe precise EKFAC Generalized Gauss Newton Hessian approximation $\\\\mathbf{G}\\_{\\\\text{EKFAC}}$ is laid out in Appendix B.5 (Lines 1045-1087), computed in the same way as previous works [1,8]. We have not made any changes to this fairly standard computation procedure.  Unfortunately, our allowed rebuttal is limited in length, so we cannot lay it out here. However, to make things clear, we will add a line in the body of the paper that $\\\\mathbf{G}_{\\\\text{EKFAC}}$ is the block diagonal matrix, whose blocks correspond to layers in the neural network, and whose entries are $\\\\hat{\\\\mathbf{F}}_l$ described in Appendix B5.\n\n> It is not very surprising that this methods performs better than the baseline.\n\nWe view this fact as a positive — our method is a principled one based on a large body of existing work and thus in principle is expected to work. Preconditioning gradient descent with EKFAC is approximate natural gradient descent [5, 6] where we approximate the Generalized Gauss Newton Hessian using EKFAC.\n\n> Does the modified NSI provably converge?\n\nOur algorithm is guaranteed to converge for a sufficiently small learning rate, an appropriate learning rate schedule, and the choice of $\\\\mathbf{P} := \\\\mathbf{G}_{\\\\text{EKFAC}}$ as is done in our paper.\n\nTo see this, note that our update rule in Equation 5 is the mini-batch gradient descent rule for the quadratic:\n\n$$\\\\frac{1}{2} \\\\boldsymbol{\\\\theta}^\\\\top (\\\\mathbf{P} + \\\\tilde{\\\\lambda} \\\\mathbf{I})^{-1}(\\\\mathbf{G}+ \\\\lambda \\\\mathbf{I})\\\\boldsymbol{\\\\theta} - \\\\boldsymbol{\\\\theta}^\\\\top (\\\\mathbf{P} + \\\\tilde{\\\\lambda} \\\\mathbf{I})^{-1}\\\\nabla_{\\\\boldsymbol{\\\\theta}} f_{z_q} (\\\\boldsymbol{\\\\theta}^s)$$\n\nand for our choice of $\\\\mathbf{P} := \\\\mathbf{G}_{\\\\text{EKFAC}}$, $(\\\\mathbf{P} + \\\\tilde{\\\\lambda} \\\\mathbf{I})^{-1}(\\\\mathbf{G}+ \\\\lambda \\\\mathbf{I})$ is positive definite. We thus inherit properties of mini-batch gradient descent on strongly convex quadratics.\n\n> How does the performance change as the models are scaled? A \"scaling law\" experiment would be beneficial.\n\nWe are at our character length, so for this question we point you to our response to scaling in our answer to reviewer ZrKT.\n\n> I am not very familiar with the prior work in this exact area, so I can not comment fully on the novelty of the proposed method for this problem.\n\nThank you for your candid remarks. The novelty here is that if we were to compute influence using EKFAC, we might as well use this as a free preconditioner if we want to then apply an iterative algorithm such as SNI. The benefit of SNI is that as we apply more compute, the iHVP tends to improve -- it is a consistent estimator of the iHVP. In TDA, this combination of EKFAC influence and iterative methods has not been explored.\n\n> The actual choice of the preconditioner $\\mathbf{P}$ seems to be arbitrary. How does this choice affect the performance? Is the preconditioner that you choose exactly the E-KFAC preconditioner? What happens if diagonal preconditioners are used instead of E-KFAC?\n\nYes. The preconditioner we choose for all our experiments is exactly the EKFAC preconditioner, stated on line 184 and line 198. While in principle one can choose any preconditioner $\\\\mathbf{P}$ to minimize the quadratic objective in Equation 4, it makes a lot of sense to choose the EKFAC preconditioner for the following reason: To compute influence, many would by default adopt EKFAC-IF since it is a state-of-the-art method, which requires computing the EKFAC decomposition. Our insight is that we could repurpose this EKFAC decomposition as a preconditioner for Stochastic Neumann Series Iterations (SNI) to get a more accurate iHVP. Rather than recomputing a preconditioner, we might as well use the already computed EKFAC decomposition at no additional cost.\n\nWe can consider what would happen if we use another preconditioner. First, note that by Equation 5, if we initialize the weights to zero, then after the first step, the solution obtained is the solution had we computed the iHVP using the (damped) preconditioner as the curvature matrix (we also explain this on line 189 onwards). Since EKFAC-based influence functions is a state-of-the-art TDA method, had we used another preconditioner (say a diagonal one), our iterative method would take many iterations just “catching up” to EKFAC-IF’s TDA performance.\n\n> How sensitive is the results to the choice of learning rate? can the authors provide a plot for one of their experiments that shows the loss at a fixed # of iterations, for different values of learning rate?\n\nWhile we cannot provide plots or links to external sites due to NeurIPS rules, we can comment on the general behavior. Fixing everything else (e.g. batch size to sample $\\\\tilde{\\\\mathbf{G}}_k$), since both SNI and our method provably converges for a small enough learning rate, in both cases, if one were to choose too small learning rate, more iterations would be needed to obtain an iHVP solution with the same loss as defined in Equation 4, and too large a learning rate causes divergence.\n\n```\nReferences:\n\n[1]: Bae, Juhan et al. “Training Data Attribution via Approximate Unrolled Differentiation”.\n[2]: Choe, Sang Keun et al. “What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions”.\n[3]: Park, Sung Min et al. “TRAK: Attributing Model Behavior at Scale\".\n[4]: Mlodozeniec, Bruno et. al. \"Influence Functions for Scalable Data Attribution in Diffusion Models\".\n[5]: George, Thomas et. al. \"Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis\".\n[6]: Amari, Shun-ichi et. al. \"Natural Gradient Works Efficiently in Learning\".\n[7]: Pruthi, Garima et. al. \"Estimating Training Data Influence by Tracing Gradient Descent\".\n[8]: Grosse, Roger et al. “Studying Large Language Model Generalization with Influence Functions”.\n```"}, {"review_id": "UJtwhwmBih", "reviewer": "Reviewer_CuEj", "summary": "The paper introduces ASTRA, a solver that pre-conditions Stochastic Neumann Iterations (SNI) with the cached eigenbasis from Eigenvalue-Corrected K-FAC (EKFAC). One EKFAC step supplies an inexpensive initial iHVP; a few hundred extra Neumann steps refine it toward the unbiased solution. On seven benchmarks—MLPs, ResNet-9, and GPT-2 fine-tuning—ASTRA consistently raises the Linear Data-modeling Score (LDS), e.g. CIFAR-10 jumps from 0.25→0.60, while adding <10 % runtime over EKFAC. Eigen-space analysis shows that rapid progress in low-curvature directions underlies these gains.", "questions": "- How does ASTRA scale when EKFAC eigenupdates dominate cost on LLMs? Can ASTRA refine low-rank curvature methods (e.g. LoGRA) to combine memory savings with higher fidelity?", "limitations": "None", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 2, "significance": 4, "originality": 4, "strengths_and_weaknesses": "**Strengths**\n- Novelty: Simple yet novel fusion of EKFAC and SNI that had not been explored for TDA.\n- Principled: Clear derivations (truncation ⇔ implicit damping) and principled pre-conditioning.\n- Convincing experiments: Seven datasets, timing curves, eigen-space breakdowns, and ensemble analysis.\n- Practical impact: Drop-in replacement for existing influence-function and unrolled-differentiation pipelines with minimal compute overhead.\n- Insightful analysis: Demonstrates that low-curvature directions dominate influence quality, guiding future work.\n\n**Weaknesses**\n- Dense writing: Readers unfamiliar with EKFAC or SNI may struggle; a brief background / preliminaries section would help.\n- Limited model scale: The largest test model is GPT-2 (124 M params); behavior on larger-scale vision and language models unclear.\n- Baseline coverage: No empirical comparison to recent TDA methods such as LoGRA, TRAK, or MAGIC.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your positive review. Based on your comments, we are very happy that you appreciate the subtleties in our work and we notably appreciate the time you put into reading our paper carefully. In particular, we are happy that you find the method “novel”,  the analysis “insightful”, the derivations “clear”, and the experiments “convincing”.  We now address your questions and concerns line by line.\n\n> Dense writing: Readers unfamiliar with EKFAC or SNI may struggle; a brief background / preliminaries section would help.\n\nThank you for this feedback. There are extended preliminaries which cover this background material in Appendix B. We are happy to move some of the preliminaries to the body to appeal to those without a background in EKFAC or SNI.\n\n> Limited model scale: The largest test model is GPT-2 (124 M params); behavior on larger-scale vision and language models unclear.\n> \n\nEvaluation for TDA is very expensive, since the very purpose of TDA algorithms is to substitute for repeated retraining, which is the ground-truth we measure against. Thus, the scale of our experiments is not bottlenecked by our TDA method, but rather by our ability to evaluate the outcome. In our main results, each setting required us to train a model over 1,000 times.\n\n**New Experiment 1:**\n\nHowever, we ran a larger experiment involving LDS, in which we trained a subset of CIFAR-100 on ResNet18. The LDS scores are:\n\nEKFAC IF 1 model, EKFAC IF 10 model: 0.143, 0.244\n\nASTRA IF 1 model, ASTRA IF 10 model: 0.232, 0.478\n\nAgain, our method ASTRA-IF provides a significant boost compared to EKFAC-IF.\n\n**New Experiment 2:**\n\nDespite not being able to estimate ground truth for large models to compute LDS, we also reran our method ASTRA-IF on a 1.5B parameter model GPT2-XL to see how the loss in Equation 4 behaves compared to influence computed with SNI. While we cannot paste the plots here due to NeurIPS rebuttal rules, we will describe the result:\n\nStarting from the same initialization, so that the loss at iteration 0 are the same, we did a learning rate sweep for both methods, increasing the learning rate until both methods diverged. We chose the best learning rates at a fixed iteration for each method and plotted the values of Equation 4 for both our method ASTRA-IF and SNI. The loss curve as defined in Equation 4 is substantially lower for our method compared to SNI for the same number of iterations: the best SNI result achieved a loss of approximately $-10^8$ , while EKFAC achieved $-5 \\\\times 10^8$ after 40 iterations. The loss of EKFAC preconditioning was lower than SNI throughout the entire 40 iterations despite starting out the same at iteration 0.\n\n> Baseline coverage: No empirical comparison to recent TDA methods such as LoGRA, TRAK, or MAGIC.\n\nWe have added comparisons to TracIn [2] and TRAK [3] to our existing experiments for completeness. Since we are not allowed to add plots, we describe the results here:\n\n**New Baseline 1**:\n\nWe ran TracIn for all settings. 1 model and 10 model ensemble TracIn scored below 0.3 LDS for all settings except FashionMNIST-N. TracIn's LDS scores are substantially lower than ASTRA-IF. \n\nWe list the 1 model and 10 model scores for TracIn as a tuple (1 model LDS, 10 model LDS), in the following order: MNIST, FashionMNIST, CIFAR-10, UCI Concrete, UCI Parkinsons, WikiText2, FashionMNIST-N.\n(0.154, 0.191), (0.201, 0.223), (0.107, 0.121), (0.191, 0.227), (0.151, 0.165), (0.1580, 0.1587), (0.577, 0.586).\n\nNote: TracIn uses the same checkpoints as ASTRA-SOURCE for comparability.\n\n**New Baseline 2**:\n\nWe also ran TRAK for all settings except GPT-2, since there does not exist a public implementation for language modeling tasks. Both 1 model and 10 model ensemble TRAK score below 0.4 for all settings we ran.\n\nWe list the 1 model and 10 model scores for TRAK as a tuple (1 model LDS, 10 model LDS), in the following order: MNIST, FashionMNIST, CIFAR-10, UCI Concrete, UCI Parkinsons, FashionMNIST-N.\n(0.132, 0.340), (0.102, 0.246), (0.145, 0.342), (0.218, 0.392), (0.073, 0.194), (0.087, 0.213).\n\nNote: TRAK uses the same projection dimensions as in [4] except CIFAR-10, in which we conduct a hyperparameter sweep among [128, 256, 512, 1024, 2048, 4096, 8192, 20480] and use the best dimension (512).\n\n\n**LoGRA discussion**\nWe would *expect* LoGRA to underperform EKFAC influence, since LoGRA computes EKFAC in lower dimensions either after 1) random projections or 2) projection to the PCA eigenbasis using a Kronecker structured projection matrix. In their paper [1], LoGRA underperforms EKFAC, with the benefit of reducing the compute and memory footprint.\n\n**MAGIC discussion**\nWe opted to not compare against MAGIC since its publication occurred well past the cutoff of March 1st, 2025 — the NeurIPS 2025 cutoff for contemperaneous work, and we were not able to find a public implementation to use. \n\n> How does ASTRA scale when EKFAC eigenupdates dominate cost on LLMs? Can ASTRA refine low-rank curvature methods (e.g. LoGRA) to combine memory savings with higher fidelity?\n\nIn Appendix B, lines 1100-1106, we lay out the main components of computing EKFAC for influence functions, which can give the reader a sense as to how EKFAC scales. But we think your idea here is great — in principle, we can do all our computation in a projected space, which may give us an even more favorable compute vs. accuracy tradeoff, and we think this is a promising direction for future work.\n\n```\nReferences:\n[1]: Choe, Sang Keun et al. “What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions”.\n[2]: Pruthi, Garima et. al. \"Estimating Training Data Influence by Tracing Gradient Descent\".\n[3]: Park, Sung Min et al. “TRAK: Attributing Model Behavior at Scale\".\n[4]: Bae, Juhan et al. “Training Data Attribution via Approximate Unrolled Differentiation”.\n```"}, {"review_id": "WWWdYUCN4l", "reviewer": "Reviewer_LxJj", "summary": "The paper introduces ASTRA, an algorithm that uses EKFAC-preconditioner on Neumann series iterations, to arrive at iHVP approximation for training data attribution more accurately and efficiently. On tabular, vision, and laguage tasks, ASTRA speeds up iHVP approximation and beats baseline methods on accuracy. The paper also provides an exploration on how curvature directions affect influence functions performance.", "questions": "How would the method scale with larger models (billion parameters)? If empirical analysis cannot be done due to compute constraint, is there estimate / small-scale scaling behavior experiment that can be done?", "limitations": "Yes.", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strength:\n- Extensive evaluation on a lot of tasks, modality, and architecture. The authors do not restrict themselves to one data modality.  They benchmark on small regression MLPs, classic image MLPs, a convolutional ResNet-9, and an autoregressive GPT-2 fine-tune. \n- Rigorous mathematical derivation. Section 3 exposes derivation step by step clearly.\n- Beside the method itself, the paper also provides a further exploration of low curvature impacts that underpins the paper's hypothesis. This analysis not only validates the method but also gives practitioners guidance on where to focus optimisation effort.\n\n\nWeakness:\n- Lacks discussion on how the method would scale with larger models. The introduction acknowledges that EKFAC can scale to billion-parameter language models, yet the largest empirical target is a 124 M-parameter GPT-2 fine-tune. Modern language and vision models are usually billion parameter levels. It's unclear how the method would scale to larger models.\n- For some evaluations, performance depends a lot on ensemble. This weakens the out-of-the-box appeal.", "was_revised": true, "final_justification": "Author's rebuttal addresses some of my concerns. I decide to maintain my score since points I raised are still valid.", "rebuttal": "Thank you for your positive review. We appreciate that you found our evaluation “extensive”, and our mathmatical derivation “rigorous”. In particular, we are glad that you found that Section 6 (which shows that low curvature directions really matter for TDA) useful for practitioners. We believe it is quite important that practitioners understand that directions of low curvature carry a lot of predictive information for TDA. We now address your questions and concerns line by line:\n\n> Lacks discussion on how the method would scale with larger models\n\nEvaluation for TDA is very expensive, since the very purpose of TDA algorithms is to substitute for repeated retraining, which is the ground truth we measure against. Thus, the scale of our experiments is not bottlenecked by our TDA method, but rather by our ability to compute ground truth. In our main results, each setting required us to train a model over 1,000 times to obtain a stable estimate of ground truth.\n\n**New Experiment 1:**\n\nHowever, we ran a larger experiment involving LDS, in which we trained CIFAR-100 on ResNet18. The LDS scores are:\n\nEKFAC IF 1 model, EKFAC IF 10 model: 0.143, 0.244\n\nASTRA IF 1 model, ASTRA IF 10 model: 0.232, 0.478\n\nAgain, our method ASTRA-IF provides a significant boost compared to EKFAC-IF.\n\n**New Experiment 2:**\nDespite not being able to estimate ground truth for large models to compute LDS, we reran our method ASTRA-IF on a 1.5B parameter model GPT2-XL to see how the loss in Equation 4 behaves compared to influence computed with SNI. While we cannot paste the plots here due to NeurIPS rebuttal rules, we will describe the result:\n\nStarting from the same initialization, so that the loss at iteration 0 are the same, we did a learning rate sweep for both methods, increasing the learning rate until both methods diverged. We chose the best learning rates at a fixed iteration for each method and plotted the values of Equation 4 for both our method ASTRA-IF and SNI. The loss curve as defined in Equation 4 is substantially lower for our method compared to SNI for the same number of iterations: the best SNI result achieved a loss of approximately $-10^8$ , while EKFAC achieved $-5 \\\\times 10^8$ after 40 iterations. The loss of EKFAC preconditioning was lower than SNI throughout the entire 40 iterations despite starting out the same at iteration 0.\n\nFinally, we can provide some color on the costs of our algorithm. Recall that our algorithm optimizes Equation 4 using preconditioned SNI. Compared with regular SNI, our method ASTRA requires an additional matrix vector product to compute the product of $(\\\\mathbf{P}+ \\\\lambda \\\\mathbf{I})^{-1}$ with $(\\\\tilde{\\\\mathbf{G}}_k+ \\\\lambda \\\\mathbf{I})\\\\boldsymbol{\\\\theta}_k$, whose computation costs slightly more than one backward pass. In terms of memory, compared with SNI, we need to additionally keep the EKFAC statistics in memory, and the exact memory cost of EKFAC depends on model architecture. We have a simplified illustration on line 113 in which the memory cost scales linearly with the dimension of the network. This is a large reduction compared to storing the Hessian naively, which scales quadratically with the dimension of the network.\n\n\n> For some evaluations, performance depends a lot on ensemble. This weakens the out-of-the-box appeal.\n\nFigure 2 compares 1 model performance, as well as 10 model ensembled performance. Relative to EKFAC-baselines, in both cases, we find improvements using our method."}, {"review_id": "CV69RDYlBf", "reviewer": "Reviewer_ZrKT", "summary": "This paper focuses on improving the accuracy and reliability of influence function-based training data attribution (TDA) methods, particularly in architectures involving convolutions, where existing EKFAC-based approaches tend to underperform. The authors propose ASTRA, a new algorithm that leverages EKFAC as a preconditioner for Stochastic Neumann Iteration (SNI) to compute inverse Hessian-vector products (iHVPs) more accurately and efficiently. By addressing the computational and numerical challenges in iHVP approximation, ASTRA enhances the quality of influence estimation.", "questions": "1. While ASTRA is empirically shown to be effective, does the method come with any theoretical guarantees or bounds on convergence or approximation error, especially compared to standard Neumann series iteration?\n\n2. EKFAC and second-order approximations can be computationally demanding. Could the authors comment on how ASTRA scales to large-scale architectures, such as modern transformers or deep convolutional networks? Are there specific memory or runtime trade-offs?\n\n3. The paper claims improved performance in convolutional settings. Are there particular architectural patterns (e.g., depth, residual connections, attention) where ASTRA struggles or underperforms? Additional results or qualitative insights would be helpful.", "limitations": "I do not see any potential negative societal impact.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n\n1. The paper proposes ASTRA, a well-motivated algorithm that uses EKFAC as a preconditioner for Stochastic Neumann Iteration (SNI), offering a more accurate and cost-effective way to approximate inverse Hessian-vector products (iHVPs), which are central to influence-function-based TDA.\n\n2. While prior work has questioned the practical value of influence functions due to poor iHVP approximations, this work shows that with improved estimation (via ASTRA), influence scores can regain predictive utility, even in settings that violate the original theoretical assumptions.\n\n3. The paper identifies and targets a key weakness in existing EKFAC-based methods—namely their poor performance on convolutional networks—and proposes a method that generalizes better to such architectures.\n\nWeaknesses:\n\n1. Although ASTRA is more efficient than full Hessian inversion, it still requires access to second-order statistics and matrix-vector products. The paper does not fully explore how well ASTRA scales to large-scale models (e.g. ResNet-152 on Imagenet), where such operations may be prohibitively expensive.", "was_revised": true, "final_justification": "After reading the rebuttal, I find that the authors have provided some reasonable responses, and they appear to have a clear understanding of both the strengths and the limitations of their method. While I am not deeply familiar with the full body of literature in this specific subfield and have not reviewed all the cited works, my evaluation is based on general research experience and standard expectations in the field. It should be viewed as a perspective from outside the core community, and I hope it offers a constructive and useful reference for the judgment of other reviewers.", "rebuttal": "Thank you for your positive review. We appreciate your comments about the algorithm being “well-motivated”. We are also glad you appreciate the unique improvement ASTRA provides for convolution architectures. We will address your questions and concerns below line by line:\n\n> it still requires access to second-order statistics and matrix-vector products\n\nThe reason for second order statistics being necessary is that the classic influence functions derivation contains the Hessian [1]. Past work has shown that inclusion of second-order information significantly improves efficacy of TDA — we can compare TDA outcomes measured via LDS for gradient dot products with methods that use second order information: (e.g. EKFAC influence and ASTRA influence), and in general we observe that better second order approximation results in a better TDA outcome. Both our results and past results in [2, 3, 4] show that this is the case. \n\nYou are correct in implying that incorporating second order information increases the computational burden. However, strong TDA methods such as EKFAC influence already computes second order information. In our paper, we show that at a limited additional cost, by repurposing the EKFAC approximation to the Hessian as a preconditioner and preconditioning the SNI update, we can achieve better TDA outcomes.\n\n> The paper does not fully explore how well ASTRA scales to large-scale models (e.g. ResNet-152 on Imagenet)\n\nEvaluation for TDA is very expensive, since the very purpose of TDA algorithms is to substitute for repeated retraining, which is the ground truth we measure against. Thus, the scale of our experiments is not bottlenecked by our TDA method, but rather by our ability to compute ground truth to compute LDS. In our results in Figure 2, each setting required us to train a model over 1,000 times to obtain a stable estimate of ground truth.\n\n**New Experiment 1:**\n\nHowever, we ran a larger experiment involving LDS, in which we trained a subset of CIFAR-100 on ResNet18. The LDS scores are:\n\nEKFAC IF 1 model, EKFAC IF 10 model: 0.143, 0.244\n\nASTRA IF 1 model, ASTRA IF 10 model: 0.232, 0.478\n\nAgain, our method ASTRA-IF provides a significant boost compared to EKFAC-IF.\n\n**New Experiment 2:**\n\nDespite not being able to estimate ground truth for large models to compute LDS, we reran our method ASTRA-IF on a 1.5B parameter model GPT2-XL to see how the loss in Equation 4 behaves compared to influence computed with SNI. While we cannot paste the plots here due to NeurIPS rebuttal rules, we will describe the result:\n\nStarting from the same initialization, so that the loss at iteration 0 are the same for both methods, we did a learning rate sweep for both methods, increasing the learning rate until both methods diverged. We chose the best learning rates at a fixed iteration for each method and plotted the values of Equation 4 for both our method ASTRA-IF and SNI. The loss curve as defined in Equation 4 is substantially lower for our method compared to SNI for the same number of iterations: the best SNI result achieved a loss of approximately $-10^8$ , while EKFAC achieved $-5 \\\\times 10^8$ after 40 iterations. The loss of our method was lower than SNI throughout the entire 40 iterations despite starting out the same at iteration 0.\n\n> does the method come with any theoretical guarantees or bounds on convergence or approximation error\n\nOur algorithm is guaranteed to converge for a sufficiently small learning rate, an appropriate learning rate schedule, and the choice of $\\\\mathbf{P} := \\\\mathbf{G}_{\\\\text{EKFAC}}$ as is done in our paper.\n\nTo see this, note that our update rule in Equation 5 is the mini-batch gradient descent rule for the quadratic:\n\n$$\\\\frac{1}{2} \\\\boldsymbol{\\\\theta}^\\\\top (\\\\mathbf{P} + \\\\tilde{\\\\lambda} \\\\mathbf{I})^{-1}(\\\\mathbf{G}+ \\\\lambda \\\\mathbf{I})\\\\boldsymbol{\\\\theta} - \\\\boldsymbol{\\\\theta}^\\\\top (\\\\mathbf{P} + \\\\tilde{\\\\lambda} \\\\mathbf{I})^{-1}\\\\nabla_{\\\\boldsymbol{\\\\theta}} f_{z_q} (\\\\boldsymbol{\\\\theta}^s)$$\n\nand for our choice of $\\\\mathbf{P} := \\\\mathbf{G}_{\\\\text{EKFAC}}$, $(\\\\mathbf{P} + \\\\tilde{\\\\lambda} \\\\mathbf{I})^{-1}(\\\\mathbf{G}+\\\\lambda \\\\mathbf{I})$ is positive definite. We thus inherit properties of mini-batch gradient descent on strongly convex quadratics.\n\nHowever, we cannot make any theoretical claims about the converge rate, since the condition number after preconditioning depends on the quality of the approximation $\\\\mathbf{G}_{\\\\text{EKFAC}}$, which does not carry theoretical guarantees, but is empirically known to work well [5, 6].\n\n> Could the authors comment on how ASTRA scales to large-scale architectures, such as modern transformers or deep convolutional networks. Are there specific memory or runtime trade-offs?\n\nCompared with SNI, our method ASTRA requires an additional matrix vector product to compute the product of $(\\\\mathbf{P}+ \\\\lambda \\\\mathbf{I})^{-1}$ with $(\\\\tilde{\\\\mathbf{G}}_k+ \\\\lambda \\\\mathbf{I})\\\\boldsymbol{\\\\theta}_k$, whose computation costs slightly more than one backward pass. In terms of memory, compared with SNI, we need to additionally keep the EKFAC statistics in memory, and the exact memory cost of EKFAC depends on the model architecture. We have a simplified illustration on line 113 in which the memory cost scales linearly with the dimension of the network. This is a large reduction compared to storing the Hessian naively, which scales quadratically with the dimension of the network.\n\n> Are there particular architectural patterns (e.g., depth, residual connections, attention) where ASTRA struggles or underperforms?\n\nFor influence functions, typically a decision is made as to which layers/modules we compute influence over. In our paper, we compute influence only over fully connected and convolution layers. For LLMs, we only consider influence of the fully connected layers, in line with [4]. We explain this in lines 1244-1247 in the Appendix. The experiments involving GPT2 and ResNet include a residual connection. However, we do not observe any reliable pattern with respect to depth and residual connections. \n\n\n```\nReferences\n[1]: Koh, Pang Wei et al. “Understanding Black-box Predictions via Influence Functions”.\n[2]: Bae, Juhan et al. “Training Data Attribution via Approximate Unrolled Differentiation”.\n[3]: Choe, Sang Keun et al. “What is Your Data Worth to GPT? LLM-Scale Data Valuation with Influence Functions”.\n[4]: Grosse, Roger et al. “Studying Large Language Model Generalization with Influence Functions”.\n[5]: Martens, James et al. “Optimizing Neural Networks with Kronecker-factored Approximate Curvature”.\n[6]: George, Thomas et. al. \"Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis\".\n```"}]

Author rebuttal:

Meta-review: All four reviews leaned towards acceptance with one accept recommendation [CuEj] and three borderline accept recommendations [Kz7b,LxJj,ZrKT].

Reviewers appreciated several aspects of the work:

+ The algorithm was considered well motivated [ZrKT]
+ The method was considered simple but novel [CuEj]
+ Showing that influence scores can regain predictive utility was appreciated [ZrKT]
+ The method was considered to have practical impactt as a drop-in replacement [CuEj]
+ Improving generalisation of EKFAC-based methods to convolutional network architectures was appreciated [ZrKT]
+ Exploration of the impact of low curvature was appreciated [LxJj] and considered insightful [CuEj]
+ The method was considered principled with clear derivation and principled preconditioning [CuEj]; similarly, the mathematical derivation was considered rigorous [LxJj]
+ The evaluations were considered extensive [LxJj] and convincing [CuEj]
+ The accuracy improvement and cost-effectiveness were appreciated [ZrKT]


However, several concerns were raised, which the authors responded to at the rebuttal stage:

- Requiring second-order statistics and matrix-vector products was criticised and scalability to large-scale models was questionsd [ZrKT]; authors argued for necessity of second-order statistics and that strong TDA methods already compute second order information. 
- One reviewer considered the studied problem limited [Kz7b]; authors argued they are addressing efficacy of the TDA algorithm itself.
- One reviewer was not convinced about what the main bottleneck in influence functions is [Kz7b]; authors discussed that their experiments show improving inverse Hessian approximation improves the TDA outcome.
- One reviewer considered the idea simple [Kz7b] and not surprising that it outperforms the baseline [Kz7b]; authors argued this to be a positive.
- Scalability to large models was considered unclear [LxJj,CuEj] and the largest model in experiments was considered small compared to scalability of EKFAC [LxJj]. Discussion on scalability to large architectures was desired [ZrKT,LxJj] and similarly, a "scaling law" experiment was desired [Kz7b]. Authors ran two larger experiments and commented about specific computation and memory costs of the steps.
- A further scaling question involved when eigenupdates dominate the cost; authors responded with some discussion.
- Lack of comparison to recent TDA methods was criticised [CuEj]; authors provided a comparison to two more methods and brief discussion of two others.
- Availability of theoretical guarantees/bounds on convergence or approximation error were desired [ZrKT,Kz7b]; authors discussed a convergence guarantee but noted they cannot guarantee the rate.
- Discussion of the impact of the preconditioner choice was desired, including if diagonal preconditions are used [Kz7b]; authors argued using EKFAC preconditioned to be natural, and argued using a different one would reduce performance.
- Discussion of sensitivity to learning rate was desired [Kz7b]; authors commented overly small learning rate would delay convergence and large would cause divergence.
- Results and discussion on settings where the method underperforms were requested [ZrKT]; authors provided some discussion stating they did not observe a reliable pattern.
- Dependence of performance on ensemble in some evaluations was criticised [LxJj]; authors argued they get improvements without ensemble too.
- The writing was considered dense [CuEj] and another reviewer considered the paper very hard to read and lacking definitions [Kz7b]; authors aimed to do some reorganisation of preliminary material and improve visibility of some equations and clarify some definitions.

After the rebuttals, [ZrKT] seemed to be mostly satisfied, [Kz7b] felt most of their concerns were resolved, whereas [LxJj] felt some of their concerns were addressed.

Overall, it seems if the authors incorporate well the planned revisions, and especially pay attention to readability which may still be a concern, the work may ultimately be at a sufficiently good state to be presented at NeurIPS.
Decision: Accept (poster)