Distributional Training Data Attribution: What do Influence Functions Sample?
Authors: Bruno Kacper Mlodozeniec
Isaac Reid
Samuel Power
David Krueger
Murat A Erdogdu
Richard E. Turner
Roger Baker Grosse
Venue: NeurIPS 2025 spotlight ()
Keywords: influence functions
data attribution
training data attribution
unrolled differentiation
deep learning theory
Primary area: theory
TL;DR: This paper introduces distributional training data attribution, a data attribution framework that accounts for stochasticity in deep learning training, enabling a mathematical justification for why influence functions work in this setting.

Abstract: Randomness is an unavoidable part of training deep learning models, yet something that traditional training data attribution algorithms fail to rigorously account for. They ignore the fact that, due to stochasticity in the initialisation and batching, training on the same dataset can yield different models. In this paper, we address this shortcoming through introducing _distributional_ training data attribution (d-TDA), the goal of which is to predict how the distribution of model outputs (over training runs) depends upon the dataset. Intriguingly, we find that _influence functions_ (IFs), a popular data attribution tool, are 'secretly distributional': they emerge from our framework as the limit to unrolled differentiation, without requiring restrictive convexity assumptions. This provides a new perspective on the effectiveness of IFs in deep learning. We demonstrate the practical utility of d-TDA in experiments, including improving data pruning for vision transformers and identifying influential examples with diffusion models.

Distributional Training Data Attribution: What do Influence Functions Sample?
Bruno Mlodozeniec 112
Isaac Reid
Sam Power David S. Krueger*
Murat A. Erdogdu+se
Richard E. Turner Roger B. Grosse
'University of Cambridge
*Max Planck Institute for Intelligent Systems "Mila Quebec Al Institute UUniversity of Toronto
University of Bristol
'Vector Institute
"Alan Turing Institute
bkm28@cam.ace aaticammacc.u. uk
ir337@cam.ac. ac.
Abstract
Randomness is unavoidable part of training deep learning models, yet some- thing that traditional training data attribution algorithms fail to rigorously account for. They ignore the fact that, due stochasticity in the initialisation and batching, training on the same dataset can yield different models. In this paper, we address this shortcoming through introducing distributional training data attribution (d- TDA), the goal of which is to predict how the distribution of model outputs (over training runs) depends upon the dataset. Intriguingly, we find that influence functions (IFs), a popular data attribution tool, are 'secretly distributional:: they emerge from our framework as the limit to unrolled differentiation, without requiring restrictive convexity assumptions. This provides a new perspective on the effectiveness of IFs deep learning. We demonstrate the practical utility of d-TDA in experiments, including improving data pruning for vision transformers and identifying influential examples with diffusion models.
Introduction
Training data attribution (TDA) techniques are fundamental interest machine learning, shedding light on the relationship between a model's properties and its training data. TDA is typically framed as a counterfactual prediction problem: estimating how a model's behaviour would change upon removal of particular examples from the training dataset [1, 2]. This invites the concept of influence, Training examples are deemed 'influential' if the model's behaviour would change significantly upon their exclusion. The practical utility of TDA has been demonstrated in applications including interpreting, debugging and improving models [2. 3], dataset curation [4], and data valuation [1. 5].
Influence Functions. It is typically prohibitively expensive to compute influence by retraining with different datapoints removed. This has motivated number of TDA methods designed approx- imate influence, but without actually retraining. Amongst such TDA methods, a leading example influence functions (IFs) [2, 6). This classical technique from robust statistics function theorem to estimate the optimal model parameters' sensitivity to downweighting a training datapoint. IFs have been deployed to investigate the generalisation patterns of 52 billion parameter large language models [3], and for data attribution of difffusion models [7]. Separately, researchers have proposed an alternative TDA method called unrolled differentiation [8, 9, 10]. Here, one differ-
the implicit
Equal contribution first authors. Order decided by can swim furthest underwater. Shared senior authors.
39th Conference on Neural Information Processing Systems (NeurIPS 2025).

---

"Standard" influence functions
"Distributional" influence functions
Training loss: (()
Perturbed training loss: L() -(4))))a(0)
Training Losses
Initialisation distribution
(c) Training Trajectories
50
100
(d) Final Models
Predict change in final distribution
Predict shift
in optimum
Figure 1: Distributional training data attribution. Classical data attribution methods like influence functions typically motivated using convex loss functions, predicting deterministic shift the unique optimal model weights (left). In contrast, this paper advocates for a distributional perspective, approximating the probability distribution over model parameters/outputs after removal of training examples (interpreted as perturbation of the training loss). This includes for non-convex loss functions (right).
entiates through a particular training trajectory directly obtain the sensitivity of the final model parameters to the weighting of a particular example in the loss function. Unrolled differentiation tends work better than IFs experiments, but more expensive compute.
Randomness in training. The success of IFs in deep learning perhaps surprising because the classical foundations of both TDA and IFs fail to account for core property of modern training: stochasticity [11]. Given the randomness inherent weight initialisation and mini- mini-batchinng, training can be understood as sampling from distribution over final models. Each training run corresponds to drawing single sample from this distribution. Yet classical TDA is only defined for deterministic training algorithms, and IFs are primarily understood for convex objectives (or by finding convex proxies [11]). Stochasticity is usually dismissed as a nuisance for TDA methods, glossed over in method derivations [2]. At best, it is sometimes heuristically managed by ensembling averaging [12]. In practice, stochasticity makes it difficult to diagnose which changes to model behaviour are attributable changes in the training dataset, and which are due sampling randomness.
Introducing distributional training data attribution. In this paper, we argue that the randomness in model training is not nuisance. Conversely, it ought to play central role in our understanding of influence, and deserves proper mathematical treatment. Viewing training as sampling from distribution over final model weights (or outputs), the goal of TDA should be to efficiently predict changes to this distribution under modifications the training dataset: a novel perspective that we coin distributional training data attribution (d-TDA). Figure provides visual schematic.
Influence functions are distributional. We show that unrolled differentiation is natively a d-TDA method (Section 3.1). Subsequently, in Section 4.1, we rigorously show that IFs approximate unrolled differentiation for long enough training times, and hence IFs are already inherently distributional.
Core contributions. (1) We introduce distributional training data attribution (d-TDA), a framework for studying data attribution in stochastic deep learning settings (Section 3). (2) We show that influence functions (IFs) are 'secretly distributional', solving special limiting cases of a d-TDA task (Section 4). This may help explain the effectiveness setting in which they were originally proposed. (3) We propose distributional inffuence, which quan- affects the distribution over
IFs in deep learning, far from the convex
tifies the importance of examples by how much their inclusion/exclun model weights and outputs. We show that distributional influence captures interesting information missing from its regular predecessor, and leads to more effective data pruning (Section 5).
2

---

2 Background
'Classical' Training Data Attribution (TDA). (TDA). Consider the space D == N=1 2N of possible finite training datasets D == (=)) In classical TDA, one is concerned with deterministic training algo- rithms 0- D + pare,, which take a dataset as their input and return 'trained' model parameters 0 (2). The goal of TDA is to predict how the output of the training algorithm 0 would change if were run using a perturbed training dataset D', with some examples removed. Concretely, given some trained model 0 (D), TDA methods (2' aim to approximate 0 (2") z (2') without actually retraining the model. Of course, in practice, one is typically interested in the change in some measurement function
R + R& when the dataset is modified for instance, the loss on a particular test example. Therefore, TDA methods (.) are typically evaluated on their ability to approximate m m00 (2)) ss m (D)).
'Classical' influence. The discussion above invites the concept of inftuence. The influence of an example is the change in the measurement m o 0" when the example is removed from the training dataset. Influential samples change the measurement by large amount. The influence of training datapoint z with respect to a measurement function m is given by:
In,) == (2)) 10-mmm0mmm (
(1)
This is extended to groups of examples (z)) k=1 C in the obvious way. To approximate Inf Inf(2x)
without actually retraining, one uses a TDA method to approximate 0" (2) 2x).
Response. A practical difficulty posed by the formulation of influence in Eq. (I) that D, the domain of the training algorithm (-), is discontinuous. Datapoints 24 are either included or not included. The binary nature of this choice makes it difficult to analyse Inf In,) directly using muddeett-based methods. Hence, it is typical to instead consider a continuous relaxation to the training algorithm.
Let us introduce a scalar e E [0, } which controls the weighting of a particular example in the training algorithm. Suppose E e= corresponds to inclusion and = n' corresponds to exclusion, with intermediate values meaning the example is still present but downweighted. The precise setup will depend on the training algorithm of interest. Let 0D =2 (s) denote the (assumed deterministic) outcome of the training algorithm with loss 12) (E). Provided 0% -2) (e) is continuous and twice-differentiable at c=0, 0, we have that
d0 -0 (s)
OD DD (e) = (2) +
+ Oe+2)
as s - O.
(2)
de
Response r(s,)
d0 Dii (s)
Hence, we define the response r(zg)) =
de le-0> such that m =x)) = (00(()) +
eVm r(z)) + O(E2). Intuitively, response measures the sensitivity of the training algorithm output with respect the weighting = of the example 2g E D. To make this more explicit, we will now give two concrete examples: influence functions and unrolled differentiation.
1. Influence functions. Many classical algorithms only depend on the data through a loss function 0)== Σm= (0) with l, Ru - R some per-example loss. One natural way codify downweighting in that case define an interpolated loss LD- the loss function Lo has a single unique minimum and that the training algorithm successfully locates it. Mathematically, this can be written as (*(2) = argmin Ln)()) Minimising the
(e):= (+):
- g. Suppose that
interpolated loss ED=Di,, (e) and applying the implicit function theorem, it is straightforward to prove that in this special case:
r ((x) = V (o (V (O+ =: TIP
(3)
We derive this result in detail in Section B. TIF is referred to as influence function (IF) a popular TDA tool. The effectiveness of IFs for deep learning is perhaps surprising given the unrealistic assumptions made during their derivation.
2. Unrolled differentiation. Suppose instead that the model is trained using stochastic gradient descent (SGD). Consider the weight update rule 0 0 = O, B -n 8% Ve, (V, (0,), where (0,) EE
3

---

denotes the trajectory of model parameters and 0, is some random initialisation. 8 with t E N are
independently and identically distributed batching variables {0, 13., with mean E[ N. In close analogy to the interpolated loss LD-D) (e), consider the interpolated update step:
N - =in=k)
(2)1 0, Th
(,
8 Ve,
(4)
where 1n= is the indicator function. If we train for T timesteps, one can directly differentiate through the training trajectory to obtain the sensitivity of the final model weights 0r with respect to the weighting Applying the chain rule, one obtains rather cumbersome expression (Eq. (57) in Section D). In this setting, we call r'm = de d0rl the unrolled differentiation response. T'n can be used classical TDA method if we consider all sources of randomness be fixed. This algorithm tends to work better than IFs in experiments, but the repeated computation and caching of Hessians makes its naive implementation expensive for long training runs. This has prompted work on approximate unrolled differentiation [9].
3 Distributional Training Data Attribution
obvious problem with the classical TDA formulation described Section is that reality training is stochastic: the randomness in model initialisation and SGD precludes defining a determin- istic map 0 2 - Re Even retraining with an identical dataset will in general give a different model; 0* (2) is better thought of as random variable. Previous work has dealt with this randomness heuristically by averaging over training ensembles [2, 9]. In contrast, in this paper we advocate for more rigorous distributional perspective. Taking the model initialisation 6, and the batch selections (5*) EeN to be random variables on some probability space (9, F, P), we frame distributional training data attribution follows.
Distributional training data attribution (d-TDA), Let (2) be the outcome of (stochastic) training with some dataset D D. Let (A)(A) == ple (D) E A] (for A E F) denote its proba- Ru +
bility distribution. Let m MoD be the distribution of some measurement function
Rd- of the trained model. The goal of distributional TDA is to reason about the behaviour of HD and m uD with respect changing especially, removing examples by taking D
Zz--
Rather than considering randomness be a nuisance, d-TDA acknowledges that the training dataset determines the distribution over trained models. Effective d-TDA methods answer questions like:
1. Given samples from Plp, how can approximately sample from M'Di, ?
2. If removed from the training dataset, which example 2g 2) would most drastically change pp?
3. Which examples should remove to change the variance of my upon retraining?
The fact that d-TDA predicts changes in distributions over measurements leads us to reevaluate the notion of influence. In particular, removing influential samples ought to substantially modify Mz With this in mind, we define distributional influence (c.f. Eq. (3)) as follows:
Definition 1. (Distributional influence). The distributional influence of a training example 2g D with respect to a measurement function m R + R& is given by:
DistInf (2x) A My I )
(5)
where A(H (2) is some 'difference function' between μ1 and U2.
There exist many possible instantiations of distributional influence, depending on the choice of A.
Letting X Mn,Y M2 denote the final measurement random variables, one could consider:
Mean influence
Variance increase influence
Wasserstein infinence
A(4 B(X) - B(Y)
Var(Y) Var(X)
3.1 Distributional influence with unrolled differentiation
To compute DistInf (2+), we need to (approximately) sample from PD,,, without retraining the model. This can be achieved using unrolled differentiation, described by the pseudocode below.
This can be roughly thought of as SGD updates with the interpolated loss function Lo- (c).
4

---

Alg. 1. Unrolled differentiation for d-TDA.
1. Sample 0 (D) Or from H4D by training the model with stochastic updates (Eq. (4)).
2. Obtain approximate samples from ap,,, without retraining by taking 0 (2) \ ( = (2) +
TTw. with n uup - I le=o the unrolled differentiation response. If interested in the distribution over some measurement, compute m(0"(2) =x)) (2) # Vm (F)(0(2)) (2)) TwD
3. Using these two sets of (correlated) samples, compute the difference function A between the empirical distributions to efficiently approximate DistInf(2;)
When computing Tun. the following observation simplifies differentiating through long training trajectories.
Remark (Unrolled differentiation is Markoy chain). Applying the chain rule of differentiation de,
to Eq. (4) gives the following recursive formula for (0,, re) = (0c, de ls-
0--1 re+1
(11 Σ
0
1h
N B
8 (
(6)
2 B
54
2 (8)) + 6 Ve (0,)
Intuitively, Eq. (6) shows that the responseT:rt depends on the response at the previous timestep Te modulated by the loss function curvature. If datapoint 2N is present in the batch sampled at timestep r+t1 also depends on the corresponding loss gradient Vl, (0.). Practically, Eq. (6) permits to compute the final response Fy at linear time and constant space complexity with respect to training duration, without caching or explicitly computing the batch Hessians (c.f. Eq. (57)). This akin to forward-mode automatic differentiation for meta-learning [13]. To the best of our knowledge, this the first application of such techniques to efficient computation of the response. Crucially, the batch selection 8" is i.i.d., Eq. (6) defines Markov Chain observation that unlocks well-studie mathematical machinery and invites us to analyse its limiting distribution (see Section 4.
Empirical demonstration. Figure 2 showcases the application of unrolled differentiation as d- TDA method, successfully predicting changes in the distribution of measurements when select subset of the dataset removed.
Samples
Means
Measurement distribution on
0.5
Prodicted with full dataset meaaurement
Sample Wy 0.00 Wy 0:22
Predicted influence
Predicted
0.0
Sample
Wa 0.24
-0.5
Query sample
Sample 0.09 Wy 0.08
-1.0
dRtee
Subeet
Full dataret
eRte 5 oñte
prodicted
Measurements by models trained on subset
Figure 2: d-TDA demo for a neural network trained UCI Concrete. d-TDA (using unrolled differentiation) gives approximate samples from the distribution of models re-trained on some fixed subset D' C D (left). Actual samples from Hny (obtained by expensive retraining) are closer to predicted samples from JaDy (obtained by efficient d-TDA methods) than they are samples from the original model μz, both in terms of their means (centre) and Wasserstein distance (right). The distributions
measurements query samples for different stochastic training runs.
Why not use regular TDA with a fixed seed? A natural question raised by the challenge of stochasticity is: instead of treating the outcome of training random variable, why not simply fix all sources randomness? Why not just stratify by the random choices like initialisation and data ordering? Naively, this seems to recover deterministic training algorithmm, to which one may apply regular (non- tnon-distributional) TDA methods. We refer to this as fixed-seed TDA'.
5

---

Distributional TDA is often preferable to fixed-seed TDA because many methods, like IFs, more accurately perform the d-TDA task, even when failing at the fixed-seed one. For instance, IFs and unrolled differentiation find local perturbations around the current optimum to predict outcomes of counterfactual retraining. However, even fixing all randomness, chaotic training dynamics can push training into completely different region of the parameter space. Moreover, when using a fixed batch-size, even tiny changes to the training set size can offset at what iteration each datum appears. This means even fixed-seed trajectories can converge to widely different 'optima" under small dataset perturbations, rendering the shift in the original local optimum inadequate. Later Section 5, we demonstrate IFs perform better downstream tasks distributional TDA method compared to as a fixed-seed TDA method. This suggests the distributional perspective provides a more accurate picture of how influence functions work.
What do influence functions sample?
In Section 3, we introduced distributional TDA, adopting rigorous mathematical perspective that accounts for stochasticity training. Here, we demonstrate how d-TDA relates to classical influence functions (IFs; Eq. (3)). From two complementary perspectives, we find that IF's are actually 'secretly distributional' appearing as asymptotic samples in specific d-TDA settings. stark contrast to usual derivations of IFs, which rely assumptions that are unrealistic for deep learning [2, 14), we place only mild constraints on (()). All proofs are in Section A.
4.1 Perspective 1: unrolled differentiation converges a.s. to influence functions
Adopting a distributional perspective, a natural question is: what is the limiting distribution of the random variable ry), updated according to Eq. (6)? Begin by considering the model weights (0,), which are updated by SGD with i.i.d. batch selection. We make the following assumptions.
A1 V L and Vc, are Lipschitz continuous and bounded. 2
A2. The step sizes (mt) z- are positive scalars satisfying Σ, T) = 0o and Σ, π < cc.
A3. The iterates of Eq. (12) remain bounded a.s., i.c. sup,ll(0,,))) r.)ll < 00 a.s.
Standard results due to e.g. H. J. Kushner and G. G. Yin [15] give us the following result:
Theorem (SGD converges stationary points [15, Theorem 2.1, Chapter 5D. Provided
assumptions A1-A3 hold, the sequence of SGD iterates (0,) -- as defined by Eq. (6) converges almost surely to the set Sc of stationary points of the corresponding ODE: 0 = - (0)) namely, 82 = {0 - (0)) = 0].
Theorem demonstrates that, with suitably decaying learning rate, SGD converges to critical points namely, saddle points or local minima. Define the set local minima as follows:
9T ( (() 0, 2 (()) = 0} C 8.
If the weights converge to a saddle point in Sz \ 8% it is intuitive that the response r, == de c- will diverge. This is because the final model parameters will become sensitive to any infinitesimal perturbation of the loss function.2 Conversely, if the weights converge to local minimum, the limiting behaviour of r, becomes tractable. Consider the following additional assumptions.
d0g
A4. Ve (0) E Span (()) for all 87.
AS. The nonzero eigenvalues of V2 ((0) for 0 E S% are uniformly bounded away from
A6. There exists some compact neighborhood (( around 8 such that gradient flow trajectories
((t) initialised therein converge uniformly over initialisations to points in ST. Moreover, their
lengths are bounded a.s., that: sup)(o)) y ((s)) lim 0(s ds' 6o.
Theorem 2. (Unrolled differentiation converges to IFs). Suppose that A1-A6 hold, and consider SGD trajectory in the set that converges to 8th (c.f. 8. \ 8T)). The sequence of iterates ((0,,) r, )), generated by Eq. (6) converges almost surely to the set R { rir((") + rus (0*) 0 E 8% TIF ()) = V L(0)+ (0), Txy (f) E (ull(v²<(())) that is, pointwise IFs, plus compo- nent in the Hessian nullspace. (-)* is the pseudoinverse.
7This could be interpreted limitation of the conventional notions of response and influence.

---

sketch. sketch. We consider the ODE which is the continuous time relaxation of Eq. (6). We prove that the solution to this ODE r(t) converges to an influence function, plus a component in the flat directions of a minimum manifold. Under the learning rate assumptions above, the SGD updates asymptotically track this ODE, which allows us to prove the final result.
Commentary on Theorem 2. The response iterate T either diverges (in the case that the weights 0, converge to saddle), converges to R" Hence, at late times, one approximately sample from M2DA, by sampling from (uD and offsetting by N'IF. Using "'I instead of T' means that 1) we assume we have trained for long enough, and 2) we neglect components of response in the Hessian nullspace. The latter may not converge and will general depend on the history of SGD iterates Q..
Assumptions. A1-A3 are standard assumptions, needed to ensure that SGD converges. A4 guaran- tees the perturbation l doesn't 't have component in the flat directions of the minimum manifold, or else unrolled differentiation diverges. Note that A4 automatically holds for any symmetries shared by L and 4. c.g. due to neural network parameterisation. AS ensures that IFs remain bounded; Hessian eigenvalues on 8 can be zero, but not nonzero and arbitrarily small. The most technically meaningful assumption is A6, which assumes gradient flow converges sufficiently fast local minima. We stress that it is much less restrictive than the requirements usually cited for IFs to apply, such strong convexity [2, 9, 14].
Remark 2. For Generalised Linear Models (GLMs), the component of the unrolled response r, in the nullspace of the Hessian is throughout training. Hence, in this setting, Theorem 2 gives exact convergence of the unrolled response to the influence functions formula.
Empirical validation. In Figure 3, we test our theoretical results on a regression task with UCI Concrete. As predicted by Theorem 2, the strength of correlation becomes very high (90%) at late times as the unrolled differentiation Markov Chain converges to IFs. As expected, the correlation is better for lower step sizes, for which the SGD iterates (normalised by learning rate) track gradient flow more closely.
0.9
Prachion of mrraerrement
Learning Vate
gradivet in the span
0.8
0.01
0.005
of the Ulssson
0.001
0.0000
0.7
0.0001
between unrolled
Pourson corvdation
0,0895
iiltwaree
0.9990
We also experimentally confirm that the compo- nent of the unrolled response т₁ in the null space of the Hessian that is, the error term that IFs cannot capture is insignificant for the practical tasks we test. In the lower panel of Figure 3, we see that the measurement gradient Vm only has tiny component in the null- space of the Hessian, living almost entirely the column space. Recalling that m (D \ 2x)) z m ( (2)) + N Vm Tup- this means it barely contributes to predicted changes in m. As such, the fact that IFs do not capture this part of the limiting distribution of r, does not appear to be of substantial concern for downstream tasks.
Training time (iterations x learning rate)
Figure 3: Validating Theorem Top: Correla- tion between changes in measurement predicted by unrolled differentiation and changes predicted by IFs. plotted against training time. The co- efficient becomes high as the Markov chain converges. The correlation is stronger for small ?) where SGD is closer to gradient flow [15]. Bottom: Norm of the measurement gradient component in the span of the Hessian divided by norm of the measurement gradient. The null- space component of Vm remains tiny.
4.2 Perspective 2: transport maps between Boltzmann distributions
Departing from unrolled differentiation, we now instead model the final weights by a Boltzmann distribution. This is motivated by the fact that it is the limiting distribution of Stochastic Gradient Langevin Dynamics [16, 17], which closely resembles SGD. Given an energy function L(0) elx(0) (0) and an inverse temperature parameter 8 E Rt, the Boltzmann distribution is:
p (0) e - -(,,,(0)e0,1100))
Z(B,5) = 1 -(((()) dO.
(8)
Z(B,e)
Let P. denote the corresponding measure. Let 8 = (0: (()) inføEr ((0)} denote the set of global minima of the loss function (c.f. ST above). Consider the following set of assumptions.
7

---

A1. The derivatives ddn are bounded for E (1, and [I.
10)
A2. The nonzero eigenvalues of the Hessian (0) are uniformly bounded away from zero on 5%
A3. The perturbation & (0) is constant on 82
A4. 0 + 2()) el (0) is Lebesgue integrable for all & in some neighbourhood of 0.
A5. (() attains its minima, i.c. 5% = (0: : ( = i }(0}} is not empty.
Theorem 3. (Asymptotic optimality of IFs with Boltzmann distributions). Let R C C* (R4, R4) denote the class of bounded vector fields r such that Tg (0) == 0 + sr(0) is a Ca diffeomorphism for all sufficiently small
0. Define the functional
I F(r, = lim 1 8-000 3 D (Te Po P
equal to the asymptotic KL divergence between the transformed base measure T Po and the true perturbed measure p Consider the subset of maps R := {rr E r(0) = TI (0) for E
} C R, for which the map is equal to influence functions the minimum manifold. Then, given
any r E RI and any r' E R \ RIF= there exists some a E I++ such that
I F(r, 8) S F
A
|e| a
(10)
Moreover, the set RI non-empty, so such diffeomorphisms do indeed exist.
Proof sketch. We start by showing that, for small enough e, there do indeed exist continuously differentiable bijections in the class (0) such that r(0) = (II(O) when 0 E S% At low tempera- tures, only the behaviour at 8% matters because the probability mass concentrates where the loss is minimised. Taking 3 - oo and using the Laplace approximation, we analyse the low-temperature KL divergence between T Po (the transformed measure, without loss perturbation) and P (the measure with loss perturbation). Among the class Tg (0). this is minimised at O(E²) terms by IFs.
Commentary on Theorem 3. For Boltzmann distributions, IFs provide exactly the transport map in (0) required transform the low-temperature (weak limit) Boltzmann distribution with loss C()) onto the Boltzmann distribution with a perturbed loss <()) ely (0), up to (s²)) terms. This provides a very explicit distributional motivation for IFs: they map samples from Po (read: (LD)) onto approximate samples from P. (read: (D))=;)) ) and do so approximately optimally in the KL sense. We also remark that, since the KL divergence is invariant under parameter transformation, this notion of optimality does not depend on the specific choice of coordinate system. Minimal assumptions are made on (0) throughout for Theorem 3 to hold.
Key takeaways from Section 4. IFss are implicitly distributional. Supposing (2) ~ HDD then the sample (2) + N' IF is approximately distributed according to H.D) in two precise mathe- matical senses: (1) as an asymptotic limit of unrolled differentiation, and (2) minimising KL divergence if the final weights follow low-temperature Boltzmann distributions. This means that we use II instead of Yup in Alg. 1 cheaper yet principled proxy, unlocking d-TDA scale. It may also help explain why IFs are effective in deep learning, far from the convexity assumptions relied upon during typical derivations from robust statistics.
Distributional Training Data Attribution in Practice
Having demonstrated how d-TDA can be operationalised using IFs (Section 4), we now discuss its practical utility. We begin by demonstrating that distributional influence captures interesting infor- mation missing from its classical counterpart.
Rethinking influence. Previous papers have heuristically considered what amounts to mean influ- ence [12] if from the training dataset, which example would change a model measurement most on average? In synthetic ID regression task shown in Figure 5, this criterion identifies 20 as the most influential datapoint. As discussed above, we could quantify the difference in distributions after retraining in different way, e.g. with Wasserstein influence. Here, in contrast, Z31 is deemed the most influential. Note that T31 is not very influential by conventional measures since its mean shift modest, yet its removal drastically changes the behaviour after training, sharply increasing uncer- tainty. This demonstrates that different notions of distributional influence can capture meaningful information about the training data missed by e.g. heuristic ensembling. As second demonstration, in Figure 11 (App. C.2.3) we use d-TDA to identify MNIST examples which lead to large change
8

---

Seed 0
Trained full dataset
Train on all data
Remove 230
Remove 231
Remove with fixed seed influence Remove with distributional influence
Training samples
Predictions
4
Most indiermtial sluf
0.85
Test Accuracy
0
Figure 6: d-TDA > fixed-seed TDA. Test accu- racy improvements on CIFAR-18 with SWIN Vision Transformer from IF data pruning. We compare two approaches subset selection: 1) traditional TDA with a fixed seed, where for each random seed we remove 5000 datapoints that are predicted to decrease the validation loss the most for the model trained with that specific fixed seed; and 2) distributional-TD.. where for each model we remove 5000 datapoints pre- dicted to decrease the validation loss the most on average. Both methods lead to test accuracy improvements upon the baseline trained with all data. However, d-TDA leads to greater over- all improvements on average. Black dots show accuracies for individual models (with seeds indicated in gray). whereas coloured diamonds indicate the average result for each method.
Predictions at steet
bdean abift: 0.14
Meenn aiatt: 0.11 Wy changes 0.56
Wy change: 0.14
2-
Figure 5: Distributional influence on 1D regres- sion. Top: Samples model functions trained on the full dataset, as well as subsets without the most influential example by mean shift (Xan) and by Wasserstein shift (Eig)). The 90th percentile of each distribution is indicated by shading. Bor- tom: Histograms of model outputs at teat after removing Zo 231 and corresponding mean and Wasserstein shifts c.f. the original model.
in variance but small change in mean. The right panel of the figure confirms that retraining does actually lead to the changes our methods predict.
Distributional influence on diffusion models. illustrate this concept at scale, we apply distrib- utional influence to identify the most influential training examples for latent diffusion model [7]. Figure ranks the most and least influential examples on ArtBench, comparing Wasserstein influ- ence, mean influence, and classical fixed-seed influence. The identified examples vary in each case.
The ability predict how different training examples impact the distribution be practically useful. For example, one could identify examples to add or remove to most reduce variance, hence reducing model (epistemic) uncertainty. Further, d-TDA methods could also be used to operationalise criteria like information gain [18, 19] or marginal likelihood [20, 21, 22, 23].
training runs may
Most
Least
influential
influential
Generated sample
Figure 7: d-TDA highlights different influences for diffusion models The figure shows most and least influential datapoints on generations of shown samples from a latent diffusion model trained on ArtBench-10. The "most influential" examples are those that change the DDPM loss (a proxy for the ligeliie) of the generated sample the most, following [7].

---

Data pruning with d-TDA. Next, we apply distributional (mean) influence to a data pruning task, where the goal to remove datapoints from the training set to improve the performance of the final trained model. We consider a SWIN transformer [24] trained on the full CIFAR- 10 dataset (see for details). For the baseline, we remove 5000 datapoints deemed to be most influential that is, estimated to decrease the validation Ioss the most when ablated using regular TDA a single model. For d-TDA, we remove 5000 datapoints estimated to decrease the validation loss the most on average, for 10 models trained using different random seeds. We then compare the final accuracies and losses for individual models trained with those examples ablated, using the same random seeds. Figure 6 shows the results. The distributional variant unlocks accuracy gains c.f. fixed-seed TDA.
Evaluating TDA methods. The observations above also invite to rethink how evaluate data attribution methods for stochastic training algorithms: d-TDA methods ought to be effective at identifying examples responsible for large changes in distribution. These changes are often missed when one only looks at the change in mean. Note that the Linear Datamodelling Score (LDS) [12], a common evaluation metric, can already be interpreted as a d-TDA evaluation metric. measures how accurately attribution methods rank training datapoints by mean influence:
LDS spearman (Dist (2))t (2))) (DistInf & (2)
M
(11)
where spearman denotes the Spearman rank correlation, Dist.Infg.. DistInfg- are distributional (mean) influence scores computed using exact retraining and a d-TDA method respectively, and D'' randomly subsampled subsets of the training data. In light of our discussion, it is natural to gener- alise Eq. (11) using other notions of distributional influence. We term such metrics distributional LDS, of which regular LDS is a special case. We show a preliminary benchmark in Figure 10 (App. C.2), showcasing that distributional LDS can better flesh out differences between d-TDA methods. Distributional LDS (e.g. Wasserstein) can be computed at virtually no additional cost over standard LDS, and we argue should become the default for benchmarking data attribution in deep learning.
Leave-one-out is not broken, just noisy. Prior works have reported that TDA methods such as IFs are incapable of accurately predicting the outcome of leave-one-out (LOO)
Predicted measuroment
Samples
0.2
retraining, often obtaining near 0% correlation to ground- truth measurements after actual retraining [14). Adopting a distributional perspective, we view this differently. For
-0.4
Mcand
2,30
big datasets, removing a training example 25 only leads to a tiny change in distribution 1A > HDD)a We have seen
-2332
al Removed datapoint
that IFs allow us to approximately sample from MD
-0.8
but we may need many empirical samples to detect such minor distributional shift empirically. In other words, real- world training is noisy: TDA methods struggle with LOO primarily because of low signal-to-noiss ratio, rather than any fundamental incompatibility. In Figure 8, we verify that common attribution methods are capable of accurately approximating the LOO distribution with enough samples the means of the predicted and ground-truth distributions correlate extremely well.
-0.6 -0.4
-0.2
True measurement
Figure 8: There is signal in leave- one-out. Measurements (model output) on a fixed query input when different examples are removed from the training set. True measurements on the z-axis against measurements predicted with unrolled differentiation the y-axis for individual models (with different ran- dom seeds) are shown with low-opacity. The distributions of measurements are noisy, and similar for each removed ex- ample, hence the LOO correlation is close to O. The empirical means of the distribution over random seeds are shown in full color in the inset axis. There is clear correlation between the means of the true and predicted mea- surements. See App. C. for full details.
Conclusion
This paper introduced distributional training data attribution (d-TDA): a new paradigm for data attribution when training algorithms are stochastic. To demonstrate its utility, we used d-TDA to more effectively identify training examples whose removal improves test loss and accuracy, and proposed novel ways to evaluate d-TDA methods. Rigorously tackling distributional questions also yielded new mathematical mo-
tivations for influence functions for deep learning.
10

---

References
[]] Ghorbani Zou, "Data Shapley: Equitable valuation of data for machine learning," International conference on machine learning, 2019, pp. 2242-2251.
[2] P. W. Koh and P. Liang, "Understanding black-box predictions via influence functions," in International conference on machine learning, 2017. pp. 1885-1894.
[3] Grosse et al., "Studying large language model generalization with influence functions," arXiv preprint arXiv:2308.032000 2023.
[4] Z. Liu. H. Ding, H. Zhong. W. Li. Dai, and C. He, "Influence selection for active learning," in Proceedings of the IEEE/CVF international conference on computer vision, 2021, pp. 9274-9283.
[5] R. Jia et al., "Towards efficient data valuation based on the shapley value," in The 22nd International Conference Artificial Intelligence and Statistics, 2019, pp. 1167-1176.
[6] F. Hampel. "The influence curve and role in robust estimation," " Journal of the american statistical association, vol. 69, no. 346, pp. 383-393, 1974.
[7] B. K. Mlodozeniec, R. Eschenhagen, J. Bae, A. Immer, D. Krueger, and E. Turner, "Influence Functions for Scalable Data Attribution in Diffusion Models," in The Thirteenth International Conference Learning Representations,
[8] S. Hara, A. Nitanda, and T. Maehara, "Data Cleansing for Models Trained with SGD." [Online]. Available: http://aaxiv.orgabbb///6..0
[9] J. Bac, W. Lin, J. Lorraine, and Grosse, "Training data attribution approximate unrolled differenti- ation," " arXiv preprint arXiv:2405. 12186, 2024.
[10] A. Ilyas and L. Engstrom, "MAGIC: Near-Optimal Data Attribution for Deep Learning," arXiv preprint 2504. 16430, 2025.
[11] J. Bac, N. Ng, A. Lo, M. Ghassemi, and R. B. Grosse, "If influence functions are the answer, then what is the question?," Advances in Neurul Information Processing Syxtems, vol. 35, pp. 17953-17967, 2022.
[12] S. M. Park, K. Georgiev, A. Ilyas, G. Leclere, and A. Madry, "Trak: Attributing model behavior at scale," arXiv preprint arXiv:2303. 230.,1603., 2023.
[13] L. Franceschi, M. Donini, P. Frasconi, and M. Pontil, "Forward and Reverse Gradient- Based Hyperpa- rameter Optimization." in Proceedings of the 34th International Conference on Machine Learning. D. Precup and Y. W. Teh, Eds., in Proceedings of Machine Learning Research, vol. 70. PMLR, 2017, pp. 1165-1173, [Online]. Available:
[14] S. Basu, P. Pope. and S. Feizi, "Influence functions in deep learning are fragile," arXiv preprint arXiv:2006. 14651, 2020.
[15] H. J. Kushner and G. G. Yin, "Stochastic approximation and recursive algorithm and applications," Application of Mahematics, vol. 35, no. 10, 1997.
[16] M. Welling and Y. W. Teh, "Bayesian learning via stochastic gradient Langevin dynamics," " in Proceedings of the 28th international conference on machine learning (ICML-II), 201 pp. 681 681-688.
[17] S. Mandt, M. D. Hoffman, D. M. Blei, "Stochastic gradient descent as approximate bayesian infer- ence," Journal of Machine Learning Research, vol. 18, no. 134, pp. 1-35, 2017.
[18] D. V. Lindley, "On measure of the information provided by experiment." The Annals of Mathematical Statistics, vol. 27, no. 4, pp. 986 1006, 1956.
[19] F. B. Smith, A. Kirsch, Farquhar, Y. Gal, A. Foster, and Rainforth, "Prediction-oriented Bayesian active learning," in International Conference on Artificial Intelligence and Statistics, 2023, pp. 7331- 7348.
[20] J. MacKay, Information theory, inference and learning algorithms. Cambridge university press, 2003.
[21] F. Fong and C. C. Holmes, "On the marginal likelihood and cross-validation," Biometrika, vol. 107, no.
2, pp. 489. 2020.
[22] K. Mlodozeniec, M. Reisser, and Louizos, "Hyperparameter Optimization through Neural Network Partitioning," " The Eleventh International Conference on Learning Representations, 2023.
11

---

[23] A. Immer. M. Bauer. V. Fortuin, G. Rätsch, and K. M. Emtiyaz, "Scalable marginal likelihood estimation for model selection in deep learning," in International Conference on Machine Learning, 2021, pp. 4563-4573
[24] et al. "Swin transtormer: Hierarchical vision transformer using shifted windows," in Proceedings of the IEEE/CVF international conference on computer vision, 2021, pp. 10012-10022.
[25] V. S. Borkar, Stochastic approximation: dynamical systems viewpoint, vol. 9. Springer, 2008.
[26] G. Krantz and R. Parks, The Implicit Function Theorem. Boston, MA: Birkhäuser, 2003. doi:
10.
[27] V. Noferini, "A Daleckii-Krein formula for the Frechet derivative of generalized matrix function," 2016.
[28] J. L. Daletskii applications to the theory of perturbations," AMS Translations (2), vol. 47, no. 1-30, pp. 10-1090, 1965.
G. Krein, "Integration differentiation of functions of Hermitian operators and
[29] P. W. W. Koh, K.-S. Ang. H. Teo, and P. S. Liang. "On the accuracy of influence functions for measuring group effects." Advances in neural information processing systems, vol. 32, 2019.
[30] S. Basu, X. You, and S. Feizi, "On second-order group influence functions for black- box predictions," in International Conference on Machine Learning, 2020, pp. 715-724.
[31] E. Barshan, M.-E. Brunet, and G. K. Dziugaite, "Relatif: Identifying explanatory training samples via relative influence," in International Conference on Artificial Intelligence and Statistics, 2020, pp. 1899- I 1909.
[32] J. Martens and R. Grosse, "Optimizing neural networks with krotecker-factored approximate curvature," in International conference on machine learning. 2015, pp. 2408 2417.
[33] R. Eschenhagen, A. Immer, R. Turner, F. Schneider, and P. Hennig, "Kroneckerffactrred approximate curvature for modern neural network architectures," Advances in Neural Information Processing Systems, vol. 36, pp. 33624-33333333 33655, 2023.
[34] T. George, C. Laurent, Bouthillier, Ballas, and Vincent, "Fast approximate natural gradient descent a kronecker factored eigenbasis," " Advances in Neural Information Processing Systems, vol. 31, 2018.
[35] L-C. Yeh, "Concrete Compressive Strength." 1998.
[36] L. Deng, "The MNIST Database of Handwritten Digit Images for Machine Learning Research," IEEE Signal Processing Magazine, vol. 29, no. pp. 141-142, 2012, doi: 10.1100//SP22222222111
[37] F. Dangel. R. Eschenhagen, W. Ormaniec, A. Fernandez. L. Tatzel, and A. Kristiadi, "Position: Curvature Matrices Should Be Democratized via Linear Operators." [Online], Available: http://aaxiivoor4abs/55501 19183
[38] J. Bae, G. Zhang, and R. Grosse, "Eigenvalue corrected noisy natural gradient," arXiv preprint arXiv: 12565, 2018.
12

---

A Proofs
A.1 Proof of Theorem 2: Unrolled differentiation converges to IFs
This appendix provides a proof of Theorem 2, repeated below for the reader's convenience.
Consider stochastic gradient descent updates given by
00-1 re+1
O, Σ v ( (0))):: +
M E ( B
(12)
with random variables 8' in {0, 10 for t N. independently and identically distributed with mean E[,%] = N.
B and for some 0g independent of (8') (8'),E) LEN"
Consider the following assumptions.
A1. V and Vlg are Lipschitz continuous and bounded.
A2. The step sizes (1)+ 1=g are positive sealars satisfying E, 1) = 0o and Σ, rh < oc.
A3. The iterates of Eq. (12) remain bounded a.s., i.c. sup, r,)l < oo a.s.
A4. (0) E Span V (())) for all 0 87.
AS. The nonzero eigenvalues of V (()) for 0 E 5 are uniformly bounded away from O.
A6. There exists some compact neighborhood ( around S such that gradient flow trajectories
(()) initialised therein converge uniformly over initialisations to points in 8% Moreover, their lengths are bounded a.s., so that: suPorojex((8"> S:= (0(s) - lim 0 )ll ds < co.
Define 8 = {0: ((( 0}, the of model parameters where the loss function has zero gradient. Also define the set of local minima, 8 = (0: V = 0, V (()) = 0) C Sc. We denote the pseudoinverse of a matrix with (.) The following is true.
Theorem A.1. (Unrolled differentiation converges to IFs). Suppose that A1-A6 hold, and consider an SGD trajectory in the set that converges to S (c.f. Sz ST). The sequence of iterates ((0,, r.)) t-o generated by Eq. (12) converges almost surely to the set R" == " (i"((") + Tus (0)) 0 8m rir V [²z( (0), rys (0) Null(V><(())) -haa is, pointwise IFs, plus component in the Hessian nullspace.
Proof.
We interested the behaviour of response for trajectories where 0, converges Sm rather than Sz \ 5% (see Theorem 1). Consider the following ordinary differential equation:
-V -=
(13)
This can be considered to be continuous time analogue to the SGD updates in Eq. (12). Ô - -V) corresponds to gradient flow. We can solve this analytically given the initial model weights e(0), obtaining a gradient flow trajectory ((t). Note that, by assumption A1 (Lipschitz continuity) and the Picard-Lindelöf theorem, for any initialisation (0(0), r()) there exists unique solution to the ODE in Eq. (13). The following is true.
Lemma A.2. (Gradient Flow ODE converges influence functions) Given assumptions AI, A4, consider any initialisation (0(0), r()) of Eq. (13) for which 1) the ODE con- verges some limiting weights 0" = lim, O()), 2) the trajectory length is bounded, i.e. not (t) - 0 dt < co, and 3) the limiting weights are (possibly degenerate) local minimum. Le. (²"((0) > O. Then lim r(t) exists and lim r(t) {Tir(0") + rxs("") (i) < Vl, (0), rus(( ()) E Null(V² (())))
Proof. Inserting the computed flow trajectory O(t), consider the ODE for influence + ~V ((0())) (0(()())) For notational simplicity and consistency with the dynamical systems literature, will write this in shorthand as:
(tt) = (t)r((t) + b
(14)
13

---

where -V2 and b(t) == (4 Since O(t) converges and the Hessian and gradients are Lipschitz continuous (assumption A1), (t) + A. and b(t) + bo converge as well.
We will now study convergence to the limiting influence. To == lim r(t). For a particular flow trajectory with limiting negative Hessian A == lim (t), let P == A A denote the projection operator onto the column space of At)). Define the variable xxt) == r(t) - ( A ).)
Rewriting Eq. (14), we have that
a()= (l)aa(() + (b(1) - b()
(15)
Note that y() - 0 if and only if 2o is in the column space of A0 Assumption A4 ensures that this indeed the case. Denote (t) Ag + A (t), with (t) + 0. Let x (t) := Px(() and æ (t) =
(I P)x(() be the components of ax()) in the column and null- sall-space of A respectively. Our goal
will be to show that x (L) + 0. Premultiplying Eq. (15) by P. we have that
d (t) = ()) (z
(t) (t)+
+ x (t)) + t(t)).
(16)
This implies that
(x ) xt = - Ao xt + (at)) (ax!l
(zll + x x +))
+
(17)
where suppressed dependence for compactness. Note that (a+)
1d (a+t)
2dl
1 d 12 11 | Also, since x=1 is in the column space of A and A is negative 2 at semidefinite, we have that (a+)) A Ac <- |0-|2 with - < 0 the greatest nonzero eigenvalue of A Combining the above.
d
x || = 1(x=1)) T 2x (xt) P(Ax + 9)
dl
< <1((=t) _ A xt (1) + (x+)) P ar v + ツ) S - | + | |PP| Ax
(18)
S A x 1| + | |||| lyl).
We used the triangle and Cauchy-Sehwar inequalities. By the assumptions in the theorem statement, | [t] is bounded by constant independent of which is guaranteed as [x(0)] is bounded and L (()) ((00)]l dl < ou; see Lemma A. below. this case, we have that
at
|
- x
1 ||> |6>/|| + llvil).
(19)
Divide through by |a*|| J|. Since (t) + 0 and yt) + 0, for any E > 0 there exists a time T. such that > (t)l + (t)
for all t Tg. At such times, di (Ix (t)lle²r) S Eetr, whereupon
E x - (it-e-ii-T)) T.) < ye x(t-T.
< (t) < | (T
(20)
For any > 0, choosing such that x < E', one can find some T > T so that x (t) < s' fort > T.. This proves that x 1 (t) + 0, whereupon we can conclude that, for such initialisations, Pr(t) -
At b
As brief digression: in Lemma A. we used that [t) is bounded by some constant 7. We stated that this is guaranteed if [z( is bounded and 1 (0(s) 0012 ds < cc. This is seen as follows.
Lemma A.3. Under assumptions A1-A7 especially. Lipschitz smoothness of the Hessian and gradients, and the convergence condition 1 Ott) 0 0- dt < oo with A_ 3 0 - the response r(t) remains bounded.
Proof. Take ætt) := rtt) (- The nonzero eigenvalues of A are uniformly bounded away from 0 on 8th and b 00 is bounded (A4), so bounded ax()) implies bounded r(t). Consider that, for (+t) = (t))((t) + (t), we have:
Par(t) can be interpreted the error between r(t) and the asymptotic influence functions formula.
14

---

d dt
| -=||
x
Ax + + Ax x S ||||=| + |
(21)
50
We used the assumption that A. 3 0. since we considering flow trajectories that converge to local minima. It follows that
S A || + levll so
27
[x(t]]| S - (e))lde
[|>(0]]| + L ly((') efiio-111 de'
P==
(22)
S [lu(t")) 2'=0
[[x(0]|+ +
This is bounded if 10 ( ds < oc and ( [1)(()) ds <o. If the first condition holds and doo is in the column space of A, (A4), then from the definition of y(t) the second condition simplifies to - lb(t) bx ds < co. Under Lipschitz smoothness assumptions for the Hessian and perturbation, these conditions
clearly guaranteed by the convergence rate condition the model weights f 10(t) - 0~ ds < 0o as claimed.
Lemma A. 2 proved that, provided r(0) is bounded and the model weights ((t) converge to a local minimum, the response vector field r(t) evolving according to the flow ODE converges to influence functions. In particular, we found that | S ye- - + &. with À the infimum over nonzero eigenvalues of the Hessian at points in 8% (assumed to bounded away from and the maximum possible [zx( (also bounded given Lemma A. 3). Assuming Lipschitz smoothness and bounded (t) (assumption A1) A (t)ll I I, (t)) - 0 and [b(t)] < Lng (0())(t)))))))))))) - 0x with L 42 bounded constants. Hence, T. is upper bounded by constant multipled by the maximum time required to guarantee that [0) 0 ooll < E'. Therefore, provided A6 holds i.e. flow trajec- tories converge uniformly in the neighborhood of the local minimum - Pr(() initialised therein also converges uniformly. This property will be important later in the proof.
We have seen that the influence ODE converges under mild conditions. Our next task is to use this result to prove the convergence of the influence SGD iterates described by Eq. (12). To do this, we invoke classic arguments made (among others) by V. S. Borkar [25].
We can rewrite Equation (12) in the following way:
(0+1 Te+1
0, Th E Vln + n.M,
=
(1-
The N n-1
V2L r. + (0,)
+
N ( - Σ,_, 8 (e,
M [HE v'e, - - Σ 8 V2,, (0
N₂
=
- -
Since the batching variables are i.i.d. and the dataset is fixed, (M,, N₁) is Martingale difference sequence with respect to the increasing family of -ffildss
g. == ( m: m m n)
(23)
That is, E((M, n+1 N 1 = a.s., O. Furthermore, (M N₂,) square integrable with E(IMM 1 12 + II 1l (F,,) < K(1+ 10 + |r,l) a.s., n 2 O, for some constant K 2 0. This allows us to use standard martingale convergence results to connect the SGD iterates the ODE solution as f + oc.
We can think of the SGD trajectories as a noisy discretisation of the corresponding gradient flow ODE, with ', = Σn
representing the amount of time that the process has been running for. Let = (t, n E N} be the corresponding to SGD steps. Let (0,, r.)) neN denote the SGD iterates, generated by Eq. (12), Define (t) == r, for t E [t, tn+1). Finally, let (0TH (t), (T)) for E Im, 00) be the solution to the gradient flow ODE in Eq. (13), initialised at (tm"(t,) ) r'" (tm = (0 m. ]. By [25, Lemma 2. since the noise is martingale difference sequence and given assumptions A5-A6 for any finite E R
15

---

lim
sup
ISSGD (t) - (t) 0 a.S.,
(24)
self
so there exists m E N such that te -T 22 2Tyl''GD (!) r' (1) Se for all > m [25]. This remains true if we can increase my to be big enough that the time interval 4m m < TV m 2 m, whereupon we have that
U W
Il + TT + 27) = Imm + T, +T,,)) oc).
(25)
Since Om - 5 we can make m² yet greater to guarantee that the SGD iterates (0,,. are in the tubular neighhorhood where convergence of gradient flow, and therefore the response, is uniform (assumption A6). Recall our earlier definition of the set R {rIp((') Tys "xx((') l E 882, r == (()) V&, (0), Tus (0) E (ull((><(())) This corresponds to influence functions at the loss function minima, plus unspecified component parallel in any degenerate directions. Let P.. denote the unique asymptotic projection operator when gradient flow is initialised
m2m;
(6.. T' and run for infinite time. From the uniform convergence of gradient flow, we have that
3 3Tg s.t. V T.
inf (t) - r"1l S inf MER
[P rm( pmrl1 + I[[I - P )r (t) (I P )r'l
rER
(26)
=0
Pm Pmr SE. inf P T MER
The second term vanishes because within the set R the null space component is unconstrained; we make no claims about convergence. Hence, can always be exactly fitted (I Pm (t). Meanwhile, the first term is can be made less than e due to convergence of r' (L) perpendicular to flat directions, which we already proved.
Choose any n such that f > 'mm- + Te. Then choose some corresponding m 2 m such that t, E [t., + T m + 2Tg], which is always possible due to Eq. (25). Combining the previous inequalities,
inf Ir, r'1 < |rr, T. (t,)) + rER
inf |rr r'll < 2E. rER
(27)
SGD ODE
ODE
influence functione
Take the union over all I, Im + Tg we can finally conclude that
R
(28)
as claimed, This completes the proof.
A.2 Proof of Theorem 3: Asymptotic optimality of IF's with Boltzmann distributions
This appendix provides a proof of Theorem 3, restated below for convenience.
Given an energy function L() - (0) and an inverse temperature parameter 8 E R+ the Boltz- mann distribution is:
( e (((
Z(B,e) ((B,E)= -
-((((())((((0)) de.
(29)
z(B,E)
Let 8. denote the corresponding measure. Let 8 = (0: ((0) inf I(0)} denote the set of global minima the loss function (c.f. 8T above). Consider the following set of assumptions.
A1. The derivatives d"t,10) ddm are bounded for E 2,,2,3) 3} and i E [1.
A2. The nonzero eigenvalues of the Hessian V (²(((e) are uniformly bounded away from zero on 89.
A3. The perturbation (0) is constant on 8%
A4. 0 + (() Elg (0) is Lebesgue integrable for all E in some neighbourhood of 0.
A5. L(0) attains minima, i.e. 8% = (0. (()) infg ((0)) is not empty.
16

---

Theorem A.4. (Asymptotic optimality of IFs with Boltzmann distributions.) Let R c C (RR Rd) denote the class of bounded vector fields such that (0) == + er(0) C
diffeomorphism for all sufficiently small & > O. Define the functional
1 I fr, lim 8-000 3 DKT. Po (Pe).
(30)
equal to the asymptotic KL divergence between the transformed base measure T Po and the true perturbed measure P Consider the subset of maps IF := {r E R: r(0) TIp (0) for E
8% } C R, for which the map is equal to influence functions on the minimum manifold. Then, given any r E RI and any E R \ Hpp, there exists some a E I+ such that
F F(r, S Fr'
A
Sa. a.
(31)
Moreover, the set R non-empty, so such diffeomorphisms do indeed exist.
Proof. We begin with the following lemma.
Lemma A.5. For small enough e, there exist continuously differentiable bijections in the class Tg such that T- (0) = 0 + (0) for 0 E Sz.
Proof. We start by defining function Te that we will show has the claimed properties. Let A E E denote the smallest nonzero eigenvalue of the Hessian V (()) on the minimum manifold S which bounded away from (assumption A2). Define the following scalar transformation f + R,
2 if A min:
f(z) =
Amm
Amn
(32)
otherwise.
Note that f Lipschitz continuous with constant 1/x2 For compactness, let A v2 Denote the operation of f on symmetric matrix A by f(A) Q f(A) where f is understood to act separately on each of the eigenvalues on the diagonal of A. Observe that, if A is positive definite and all its eigenvalues are greater than equal to Annin then f(A) = A and we recover the regular matrix inverse. Similarly, if A is positive semi-definite with all non-zero eigenvalues greater than or equal Amia and v is a vector in Span(A), then A f(A)u. Hence, if we take T = 0 + er(e) with r0)) / 2 (())) V&, this clearly satisfies r()) Tie ( V (0) Vl for E Sc (Assumption A.3). Hence, we only need to show that it's a continuous bijection. To this end, we will use the following theorem*:
Lemma A.6. (Hadamard's Global Inverse Function Theorem S. G. Krantz and H. R. Parks [26, Theorem 6.2.8)) Let Rdd Rd be continuously differentiable function. If:
1. h is proper (for every compact set K C Rd h (KK) is compact), and the Jacobian of h vanishes nowhere,
then h is homeomorphism (continuous bijection with a continuous inverse).
We will first show that the Jacobian vanishes nowhere. The Daleckii-Krein Theorem [27, 28] gives us the following:
VS(A) = Q QVAQ)
(33)
where O denotes the Hadamard matrix product, (A O
B
A.
B,
and
(AHA 2
if A *
Rij
(34)
(())
otherwise.
Note that sup,, R, < 1.
The following is true:
x
IV,A
R Q S Vdparan Rup if
=
|
S
Vd
paran
(35)
Amn
'Presented for the specific case of the standard topology R
17

---

Here, |I, denotes the Frobenius norm of 8 (()) which is bounded by a constant if the 30;
third derivative of the loss is bounded (assumption A.I).
The Jacobian of this transformation is:
(0) = I - -V+ + +
(36)
where I denotes the dparan x daran identity matrix. Given the previous, the spectral radius of the term in square brackets is bounded under assumptions A1-A2, so the Jacobian is positive definite at small enough E. This means that Te is locally invertible everywhere for small enough
To show that T, is proper, note that Tg R - R is proper iff lim IT (2r,, lll for every sequence (0x,), n=1 s.t. |2r,122 + oo as + 00. To show the latter, note that V L and V4 are bounded (Assumption A.1), and so is f(v) (since f is Lipschitz), and hence r(0)= f(V? V is also bounded. Hence, lim F + er ll co as I| l2 oo as required, and by the Hadamard's Theorem, T. a homeomorphism.
Equipped with Lemma A. 5, for small enough & we can apply the change of variables formula. Since the KL-divergence is reparameteriiatoonnnnvaann for any continuously differentiable bijection T.
we have that DKL. PT p
DKL. (T Tex Po T P. = DxL. [P81 Ta P!l
p8 has density given by Note that c#
p (T, det ()) --B(01T,10))44111.000
det VT. (0),,
(,(,,)
Hence, we have that:
8 DKL T 1o1 p. = 3 DKL [ Tes ps]
3 ʃ () log (0)) det VT. (0)
1
Fo (0)
d0
(37)
= / 1) ((TT,())) - =( + log det VI. d0 + ( [pp] + log (p(e, B)).
Here, H6 [P]] == - p (0) log p (0) d0 denotes the entropy of 8 Note, the terms 1/2 [ [6] +
log (e, B) do not depend on T. Moreover, we have that lim g j [Po] O. Given assump- tions A1-A2 used to keep the transformation locally invertible, the eigenvalues of VT. (0) are
bounded by constant independent of 0 so lim 3 P (0) log det (T)( = 0.
Putting in ( = 0 + cr))), let us now consider the expectation E E [C(0 Er(0))]. Below, will use Einstein summation notation, with the implicit understanding that one should sum over repeated indices. Taylor expanding in e with f fixed,
( + =r())))))))))))))) - el, (0 +
L(0 = +10)+
+ -- - + (0,
(38)
2
Here, R6 (0,,,)) = (00 is the error term, for some e (0, Similarly,
El 1 er) = =4 (0)+ 0,4 (0), + R (0,,,))
((,
(39)
where this time 1,000,.0) = 0++ Since the first three derivatives of L and Ey- as well as r and r', are bounded a.e. (assumption A1), lim bounded by constants independent of follows that
R and lim (0,00,)))) ga
are
'iith a slight abuse of notation, we included r as an argument to emphasise that the remainder term will depend the particular choice of function Once the function is fixed, the arguments of Ry are of course 0 and e.
18

---

=0+10++0 = ppp <(() =0,,)rr)(())r; + =² 0,0, =4x(0) - O (40)
B
+ sr) |0
[4(0
Next, consider the log partition function, log (,,,, B3 == log e (000000))- Applying the Laplace approximation [20], we find that:
lim log (p(5, =- inf aa
(<(0) - =l
(41)
Assembling the various pieces, we then have that:
F f(r,E) = lim Eg pp (2(0) - el, PER
inf [L(00 cl2 +
8
(42)
2² lim E 90000
p (,0, 0,4 + O(5))
We dropped the 8-+0 E 0-pb (0)r term since the the weak limit Poe has support on the minimum manifold S (e ( inf L())). where 0,( 0 by definition. Since (,)rr is continuous and bounded, the limit of the expectations is the expectation under the weak limit.
Let us now consider some r E RiF and some r' E R' \ RyF- Since TIF (0) = (2 (4) directly minimises the square bracket the second line of Eq. (42) for each Se, we have that
(2
-(r,s)-F(r' - F((',
2 lim E. o 8000
-0,4000 - -000)10000(5500(000))
(43)
1m-Ac0
+ lim E en ( R (0,,,)) R2' ( (0,,,r))) 8-000
Every remainder term bounded by constant independent of 0 and e, the magnitude of the expectation on the bottom line is bounded by a constant C' E R Hence,
F F(r - F s) v A Cs
(44)
whereupon F(r,e) S Fr'', e) is guaranteed for |e| < This completes the proof.
Extra remark. In the special case that the perturbation does introduce any new global minima/ break the degeneracy of the minimum manifold, then the KI. divergence actually vanishes up to (ss)) so we have even stronger result. Assume that the perturbation is constant on Se (assumption A3). Consider the following:
inf [C(0) ely (0) 6
(((0) - =l,( + inf 0-es,
ri - 0,0,,0000 1 ( (0) - (0') + 0(E)) (45)
=((0)) =E s² inf (0r) T ((0) (0)) +0(=)),
1
r'e
where Taylor expanded in & and used the implicit function theorem. The infimum will cancel with the expectation on the lower line of Eq. (42) if its argument (l₂(0") V ((0') Ve, (0*) is identical for all 0* 5.. This will not true for general perturbations - just a specific class that does not break the manifold symmetry at (e²2).
To provide intuitive summary: influence functions are always the best local transport map at O(=2), parameterised by 0 - 0 + (#)) But they are also the best non-local map, making ((=2) terms vanish so that the KI. divergence is truly (ee3), in the special case that the perturbation does not induce symmetry breaking of the minimum manifold at (s²2). This condition is formalised by
19

---

V4,(0") V'= being constant for all 0" E S which is also intuitive because it gives the change l at the new minima.
B Derivation of Influence Functions
The purpose of this appendix section is to provide a standalone "classical" derivation of the influence functions framework for the "classical" training data attribution task. We state the Implicit Function Theorem (Section B.1); then, in Section B.2 we introduce the details of how can be applied to predict local changes in the minima of loss function L(e, 0) parameterised by continuous hyperparameter E (c.g. (E,, 0) = Lp (0) El2 (0)), so that e controls how down-weighted the loss terms on some examples are). This derivation largely mirrors that in [7, Appendix A].
There appears be prevalent misconception that influence functions can only be applied convex loss functions [14]. This appendix hopefully makes clear that they can be applied to a loss function with multiple minima, as long each minimum is strict local minimum; influence functions in that case will simply predict the change in the corresponding local minimum. The rest of this paper then makes formal what influence functions do in the more general and complex setting of stochastic optimisation for general loss functions with possibly degenerate minima.
B.1 The Implicit Function Theorem
Theorem (Implicit Function Theorem G. Krantz and H. R. Parks): Let F: Rd x 1T R be continuously differentiable function, and let R" x R have coordinates (æ,, y). Fix point (a, b) = (a1; ..., an b1..... m) with F(a, b) = O, where Rm the zero vector, If the Jacobian matrix V Faa b) E Rmmm of y F(a, y). defined as
(y,( = OF dy,
b),b),
(46)
is invertible, then there exists an open set U C R containing such that there exists unique function g U - Rm satisfying g(a) = b, and F(x,g(z)) = for all æ E U. Moreover, g is continuously differentiable.
Remark (Derivative of the Implicit Function): Denoting the Jacobian matrix of x H> F(x, y) as V F(F y), the derivative 0g U + mxx of g given by Theorem I can be written as:
0g IV F(a, ((x)) 'v
FFa F(a,g(=))
(47)
dr
This can readily be seen by noting that, for w E U:
dF dF (()) 0.
F(a', gg(((()))))))))))))))))) = Ya' E U
(48)
dx
Since g is differentiable (by Theorem 1), can apply the chain rule of differentiation get:
0 dF(a,g(a)) V F(x,g(z)) V FyF.
0g(r)
(49)
dx
dx
Rearranging gives equation (47).
B.2 Applying the implicit function theorem to quantify the change in the optimum of a loss
Consider a loss function L: Rd R + R that depends on some hyperparameter E R" (e.g. the scalar by which certain loss terms are down-weighted and some parameters E R" At the minimum of the loss function L(e, A), the derivative with respect to the parameters 0 will be zero. ØL Hence, assuming that the loss function is twice continuously differentiable (hence de is continuously differentiable), and assuming that for some e' E R" we have a set of parameters O" such that 0² a. (s', 0*)) and the Hessian 3 (E', 0*) is invertible, we can apply the implicit function theorem de
20

---

to the derivative of the loss function 04 _R" x R + _Rm to get the existence of a continuously differentiable function such that ac 85
de (E)) = for e in some neighbourhood of
Now g(E) might not necessarily be minimum of 0 - Liee 0) However, by making the further assumption that & strictly convex we can ensure that whenever 8e 0) = 0, 0 a unique minimum, and g( represents the change in the minimum
vary bold(e). Alternatively, if 0" = g(=") is local minimum, then g(E) will give the shift in this particular local minimum we vary e some neighbourhood around e'.
We can make this more precise with the following lemma:
Lemma 1: Let L: IR" x I + R be twice continuously differentiable function, with coordinates denoted by (E,0) E R" x R such that 0 + (E,, strictly convex YE E R". Fix a point (s', 0*) such that (E' 0*) =0. Then, by the Implicit Function Theorem applied to dc exists an open set U R" containing 0* and unique function
30 there
U - Rm satisfying:
g(E') = 0", and
g(e) is the unique minimum of 0 + [(e, 0) for all e U.
Moreover, g is continuously differentiable with derivative:
0g(E) L 002 9(E))
02
0:00 (e, g())
(50)
de
Again, dropping the assumption of strict convexity, and replacing it with the assumption that (E', 0) merely yield a local minimum, gives a similar conclusion, but only guarantees existence of a function such that g( is a local minimum for all E E U.
Equation (50) might still look a bit distinet from the influence function formula. The one missing piece is restricting ourselves to look at L of the form (E,,, 0) Ln, - =É(O), matching the loss interpolations we consider in the main paper body.
Remark 2: For a loss function IR R of the form L(s, Lup (0) =l(O), 8e80
8
(c, g()))
in the equation above simplifies to:
Lº L
ae (g(e))
5 g()))
(51)
0=00
ae
The above give the final influence functions formula. Namely, for the loss of the form:
1 M
[(e, 0)
(0)
c (0)s
(52)
N
M
2-1
Am
we can substitute 026 = M 2ك=1 8:00
(0),(() into (50) to get the existence a function g with the properties given by Lemma with the derivative taking the following familiar form:
0g(s)
02
M
0 1 00)
(e (s))
(53)
de
002
M
00
and, at 0:
0g 0
(0) = 05
802 (g(0))
4 (0). 00
(54)
M 1=1
C Additional experimental results
C.1 Investigating leave-one-out
21

---

In Figure 9, we show that d-TDA methods are able to approximate the leave-one-out (LOO) distribution over measurements rather well. We simply need very large number of samples from each distribution to get good empirical estimate of the distribution in order to observe this. The setting for the LOO experiment was training an MLP on the UCI Concrete dataset (see Section E. for architectural and training details). We plot the measurement (model output) for a fixed query input for 500 models trained with different random seeds. For each one of removed (leave-one- out) training examples, we retrain model with each random seed without that example to obtain the ground -ruth measurement shown on the r-axis (left & middle plots in Figure 9). To obtain the 'predicted" measurement, we compute the predicted change the measurement of the model trained on all the data using (exact) unrolled differentiation; this is shown on the y-axis (left & middle plots in Figure 9).
Samples
Means
Fearson Correlation
-0.2
MSE
10-2
Predicted
0.4
D.3000
10
-0.6
Retuoved exomple
0.3009
0.2
10-4
-0.8
9975
0.0
006
0.4
0.2
-.6
0.4
0.2
100
200
300
400
500
Num. samples (scods)
Measurement by models trained on different leave one out subsets
Figure 9: There is signal in leave-one-out. Left: Measurements (model output) on a fixed query example predicted by a d-TDA method (unrolled differentiation) when different singular examples are be removed from the training set against the actual measurements on models retrained without those examples. The distributions of measurements are noisy, and very similar for each removed example, hence the LOO correlation close to Middle: we look at the means of the distributions, we see that the measurement distributions are subtly different, and the d-TDA method is able to pick up on the shift in mean. The differences are tiny, however; note the scale on the zoomed-in plot. For reference, the mean of the mcasurement distribution with model trained on the full dataset is shown (on the y-axis) with rectangles; we see that the d-TDA method improves upon using the mcan of the original distribution. Right: Correlation between the means of the true measurement distributions after retraining. and the means of the predicted distributions, against the number of seeds we use to empirically estimate each distribution. As the number of seeds goes up into the hundreds, the correlation approaches 90%. The seeds (determining data ordering and initialisation) were chosen independently for the fully trained model and the retrained models, indicating that we don't need to correlate the retraining trajectories with the fully trained models to get good LOO scores [10].
C.2 Distributional Linear Datamodelling Score (LDS)
In Figure 10, show distributional LDS scores using different notions of distributional influence. can be seen that different distributional LDS metrics do reveal differences between methods that can't t be seen when only using the mean influence. For instance, the performance difference when using the full Hessian vs. block-diagonal Hessian for influence functions only becomes apparent when using the Wasserstein LDS metric. Similarly, EK-FAC with and without score normalisation (described in Section C.2.2) perform identically (up to numerical accuracy) mean influence LDS, but we see that the normalisation helps slightly when using the Wasserstein and variance change influence LDS metrics. Lastly, it's clear that all methods except for unrolled differentiation fall short of being able to capture the variance change in the measurement distributions in the settings considered.
For future work, we would strongly recommend using the Wasserstein LDS metric benchmarks. It's natural choice from theoretical standpoint- Wasserstein distance is able to capture differences in distributions that go beyond changes in mean. also makes sense intuitively as a notion of influence in stochastic training settings. Lastly, it's easy to implement it differs only marginally
22

---

from the mean LDS metric - and empirically seems to capture interesting information missed by mean LDS.
Mean
Marginal Wasserstein Variance Change
Concrrtte
79.3%
56.2%122
56.2
2.5%
EK-FAC Influence
MLP
79. 1..1
15.88%+2.2
+.5
EK-FAC Influence (Normalised)
4
13.6%+1..
80.2%==o 0
Uurolled Differentiation
87.1%++o.7
66.8%=+1.7
3.4%++2.
Exact Hessian Infiuence
Concrete Tlyy MLP
84.7%6+0.7
07.7%
#9.0%+2..
Block-diagonal Hessian Influencce
78.0% 03:
53.8%
#.4%%+2
EK FAC Influence
77.7%++12
54.4%+2.11
#4.4%+2.
EK-FAC Influence (Normalised)
96.0%=2..
85.4%%=1:0
04.9%%=0.-
Uurolled Differentiation
100
Rank Correlation % (Distributional LDS)
Figure 10: Distributional LDS. The distributional LDS scores using different notions of distribu- tional inffuence (mean, variance change and Wasserstein). Each axis row corresponds to different training setting, and each nested row to a different d-TDA method.
C.2.1 Influence rankings are different according to different notions of influence
Using distributional LDS with notions of influence other than mean influence would be in vain if they produced the same orderings over (groups of) datapoints. We observe, however, that this is not the case: Wasserstein influence and variance influence produce meaningfully different rankings, presenting different challenge for d-TDA methods. This is demonstrated in Table 1 below.
Table 1: Similarity between influence rankings for random subsets of the training dataset when using different notions of distributional influence The distributional influence (and the corresponding rankings) are empirically estimated by retraining. "Top 10% overlap" refers to the fraction of the 10% most influential subsets that is shared by rankings according to different notions of influence. "Footrule distance" represents the total number of places each element in ordering would have to shifted by to match the other ordering. The reported footrule distances and top 10% overlaps the average over all query points
[Mean] vs. Wasserstein influence
[Mann vs. Variance influence
Setting
Footrule distance
Top 10% overlap
Footrule distance
Top 10% overlap
Concrete MLP
27.4 (max 200)
47%
129.7 (max 200)
8%
MNIST MLP
32.2 (max 200)
54%
106.3 (max 200)
11%
C.2.2 Normalised Hessian-pproxiimations for influence functions
One previously observed issue when using Hessian approximations such K-FAC with influence functions is that, although the correlation to ground-truth measurements is good, the scale in the predicted change is often off by large factor [7]. This deficiency not captured when looking at classic (mean) LDS metric, as the metric is invariant to the scale of the predicted change in measurement. However, the scale of the predicted change matters when we rank subsets according to influence using other notions of difference in the distribution. Henee, distributional LDS metric can detect methods that are off by large scale factor in their predictions.
To alleviate this limitation of influence functions on distributional influence tasks, we propose a method empirically normalise the Hessian approximation. Concretely, we do in way that doesn't require any retraining, unlike hyperparameter sweeps done to maximise an LDS score.
Concretely, we note that for Hessian approximation H to the Hessian H, for any vector v in column space of the Hessian, we would want:
23

---

H H
(55)
If the Hessian approximation is good approximation the Hessian, but is off by some scale factor, ie. aH av for some a, we can find a by trying to minimise:
Σ aft Ho; Ho, v, 12
(56)
for a set of vectors v; that we expect be in the column space of the Hessian. This is exactly our proposed normalisation method. For the set of vectors Ut. gradients, as we would expect them to be in the column space the true Hessian (otherwise, the response would diverge training goes on, shown in our theory section). We can compute Hv; a Hessian-vector product- - relatively cheaply even for large models, at roughly the cost of forward- backward pass, by using torch. hvp. Lastly, Eq. (55) so can be solved analytically. Hence, we don't I need to run optimisation to find the normalisation factor a. At the end, we simply multiply the Hessian approximation by the normalisation factor a to get the normalised Hessian approximation. We see minor improvements in the distributional LDS scores from using the normalisation factor, but we observe that the normalisation factor is necessary for the predicted changes in distribution by EK-FAC influence to look visually reasonable.
the per-datapoint training loss
second-degree polynomial in a, and
C.2.3 Identifying examples responsible for high-variance on MNIST
Query point
Change in mean
5
Remove top-2 by variance decrease
Top by variance derrease:
3
H
9
1
3 1op
Remuove top-2 by mean change
Top by
Log
"up
1op
1bp
20p
/
5
3
3
Top
Change
variance
Figure 11: d-TDA for MNIST. d-TDA with influence functions (see Section 4) successfully deter- mines which training examples to remove for decrease in measurement variance. These differ to examples identified for a change in mean. Different d-TDA variants capture diverse information about the training data. This experiment uses a multi-layer perceptron (MLP) trained on MNIST.
C.2.4 Data Pruning Results
24

---

Train on data
Remove 5000 that decrease loss for this model
Distributional: Remove 5000 that decrease loss on average
0.50
0.55
0.60
0.65
0.70
0.75
0.80
Validation Loss
Validation loss improvements on CIFAR- 10 with SWIN Vision Transformer from IF data pruning. For matching results for accuracy see Figure 6. We compare two approaches subset selection: 1) traditional TDA with fixed seed, where for each random seed we remove 5000 datapoints that are predicted to decrease the validation loss the most for the model trained with that fixed seed; and 2) distributional-TDA. where for each model remove 5000 datapoints predicted to decrease the validation loss the most on average. Both methods lead to validation loss improvements upon the baseline trained with all data, but d-TDA leads greater average improvement. Black dots show accuracies for individual models (with seeds indicated in gray), whereas coloured diamonds indicate the average result for each method.
Figure 12:
D Related Work
Influence Functions Influence functions originally proposed as method for data attribution in deep learning by [2]. Later, [29] explored influence functions for investigating the effect of removing adding groups of data points. Further extensions proposed by [30] who explored utilising higher-order information
and [31], who aimed to improve influence ranking via re- normalisation (different from our normalisation). Initial works influence functions [2, 29] relied on using iterative solvers to compute the required wwverse-Hessian-veettee products. [3] later explored using EK-FAC
alternative solution efficiently approximate calculations with the inverse Hessian. [14] investigated the empirical limitations of influence functions for predicting changes in measurements in the leave-one-out setting, without taking into consideration the distributional aspects of the training algorithm. [11] also investigated the limitations of influence functions, and propose perspectives on what not counterfactual retraining they might mightatuually approximate. this paper, we propose alternative perspectives, which are truthful to the underlying goal of predicting outcomes of coutnerfactual retraining with data removed.
Unrolled differentiation Orthogonally, pointing out the limitations of influence functions, [8, 9, 10] have proposed to unrolled differentiation for computing influence instead. For SGD trajectories, one can apply the chain rule differentiation to obtain a closed-form formula for the unrolled differentiation response:
d0,(e) (E)
o e (0)) ..))) ve, (0.) =: TuD (57)
de
B
==
t=
K-FAC and EK-FAC The need for approximate computation with the training loss Hessian in deep learning is evident, and cronecker-Factoree Approximate Curvature (K-FAC) has been one of the best performing Hessian approximations in TDA that can be run on a large scale. K-FAC was origi- nally proposed by [32] to approximate the Fisher Information matrix for natural gradient descent. was initially formulated only for multi-layer perceptrons, but has since been generalised any architecture with linear layers with weight-sharing (which includes convolutional neural networks, recurrent neural networks, and transformers) by [33]. [34] introduced aigenvalue-correate K-FAC (EK-FAC), which corrects K- FAF by using the optimal diagonal approximation in the Kronecker- factored eigenbasis. This was originally done in the context of approximate natural gradient descent, but [3] later used EK-FAC in the influence function approximation study generalisation in large language models.
25

---

E Experimental Details
E.1 Settings
We work with the following training settings in the empirical investigations in this paper:
Concrete I MLP. In this setting, we train a multi-layer perceptron (MLP) on a (ID target) regression setting on the UCI Concrete dataset [35]. The MLP with an input size of 8. hidden dimensions of [128, 128, 1281, and GeLU activation functions, was trained using Stochastic Gradient Descent (SGD) with a learning rate of 0.03 and momentum of 0.9. We applied a weight decay of 10-5 and gradient clipping at 1.0. The model was trained for 580 iterations using mean squared error (MSE) loss function and a batch size of 32. The initial 58 iterations (10% of the total) are dedicated to a linear learning rate warmup from O. For all the retrained models with the data removed, we keep the same number of training iterations as the original model, no matter how much data is removed.
Concrete Tiny MIP. This setting is the same as the previous one, but we use smaller MLP with hidden dimensions of [64, 64], which enables exact Hessian inversion.
MNIST MLP. For the MNIST MLP setting, we train a multi- qulti-layer perceptron (MLP) on the MNIST dataset [36]. The MLP takes flattened 28 x 28 images (input size 784), has hidden dimen- sions of [512, 256, 128], and an output size of 10. The model was trained using SGD with a learning rate of 0.03 and momentum of 0.9. We applied weight decay of 10-3 3, The model was trained for 1560 iterations with a cross-entropy loss function and a batch size of 64, A linear learning rate warmup from 0 was applied for the initial 5% of the total iterations.
SWIN Vision Transformer CIFAR-10. We train SWIN Transformer described in [24] with 2 sets of blocks with 4 layers each, with channel dimensionality 128 + 128 and 128 + 256 respec- tively. We use attention head dimension of 32, with patch-size of 2, window-size of 4, and 30% dropout applied to the final layer (head). We train the model with Adam with learning rate of 10-4, weight decay of 10-1, linear warmup for the first 5% of training iterations, a cosine schedule, and total of 200000 training iterations with satch-size of 64. For the dataset, we use the full 50000 images from the train set of CIFAR-10.
Latent Diffusion Model | ArtBench. For training the model, we follow the ArtBench setting with Latent Diffusion Model as described in Appendix J in B. K. Mlodozeniec, R. Eschenhagen, J. Bae, A. Immer, D. Krueger, and R. E. Turner [7]. The only difference is that we train 5 models with different random seeds on the full dataset.
E.2 Inffuence computation
For influence functions computation, we rely on the following methods:
Exact Hessian We use curvlinops [37] to compute the exact Hessian for the training loss. We add a damping factor equivalent to weight-decay, so that the Hessian corresponds to the actual training loss with the l penalty.7
Block-diagonal Hessian We compute the full exact Hessian described above, but then extract the per-peraaee parameter (weights and biases of each layer) block-diagonal parts of the Hessian. The inverse of block-diagonal matrix is the block-diagonal matrix of inverses of each block, which allows to invert the Hessian block for each parameter separately. For this, we use the same solver and settings as for the exact Hessian.
Curvature (EK-FAC) [32, 33. can be viewed as a Kronecker-factonne factored approximation to the Hessian. We use the curvlinops [37] implementation of EK-FAC, with the slight modification to compute the pseudo-inverse rather than regular inverse, as our theory suggests we ought to do. This amounts to thresholding eigenvalues, and only inverting
EK-FAC Eigenvalue correctee
"This is so that the "trajectory length" will be roughly equivalent for trained and retrained models.
"Note that, without the l penalty term, the Hessian not in general be positive semi-definite when training has converged. Dealing with deeght-decay is a tacit detail that is not often mentioned in the literature. For inversion, we use the default pytorch mecudo-nnvers, solver torch. linalg. piny with relative tolerance of 1 and absolute tolerance of O.
26

---

the ones above certain threshold, while setting the ones below it to zero. This is because, due to numerical errors, 0 eigenvalues might actually be compute as very small values, which after inversion will dominate the inverse matrix spectrum. The threshold is set relative to the largest eigenvalue for each layer, and we set it to 10f times the largest eigenvalue by default. Just as for the exact Hessian, before taking the pseudo-inverse, we add a damping factor equivalent to weight-decay, so that the Hessian corresponds to the actual training loss with the l penalty.
Unrolled differentiation We compute exact unrolled differentiation with forward-mode automatic differentiation, by keeping track of the d0, terms during training, and computing the forward derivative through the optimiser update using foreard-mode autodiff (torch. fuwc. There is one such term for every datapoint (or group of datapoints considered) to be removed. In the case of stateful optimisers (like SGD with momentum Adam), also need to keep track of the derivative of the optimisier state at iteration with respect to e.
ar
E.3 Individual experiment details
Figure 2. For this figure, we train 50 MLP models on the UCI Concrete (see Section E.1 for architecture and training details). We remove a fixed randomly sampled subset of 10% of the training datapoints to obtain 2 for the 'retrained' models. We measure and plot the ID model output on the left and center plot. The 'predicted' measurements are computed using exact unrolled differentiation applied to the models trained on the full dataset.
To investigate the correlation between unrolled differentiation and exact Hessian influence functions, we restrict ourselves to the 'Tiny MLP' UCI Concrete setting described in Section E.1. This smaller setting allows us to compute the exact Hessian. For the top axis, we compute changes in measurement (model output) on 103 test points from UCI Concrete, and plot the Pearson correlation between the changes predicted by unrolled differentiation and exact Hessian influence. We also computed the correlation between exact Hessian influence functions and unrolled differentiation with changes to parameters project to lie within the span of the Hessian (assuming eigenvalues below relative tolerance are 0). The lines were virtually overlapping with the original correlation unrolled differentiation without the projection, and hence removing the null-space component doesn't t affect the results significantly.
Figure 3.
For the bottom axis, we compute the predicted changes in measurement (model output) for 103 different test points using (I) unrolled differentiation and (2) unrolled differentiation, but projecting the predicted change parameters onto the span of the Hessian (again, assuming anything below the relative tolerance of 10 is O eigenvalue). We report the Pearson correlation across the 103 test points between measurement changes computed with (1) and (2).
Transformer Data Pruning (Figure 6 & Figure 12). For the data pruning task, we train 10 models with 10 different random seeds on the full CIFAR- 10. We compute influence functions using EK-FAC [3, 38] with an adaptation to the EK-FAC implementation in the curvlinops library [37]. Concretely, use numerically stable pseudo-inverse as described in Section E.2. We then find 5000 datapoints to remove for each method in the following way:
Fixed-seed TDA. For each of the 10 models trained with different random seeds, we influence functions identify different 5000 datapoints to remove for each model. We select the 5000 datapoints that are predicted to reduce the validation loss when removed the most for that model. When we retrain with the datapoints removed, we use the same random seed as for training the original full model.
Distributional (mean influence) TDA.
We identify one subset of 5000 datapoints to remove for all 10 models. Concretely, we find 5000 datapoints that are predicted to reduce the the validation loss the most on average over the 10 models. We then retrain each of the 10 models with that same subset removed.
influence functions performed perfectly as classical TDA methods, identifying a separate subset to remove for each random seed should perform at least as well as picking one shared subset for all seeds. The fact that the latter performs better implies that influence functions are indeed better understood (and more accurate) d-TDA method.
27

---

As baseline, we compare against the 10 models trained on the full dataset. Naturally, this outper- forms removing 5000 datapoints random, and hence stronger baseline.
Distributional Inftuence for Latent Diffusion Models (Figure 7). To identify top influences for the latent diffusion model, we again use EK-FAC influence, this time without the numerically stable pseudo-inverse. Instead, we directly use the reference implementation open-sourced in [7], and use the same IF settings with damping as described in the appendix of [7] for the ArtBench setting. We compute distributional influence with 5 models trained with 5 different random seeds. For the fixed- seed TDA reference, apply influence functions to model only (with seed 0) and report the top influences for that seed.
E.4 Distributional LDS experiments
For the distributional LDS experiments in Section C.2, we subsample 20 datasets from the original dataset uniformly at random without replacement, each with 10% of the datapoints removed (c.f. 100 datasets and 50% examples removed in [12]). For each subdataset, we train 50 models with different seeds. To estimate the distribution of the fully trained model measurements, we also use 50 seeds, and we apply the d-TDA method produce an approximate sample from the models trained on subdatasets to each of the 50 fully trained models. To estimate mean, variance change and Wasserstein influences, we compute the mean differences, variance differences and Wasserstein distances for the empirical distributions of the measurements using the 50 seeds.
28

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
The abstract promises new lens training data attribution for stochastic algo- rithms, which Section 3 thoroughly discusses. Secondly, the abstract indicates that influence functions can be derived for deep learning in principled fashion using our framework, a promise we deliver on in Section 4.1.
Justification:
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
is fine to include aspirational goals as motivation as long 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]
The paper provides theoretical results, and explicitly states the assumptions on which those rely. The limitations and reach of these assumptions is discussed, as well as empir- ically evaluated.
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 are encouraged create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-spccification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover
29

---

limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
The key results of the paper presented in Theorem and Theorem 6 are accom- panied by an explicit list of assumptions, well complete proofs in the Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof' can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof' provided 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 exper- imental results of 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: Yes, the experimental details are provided in the appendix, and the code for all the experiments will be released upon acceptance.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is a dataset and/or model, the authors should describe the steps taken 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 it may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access the model. 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, other means that are appropriate to the research performed.
specific model and empirical evaluation,
While NeurIPS does not require releasing code, the conference does require all submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
1. If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
30

---

2. If the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
3. If the contribution is new model (e.g.. a large language model), then there should either be a way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-souree dataset or instructions for how to construct the dataset).
4. We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but 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 instructions faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
The code will be open-sourced upon acceptance, together with the instructions on how to reproduce the main experiments. Anonymised source-code is included in the supple- mentary for the review.
Justification:
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https//nips..c/pubiic
quiddssCooeSubmissiooooiic for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (htips:// nips.co
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 a subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify the training and test details (e.g., data splits, hyperpara- meters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: The key information critical to understanding the results is provided in the main paper body. The technical details pertaining to the exact setup are given in the appendix, and the code with instructions on how to run the experiments will be open-sourced upon acceptance.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
31

---

The full details can be provided either with the code, in appendix, or as supplemental material. 7. Experiment statistical significance
Question:
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yss]
Justification:
Confidence intervals and error bars are given and described where relevant.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or 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, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call 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 --sigma error bars, but one should state it. The authors should prefer- ably 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 of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text. 8. Experiments compute resources
Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Answer: [Yes]
Justification:
The appendix briefly describes the computational resources used for the experi- ments, although novelty of scalability not main focus of the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, 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 full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question:
Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics http:///euun...ccpoooocchhhiiitt
Answer: [Yes]
32

---

All the guidelines in the NeurIPS Code of Ethics were followed throughout the submission and writing process.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
The authors should make
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification:
Although the work is foundational in nature, the work discusses applications of
the kind methods we introduce a new perspective for in the introduction. For each application, we reference the papers that include the discussion of the societal impact of the corresponding application.
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 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 direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: The models used in the paper do not pose high risk for misuse. No novel datasets are released.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have a high risk for misuse or dual-use should be released
with necessary safeguards to allow for controlled use of the model, for example by
33

---

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 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 are the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
The experiments in the paper do rely on existing datasets, with the creators being adequately credited. The license terms for the use of the datasets are being respected. The paper does not use any models that are not publicly available.
Justification:
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 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, paperswitheode..com cot.cccaaccco/mmtatttttttsss has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of derived asset (if has changed) should be provided.
If this information the asset's 's creators.
available online, the authors are encouraged to reach out to
13. New assets
Question: Are new assets introduced the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: No new assets introduced.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details the dataset/code/modell as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file. 14. Crowdsourcing and research with human subjects
34

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, applicable, well as details about compensation (if any)?
Answer:
[NA]
Justification:
No crowdsourcing nor research with huma subjects was conducted in this paper.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with
human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, cura- tion, or other labor should be paid at least the minimum wage 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 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, non- standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification: Use of LLMs was not an important, original, or non-standard component of this research.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer LLM policy (ttp:///eerrp..cccooomteeees///f for what should or should not be described.
35


Reviews:
[{"review_id": "IsdkknoLvL", "reviewer": "Reviewer_gp8r", "summary": "The paper introduces Distributional Training Data Attribution (d-TDA), a new framework for attributing model behavior to training examples under stochastic training dynamics. Unlike classical methods like influence functions that assume deterministic training and focus on mean effects, d-TDA models the full distribution over model outputs induced by random initialization and batch sampling. The authors define distributional influence as the change in output distribution (e.g., mean, variance, Wasserstein distance) when a training point is removed. They propose an efficient approximation using unrolled differentiation, avoiding retraining by computing a first-order response that estimates distributional shifts. The paper also provides new mathematical justifications for influence functions, showing they arise as limiting cases of d-TDA under certain conditions. Empirical results demonstrate that d-TDA captures richer signals—such as uncertainty changes—that traditional methods miss. This work redefines data attribution by treating randomness as a feature, not a nuisance, and opens up scalable and principled analyses for deep learning models.", "questions": "Q1: Do you think the actual distribution of influence function is achievable in practice? \n\nQ2: Whether you think the estimation of $\\theta^*(D\\z_k)$ is always necessary, since we sometimes just need to estimate the training samples contribution to a specific prediction on a testing sample. For example, I am aware that paper which bypass such estimation for empirical convenience, which can locate the real problems in public data sets.\n\n> Debugging and Explaining Metric Learning Approaches: An Influence Function Based Perspective (NeurIPS 2022)\n\nIn their approach, they just find a $\\theta$ where a training sample and a testing sample are mostly likely to be correlated. \n\nQ3: I like your interpretation on the leave-one-out training, but what is the useful measurement for real-world experiment of data attribution?\n\nQ4: In practice, how your d-TDA can be significantly different from classical TDA to identify problematic training data samples?", "limitations": "It is overall a wonderful paper with solid theoretical contribution and insights. However, it seems to be limited in answering how this work can empirically support real-world application.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "# Strengths and Weakness\nStrength:\n+ highly theoretically interesting\n+ theoretically sound\n+ new insights into many aspects of influence function\n\n\nWeakness\n- the IF calculation to become even more computationally expensive\n- it is not clear how could we sample seeds for *actual* distribution, which leads to concern whether the ground-truth distribution is achievable\n- little evidence to show the distribution is really necessary for data attribution in practice", "was_revised": true, "final_justification": "I am still supportive for this work, which is interesting and insightful.", "rebuttal": "We thank the reviewer for their favourable comments on the text. We appreciate they describe our perspective on IFs and data attribution in deep learning as “highly theoretically interesting” and “theoretically sound” -- key goals of the paper. We are happy to address their concerns and some points of minor misunderstanding below, also emphasising the **extra transformer and diffusion experiments** added in response to their helpful feedback. \n\n> 1. the IF calculation to become even more computationally expensive\n\nWe stress that a core contribution of the paper is showing how standard influence functions (IFs) can already be interpreted as resampling within a distributional framework (Thms 2 and 3). Concretely, offsetting trained model weights by an IF enables one to (approximately) sample from the distribution of models trained with a modified dataset. Our paper makes this claim mathematically rigorous. Hence, our work shows that standard influence functions applied to a single model are *already distributional*. In this sense, d-TDA is not inherently more expensive than IFs: d-TDA simply provides a new interpretation of what IFs actually do in deep learning. Without this distributional interpretation, IFs must be derived with unrealistic assumptions on the loss function: e.g. convexity or the presence of a unique minimum (e.g. the PBRF approximation in (Bae et al., 2022)).\n\nThe distributional reinterpretation allows us to *in addition* define novel measures of distributional influence. In our minds, the primary utility of of distributional influence is for *evaluating* methods like IFs in deep learning. Since IFs are inherently distributional, we think distributional evaluations are appropriate to verify how well they approximate actual resampling. As the reviewer correctly notes, calculating e.g. Wasserstein influence does indeed require an ensemble of samples – as few as 5 suffices in some of our experiments. However, when evaluating TDA methods, it is already standard practice to average TDA results over a small number of models (Park et al., 2023). Hence, our algorithm does not actually substantially increase the cost compared to usual in these cases. \n\n> 2. It is not clear how could we sample seeds for actual distribution, which leads to concern whether the ground-truth distribution is achievable\n\nCould the reviewer kindly clarify what they mean by “how could we sample seeds for actual distribution”? If you are asking how we sample from the “ground-truth” training distribution trained on a given dataset: to obtain a sample, we change the seeds for the training sources of randomness – 1) the data batch ordering, 2) model initialisation, and 3) any other sources of stochasticity (e.g. dropout masks). This is what we report as grount-truth samples in e.g. Fig. 2.\n\n> 3. ‘Little evidence to show the distribution is really necessary for data attribution in practice.’\n\nThe reviewer is right to note that this initial work focuses on introducing the conceptual foundations of d-TDA and understanding its mathematical properties. Whilst we provide some empirical demonstrations, we do indeed defer a full–scale experimental evaluation of downstream applications as important future work. We again stress that the crux of our paper is about showing that influence functions inherently answer a distributional question; we do not argue that everyone should always estimate the whole distribution for every downstream task. \n\nThat said, we agree that the paper might benefit from a clearer demonstration of the downstream utility of distributional influence. Therefore, in response to the reviewer's feedback, **we added an extra dataset pruning experiment comparing distributional influence against regular influence on a subset selection task with SWIN Vision Transformers on CIFAR-10**. Our results show that, by using distributional (mean) influence, we can select a better subset of examples to reduce the test loss than when we compute influence on a single model with a fixed seed. Our distributional subset selection  removing 10% of data improves the final test loss from 0.725 to 0.509 (compared to 0.528 with non-distributional variant). We use 10 samples to estimate the distribution.\n\n## Questions:\n>Do you think the actual distribution of influence function is achievable in practice?\n\nYes. Our novel theoretical results show that you can approximately sample from the distribution over retrained models with as few as 1 training run and its corresponding IF (see Thm 2). Moreover, taking a small ensemble of models (~5 models), one can effectively capture changes to the distribution over model measurements (Fig 2 right). We again stress that training with multiple seeds is the default practice when evaluating methods like IFs (see Appendix C.2).\n\n> Whether you think the estimation of [influence on parameters] is always necessary, since we sometimes just need to estimate the training samples contribution to a specific prediction on a testing sample.\n\nThanks for the comment. We certainly agree that there is an important place in the literature for empirical debugging and explanation tools, like those referenced by the reviewer. Such techniques are sometimes framed as analogous to IFs when used for TDA. However, our paper is more concerned with the following mathematical question: how does the distribution over final model predictions change depending on the dataset? Whilst we agree that fast, heuristic alternatives to unrolled differentiation and IFs might suffice in certain down-stream applications, these techniques do not necessarily address this question. There are many cases where we really care about understanding how the model would behave if retrained, e.g. copyright concerns (Mlodozeniec et al., 2024) or understanding LLMs (Grosse et al., 2023). As such, we need to develop principled methods that aim to answer this question. In this sense, we consider these works to be orthogonal to our work.\n\n> I like your interpretation on the leave-one-out training, but what is the useful measurement for real-world experiment of data attribution?\n\nWe are pleased the reviewer appreciates our contribution that LOO is not inherently incompatible with IFs:  a distributional interpretation reveals that the problem is just a low signal to noise ratio (SNR), and with enough samples IFs can indeed approximate the LOO distribution. For evaluation purposes, we believe leave-many-out retraining (as is standard practice in the literature (Park et al. 2023)) is a useful alternative for real-world experiments that has a sufficient SNR. We report such experiments in Appendix C.2.\n\n> In practice, how your d-TDA can be significantly different from classical TDA to identify problematic training data samples?\n\nFig 3 already demonstrates an example of a datapoint whose exclusion does not substantially modify the mean model behaviour, but which does drastically modify the distribution over model outputs in e.g. the Wasserstein sense. We provide further examples for MNIST in Fig 4. Training examples that are visually and intuitively ‘influential’ are sometimes missed by mean TDA. **In response to this comment, we have now also added experiments comparing top Wasserstein shift vs mean shift samples on Latent Diffusion Models on 50000 Artbench images (256x256 resolution)**. Once again, we find that most important examples do differ between Wasserstein and mean influence.\n\n\n---\n\nWe again warmly thank the reviewer for their feedback. We hope we have resolved all their questions and concerns, but invite them to let us know if any remain. If satisfied with our clarifications and the extra experiments, we respectfully ask them to consider raising their score.\n\n## References\n(Park et al., 2023) TRAK: Attributing Model Behavior at Scale, https://arxiv.org/abs/2303.14186\n(Bae et al., 2022) If influence functions are the answer, then what is the question?, https://arxiv.org/abs/2209.05364\n(Mlodozeniec et al., 2024) Influence Functions for Scalable Data Attribution in Diffusion Models, https://arxiv.org/abs/2410.13850\n(Grosse et al., 2023) Studying Large Language Model Generalization with Influence Functions, https://arxiv.org/abs/2308.03296"}, {"review_id": "2DuiHuwXrz", "reviewer": "Reviewer_wFZD", "summary": "This paper focuses on the problem of training data attribution (TDA), often also called \"influence estimation.\" The goal of TDA is to estimate the impact of removing a training datapoint on the model parameters or some function of the model parameters. The paper notes that previous methods have ignored or glossed over the randomness present in training modern machine learning models via the random initialization and batch ordering. In particular, existing TDA methods focus on providing a point estimate that estimates the influence for a datapoint given a particular random seed. The paper proposes to instead compute *distributional* TDA (d-TDA or distributional influence), which computes the influence under all possible random seeds. In synthetic experiments, the paper shows that distributional influence can provide different answers than does regular influence. The paper then proves that influence functions (the typical method for estimating regular influence) are in fact a limit of distributional influence in that they approximate the d-TDA distribution.", "questions": "My main questions are:\n1. When is it practically meaningful to compute distributional influence?\n2. What is the paper's position on the use of IFs?\n3. How is the paper proposing to compute IFs? If the answer is the algorithm on lines 150-155, what exactly is going on in this algorithm?\n\nThank you!", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "I found this to be a technically sound paper that computes an interesting quantity: the distributional influence of a datapoint. As far as I'm aware, this is a novel contribution that hasn't been considered before. And the paper successfully demonstrates an example where the use of distributional influence results in different estimates than the use of regular influence, which is an important step for this type of algorithm being practically useful.\n\n**Overall motivation**\n\nOne of the biggest issues I see with the paper is that it's unclear why one would want to compute distributional influence. As far as I'm aware, one is generally interested in influence of a datapoint over the finished end model so that, e.g., one can wonder about whether a datapoint is negatively contributing to its performance (for the purposes of re-running the training) or if a datapoint was valuable to its performance (for the purposes of compensating data providers). I don't see a clear argument for anyone being interested in knowing how valuable a datapoint was for a hypothetical model with different number of layers, order of layers, activation functions, number of training epochs, or other hyperparameters. And I see random seed as just another hyperparameter that defines the particular model one gets after training. So why is it of interest to get results under some distribution of this hyperparameter? I don't think the paper currently makes this case, but instead points out that distributional influence is different, and then concludes that this means that it is meaningful. For example:\n\n- Figure 3 shows that regular influence functions capture different information than d-TDA. The paper concludes that \"This showcases that d-TDA captures meaningful information about the training data missed by previous notions of influence.\" I agree Fig 3 shows it captures different information, but I don't see why Fig 3 shows this information is *meaningful* for any tasks.\n- \"The ability to predict how different training examples impact the distribution over training runs may be practically useful. For example, one could identify examples to remove to reduce variance, improving model repeatability over training runs.\" Why is this practically meaningful? At first thought, variance over training runs sounds like a good thing because it allows us to find more local optima, some of which may be better than our initial training run. So why should a practitioner care if they can consistently get to the same local optimum using different random seeds? They already know how to get to that local optimum.\n\nI also have two other points about the motivation of the paper:\n\n- I'm not sure exactly what the takeaways from Theorem 3 are. I would say that Thm 3 shows that, for Boltzmann distributions, IFs are optimal to use. But why do I care about Boltzmann distributions? How does this relate to a user thinking of using IFs for some model in practice?\n- I can't tell what the paper's position on IFs is. Earlier in the paper, it seems to be that they're not good because they're not measuring distirbutional influence. But then, in the final paragraph before the conclusion, the paper seems to turn around and say that the whole paper motivates the use of IFs. So is the paper providing anything besides suggesting the use of IFs? Maybe its algorithmic contribution is to use unrolled differentiation, but I believe previous work has covered this (see Luo et al. (2023)).\n\n**Related work**\n\nThe work of Luo et al. (2023) is highly related to this paper and its discussion of unrolled differentiation. Luo et al. (2023) study the approximation of LOO in models fit with SGD, and essentially use similar unrolled differentiation techniques as proposed in Eq (6). So I think the quote from line 164 of the current paper, \"To the best of our knowledge, this is the first application of such techniques\" should be updated.\n\nThe paper notes that previous works like [20] argued that IFs are not good for predicting leave-one-out retraining for deep learning models. The paper claims that this is incorrect, and LOO is in fact just noisy, rather than IFs being bad. Is the claim here that d-TDA applied to the situations in [20] would work well? I think this is a major empirical claim that would be pretty impactful if true. But big claims need big evidence, and there's no discussion or empirical evidence backing this up; it feels like investigating this claim could be an entire paper on its own. I think this paragraph either needs to be removed or needs significant expansion.\n\n\n**Unclear things**\n\nI was unclear about what some of the experiments were trying to show, as well as found a few other places where things were unclear; I've detailed the issues below:\n\n- The algorithm proposed to compute d-TDA is on lines 150-155. I'm not clear on what the algorithm is. The key Step 3 says \"using these two sets of (correlated) samples...\", but what are the two sets? It seems like only step 2 describes a set of samples.\n- Figure 2 is supposed to be an initial empirical explanation of d-TDA, but I wasn't clear what Figure 2 is displaying. First, what is \"query sample\" in the legend? Are the models one-dimensional here? Because it seems like the leftmost figure is supposed to be plotting models. Next, the caption states that the middle figure shows that the exact retraining samples are closer to predicted samples than samples from the original model. But the center figure seems to only be showing two datapoints per color here. To show thing $a$ is closer to thing $b$ than is thing $c$, I would think we need three points on the plot ($a,b,c$) for each color. Finally, on the right, I'm not sure what each of the terms in the legend means.\n- Figure 4 also uses the term \"query point,\" and I'm not sure what it means here either. I tried searching the text for the word \"query\", but it's only used once: in the caption of Fig 2.\n- Figure 5 (top) claims to show measurements predicted by unrolled differentiation versus changes predicted by IFs. But I only see one set of lines here. Are the dashed lines supposed to be one unrolled differentiation or IFs? But then why do only some of the lines have corresponding dashed lines? Also why is the top line so thick?  Finally, I see only three lines, but there are five lines listed in the legend.\n- Figure 5 bottom looks like it has a bunch of things overlapping. Are these supposed to correspond to the different learning rates?\n- Sometimes $k$ seems to index datapoints and sometimes $i$ seems to index datapoints.\n\n\n**Theoretical development**\n\nI thought there were a few small issues with the theoretical development:\n\n- Assumption A1 says that the loss function gradients $\\nabla \\ell_k$ are \"Lipschitz continuous and bounded\". Bounded for each $k$? Or uniformly? Also, is it true for neural networks that their Hessian is bounded and Lipschitz?\n- \"If the [model] weights converge to a saddle point ... the final parameters will become sensitive to any infinitesimal perturbation of the loss function.\" I'm pretty sure this isn't true. Consider the loss $L(\\theta) = \\sum_{n=1}^N (x_n - \\theta)^3$. If $x_1 = \\dots = x_N = 0$, then there is a saddle point at $\\theta = 0$, and this is still a saddle point even after you remove any given $x_n$.\n- I didn't see any motivation for assumptions A4 and A5. Should we expect these to hold up in cases where we're typically using SGD? I see the discussion in 291-298 describes why they are useful for the proof, but why are they reasonable?\n- Duplicate assumption numbers are used -- there are two A1, A2, A3, A4, and A5's in the paper.\n- The second A2 states \"the nonzero eigenvalues of the Hessian are uniformly bounded away from zero.\" How could this be violated? The Hessian is of finite dimension, so the minimum non-zero eigenvalue is bigger than zero.\n- Theorem 3 uses the term \"minimum manifold,\" but I don't think this was defined.\n\n\n**References**\n\nYuetian Luo, Zhimei Ren, and Rina Barber. Iterative Approximate Cross-Validation. ICML. 2023", "was_revised": true, "final_justification": "The authors have, I think, addressed a lot of the issues brought up in the reviews. I think the re-framing of their work, as well as the new experiment showing the potential for d-TDA to provide practically meaningful guidance have really made the paper's contribution a lot more clear. So, I've updated my score to vote for an accept.", "rebuttal": "We thank the reviewer for thoughtful comments and for finding our work technically sound and novel. We address their main concerns below and highlight our **new diffusion and transformer experiments**, added in response to their feedback.\n# Motivation for d-TDA\n**Why is distributional influence interesting?** It is widely acknowledged that conventional TDA methods are limited by their inability to account for stochasticity in training (K et al., 2021, _Revisiting Methods for Finding Influential Examples_). It is difficult to distinguish changes to model behaviour due to training data modifications from changes due to sampling randomness. Furthermore, usual derivations of IFs do not account for non-convex losses with multiple minima (Bae et al., 2022, _If Influence Functions are the Answer, then What is the Question?_) which are typical in deep learning. In contrast, in the d-TDA setting IFs can be naturally derived without these restrictive assumptions.\n\n**Why not regular TDA with a fixed seed?** This is an important question that we will dedicate a discussion to in the paper. Even with a fixed seed, chaotic training dynamics can push training to a different parameter space region when perturbing a dataset. To aggrevate this, in standard deep learning training with a fixed batch-size, training on a subset completely offsets at what iteration each datum is presented, even with a fixed seed. There are no guarantees that in these settings IFs would give us an answer to the counterfactual retraining question with a fixed seed. This is why prior work relies on convexity assumptions, where everything converges to the same point. In contrast, our work rigorously shows that IFs do something meaningful _even in this challenging setting_. We just have to interpret IFs as approximate resampling instead.\n\nTo illustrate this, we will add an experiment where we vary the peak learning rate at the beginning of training, showing the break-down in correlation between IF prediction and retraining _with the same seed_. Nonetheless, we illustrate that IFs are good at approximately resampling from the target distribution in all cases. Furthermore, *our new vision transformer data pruning experiments* illustrate that distributional (mean) influence does better in practice than applying IFs to a single model with a fixed seed.\n\nWe also think that the distributional properties of training may be of *independent interest*. Practitioners might be interested in selecting the data to reduce the variance for the output for a given unlabelled input, or they might want to do data selection to ensure it leads to improvements in validation loss for 90% of training runs. Our d-TDA methods allow practitioners to curate the dataset to target both these desiderata, or yet more complicated distributional properties. **We demonstrate this with new Latent Diffusion and vision transformer experiments on 50000 Artbench (256x256 resolution) and CIFAR-10 datasets respectively.**\n\nWe agree that exploring the full scope d-TDA applications is an important future research direction. However, our paper's chief goal is to build the mathematical foundation of influence functions for deep learning through the lens of d-TDA.\n# Takeaways from theorems and position on IFs\n**Purpose of Thm 3.** Thm 3 shows that IFs are the optimal KL-divergence approximation to the d-TDA question if the fully trained model weights obey a Boltzmann distribution with the energy given by the loss – a reasonable approximation which arises from approximating the batching noise in SGD (Mandt et al., 2017, _Stochastic Gradient Descent as Approximate Bayesian Inference_). In this setting, IFs provide precisely the KL-divergence-optimal asymptotic transformation to map samples trained with the original dataset to the samples trained with the perturbed dataset. The isotropic Gaussian noise is a simplifying relaxation, akin to how prior works have studied convex losses (albeit arguably a weaker one). We include Thm 3 to demonstrate that IFs have another distributional motivation (similarly to Thm 2). By slightly simplyfing the setting compared to Thm 2, we obtain a stronger result: IFs are the optimal distributional approximation in terms of KL-divergence. It is not intended to be interpreted as practical guidance for training models.\n\n**Position on IFs (Q2).** The core motivation for this work is that IFs appear to work well in downstream tasks, but lack theoretical and conceptual motivation in deep learning due to stochastic training and non-convex losses. Our paper provides a justification for IFs in these settings by showing that they can be interpreted as _resampling_. Hence, our paper is pro-IFs in that *it provides new strong motivation for their use through links to unrolled differentiation*.\n\n## Minor points\n_Related work of Luo et al. (2023)._ Forward-mode unrolled differentiation in deep learning was proposed even earlier by (Franceschi et al., 2017, _Forward and reverse gradient-based hyperparameter optimization_), whom we cite. We claim that our paper is the first such application _to TDA_ (hence ‘first application [...] for the computation of the response’). We have clarified this; thanks.\n\n_Empirical evidence that LOO is just noisy._ Please see Fig. 6 and App. C.1, in which provide the suggested experiments. \n\n_Figures 2, 4 and 5_. We apologise that parts of the figures were unclear to the reviewer. The captions are dense due to the intial 9-page limit.\n\n_Fig 2_. We consider neural networks trained on UCI concrete, with either the full dataset or a subset. We consider 4 measurements: the 1D neural network output on 4 chosen ‘query samples’. Each is represented by a different colour. **Left:** we show actual outputs trained on a subset (x axis) plotted against against predicted outputs (y axis) for each measurement (colour). The prediction for the outputs trained on a subset (“retrained”) are made by applying unrolled differentiation to models trained on the full dataset. Each sample corresponds to a single random seed. Here, a perfect prediction would lie on the black $x=y$ line – meaning predicted matches ground-truth retraining. Most samples lie far away from the line. **Centre:** Repeats the left panel, but shows the mean over all random seeds. We see that the _mean of each measurement distribution predicted with unrolled differentiation_ (circles) actually lies close to the $x=y$ line, certainly closer than the mean of the original measurement distribution of a model trained on all the data (diamonds). **Right:** we show that the true retrained distribution (colour, line) is closer to the predicted retrained distribution (colour, shaded) than the distribution trained on the full dataset (black). We will make sure this is explained more thoroughly in the caption.\n\n_Fig 4_. As in prior work, ‘query point’ refers to an example (input) for which we compute the measurement function. We now introduce the term explicitly. \n\n_Fig 5_. As the y-label indicates, the top plot shows the Pearson correlation between IF-predicted and UD-predicted measurement changes (y-axis) against training time (x-axis). 5 differently-colored lines correspond to runs with different learning rates. Three smallest learning rates give similar correlations, so they overlap. We have made this visually clearer. The dashed lines are spurious and have been removed.\n\n> Also, is it true for neural networks that their Hessian is bounded and Lipschitz?\n\nThis is a common assumption in deep learning theory (_far_ weaker than convexity). For instance, it holds if we consider optimisation on a bounded domain (which is virtually always true in practice due to finite-precision arithmetic).\n\n> \"If the [model] weights converge to a saddle point ... the final parameters will become sensitive to any infinitesimal perturbation of the loss function.\" I'm pretty sure this isn't true.\n\nWe mean that a gradient flow trajectory converging to a saddle point will *in general* (not always) be unstable under infinitesimal perturbations to the loss function. One can construct careful perturbations for which this is not the case. To illustrate, consider $L(\\theta) = \\theta_1^2$. The set $\\{\\theta \\in \\mathbb{R}^2: \\theta_1=0\\}$ are all minima, but any infinitesimal perturbation $L(\\theta) + \\epsilon \\theta_2^2$ collapses the minimum set to a single point. We have now phrased this more carefully in the text.\n\n_Motivation for assumptions A4 and A5_. A4 is intuitive as it's necessary for unrolled differentiation to not diverge; if unrolled differentiation diverges then we certainly cannot expect it converges to IFs. Likewise, A5 is needed to make sure that minima cannot be arbitrarily flat so IFs themselves do not diverge. We stress A1-A5 are intended to describe sufficient conditions for unrolled differentiation to converge to IFs. Whilst reasonable for many deep learning applications, we certainly do not claim that they will hold universally. Indeed, situations where they do not hold may be of independent interest for understanding when IFs break down. We'll make sure to expand the discussion of the motivation of each assumption.\n\n_Bounding nonzero Hessian eigenvalues away from zero._ A5 refers to the nonzero Hessian eigenvalues on the set of critical points (‘minimum manifold’), rather than at a single location $\\theta$ – hence why the assumption is nontrivial.\n\n_Missing definition._ The set $\\mathcal{S}^g_\\mathcal{L}$, to which ‘minimum manifold’ refers, is defined in line 339; we made the connection clearer.\n\nWe thank you for your thorough feedback. If anything remains unclear, please let us know. We hope our clarifications and new experiments warrant recommending acceptance."}, {"review_id": "b5aunZwlTZ", "reviewer": "Reviewer_sTvu", "summary": "This paper introduces the idea of distributional data influence (d-TDA), which characterizes the potential distribution of the result model when removing some training samples from a stochastic training algorithm. The authors present an efficient d-TDA method based on unrolled differentiation and evaluate its ability in estimating the ground truth distribution change. Finally, the paper presents a lot of empirical analysis of the estimation results and provide a rigorous analysis on how the well-known influence function naturally arose in this setting.", "questions": "See weakness and some more questions:\n\n1. In line 150, typically how many trained models are needed to accurate estimate the overall distribution? I would like to clearly understand the estimator's cost.\n2. The current unrolled differentiation framework assumes SGD training, will the faithfulness of the result be severely damaged if trained with other optimizers such as AdamW or Muon?\n\nThis is a very interesting work, but my main concern still revolves around the practical usability of this framework comparing to the traditional non-distributional ones. If the author could add some experiments or comparisons to better illustrate this point I am happy to raise my score.", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n1. The paper generalizes training-data attribution from single-model effects to shifts in the entire distribution of models induced by training randomness, which is a natural and interesting setting.\n2. The exposition is clear and technically rigorous, with some insightful theoretical analysis.\n\n**Weakness:**\n1. Lack of numerical results: Most experiments are conducted only on toy regression tasks, and the biggest dataset/model tested is on the MNIST dataset. This raises question on how practical is the current methodology to the current large machine learning models.\n2. The paper currently does not make any experimental comparisons of the distributional method to standard (non-distributional) attribution (perhaps some downstream tasks such as dataset pruning or noisy-label detection?), so the claimed advantage remains speculative. If the new distributional setting offers little improvement while significantly increasing computational cost, its practical value remains doubtful despite its theoretical appeal.", "was_revised": true, "final_justification": "The paper introduces a novel data-influence framework called distributional data influence. The work is overall technically sound and rigorous. The evaluation can successfully demonstrate the advantage of d-TDA over traditional TDA with the promised experiments on ViT, though they are still not very comprehensive. The presentation of the paper can also be further improved, as I see many reviewers share the same confusion which is not well-answered in the current draft. Overall I think the merit of this paper is bigger than its limitations, and I thus give a score of 4.", "rebuttal": "We thank the reviewer for their thoughtful feedback, and are pleased that they identify the work as interesting, clear and technically rigorous. We address all their comments and questions in detail below. \n\n**Few numerical results.** The reviewer is right to note that our chief contributions are conceptual and theoretical, offering some preliminary empirical evidence via UCI concrete and MNIST but deferring more exhaustive experiments to future work. _Our main intention is to introduce and explore the mathematical foundations of d-TDA – especially, as a lens to understand the unreasonable effectiveness of IFs in deep learning_. Given how widely influence functions are being used in deep learning, and the limitations of current interpretations, we believe this is a major theoretical contribution. Future papers can further explore applications and improvements to IFs based on insights and evaluations suggested by framework.\n\nThat said, we do agree that the paper could benefit from some additional empirical results. For this reason, following the reviewer’s comments, we have added **additional experiments applying our d-TDA methods to transformers and diffusion models**, including the suggested **data pruning** task with vision transformers. For the vision transformer experiments, we show that distributional (mean) influence can better identify training datapoints that will improve the final test loss _fo a single model_ when removed from the training set. By removing 10% of datapoints, we can improve the test loss from 0.715 to 0.509 (compare to 0.528 for regular TDA with a fixed seed). For latent diffusion trained from scratch on a dataset of 50000 Artbench images (256x256 resolution), we compare top predicted Wasserstein shift and mean shift samples, illustrating that they affect the ground-truth retraining distribution in the predicted manner.\n\n**Comparison to standard (non-distributional) TDA**. Thanks for the comments. We again emphasise that our chief goals are to 1) introduce a new training data attribution framework that properly takes into account stochasticity, and 2) show how existing standard methods such as IFs can be properly motivated and used within this mathematically rigorous framework, in contrast to the non-distributional TDA setting.  Whilst we do spend some time comparing distributional influence to its regular counterpart (see especially Fig 7 in the appendix, where ‘mean influence’, albeit distributional, is the current standard practice), our principal focus is not new data attribution methods per se, but rather a novel way of thinking about existing and future training data attribution methods.\n\nThat said, we agree showing the practical limitations of applying IFs to a single model with a fixed random seed would be beneficial. We believe **the newly added data pruning experiments with vision transformers on CIFAR-10**, showing empirical improvements from using distributional influence, illustrate this point.\n\n**Questions**. Thanks for the great questions.\n\n*How many samples are needed for d-TDA?* One can adopt a distributional perspective with only 1 sample. In this case, the original sample is interpreted as distributed according to the measure for the original dataset. Meanwhile, the modified sample, obtained by offsetting with IFs or unrolled differentiation, is (approximately) distributed according to the measure for the new dataset. Through Thms 2 and 3, we make this statement mathematically rigorous for the first time. More pragmatically, if computing e.g. Wasserstein or mean influence, we typically use 5-100 samples, with full details for each experiment reported in the Appendix. Also note that practitioners often heuristically average over a small ensemble of IFs to improve performance. Hence, one could argue that we are really just suggesting a more rigorous way to use this existing set of samples, at essentially identical computational cost. \n\n*Does d-TDA work with other optimisers?* Yes. Unrolled differentiation can be extended to arbitrary optimisers, provided their updates are differentiable. We believe our Thm 2 can be extended in kind. In fact, following the reviewer’s comments, we have verified that this is the case for SGD with momentum. However, we leave a full exploration of the theoretical details to important future work, following standard practice in the literature by initially focussing on SGD (Bae et al, 2024). \n\nWe again thank the reviewer for their time and comments. With the above clarifications in mind, as well as the extra experiments and theory added in response to their comments, we respectfully ask that they consider raising their score and recommending acceptance. \n\n**References**.\n(Bae et al., 2024) Training Data Attribution via Approximate Unrolled Differentiation, https://arxiv.org/abs/2405.12186"}, {"review_id": "w9VIycI4Md", "reviewer": "Reviewer_A7Ts", "summary": "This paper introduces a new framework called distributional training data attribution for evaluating the influence of training data on machine learning models, particularly in the context of stochastic training methods.\nThe authors argue that traditional training data attribution methods are insufficient because they do not account for the randomness inherent in deep learning model training, which arises from factors like random weight initialization and stochastic batching. \nThis randomness means training on the same dataset can produce different models.\nThe core idea of d-TDA is to shift the focus from predicting a single, deterministic change in model output to predicting how the entire distribution of model outputs changes when a training example is removed or altered. \nThis approach allows for a more comprehensive understanding of influence, capturing effects on not just the mean outcome but also its variance or other distributional properties.\nKey contributions:\n1. The paper formally defines the problem of data attribution in stochastic settings as predicting changes in the distribution of model weights or outputs. \nThis new concept quantifies a training point's importance based on how its removal affects the output distribution, using metrics like the Wasserstein distance. \nThe paper demonstrates that this can identify influential examples missed by traditional mean-based methods.\n2. The paper shows how unrolled differentiation can be used to estimate changes in the output distribution without the need for expensive retraining.\n3. The paper rigorously shows that classical influence functions, a popular but poorly understood tool in deep learning, naturally emerge from the d-TDA framework under certain assumptions. \nIt proves that IFs can be seen as a long-training-time limit of unrolled differentiation and as an optimal way to map between Boltzmann distributions representing the model before and after data removal. \nThis provides a new theoretical grounding for their effectiveness in non-convex settings.", "questions": "See weaknesses", "limitations": "yes", "rating": 5, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strengths:\n1. The primary strength of this paper is its introduction of d-TDA. This is a fundamentally new and important way to think about data attribution that directly confronts the stochasticity inherent in modern machine learning. \nInstead of treating randomness as a nuisance to be averaged away, the authors reframe it as a central object of study. \nThis conceptual shift from predicting a deterministic change to predicting a change in the distribution of model outputs is a significant and timely contribution to the field.\n\n2. The d-TDA framework allows for a much richer understanding of influence. \nA key finding is that some data points have a minimal effect on the mean of a model's output but can drastically alter its variance or other distributional properties. \nThis is demonstrated effectively in a 1D regression task where the example most influential by Wasserstein shift was different from the one most influential by mean shift, and its removal sharply increased prediction uncertainty. \nThis ability to identify variance-inducing examples is a powerful new tool for debugging models, improving their stability, and understanding generalization. \n\n3. A major contribution is providing a new, rigorous mathematical motivation for the surprising effectiveness of influence functions in deep learning. \nThe authors show from two complementary perspectives that IFs emerge naturally from the distributional framework. \nThese derivations are significant because they do not require the restrictive convexity assumptions of traditional IF theory, making them far more applicable to non-convex deep learning models. \n\nWeaknesses:\n1. The paper claims that its framework, particularly the use of IFs, unlocks d-TDA for large-scale models like diffusion and language models. However, the experiments presented are on smaller-scale problems (UCI datasets, MNIST with an MLP). \nWhile the efficient forward-mode computation of the unrolled response is a notable contribution to efficiency, the paper does not empirically demonstrate that the full d-TDA approach is computationally feasible or provides meaningful insights on the massive models where attribution methods are arguably most needed. \n\n2. The theoretical results, particularly Theorem 2 regarding the convergence to IFs, rely on a set of assumptions (A1-A6). \nWhile the authors argue these are less restrictive than prior work, their applicability to all practical deep learning scenarios could be questioned. \nFor example, A6 assumes that gradient flow trajectories converge \"sufficiently fast\" to local minima. \nThe extent to which the complex, high-dimensional, and often chaotic dynamics of large neural network training satisfy these conditions remains an open question. \nThe paper acknowledges and discusses these assumptions, which is good practice, but they remain a limitation on the proven generality of the claims. \n\n3. The paper presents two distinct mathematical motivations for IFs, one based on the limit of SGD dynamics and the other on transport maps between Boltzmann distributions. \nWhile both are elegant and converge on the same result, the paper does not explicitly discuss the relationship between them. \nAre they simply parallel justifications, or is there a deeper connection between the long-term behavior of SGD and the properties of the Boltzmann distribution that this work implicitly relies on? \nClarifying this would strengthen the theoretical narrative.\n\nMissing citations:\n1. The journey, not the destination: How data guides diffusion models. Arxiv 2024\n2. Intriguing Properties of Data Attribution on Diffusion Models. ICLR 2024", "was_revised": true, "final_justification": "I thank the authors for their comprehensive rebuttal.\n\nThe newly added experiments on Vision Transformers and diffusion models strengthen the paper.\n\nThe clarifications on the theoretical assumptions and the relationship between Theorems 2 and 3 are also well-received.\n\nI would like to maintain my positive score.", "rebuttal": "We thank the reviewer for their detailed comments and praise of the paper. We are happy to address their minor concerns below. We also politely draw their attention to our **extra experiments**, added in response to their comments.\n\n**Applicability to transformers&diffusion models**. The reviewer correctly notes that our chief contributions are conceptual and theoretical, introducing d-TDA, explaining how it can be operationalised with unrolled differentiation and IFs, and rigorously motivating IFs in a deep learning context within this framework with Thms 2 & 3. Whilst we provide a preliminary empirical exploration (see especially Figs 2, 3, 4 and 7), a full study of our methods’ applicability to larger models like LLMs must be left to future work. \n\nThat said, we do agree that our messaging might be clearer with experiments on different, bigger architectures. For this reason, we have added **additional experiments applying our d-TDA methods to transformers and diffusion models**.\nFor the vision transformer experiments, we show that distributional (mean) influence can better identify training datapoints that will improve the final test loss _fo a single model_ when removed from the training set. By removing 10% of datapoints, we can improve the test loss from 0.715 to 0.509 (compare to 0.528 for regular TDA with a fixed seed). For latent diffusion trained from scratch on a dataset of 50000 Artbench images (256x256 resolution), we compare top predicted Wasserstein shift and mean shift samples, illustrating that they affect the ground-truth retraining distribution in the predicted manner. Thanks for prompting this; we hope this addition strengthens the paper in the reviewer's eyes.\n\n**Assumptions.** The reviewer is right to observe that the assumptions used for our theoretical work are much less restrictive than those commonly used to derive IFs in the literature (namely, convexity). They are also very standard for theoretical work involving SGD. Whilst we do not claim that they will always be universally applicable for deep learning, we suggest that they are actually remarkably general. For instance, bounded loss gradient assumptions could be enforced by considering optimisation on a bounded domain, which is the case in practice when using finite-precision floating-point numbers. We also emphasise that A1-A6 should be interpreted as sufficient conditions for unrolled differentiation to converge to IFs. There could be pathological situations where e.g. SGD itself does not converge (in which cases, deep learning training wouldn't perform), in which case our theoretical results will certainly not hold. We respectfully suggest that this should not be construed as a weakness of our analysis – it is more of a reflection of the mathematical character of the problem. \n\n**Relationship between Thms 2 and 3**. Thanks for the question. These are parallel justifications with some relation to one-another. Roughly speaking, the relationship is that, after running SGD with a small learning rate for a sufficiently long time, the distribution over final model weights will converge towards a set of delta-functions at the loss function minima (~ a low temperature Boltzmann distribution, under some assumptions on the batching noise). In this case, when the loss is perturbed, the best possible transport map (in terms of KL-divergence) will simply translate these delta functions to their new locations. Thm 3 shows that this best map is given by influence functions. Thm 2 shows that unrolled differentiation converges to the same IF formula. Hence, when training for sufficiently long we get that ‘optimal transport map’ ≈ ‘influence functions’ ≈ ‘unrolled differentiation’. We will add this discussion to the text; thanks for prompting this. \n\n**Citations**. Thanks for the extra references, which we have now added. \n\nOnce again, we warmly thank the reviewer for their comments. If the reviewer thinks the newly added experiments and explanations strengthen the paper, we respectfully ask that they consider raising their score."}]

Author rebuttal:

Meta-review: The paper introduces distributional training data attribution (d-TDA), a new framework for data attribution that explicitly accounts for the stochasticity in deep learning training. Traditional methods typically ignore the randomness from initialisation and data batching, providing a single point estimate of an example's influence. In contrast, d-TDA aims to predict how the entire distribution of model outputs changes when a training example is removed. The authors demonstrate that this approach can identify influential examples missed by conventional mean-based methods, such as those that primarily affect model variance. A key theoretical contribution is a new mathematical justification for influence functions (IFs). The paper shows that IFs emerge naturally from the d-TDA framework as a long-training-time limit of unrolled differentiation, without requiring the restrictive convexity assumptions of prior work.

Strengths. The paper's primary strength is its novel and well-motivated conceptual reframing of the data attribution problem. It addresses the fundamental issue of stochasticity in modern model training. The theoretical work is rigorous and provides a significant new justification for the effectiveness of influence functions.

Weaknesses. The initial submission's primary weaknesses were a lack of empirical validation on large-scale models and some ambiguity in its core message. Reviewers found it unclear whether the paper proposed a new practical algorithm to replace influence functions or a theoretical device to justify them. These issues were substantially addressed during the rebuttal period.

Discussion. The discussion period was smooth and highly constructive. Reviewers raised concerns about the limited scale of the experiments and the clarity of the paper's contributions. The authors added new experiments on ViTs and diffusion models. They also agreed to sharpen the paper's messages. 

Why accept? New theoretical grounding for influence functions. The authors were responsive to feedback. Identifying and addressing a novel problem.

PS. "A Bayesian Approach To Analysing Training Data Attribution In Deep Learning" from NeurIPS 2023 is also highly relevant prior work -- recommend reading it.
Decision: Accept (spotlight)