Influence Functions for Scalable Data Attribution in Diffusion Models
Authors: Bruno Kacper Mlodozeniec
Runa Eschenhagen
Juhan Bae
Alexander Immer
David Krueger
Richard E. Turner
Venue: ICLR 2025 Oral (2025)
Keywords: diffusion models
influence functions
Generalised Gauss Newton
GGN
data attribution
Hessian approximation
interpretability
curvature
Kronecker-Factored Approximate Curvature
K-FAC
Primary area: interpretability and explainable AI
TL;DR: We present a method for attributing the influence of training data on diffusion model’s output by adapting influence functions and a KFAC approximation for diffusion models, and we explore what measurements we want to attribute for in the first place

Abstract: Diffusion models have led to significant advancements in generative modelling. Yet their widespread adoption poses challenges regarding data attribution and interpretability. In this paper, we aim to help address such challenges in diffusion models by extending influence functions. Influence function-based data attribution methods approximate how a model's output would have changed if some training data were removed. In supervised learning, this is usually used for predicting how the loss on a particular example would change. For diffusion models, we focus on predicting the change in the probability of generating a particular example via several proxy measurements. We show how to formulate influence functions for such quantities and how previously proposed methods can be interpreted as particular design choices in our framework. To ensure scalability of the Hessian computations in influence functions, we use a K-FAC approximation based on generalised Gauss-Newton matrices specifically tailored to diffusion models. We show that our recommended method outperforms previously proposed data attribution methods on common data attribution evaluations, such as the Linear Data-modelling Score (LDS) or retraining without top influences, without the need for method-specific hyperparameter tuning.

Published as a conference paper at ICLR 2025
INFLUENCE FUNCTIONS FOR SCALABLE DATA ATTRIBUTION IN DIFFUSION MODELS
Bruno Mlodozeniec*.4; 2,2 Alexander 1mmer2.5
Runa Eschenhagen
Juhan Bae 4
David Krueger
Richard Turner
Department of Engineering, University of Cambridge, UK
2 Max Planck Institute for Intelligent Systems, Tübingen, Germany
Department of Computer Science University of Toronto, Canada Veecoor Institute, Toronto, Canada
Department of Computer Science, ETH Zurich, Switzerland 6 Mila - Quebec AI Institute, Montreal, Canada The Alan Turing Institute, London, UK
ABSTRACT
Diffusion models have led to significant advancements in generative modelling. Yet their widespread adoption poses challenges regarding data attribution and interpretability. In this paper, we aim to help address such challenges in diffusion models by developing an influence function framework. Influence function- baaed data attribution methods approximate how model's s output would have changed if some training data were removed. In supervised learning, this is usually used for predicting how the loss a particular example would change. For diffusion models, we focus on predicting the change in the probability of generating a particular example via several proxy measurements. We show how to formulate influence functions for such quantities and how previously proposed methods can be interpreted as particular design choices in our framework. To ensure scalability of the Hessian computations in influence functions, we systematically develop K-FAC approximations based on generalised Geuss-Newton matrices specifically tailored to diffusion models. We recast previously proposed methods as specific design choices in our framework, and show that our recommended method outperforms previous data attribution approaches on common evaluations, such as the I Linear Data-modelling Score (LDS)) or retraining without top influences, without the need for method-speciff hyperparameter tuning.
INTRODUCTION
Generative modelling for continuous data modalities - like images, video, and audio - has advanced rapidly, propelled by improvements in diffusion- approaches. Many companies now offer easy access to Al Al-generated bespoke image content. However, the use of these models for commercial purposes creates a need for understanding how the training data influences their outputs. In cases where the model's 's outputs are undesirable, it is useful to be able to identify, and possibly remove, the training data instances responsible for those outputs. Furthermore, as copyrighted works often make up significant part of the training corpora of these models (Schuhmann et al., 2022), concerns about the extent to which individual copyright owners' works influence the generated samples arise. Some already characterise what these companies offer "copyright infringement as service" (Saveri & Butterick, 2023a), which has caused a flurry of high-profile lawsuits Saveri & Butterick (2023a;b). This motivates exploring tools for data attribution that might be able to quantify how each group of training data points influences the models' outputs. Influence functions (Koh & Liang, 20;7b; Bae
"Correpoondence to: bkm2 80cam .ccuk Source code available https: / / thub. whbb.ccm//rrunoow//////o ffusion-iii inniiinlleen fuence

---

Published as a conference paper at ICLR 2025
et al., 2022) offer precisely such tool. By approximating the answer to the question, "If the model was trained with some of the data excluded, what would its output be?", they can help finding data points most responsible for low loss on example, or high probability of generating a particular example. However, they have yet to be scalably adapted to the general diffusion modelling setting.
Influence functions work by locally approximating how the loss landscape would change if some of the training data points were woinhwei in the training loss (illustrated in Figure 4). Consequently, this enables prediction for how the (local) optimum of the training loss would change, and how that change in the parameters would affect a measurement of interest (e.g., loss on a particular example). By extrapolating this prediction, one can estimate what would happen if the data points were fully removed from the training set. However, to locally approximate the shape of the loss landscape, influence functions require computing and inverting the Hessian of the training loss, which is computationally expensive. One common approximation of the training loss's 's Hessian is the generalised -eeuto matrix (GGN, Schraudolph, 2002; Martens, 2020). The GGN has not been clearly formulated for the diffusion modelling objective before and cannot be uniquely determined based on its general definition. Moreover, to compute and store a GGN for large neural networks further approximations
necessary. We propose using Cronecker-Factoree Approximate Curvature (K-FAC, Heskes, 2000; Martens & Grosse, 2015) and its variant arrenvalue-correctee K K-FAC (George al., 2018, EK-FAC) approximate the GGN. It is not commonly known how apply neural network architectures used in diffusion models; for example, Kwon et al. (2023) resort to alternative Hessian approximation methods because "[K-FAC] might not be applicable to general deep neural network models as it highly depends on the model architecture". However, based on recent work, it is indeed clear that it can be applied to architectures used in diffusion models (Grosse & Martens, 20 Eschenhagen et al., 2023), which typically combine linear layers, convolutions, and attention (Ho et al., 2020).
In this work, we describe a scalable approach to influence function based approximations for data attribution in diffusion models, using FFAC-FAC approximation of GGNs as Hessian approximations. We articulate design space based on influence functions, unify previous methods for data attribution in diffusson models (Georgiev et al., 2023; Zheng et al., 2024) through our framework, and argue for the design choices that distinguish our method from previous ones. One important design choice is the GGN used as the Hessian approximation. We formulate different GGN matrices for the diffusion modelling objective and discuss their implicit assumptions. We empirically ablate variations of the GGN approximation and other design choices in our framework and show that our proposed method outperforms the existing data attribution methods for diffusion models as measured by common data attribution metrics like the Linear Datamodeling Score (Park et al., 2023) or retraining without top influences. Finally, we also discuss interesting empirical observations that challenge our current understanding of influence functions in the context of diffusion models.
BACKGROUND
This section introduces the general concepts of
diffusion models, influence functions, and the GGN.
2.1 DIFFUSION MODELS
Diffusion models a class of probabilistic generative models that fit model po(x)) parameterised
by parameters 0 E Rd to approximate a training data distribution q(x), with the primary aim being to sample new data (-(-) (Sohl- Dickstein et al., 2015; Ho et al., 2020; Turner et al., 2024). This is usually done by augmenting the original data æ with T fidelity levels as (U:T) [zx(0) (P))
with an augmentation distribution ((x:(0:T) that satisfies the following criteria: 1) the highest fidelity y²()) equals the original training data q(x:(0)) = q(=), 2) the lowest fidelity (T)) has distribution that is easy to sample from, and 3) predicting lower fidelity level from the level directly above it is simple to model and learn. To achieve the above goals, q is typically taken to be first-order Gaussian regrssegressive (diffusion) process: (z:=()) = (A,((t (1
T)) with hyperparameters d set so that the law of a!(T) approximately matches standard Gaussian distribution. (0, I). In that case, the reverse conditionals ((t ((_((t1([_'()77))
q( (()))) are first-order Markov, and if the number of fidelity levels T is high enough. they can be well approximated by diagonal Gaussian, allowing them to be modelled with a parametric model with
2

---

Published as a conference paper at ICLR 2025
Generated sample
Top influences
Most neutral
Negative influences
Figure I: Most influential training data points as identified by K-FAC Influence Functions for samples generated by a denoising diffusion probabilistic model trained on CIFAR-10 10. The top influences those whose omission from the training set is predicted to most increase the loss of the generated sample. Negative influences are those predicted to most decrease the loss, and the most neutral are those that should change the loss the least.
simple likelihood function, hence satisfying (3) (Turner et al., 2024). The marginals ((x(0) (z)( = N (2(1)) (IT== )v) x0())
(1- (I-IT== =1 x 1) also have a simple Gaussian form, allowing for the
augmented samples to be sampled as:
1/2 (00)
(0) (0)
II. X?.
with c()) ~ N).
(1)
Diffusion models are trained to approximate the reverse conditionals pa (=( = ((1((t1))/2(0)
by maximising log- log-probabilities of samples (t(-1) conditioned on (*),, for all timesteps t = I.....
We can note that (xx01- (aa(t-1)aaf (2((1) (0))) has a Gaussian distribution with mean given by:
1 de
1- X7
with (t) der (x(9) II (0)
M4 (4-1)1.(( (())
(1)
-
(1 II= 1²/2
as in Equation (1). In other words, the mean is a mixture of the sample (+!) and the noise (!)) that was applied to (0) to produce it. Hence, we can choose to analogously parameterise po (x(x(t1))(<(0))) (N(x(tt-))m-- t.o (2:(1), (c()),c(x((()) Tf.). That way, the model <( (x(()) simply predicts the noise (1) that was added to the data to produce (1)), The variances of are usually chosen as hyper- parameters (Ho et al., 2020). With that parameterisation, the negative expected log-likelihood Bycur-i,= -0000000))) [- log (z(()))))) up to scale and shift independent of 0 or (:()), can be written as (Ho et al., 2020; Turner al., 2024):²
e ~N(00T)
(0, x (0) = (000
-6
s (x(1)
10 = -II P== (o) (0)
I1-II..-
x
(2)
This leads training loss é for the diffusion model (x(()() that is sum of per-diffusion timestep training losses:?
((0, x) = Ey (0,
i~ Uniform([ZT).
The parameters are then optimised to minimise the loss averaged over a training dataset D=( 1.
8* (D) arg min (D) LD
der
((0, xn)
N
Note that the two random variables (1) Equivalently,, weighted of per-timestep negative log-likelihoods - log pe (((1))))(())
; (*) are deterministic functions of one-another.
3

---

Published as a conference paper at ICLR 2025
Other interpretations of the above procedure exist in the literature (Song & Ermon, 2020; Song et al., 202ib;a; Kingma et al., 2023).
2.2 INFLUENCE FUNCTIONS
The aim of influence functions is to answer questions of the sort "how would my model behave were trained on the training dataset with some datapoints removed". To do so, they approximate the change in the optimal model parameters in Equation (3) when some training examples (z,) (z))ezz, I= - iM) C [N], are removed from the dataset 'D. To arrive at a tractable approximation, useful consider continuous relaxation of this question: how would the optimum change were the training examples (rz)jez down-weighted by E in the training loss:
N
1 N
r r-z(E) = arg min
((0, xn) - > 2z) Σ
JEI
The function r-I R (well-defined if the optimum is unique) is the response function. Setting & to 1/4 recovers the minimum of the original objective in Equation (3) with examples (Fi removed.
Under suitable assumptions (see Appendix A), by the Implicit Function Theorem (Krantz & Parks, 2003), the response function is continuous and differentiable defined as linear approximation to the response function r-z by a first-order Taylor expansion around 8 0:
0. Influence functions can be
dr dr-i(==) de' ===00
=(=))=
+ o(())
(5)
- (D)
Σ (V))L Vo- ((0x(() (ox00(000-~"V00000000)))) IEI
+ o(s),
as 0. See Appendix A for formal derivation and conditions. The optimal parameters with examples (Ni)sez removed can be approximated by setting & to 1/3 and dropping the o(c) terms.
Usually, we are not directly interested in the change in parameters in response to removing some data, but rather the change some measurement function m0 (D), x²) 6') at a particular test input x' (e.g., per-example test loss). We can further make a first-order Taylor approximation to m(( x²) at (D) - m(0, 2') = m(f", x²) + Vi m(0 x (0) - 8*) + o (Il@ (2)) and combine it with Equation (5) to get simple linear estimate of the change in the measurement function:
m(r-z(s),z') m(0 2') +
Vi-m(0 x') (V2. CD())) Vg. (0=
(6)
JEZ
2.3 GENERALISED GAUSS-NEWTON MATRIX
Computing the influence function approximation in Equation (5) requires inverting the Hessian (0) E addeamnsdaaara In the context of neural networks, the Hessian itself is generally compu- tationally intractable and approximations are necessary. A common Hessian approximation the generalised Gauss-Newton matrix (GGN). We will first introduce the GGN in an abstract setting of approximating the Hessian for general training loss L(0) = Ez [p(0, 2)), to make it clear how different variants can be arrived at for diffusion models in the next section.
In general, if we have function p(e, 2) of the form h2 f0), with h2 convex function, the GGN for an expectation Ez [p(0, =)] is defined as
GGN(O) = E. ()) ( (h vofa())
where (8) is the Jacobian of far Whenever fa is (locally) linear, the GGN is equal to the Hessian Ez (2p((2=)) =)]. Therefore, we can consider the GGN as an approximation to the Hessian in which we "linearise" the function fz. Note that any decomposition of p(0, z) results in valid GGN as long as hy is convex (Martens, 2020). We give two examples below.
h1 typically required to be convex guarantee the resulting GGN a positive semi- cemi-defnnit (PSD) matrix. valid PSD approximation to the Hessian can be formed with non-convex ht as well; all the arguments about the exactness of the GGN approximation for a linear fa would still apply. However, the PSD property helps with numerical stability of the matrix inversion, and guarantees that the GGN will be invertible if small damping term is added the diagonal.
4

---

Published as a conference paper at ICLR 2025
Option 1. A typical choice would be for fa to be the neural network function training datapoint 2, and for h2 to be the loss function (e.g. E2-loss), with the expectation E, being taken over the empirical (training) data distribution; call the GGN for this split GGN The GGN with this split is exact for linear neural networks (or when the model has zero residuals on the training data) (Martens, 2020).
f= = mapping from parameters to model output LGGNNddel (0)
h- = loss function (e.g- f2-loss)
Option 2. Alternatively, different GGN can be defined by using a trivial split the loss p(0, x) into the log map h log and the exponentiated negated loss fz exp(-p(:, 2)), and again taking the expectation over the empirical data distribution. With this split, the resulting GGN is
f= == exp(-p(::z)
GGN loss (0) E Vop(O,2)VoP(()
(8)
h2 = log
This is also called the empirical Fisher (Kunstner et al., 2019). Note that GGN only equal to the Hessian under the arguably more stringent condition that exp( sap(-p(( p(:, 2)) the composition of the model and the exponentiated negative loss function is linear. This is in contrast to GGN for which only the mapping from the parameters to the model output needs to be (locally) linear. Hence, might prefer use GGN for Hessian approximation whenever have nonlinear loss, which is the case for diffusion models.
3 SCALABLE INFLUENCE FUNCTIONS FOR DIFFUSION MODELS
In this section, we discuss how we adapt influence functions to the diffusion modelling setting in scalable manner. We also recast data attribution methods for diffusion models proposed in prior work (Georgiev et al., 2023; Zheng et al., 2024) as the result of particular design decisions in our framework, and argue for our own choices that distinguish method from the previous ones.
3.1 APPROXIMATING THE HESSIAN
In diffusion models, we want to compute the Hessian of the loss of the form
(D Er,, [e(0, xr,)] Ex,. [Eg ( [Ile() +((=))
where En., ||] = ( Σ -) is the expectation over the empirical data distribution. "We will describe how formulate different GGN approximations for this setting.
3.1.1 GGN FOR DIFFUSION MODELS
Option 1. To arrive at GGN approximation, discussed in Section 2.3, can partition the function
0 - ([c(0) - c (1)) into the model output 0 + eg (o) and the Eo-loss function ||c(*) -11 This
results in the GGN:
Ja === (1)
+ +GGNNodel (0) = Ex,
EG
E
VJel (a(1) (a())(11)) (21) DEg (1()
(9)
h2 = (ce(0] .12
where I is the identity matrix. This correspond to "linearising" the neural network cg. For diffusion models, the dimensionality of the output of e is typically very large (e.g. 32 x23322 x for CIFAR), so computing the Jacobians Vorh explicitly is still intractable. However, we can express GGN D as
(f) = Exp
E
E
Erao (9) (0) [6]
Good ~
x
.1)
(10)
where ga (!) V TollCmod - og (11 E Rdbasss; see Appendix B for the derivation. This formulation lends itself to a Monte Carlo approximation, since we can now compute gradients using auxiliary
*Generally, Rz,, might also subsume the expectation over data augmentations applied to the training data points (see Appendix J.9 for details how this handled).
5

---

Published as a conference paper at ICLR 2025
targets Good sampled from the model's output distribution, as shown in Equation (10). FD can be interpreted as a kind of Fisher information matrix (Amari, 1998; Martens, 2020), but it is not the Fisher for the marginal model distribution po(x:).
Option 2. Analogously to Equation (8), we can also consider the trivial decomposition of ((,, x) into the log and the exponentiated loss, effectively "linearising" exp(-((;,x)) The resulting GGN is:
f= ((. hg == - log
+ -GNN = Ez, [Voe(0, Jn
(11)
where z() is the diffusion training loss defined in Equation (2). This Hessian approximation GGN D turns out to be equivalent to the ones considered in the previous works on data attribution for diffusion models (Georgiev et al., 2023; Zheng al., 2024; Kwon al., 2023). In contrast,
in this work, we opt for GGN in Equation (9), or equivalently FD, since it is arguably metter-motivated approximation of the Hessian than GGN (c.f. Section 2.3).
In Zheng et al. (2024), the authors explored substituting different (theoretically incorrect) training loss functions into the influence function approximation. In particular, they found that replacing the loss ||e(*) - )x) with the square norm loss ||e (=) (effectively replacing the "targets" e()) with 0) gave the best results. Note that the targets ( do not appear in the expression for GGN in Equation (9).5 Hence, in our method substituting different targets would not affect the Hessian approximation. In Zheng et al. (2024), replacing the targets only makes difference to the Hessian approximation because they use GGN (an empirical Fisher) to approximate the Hessian.
3.1.2 (E)K-FAC FOR DIFFUSION MODELS
While (0) and GGN do not require computing full Jacobians or the Hessian of the neural network model, they involve taking outer products of gradients in Re which is still intractable. Kronecker- -Fcctored Approximate Curvature (Heskes, 2000; Martens & Grosse, 2015, K-FAC) is a common scalable approximation of the GGN overcome this problem. It approximates the GGN with a block-diagonal matrix, where each block corresponds to one neural network layer and consists of Kronecker product of two matrices. Due to convenient properties of the Kronecker product, this makes the inversion and multiplication with vectors needed in Equation (6) efficient enough to scale large networks. K-FAC is defined for linear layers, including linear layers with weight sharing like convolutions (Grosse & Martens, 2016). This covers most layer types in the architectures typically used for diffusion models (linear, convolutions, attention). When weight sharing used, there are two variants- - K-FAC-expannd and K-FAC-reduce (Eschenhagen et al., 2023); see Appendix C. for an overview. For the parameters Or of layer the GGN FD in Equation (10) is approximated by
N
FD((1) ~
E
E
1 (D) aaa
X
E
(12)
N2
with !!! gdt being the inputs to the 4th layer for data point R and 6 R being the gradient of the E2-loss w.r.t. the output of the /th layer, and @ denoting the Kronecker product. The approximation trivially becomes an equality for single data point and also for deep linear networks with loss (Bernacchia et al., 2018; Eschenhagen et al., 2023).
For our recommended method, we choose to approximate the Hessian with K-FAC approximation of FD. akin to Grosse et al. (2023). We approximate the expectations in Equation (12) with Monte Carlo samples and use K-FAC-expand whenever weight sharing is used; in the case of convolutional layers this corresponds to Grosse & Martens (2016). See C.2 for the full derivation of K-FAC for diffusion models that also considers weight sharing. Additionally, we choose to use eigenvalue-corrected K-FAC (George et al., 2018, EK-FAC) in our experiments which improves performance notably and can be directly applied on top of our K-FAC approximation. Lastly, to ensure the Hessian approximation is cell-condi and invertible, we follow standard practice and add a damping term consisting of a small scalar damping factor times the identity matrix. We ablate all of these design choices in Appendices C.3 and G (Figures 5, 7 and 9).
suggested by Grosse et al. (2023)-
*This is because the Hessian of an fosso w.r.t. the model output is a multiple of the identity matrix. 'For the sake of simpler presentation this does not take potential weight sharing into account.

---

Published as a conference paper at ICLR 2025
3.2 GRADIENT COMPRESSION AND QUERY BATCHING
In practice, we recommend computing influence function estimates in Equation (6) by first computing and storing the approximate Hessian inverse, and then iteratively computing the preconditioned inner products Vg m( x)) (V). LD()))) Vg. (,-E((",2)) .2y)) for different training datapoints 2.. Following Grosse et al. (2023), we use query batching to avoid recomputing the gradients Va. 7=-((0" :zy) when attributing multiple samples
found that compression by quantisation works much better for diffusion models compared to the SVD-based compression used by Grosse et al. (2023) (see Appendix F), likely due to the fact that gradients Vol(0, x,,) are not low-rank in this setting.
We also use gradient compression;
3.3 WHAT TO MEASURE
For diffusion models, arguably the most natural question to ask might be, for a given sample generated from the model, how did the training samples influence the probability of generating a sample x?? For example, in the context of copyright infringement, we might want to ask removing certain copyrighted works would substantially reduce the probability of generating æ. With influence functions, these questions could be interpreted as setting the measurement function mn(0 æ) to be the (marginal) log-probability of generating from the diffusion model: log (z
Computing the marginal log-probability introduces some challenges. Diffusion models have originally been designed with the goal of tractable sampling, and not log- lillikiiiihodd evaluation. Ho et al. (2020); Sohl-Dickstein et al. (2015) only introduce a lower-bound on the marginal log-probability. Song et al. (2021b) show that exact likelliodo evaluation is possible, but it only makes sense in settings where the training data distribution has a density (e.g. uniformly dequantised data), and it only corresponds to the marginal log- likelihoodd of the model when sampling deterministically (Song et al., 2021a)." " Also, taking gradients of that measurement, as required for influence functions, is non-trivial. Hence, in most cases, we might need a proxy measurement for the marginal probability. We consider a couple of proxies in this work:
1. Loss. Approximate log (a) with the diffusion loss ((0.x)) in Equation (2) on that particular example. This corresponds to the ELBO with reweighted per-timestep loss terms (see Figure 20).
2. Probability of sampling trajectory. If the entire sampling trajectory (0:T) that gen- erated sample x is available, consider the probability of that trajectory pa (x
p(x") I t= (o(=(t)))=(())
3. ELBO. Approximate log (z) with an Evidence Lower-Boume (Ho et al., 2020, eq. (5)).
4 EXPERIMENTS
Evaluating Data Attribution. To evaluate the proposed data attribution methods, we primarily focus on two metrics: Linear Data Modelling Score (LDS) and retraining without top influences. These metrics are described below. In all experiments, we look at measurements on samples generated by the model trained D. We primarily focus Denoising Diffusion Probabilistic Models (DDPM) (Ho et al., 2020) throughout. Runtimes are reported in Appendix E.
LDS measures how well given attribution method predicts the relative change in measurement as the model is retrained on (random) subsets of the training data. For an attribution method a(D, D xx) that approximates how a measurement m00* (D), æ)) would change if model was trained on an altered dataset D', LDS measures the Spearman rank correlation between the predicted changes in output a(D, D), a(D, 2)) and the actual changes in output m (Di), x..... mn(0 (DM), after retraining on M independently subsampled versions D, of the original dataset D, each containing 50% of the points sampled without replacement. However, training on a fixed dataset can produce different models with functionally different behaviour depending on the random seed used for the initialisation and data order during stochastic optimisation. Hence, for any given dataset D', different
7 Unless the trained model satisfies very specific "consistency" constraints (Song et al., 2021b, Theorem 2). 'Higher LDS values can be obtained when looking validation examples (Zheng et al., 2024), but diffusion models used primarily for sampling, so attrihuting generated samples of primary practical interest.
7

---

Published as a conference paper at ICLR 2025
measurements could be obtained depending the random seed used. To mitigate the issue, Park et al. (2023) suggest using an ensemble average measurement after retraining as the "oracle" target:
(a)D.
M i-1 ( E
K k=1 m(O
M
LDS = spearman
(13)
where (D') E Buuee are the parameters resulting from training on D'' with particular seed k. Retraining without top influences, on the other hand, evaluates the ability of the data attribution method to surface the most influential data points - namely, those that would most negatively affect the measurement m m(""(D'), .r) under retraining from scratch on dataset DD' with these data points removed. For each method, we remove a fixed percentage of the most influential datapoints from D to create the new dataset D', and report the change in the measurement mn0((D') x) relative to (()(D)) x) (measurement by the model trained on the full dataset D).
Methods. We compare influence functions with EK-FAC and GGN (MC-Fisher; Equation (10) as the Hessian approximation (termed K-FAC Influence) to TRAK as formulated for diffusion models in Georgiev et al. (2023); Zheng et al. (2024). In our framework, their method can be tersely described as using GGN (Empirical Fisher) in Equation (11) Hessian approximation instead of G (MC-Fisher) in Equation (10), and computing the products using random projections (Dasgupta & Gupta, 2003) rather than K-FAC. We also compare to the ad-hoc changes to the measurement/rraiingg loss in the influence function approximation (D-TRAK) that were shown by Zheng et al. (2024) to give improved LDS performance. Note that, the changes in D-TRAK were directly optimised for improvements in LDS scores in the diffusion modelling setting, and lack any theoretical motivation. Hence, a direct comparison for the changes proposed in this work (K-FAC Influence) TRAK; the insights from D-TRAK are orthogonal to our work. These are the only prior works motivated by predicting the change in a model's measurements after retraining that have been applied to the general diffusion modelling setting that we are aware of. We also compare to naïvely using cosine similarity between the CLIP (Radford et al., 2021) embeddings of the training datapoints and the generated sample as proxy for influence on the generated samples. Lastly, we report LDS results for the oracle method of "Exact Retraining," where we actually retrain a single model to predict the changes in measurements.
inner
LDS. The LDS results attributing the loss and ELBO measurements are shown in Figures 2a and 2b. The LDS results attributing the marginal log-probability on dequantised data are shown in Appendix I. K-FAC Influence outperforms TRAK in all settings. K-FAC Influence using the loss measurement also outperforms the benchmark-tuned changes in D-TRAK in all settings as well. In Figures 2a, 2b and 21, we report the results for both the best damping values from sweep (see Appendix G), as well as for "default" values following recommendations in previous work (see Appendix 1.5). TRAK and D-TRAK appear be more sensitive to tuning the damping factor than K-FAC Influence. They often don't perform at all if the damping factor is too small, and take a noticeable performance hit if the damping factor is not tuned to the problem or method (see Figures and 10 in Appendix G). However, in most applications, tuning the damping factor would be infeasible, as it requires retraining the model many times over construct an LDS benchmark, so this is significant limitation. In contrast, for K-FAC Influence, we find that generally any sufficiently small value works reasonably well (see Figures 7 and 9).
Retraining without top influences. The counterfactual retraining results are shown in Figure 3 for CIFAR-2, CIF with 2% and 10% of the data removed. In this evaluation, influence functions with K-FAC consistently pick more influential training examples (i.e. those which lead to higher loss reduction) than the baselines.
4.1 POTENTIAL LIMITATIONS OF INFLUENCE FUNCTIONS FOR DIFFUSION MODELS
One peculiarity in the LDS results, similar to the findings in Zheng et al. (2024), is that substituting the loss measurement for the ELBO measurement when predicting changes in ELBO the marginal log- prebbabbbblyty actually works better than using the correct measurement (see Figure 2b "K-FAC Influence (measurement loss)")) " To try and better understand the properties of influence functions, in this section we perform multiple ablations and report different interesting phenomena that give some insight into the challenges of using influence functions in this setting.
Note that, unlike Zheng al. (2024), we only change the measurement function for proxy in the influence function approximation, keeping the Hessian approximation and training loss gradient in Equation (6) the same.
8

---

Published as a conference paper at ICLR 2025
C
CIFAR-10
ArtBench
6.37
$ 5.2%
CLIP Similarity
$10.3%
(1,000
O4 8. 88.96
(5.2%)
04 9.8%
TRAK
26
(20,99)
22.66% (20.45))
16.9%
(16.5))
K-FAC Influence
20.1553
15.44%
13.8%
D-TRAK
30.33
21.39%.0..
Exact Retraining
0
50
0
25
50
20
Rank Correlation % (LDS)
(a) LDS results on the loss measurement.
CIFAR-2
CIFAR-10
8.0%6nn
3.676am
CLIP Cosine Similarity
G.3 3.3%
3. 6
TRAK
13.6%
10.4%
K-FAC Influence
18.5%in
20
D-TRAK
27.1%
(3.790)
15.4%
(66.00)
K-FAC Influence (m. loss)
43.7
7
16.
16.4%
Exact Retraining
40
60
15 20
Rank Correlation % (LDS)
(b) LDS results on ELBO measurement.
Figure 2: Linear Data-modelling Score (LDS) for different data attribution methods. Methods that substitute incorrect measurement functions into the approximation are separated and plotted with .. Where applicable, we plot results for both the best "default" damping value with o. The numerical results are reported black for the best damping value, and for the "default" damping value in (gray). "(m. loss)" implies that the appropriate measurement function was substituted with the loss {(0, x)) measurement function in the approximation. Results for shown with Standard error in the LDS score estimate is
damping value with and
the exact retraining method (oracle), indicated with 't', where the mean is taken over different generated samples æ on which the change in measurement is being estimated.
CIFAR-2
2% removed
CIFAR-2
10% removed
H+
0.0019
#0.0005
H+
0.0022
10.005
Random
0.0025 +0.0005
0.0031 #0.0005
TRAK
0.0029 +0.0006
0.0037 +0.0005
D-TRAK
0.003 +0.0006
0.0044 -0.0006
K-FAC Influence
0.0000 0.0025 0.0050 0.0075 0.0100
0.0000
0.0025
0.0050
0.0075 0.0100 0.0125
CIFAR-10
2% removed
CIFAR-10
10% removed
0.00026
10.0003
0.00031
10.0003
Random
4-4
0.00055 10.0003
+-|
0.0008 +0.000
TRAK
0.00079
10.0003
0.0012
10.0003
D-TRAK
0.00067
+0.0005
0.0012
+0.0005
K-FAC Influence
0.000
0.002
0.004
0.006
0.000
0.002
0.004
0.006
0.008
Measurement Change
Measurement Change
Figure 3: Changes in measurements under counterfactual retraining without top influences for the loss measurement. The standard error in the estimate of the mean indicated with error bars and reported after '+', where the average is over different generated samples for which top influences are being identified.
As illustrated in Figure 20, gradients of the ELBO and training loss measurementss, up to a constant scaling, consist of the same per-diffusion- diffusion-timestep loss term gradients (ol, æ². but with different weighting. To try and break-down why approximating the change in ELBO with the training loss measurement gives higher LDS scores, first look at predicting the change in the per-diffusion-

---

Published as a conference paper at ICLR 2025
timestep losses l while substituting different per-diffusion-timeste losses into the K-FAC influence approximation. The results are shown in Figure 13, leading to the following observation:
Observation Higher-limestep losses (0 x) act as better proxies for lower-timestep losses.
More specifically, changes in losses & can in general be well approximated by substituting measure- ments by' into the influence approximation with t' > In some cases, using the incorrect timestep > even results in significantly better LDS scores than the correct timestep 1' =
Based on Observation 1, it is clear that influence function-basee approximations have limitations when being applied to predict the numerical change in loss measurements. We observe another pattern in how they can fail:
Observation 2 Influence functions predict both positive and negalive influence on loss, but, in practice, removing data points predominantly increases loss.
We show Figures 17 and 18 that influence functions tend to overestimate how often removal of group data points will lead to improvements in loss generated sample (both for aggregate diffusion training loss in Section 2.1, and the per-diffusion-timest
loss Equation (2)).
Lastly, although we have argued for estimating the effect that removing training datapoints has on the model's s marginal distribution, one property of diffusion models complicates the usefulness of pursuing this goal in practice for models trained on large amounts of data:
Observation For sufficiently large training set sizes, the diffusion model's marginal proba- bility distribution is close to constant (on generated samples), irrespective of which examples were removed from the training data.
As illustrated in Figure 19a, the exact marginal log probbabbbilty measurement is close to constant for any given sample generated from the model, no matter which 50% subset of the training data removed if the resulting training dataset is large enough. In particular, it is extremely rare that one sample is more likely to be generated than another by one model, and is less likely to be generated than another by different model trained on a different random subset of the data. Our observation mirrors that Kadkhodaic et al. (2024) who found that, diffusion models are trained on non-overlapping subsets of data of sufficient size, they generate near-identical samples when fed with the random seed. We find this observation holds not just for the generated samples, but for the marginal log- probability density, the ELBO and the training loss measurements well (see Figures 19b and 19c).
5 DISCUSSION
In this work, we extended the influence functions approach to the diffuusion modelling setting, and showed different ways in which the GGN Hessian approximation can be formulated. Our proposed method with recommended design choices improves performance compared to existing techniques across various data attribution evaluation metrics. Nonetheless, experimentally, we are met with two contrasting findings: on the one hand, influence functions in the diffuusin modelling setting appear to be able to identify important influences. The surfaced influential examples do significantly impact the training loss when retraining the model without them (Figure 3), and they appear perceptually very relevant the generated samples. On the other hand, they fall short of accurately predicting the numerical changes in measurements after retraining. This appears to be especially the case for measurement functions we would argue are most relevant in the image generative modelling setting proxies for marginal probability of sampling particular example (Section 4.1).
Despite these shortcomings, influence functions can still offer valuable insights: they can serve as a useful exploratory tool for understanding model behaviour in a diffusion modelling context, and can help guide data curation, identifying examples most responsible for certain behaviours. To make them useful in settings where numerical accuracy in the predicted behaviour after retraining is required, such as copyright infringement, we believe more work is required into 1) finding better proxies for marginal probability, and 2) even further improving the influence functions approximation.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENTS
We thank Jihao Andreas Lin for useful discussions on compression, and help with implementation of quantisation and SVD compression. We also thank Kristian Georgiev for sharing with us the diffusion model weights used for analysis in Georgiev et (2023), and Felix Dangel for help and feedback on implementing extensions to the curvlinops (Dangel et al., 2025) package used for experiments this paper. Richard Turner supported by the EPSRC Probabilistic Al Hub (EP/Y028783/1).
REFERENCES
Shun-Ichi Amari. Natural gradient works efficiently learning, Neural computation, 10(2), 1998.
Juhan Bae, Nathan Ng. Alston Lo, Marzyeh Ghassemi, and Roger Grosse. If Influence Functions are the Answer, Then What is the Question?, September 2022.
Juhan Bae, Wu Lin, Jonathan Lorraine, and Roger Grosse. Training data attribution via approximate unrolled differentiation, 2024. URL https: /atxivivvv////////////////////a/ org/ httg/bb/////aa////////////////// /abs//400 /2005 12186.
Elnaz Barshan, Mare- Etienne Brunet, and Gintare Karolina Dziugaite. Relatif: Identifying explanatory Mare-Etienne training samples via relative influence. In International Conference on Artificial Intelligence and Statistics, pp. 1899-1909. PMLR, 2020.
Samyadeep Basu, Xuchen You, and Soheil Feizi. On second-order group influence functions for black-box predictions. International Conference on Machine Learning, pp. 715-724. PMIR,, 2020.
Alberto Bernacchia, Mate Lengyel, and Guillaume Hennequin. Exact natural gradient in deep linear networks and its application to the nonlinear case. In NeurlPS, 2018.
Aleksandar Botev, Hippolyt Ritter, and David Barber. Practical Gauss-Newton optimisation for deep learning. JCML, 2017.
Felix Dangel, Runa Eschenhagen, Weronika Ormaniec, Andres Fernandez, Lukas Tatzel, and Agusti- nus Kristiadi, Position: Curvature matrices should democratized via linear operators. arXiv 2501. 2025.
Sanjoy Dasgupta and Anupam Gupta. An elementary proof of a theorem of Johnson and Lindenstrauss. Random Structures & Algorithms, 22(1):60-65 65. January 2003. ISSN 1042-9832, 1098-2418. doi: 10.1002/rsa. 100733
Runa Eschenhagen, Alexander Immer, Richard E. Turner, Frank Schneider, and Philipp Hennig. Kronecker-Factored Approximate Curvature for modern neural network architectures. In NeurIPS, 2023.
Thomas George, César Laurent, Xavier Bouthillier, Nicolas Ballas, and Pascal Vincent. Fast approximate natural gradient descent in a kronecker factored eigenbasis. Advances in Neural Information Processing Systems, 31, 2018.
Kristian Georgiev, Joshua Vendrow, Hadi Salman, Sung Min Park, and Aleksander Madry. The Journey, Not the Destination: How Data Guides Diffusion Models, December 2023.
Roger Grosse and James Martens. A Kronecker-factored approximate Fisher matrix for convolution layers. JCML, 2016.
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, Evan Hubinger, Kamilé Lukošiūté, Karina Nguyen, Nicholas Joseph, Sam McCandlish, Jared Kaplan, and Samuel R. Bowman. Studying Large Language Model Generalization with Influence Functions, August 2023.
Tom Heskes. On "natural" learning and pruning in multilayered perceptrons. Neural Computation, 12(4), 2000.
11

---

fwwriinng-eettues-20000 Published as a conference paper at ICLR 2025
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising Diffusion Probabilistic Models. In Advances in Neural Information Processing Systems, volume 33, pp. 2840-6851. Curran Associates, Inc., 2020.
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang. Lu Wang, and Weizhu Chen. Lora: Low-rank adaptation of large language models. arXiv preprint (xXiv:2106009655 2021.
Alexander Immer, Tycho F.A. van der Ouderaa, Gunnar Rätsch, Vincent Fortuin, and Mark van der Wilk. Invariance learning deep neural networks with differentiable Laplace approximations. In NeurIPS, 2022.
William Johnson, Joram Lindenstrauss, and Gideon Schechtman. Extensions of lipschitz maps into banach spaces. Israel Journal of Mathematics, (4(2):129-138 1986.
Zahra Kadkhodaie, Florentin Guth, Bero P. Simoncelli, and Stéphane Mallat. Generalization in diffusion models arises from geometry-adaptive harmonic representations, April 2024.
Diederik P. Kingma, Tim Salimans, Ben Poole, and Jonathan Ho. Variational Diffusion Models, April 2023.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pp. 1885-1894. PMLR, 2017a.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 1894. PMLR, 06-11 Aug 2017b. URL https: https::/proceee:e http///ccoednggssee/eeeeee. mlr, wtt.ssmw.....s./eeee.... pwwss//70/kohh//oooooe html.
Pang Wei Koh, Kai-Siang Ang, Hubert H. Teo, and Percy Liang. On the Accuracy of Influence Functions for Measuring Group Effects, November 2019.
Steven G. Krantz and Harold R. Parks. The Implicit Function Theorem. Birkhäuser, Boston, MA,
2003. ISBN 978-1-461 1--4661116--1111 978-1-4612-0059-8, doi: 11007/978-1-4611-000
Alex Krizhevsky. Learning Multiple Layers of Features from Tiny Images. Technical
report, University of Toronto, 2009. URL http://wwa es utoronto. /kriz/ pdf.
Frederik Kunstner, Lukas Balles, and Philipp Hennig. Limitations of the empirical Fisher approxima- tion for natural gradient descent. In NeurIPS, 2019.
Yongchan Kwon, Eric Wu, Kevin Wu, and James Zou. Datalnf: Efficiently Estimating Data Influence in LoRA-tuned LLMs and Diffusion Models. In The Twelfth International Conference on Learning Representations, October 2023.
Yann LeCun, D Touresky, G Hinton, and T Sejnowski. A theoretical framework for back-propggation. In Proceedings of the 1988 connectionist models summer school, volume 1, pp. 21 21-28, 1988.
Peiyuan Liao, Xiuyu Li, Xihui Liu, and Kurt Keutzer. The artbench dataset: Benchmarking generative models with artworks. arXiv preprint arXiv:2206.1 10006..1400 11404, 2022.
Jinxu Lin, Linwei Tao, Minjing Dong, and Chang Xu. Diffusion attribution score: Evaluating training data influence in diffusion model. arXiv preprint arXiv:2410.1866 2024.
James Martens. New insights and perspectives on the natural gradient method. JMLR, 21(146), 2020.
James Martens and Roger Grosse. Optimizing neural networks with Kronecker-factored approximate curvature. In ICML, 2015.
Kazuki Osawa, Shigang Li, and Torsten Hoefler. PipeFisher: Efficient training of large language models using pipelining and Fisher information matrices. arXiv 2.11.14133, 2022.
12

---

Published as a conference paper at ICLR 2025
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclere, and Aleksander Madry. TRAK: Attributing Model Behavior at Scale, April 2023.
J. Gregory Pauloski, Qi Huang, Lei Huang, Shivaram Venkataraman, Kyle Chard, lan T. Foster, and Zhao Zhang. KAISA: an adaptive second-order optimizer framework for deep neural networks. In International Conference for High Performance Compuling, Networking, Storage and Analysis (SC21), 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal,
Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision, 2021. URL https:
//arxiv. www./oggbb....oo.oooo.. abs/ /2003.11000 00020.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- ence on computer vision and pattern recognition, pp. 10684-10695, 2022.
Joseph Saveri and Matthew Butterick. Image generator litigation.
https://
Amnaeggeneratorltigttt.
com/, 2023a. Accessed: 2024-07-06.
Joseph Saveri and Matthew Butterick. Language model litigation. https://11mm wellllmmmiigge http:///timmiiii.i......oo
com/, 2023b. Accessed: 2024-07-06.
Nicol N Schraudolph. Fast curvature metrix-vector products for second-order gradient descent. Neural computation, 14(7), 2002.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, Patrick Schramowski, Srivatsa Kundurthy, Katherine Crowson, Ludwig Schmidt, Robert Kaczmarczyk, and Jenia Jitsev. Laion-5b: An open large-scale dataset for training next generation image-text models, 2022. URL https: //tp///aaavv////////////aa/aawaa org/ abs / 2210.200.22 08402.
Jascha Sohl-Dickstein, Eric A. Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep Unsupervised Learning using Nonequilibrium Thermodynamics, November 2015.
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising Diffusion Implicit Models, October 2022.
Yang Song and Stefano Ermon. Generative modeling by estimating gradients of the data distribution, 2020. URL https::///// arxiv orgg abs / 907.0000000000
Yang Song, Conor Durkan, lain Murray, and Stefano Ermon. Maximum Likelihood Training of Score Based Diffusion Models, October 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-Based Generative Modeling through Stochastic Differential Equations, February 2021b.
Zedong Tang, Fenlong Jiang, Maoguo Gong, Hao Li, Yue Wu, Fan Yu, Zidong Wang, and Min Wang. SKFAC: Training neural networks with faster Krotecker-factored approximate curvature. In CVPR, 2021.
Richard E. Turner, Cristiana-Diana Diaconu, Stratis Markou, Aliaksandra Shysheya, Andrew Y. K. Foong, and Bruno Mlodozeniec. Denoising diffusion probabilistic models in six simple steps, 2024.
Vaswani. Attention all you need. Advances Neural Information Processing Systems, 2017.
Guodong Zhang, Lala Li, Zachary Nado, James Martens, Sushant Sachdeva, George F. Dahl, Christopher J. Shallue, and Roger B. Grosse. Which algorithmic choices matter at which batch sizes? Insights from a noisy quadratic model. In NeurIPS, 2019.
Xiaosen Zheng, Tianyu Pang, Chao Du, Jing Jiang, and Min Lin. Intriguing Properties of Data Attribution on Diffusion Models, March 2024.
13

---

Published as a conference paper at ICLR 2025
A DERIVATION OF INFLUENCE FUNCTIONS
In this section, we state the implicit function theorem (Appendix A.I). Then, Appendix A.2, we introduce the details of how it can be applied in the context of a loss function L(e, 0) parameterised by a continuous hyperparameter (which is, e.g., controlling how down-weighted the loss terms on some examples are, as in Section 2.2).
Linear Extrapolation
Minlmam
2 En (an, 0) )((z5 0)
Cle. 0)
1V > N
n= Mann 0) ((eg,
N D> (1),))
Figure 4: Illustration of the influence function approximation for 1-dimensional parameter space E
R. Influence funcitons consider the extended loss landscape L(s, 0) der Σm (( e( , 0). where the loss f 0) for some datapoint 2, (alternatively, group of datapoints) is down-weighted by By linearly extrapolating how the optimal set of parameters & would change around we can predicted how the optimal parameters would change when the term f(zj. 0) is fully removed from the loss
(o).
A.1 IMPLICIT FUNCTION THEOREM
Theorem (Implicit Function Theorem (Krantz Parks, 2003)) F R mm + R. be continuuusly differentiable function, and let R" x IRm have coordinates (x,y). Fix point (a, b) = (a1..... On bi..... 6m) with F(a, b) = 0, where 0 1 is the zero vector. If the Jacobian matrix F(a b) pmxm of + Fa
OF,
[V b)l., dyj
(a.
invertible, then there exists open set U C Ra containing such that there exists unique function g -+ TR. such that g(a) = b, and Fxx,g(x)) = 0 for all x E U. Moreover, is continuously differentiable.
Remark (Derivative the implicit function) Denoting the Jacobian matrix of + F(x, y) as: OF,
V [x
(x,
Ox;
the derivative
0g U + Rm of U + R in Theorem 1 can be written as: 8x 0g(x) -Vxxxxx (VFF
VxF(xxg(x)).
(14)
dx
This can readily be seen by noting that, for E U: ((x'_((x))) 0 VX EU
dF g(x)) 0. dx
Hence, since g is differentiable, we can apply the chain rule of differentiation to get:
dF (F(xxg(x)) VxF(x, g(x)) + F(x g(x)) 0g(x)
0
dx Rearranging gives equation Equation (14).
dx
14

---

Published as a conference paper at ICLR 2025
A.2 APPLYING THE IMPLICIT FUNCTION THEOREM TO QUANTIFY THE CHANGE IN THE OPTIMUM OF A LOSS
Consider a loss function L R I R that depends on some hyperparameter e R (in Section 2.2, this was the scalar by which certain loss terms were down-weighted) and some parameters IR At the minimum of the loss function [(e, 0). the derivative with respect the parameters will be zero. Hence, assuming that the loss function is twice continuously differentiable (hence is continuously differentiable), and assuming that for some e' E ER" we have a set of parameters de
0 such that 0L (E', 0*) O and the Hessian 0 (1)) is invertible, function theorem to the derivative of the loss function de R x R - IR"", to get the existence of continuously differentiable function such that ac dg (E. g())){ = for e in some neighbourhood e'.
apply the implicit
Now g(E) might not necessarily be minimum of 0 Lie, 0). However, by making the further assumption that L is strictly convex can ensure that whenever a0 (E. 0) 0, is unique minimum, and so g(E) represents the change in the minimum as we vary E. This is summarised the lemma below:
Lemma Let L R Rum - T. be twice continuously differentiable function, with coordinates denoted by (e, 0) E R x Rm, such that + L(e, is strictly convex Ve E R.. Fix a point (E', (*) such that Oc (E' = =. Then, by the Implicit Function Theorem applied to ac there 29
2g exists an open set U Ih" containing O* such that there exists unique function g U U-k such that g(E')= = 0-, =0" continuously differentiable with derivative:
and g(e) is the unique minimum of 0 + Lie,
0) for all E E U. Moreover, is
Ōg(E) de
TOCC 00
0CC
Ee, ((()))
(e, gE))
(15)
aod0
Remark 2 For loss function L: R x Rm of the form Lle, 0) Li (0) + (02 (such as that in o Equation (4)), Bada (e, g(E)) the equation above simplifies to:
G d=d0 (e, (()()(())))))))))))))) = 00
OL2 (g(e))
(16)
The above lemma and remark give the result in Equation (5). Namely, section 2.2:
C(E, 0) =
1
N
E(O
1
M ((0,
E(O,
eq. (1.
0C
1 M
M
((0
N
M
d=00
a0
j=l
j=1
(15(1
0g(E)
00 (c, 9()))
1
M
x31z)
de
M
00
B DERIVATION OF THE FISHER "GGN" FORMULATION FOR DIFFUSION MODELS
As discussed in Section 2.3 partitioning the function + ||e(*) - ())((1))1 into the model output 0 + ei (x) and the & loss function is a natural choice and results in
GGN (0)
N
V2 (2())) (00) - -eb ( Tooh (z100 [at0)]
E
E
((()
Vich
(
N
n=1
N
2
Vicy (z(0)
IV Vo (2(()
E2
E (8)
(17)
N
15

---

Published as a conference paper at ICLR 2025
Note that we used
(
||000
[21
(2(0) =I.
2
We can substitute I with B, [ for any random vector 7j with identity second moment. This allows us to rewrite the expression for the GGN in Equation (17)
G
(0)
2 N ==1
1 E
Vocl (z1)) m Verl ( (x0))
where the term V (
E Rd which is vector-Jacobian product, can be efficiently evaluated roughly the cost of a single backward pass. For example, can be rewritten manner lending itself to implementation in standard autodiff library as:
V o (m:(e)
(mod <5 (2x10) (I=
Cod stopgrad
[&t (x()) +7
where st opgrad denotes the common stopping of passing of the adjoint through a certain operation common in autodiff frameworks.
C (E)K-FAC FOR DIFFUSION MODELS
C.I BACKGROUND: COONECKER-FACTOREE APPROXIMATE CURVATURE
Kronecker-Factored Approximate Curvature (Heskes, 2000; Martens & Grosse, 2015; Botev et al., 2017, K-FAC) is typically used layer-wise block-diagonal approximation of the Fisher or GGN of a neural network. Each layer-wise block can be written as a Kronecker product, hence the name.
We will first describe the K-FAC approximation for simple linear layer for deep neural network for standard regression or classification setting. We assume loss function N ΣN Eyyn. fo (2x)) N where fe is neural network parametrised by 8, D = {zn, Yo. } n=1 is the dataset with inputs on and labels Vu. and E(-,-) is loss function like the cross-entropy
mean square error. To derive the K-FAC approximation for the parameters of a linear layer with weight matrix W, 10 we first note that can write the GGN block for the flattened parameters 0 vec(Wi) as
1 Vo,ooun) (Wn (o fo(x,))) Vo,So(zn): N
GGND(O))
(18)
here we choose the split from GGN in Equation (7), but the derivation also follows similarly for other splits. Given that VJ, (2 = a x gn where ( is the input to the ith layer for the
nthh example and g the transposed Jacobian of the neural network output w.rt. to the output of the /th layer for the noth example, we have
N
GGN = N
&
900h)
(V), (vn. fo(o.))) (al) & 9
(19)
(()
x ( (V). (yn fo(xx))) 9 (!!TT
(20)
K-FAC now approximating this sum of Kronecker products with a Kronecker product of sums, i.e.
1
GG)
a
CCT
Σ In (V) elyn So (x..)) (()
(21)
N
This approximation becomes an equality in the trivial case of N = or for simple settings of deep linear networks with mean square error loss function (Bernacchia et al., 2018). After noticing
A potential hias he absorhed into the weight matrix.
16

---

Published as a conference paper at ICLR 2025
that the Hessian V1, yyr fo (2,)) is the identity matrix for the mean square error loss, the K-FAC formulation for diffusion models in Equation (12) can now be related to this derivation the only difference is the expectations from the diffusion modelling objective.
Note that this derivation assumed a simple linear layer. However, common architectures used for diffusion models consist of different layer types as well, such as convolutional layers and attention. As mentioned in Section 3.1.2, K-FAC can be more generally formulated for all linear layers with weight sharing (Eschenhagen et al., 2023).
First, note that the core building blocks of common neural network architectures can be expressed as linear layers with weight sharing. If a linear layer without weight sharing can be thought of a weight matrix W Rd x that is applied layer is applying the transposed weight matrix to right of an input matrix X E & Mxda, i.e. XW This can be thought of regular linear layer that is shared across the additional input dimension of size M. For example, the weight matrices in the attention mechanism are shared across tokens, the kernel in convolutions is shared across the spatial dimensions, and in graph neural network layer the weights might be shared across nodes or edges; see Section 2.2 in Eschenhagen et al. (2023) for more in in-depth explanation of these examples.
input vector IR dmn a linear weight sharing
Given this definition of linear weight sharing layers, we can identify two different settings in which they are used. In the expand setting, we have R x N loss terms for dataset with N data E(yn.r. (X(,)))) The additional R loss terms
1 points; we have a loss of the form N. can sometimes exactly correspond to the weight sharing dimension of size M, i..., R M; for example, when the weight sharing dimension corresponds to tokens in a sequence as it is the case for linear layers within attention. In contrast, in the reduce setting, we have a loss of the form EN E(yn: fo (Xx)). These two settings can now be used motivate two different flavours of the K-FAC approximation.
N
The first flavour, K-FAC-expand, is defined as
N
M
N
M
R
GGND(O1) x
ΣΣ
a
i m.
(T) a.m
X
Σ
Σ
Σ
Hir
Cxped
m== m=1
m=
(22)
where Gxped = N²RM, an.m is the mth row of the input to the Ith layer for the nihh example, He,r -.. (yn,r, (yn,r, fo(,)))((X.)))) and ga.rm is the transposed Jacobian of the rth row of the matrix output of the neural network w.r.t. the mtt row of the output matrix of the /th layer for the noth example. K-FAC-expand is motivated by the expand setting in the sense that for deep linear networks with mean square error as the loss function. K-FAC-expand is exactly equal to the layer- wise block- diagonal of the GGN. For convolutions, K-FAC expand corresponds to the K-FAC approximation derived in Grosse & Martens (2016) which has also been used for attention before (Zhang et al., 2019; Pauloski et al., 2021; Osawa et al., 2022).
The second variation, K-FAC-reduce, is defined as
M (T
M Σ
M
GGND(O))
Σ
Σ
anym
Σ
a
Σ
Ir.m
Hn
Σ
g
Crde
(23)
with Gss- = (22 Analogously to K-FAC-expand in the expand setting, in the reduce setting, K-FAC reduce is exactly equal to the layer-wise block-diagonal GGN for a deep linear network with mean square error loss and a scaled sum as the reduction function. With reduction function we refer to the function that is used to reduce the weight-sharing dimension in the forward pass of the model, e.g. average pooling to reduce the spatial dimension in a convolutional neural network. Similar approximations have also been proposed in different context (Tang et al., 2021; Immer et al., 2022).
Although each setting motivates a corresponding K-FAC approximation in the sense described above, can apply either K-FAC approximation in each setting. Also, note that while we only explicitly consider a single weight sharing dimension, there can in principle be an arbitrary number of them and we can choose between K- K-FAC-expand and K- K-FAC-reduce for each of them independently.
In principal, the input to the neural network does not necessarily have to have weight-sharing dimension, even when we the model contains linear weigh-sharing layers; this holds for the expand and the reduce setting.
17

---

Published as a conference paper at ICLR 2025
C.2 DERIVATION OF K-FAC FOR DIFFUSION MODELS
We will now explicitly derive K-FAC for diffusion models while also taking weight sharing into account. As described above, we will derive K-FAC-expand as well as K-FAC-reduce for FD (Equation (10)), or equivalently, GGN (Equation (9)). Note that once we have derived K-FAC, we can immediately apply an digenvalue-correctoon to get EK-FAC (George et al., 2018).
Let gan ( Vellc - (x ( I| = - - E Em= am @ byy
where dm is the math row of the input to the /th layer, by is the gradient of (er en (8) ),)² w.r.t. the mth row of the /th layer output, . is the size of the weight sharing dimension of the /th layers. and R is the size of the output dimension, c.g., height width x number of channels for image data. We have
(0:) = Exn
Emod
9
mad
(24)
Food
R
M
R
M
Eu...
Ei
E
( 2 (I r.m
Σ
6
(25)
m
FM
r=1 m==
m=1
R
M
Ei
E
(0) e0
18) Fead
Σ
am
am
(
(26)
m=1
(27)
First, we can use K-FAC-expand, which is exact for the case of a (deep) linear network:
expand
FD(
Exn
Eg
(t)))
a
(28)
M
M
M
E
Exn
Ei
i
am am
(1)
X Ez,
Ej
E
(
(
(E) Food
((C)) ((U)T
m
===1
(29)
(30)
Alternatively, K-FAC-reduce:
FD(
roduce
Ex
E
(a @ b()) )
(31)
M
1
< 800BUT]]
Er.
E;
(32)
M
an
2
1
E
39 E
E.
E
JUROT
(33)
M
with a(0)
> mm and 800 mm=l
(
Dr.m-
To approximate GGN]ess (Equation (11)) instead of GGN with K-FAC, we only have to make minimal adjustments to the derivation above. Importantly, the expectation over ened is replaced by an expectation over e(,) i.c. use targets sampled from the same distribution in training instead of the model's output distribution. Notably, for K-FAC-expand the approximation will be the same for GGN since we have to move all expectations and sums in the definition of GGN the front be able proceed (the approximation denoted
ignoring the distribution over targets
expand
with Rs above). However, for K-FAC-reduce the resulting approximation will be different:
reduce 1
GGN (0.)
Ex.
(+)T
c Ex.,
(34)
M
with g0)) = Ej E E. 10,0000 200 [a]]]] and T(()) = E2
E
[800 [3:0]]
18

---

Published as a conference paper at ICLR 2025
C.3 EMPIRICAL ABLATION
Here, we explore the impact of the Hessian approximation design choices discussed in Section 3.1 and Appendix C.2. We use K-FAC EK-FAC approximate either the GGN in Equation (10) or the D in Equation (1I). We also compare the "expand" and the "reduce" variants generally introduced in Appendix C.I and derived for our setting in Appendix C.2.
Firstly, we find that the better-motivated "MC-Fisher" estmmaar of GGN in Equation (9) does indeed perform better than the "empirical Fisher" in Equation (11) used in TRAK and D-TRAK. However, the difference is relatively small, which is maybe not too surprising given that the only difference between G and GGN FF)x-FAC-e is the distribution from which the labels are sampled (see Appendix C.2). In this sense, our GGN (E)K-FAC expand approximation already fixes a potential problem (low-rate (ow-rankness) with the "empirical Fisher" GGN D
ess used in TRAK
and D-TRAK. As described in Appendix C.2 this is not the case for GGN (E)K-FAC-reduce in Equation (34). As expected, the K-FAC-reduce approximation of GGN performs worst among all the variants." Secondly, we find that using the eigenvalue-correeted K-FAC (EK-FAC) variant, which should more closely approximate the respective GGN, improves results for all configurations. Thirdly, we find that K- K-FAC-expand noticeably outperforms K- K-FAC-reduce, which stands in contrast to some results in the second-ord -odder optimisation setting where the two are roughly on par with one another (Eschenhagen et al., 2023). This difference cannot be easily attributed since there multiple differences between the two settings: we use a square loss instead of a cross entropy loss, a full dataset estimate, different architecture, and evaluate the approximation in a different application. Notably, in some limited sense K-FAC-expand is better justified than K-FAC-reduce for our diffusion modelling setting, since it will be exact in the case of (deep) linear network.
All our results seem to imply that better Hessian approximation directly results in better downstream data attribution performance. However, we do not directly evaluate the approximation quality of the estimates.
6.2%
(0.6%)
GGN
K-FAC reduce
17.6%.2.9
(17.6%)
GGN
K-FAC expand
7.5 7.
(6.8%)
GGN
K-FAC reduce
20.6%=o.9
(20.4%)
GGN
K-FAC expand
+ 25.2%+0.9
(24.8%)
GGNloss
EK FAC-expand
18.8%+o.s
(18.6%)
GGNmodel EK FAC-reduce
(26.2%)
+ 26.7 26.7%
GGNaodel EK FAC-expand
50.3%+o.2
Exact Retraining
20
30
40
50
Rank Correlation % (LDS)
Figure 5: Ablation the different Hessian approximations introduced in Section 3.1 and Ap- pendix C.2. We ablate two versions of the GGN: the "MC" Fisher GGNmodall in Equation (9) and the "Empirical" Fisher GGN in Equation (11), as well as multiple settings for the K-FAC ap- proximation: "expand" and "reduce". and whether we use the or not (K-FAC). Same in Figure 2, we report the results for both the best damping value with and a default damping value of 10 with o. The damping value ablation for the selection of these results reported in Figure 11.
variant (EK-FAC)
Sincce the "reduce" variant for GG in Equation (34) is not supported by curvlinops (Dangel et al., 2025), we only implemented this variant for regular K-FAC. and not for EK- FAC However, we have no reason to expect the disappointing performance of the "reduce" variant with the GGN would not persist for EK-FAC.
19

---

Published as a conference paper at ICLR 2025
D RELATED WORK
Data attribution in diffusion models Data attribution in diffusion models has been tackled by extending TRAK (Park et al., 2023) in two previous works. Georgiev al. (2023) apply TRAK to diffusion models, and argue that attributing the log- probaabiiyl of sampling a latent from the sampling trajectory of an example at different diffusion timesteps will lead to attributions based on different semantic notions of similarity. In particular, directly attributing based on the diffusion loss of the generated sample leads poor LDS scores. We also observe this is the case for TRAK without damping factor tuning in this work. Zheng et al. (2024). on the other hand, show that substituting a different (incorrect) measurement function and training loss into the TRAK approximation improves performance in terms of the LDS scores, and they recommend a particular setting based on empirical performance. Significant gains in performance in Zheng et al. (2024) can observed from the tuning of damping factor. Kwon et al. (2023) propose approximation to the inverse of a Gauss-Newton matrix (specifically, the empirical Fisher) by heuristically interchanging the sum and the inverse operations; they apply influence functions with this Hessian approximation to LoRA-finetur (Hu et al., 2021) diffusion models. However, they do not consider what GGN matrix to use for approximating the Hessian in diffusion models, resorting to the Empirical Fisher matrix, and they do not explore what measuremen//los. function would be appropriate for this setting" 13, which we do in this paper. In concurrent work, Lin et al. (2024) explore alternative measurement functions for data attribution in diffusion models with TRAK.
Influence Functions Influence functions were originally proposed as a method for data attribution in deep learning by Koh & Liang (a). Later, Koh et al. (2019) explored influence functions for investigating the effect of removing or adding groups of data points. Further extensions were proposed by Basu et al. (2020) who explored utilising higher-order information - and Barshan al. (2020), who aimed to improve influence ranking via re-normalisation Initial works on influence functions (Koh & Liang, 2017a; Koh et al., 2019) relied on using iterative solvers to compute the required sivers-Hesssianveee vector products. Grosse et al. (2023) later explored using EK-FAC as an alternative solution to efficiently approximate calculations with the inverse Hessian.
For broader overview of scalable training data attribution in the context of modern deep learning, see (Grosse et al., 2023, $ 4).
E RUNTIME MEMORY AND COMPUTE
Influence functions can be implemented in different ways, caching different quantities at intermediate points, resulting in different trade-offs between memory and compute. A recommended implemen- tation will also depend on whether one just wants to find most influential training examples for selected set of query samples once, or whether one wants implement influence functions in system where new query samples to attribute come in periodically.
The procedure we follow in our implementation can roughly be summarised as informally depicted with pseudo-code in Algorithm 1.
For deployment, where new query samples might periodically come in, we might prefer to store compressed preconditioned training gradients instead. This is illustrated in Algorithm 2.
In principle, we used an empirical Fisher approximation (like in Equation (10) to approximate the GGN, we could further amortise the computation in the latter variant by caching training loss gradients during the K-FAC computation.
Note that, for applications like classification with cross-entropy loss or auto-regressive language modelling Vaswani (2017), the gradients have a Kronecker structure, which means they could be stored much more efficiently Grosse al. (2023). This is not the case for gradients of the diffusion loss in Section 2.1, since they require averaging multiple Monte-Carlo samples of the gradient.
We will primarily describe the complexities for the former variant (Algorithm 1), as that is the one we used for all experiments. The three sources of compute cost, which we will describe below, are: 1) computing and inverting the Hessian, 2) computing, pre-conditioning and compressing the query
They state that they "used negative log- likeeiiddd of a generated image a loss function", although they actually only use the diffusion training loss measurement.
20

---

Published as a conference paper at ICLR 2025
Algorithm 1 K-FAC Influence Computation (Single-Use)
1: Input: Training set { Ji query points {ă, { number of Monte-Carlo samples S E N
2: Output: Influence scores (score;j : E (1,...,N,, E ........H for all pairings of training
examples with query examples
3: H-1 <- ComputeAnndlnvertrFAa
Compute and store K-FAC inverse D Process query points Compute gradient with S samples > Precondition gradient
4: for:
to Q do
5: V +- Vem(2z;:)
6:
Vj < H-
7:
Store compressed yj
8:
end for
9:
for
1 to do U; f Ve((2,; 8)
Process training points Compute gradient with S samples
10:
11: for = to Q do
12:
scoreij y U/
Compute influence score
13: end for
14: end for
Algorithm 2 K-FAC Influence Computation (Continual Deployment Setting)
1: Input: Training set {a; J samples $
2: Output: Cached preconditioned training gradients {ur for efficient influence computation
3: !! CompuucAnndimmerKKFFAA)
Compute and store K-FAC inverse Preprocess training set Compute gradient with S samples b Precondition gradient
4: for = 1 to do
5: U/ Vol(z,;
6: y: < H H U;
7: Store compressed y/i
8: end for
9: procedure COMPUTRINHL ($,)
Called when queries arrive
10: for 1 to Q do
11:
8j < Vol(i,; for i score; < v f/
b. Compute query gradient
N do
13: end for
D Compute influence score
15:
end for
return (score,y i {L...., (L....,Q) Q}
17:
end procedure
gradients, and 3) computing the training gradients and taking inner- anerrproddcct with the preconditioned query gradients.
E. ASYMPTOTIC COMPLEXITY
Here, will describe how runtime compute and memory scale with the number of query examples to attribute Q, the number of of training examples N, the number of Monte-Carlo samples S, for standard feed-forward network with width W and depth L These variables summarised in Table 1. The number of parameters of the network P is then OWW²
Altogether, the runtime complexity of running K-FAC influence in this setting is O ((N + S NQW2L + L) and requires O(QW2 NQ) storage. We break this down below.
19This can either be a multi-layer perceptron, convolutional neural network with W denoting the channels. The feed-forward assumption is primarily chosen for illustrative purposes, but the analysis is srright-forward to extend to other architectures, and the asymptotic results do not differ for other common architectures.
21

---

Published as a conference paper at ICLR 2025
Number of query data points
N
Number of training examples
W
Maximum layer width
L
Depth of the network
Number of samples for Monte-Carlo evaluation of per-example loss measurement gradients
s
Table 1: Variables for scaling analysis.
E.1.1 K-FAC AND K-FAC INVERSION
For each training example, and each sample, the additional computation of K-FAC over a simple forward-backward pass through the network (LeCun et al., 1988) comes from computing the outer products of activaitiiaiiss and gradients of loss with respect to the pre-activations. Overall, this adds a cost of O(W²L) on top of the forward-backwart pass, and so single iteration has the same (( L.) cost scaling as forward- backbard pass. Hence, computing K-FAC for the entire training dataset with S samples has cost
Since K-FAC is block-wise diagonal approximation, computing the inverse only requires computing the per-layer inverses. For linear layer with input width Win and output width Wout computing the inverse costs O in + )uu)) due to the Kronecker-factores form of the K-FAC approximation.
Similarly, storing K-FAC (or the inverse) requires storing matrices of sizes W x WIn and Wous X Wout for each linear layer.
Hence, computing K-FAC has runtime complexity of O(NSLW?). An additional OLL'3) will be required for the inversion, which is negligible compared to the cost of computing K-FAC for larger datasets. The inverse K-FAC requires (LW3) storage. In practice, storing K-FAC (or inverse K-FAC) requires more memory than storing the network parameters, with the multiple depending on the ratios of layer widths across the network.
1.1.2 PRECONDITIONED QUERY GRADIENTS COMPUTATION
Computing single query gradient takes O(SW2 time, and preconditioning with K-FAC requires further matrix-vector product costing O(W' L). The of compressing the gradient will depend on the method, but, for quantisation (Appendix F), it's negligible compared to the other terms. Hence, computing all Q query gradients costs OQQSVL L)
Storing the Q preconditioned gradients requires 3 storage (although, in principle, this could be more efficient depending on the compression method chosen and how scales with the network size while maintaining precision).
E. TRAINING GRADIENTS AND SCORES COMPUTATION
Again, computing a single training gradient takes ((SW*L) and an inner product with all the preconditioned query gradients takes an additional (QQW2L). Hence, this part requires O NNQW"))) operations.
Storage-wise, storing the final "scores" (the preconditioned inner products between the training and query gradients) requires a further O(NQ) memory, but this is typically small (4NQ bytes for float 32 precision).
E. COMPARISON WITH TRAK
The complexity of TRAK (Park et al., 2023) additionally depends on the choice of the projection dimension R. The computational cost of running TRAK is O(NN + Q)SW2 L NQR + R²). Similarly, the memory cost of the implementation by Park et al. (2023) O(N + Q)R R²).
Note that, it is unclear how R should scale with the neural network size 2 L. Random projections do allow for constant scaling with vector size while maintaining approximation quality in some settings (Johnson al., 1986). To the best of our knowledge, it has not been shown, either empirically
22

---

Published as a conference paper at ICLR 2025
theoretically, what the expected scaling of R with network size might be in the context of influence function preconditioned inner products (Equation (6)). In the worst case, the projection dimension might be required to scale proportionally to the network size to maintain desired level of accuracy.
E.2 RUNTIME COMPLEXITY
We also report the runtimes of computing TRAK and K-FAC influence scores for the experiments reported in this paper. We discuss what additional memory requirements one might expect when running these methods. All experiments were ran single NVIDIA A100 GPU.
The runtime and memory reported for computing influence for 200 query data points. As discussed at the beginning of Appendix E, K-FAC computation and inversion costs are constant with respect to the number of query data points, and computing the training gradients can be amortised in sensible aalloymmet-ggaacde implementation at the added memory cost of storing the (compressed) training gradients.
E.2.1 RUNTIME RESULTS
Tables 2 and report the runtimes on a single NVIDIA A100 GPU of the most time-consuming parts of the influence function computation procedure.
# network param. (millions)
Dataset (size)
Dataset size
MC samples
K-FAC computation
Query gradients
Training gradients
CIPAR-2
5000
38.3
250
03:30:32
01:12
00:34:42
CIFAR-10
50000
38.3
250
35:01:33
01:12
05:43:14
ArtBench
50000
37.4 4-8..6"
125
32:48:08
04:18
18:56:57
Table 2: Runtime for K-FAC influence computation across datasets. "x" indicates parameters of pre-trained part of the model (e.g. VAE for Latent Diffusion Models).
Dataset (size)
Dataset size
network param. (millions)
# MC samples
Query gradients
Training gradients
Hessian inversion
Computing scores
CIRAR-2
5000
38.3
250
3:38
01:32:10
00:11
3:43
C1PAR-10
50000
38.3
250
3:44
15:15:04
00:57
3:54
Art:Bench
50000
37.4 +83.6"
125
5:46
23:58:44
00:28
6:24
Table 3: Runtime for TRAK score computation across datasets. "s" indicates parameters of a pre- trained part of the model (e.g. VAE for Latent Diffusion Models).
E.2.2 MEMORY USAGE
Tables 4 and 5 report the expected memory overheads due to having to manifest and store large matrices or collections of vectors in the influence function implementations of K-FAC Influence and TRAK.
network param. (millions)
Inverse K-FAC (GB)
Cached query gradients (GB)
Dataset
CIFAR-2
38.3
1.57
7.66
CIFAR-10
38.3
1.57
7.66
ArtBench
37.4 +83.6"
1.57
7.47
Table Memory usage linked K-FAC Influence. the model (e.g. VAE for Latent Diffusion Models).
"*" indicates parameters of pre-trained part of
23

---

Published as a conference paper at ICLR 2025
network param. (millions)
Projection dimension
Projected train gradients (GB)
Projected query gradients (GB)
Projected Hessian inverse (GB)
Dataset
CIPAR-2
38.3
32768
0.66
0.027
4.29
CIFAR-10
38.3
32768
6.55
4.29 0.027
ArtBench
37.4 +83.6
32768
6.55
0.027
4.29
Table Memory usage linked TRAK. "x" indicates parameters pre-trained part of the model (e.g. VAE for Latent Diffusion Models).
F GRADIENT COMPRESSION ABLATION
int8
1ot18
1.00
Rank correlation uncompressee essedd scores
Rank 100
Rank 150
Rank 50
0.95
Rank 20
Bank 10
0.90
Bank 5
0.85
Rank 2
Rank
0.75
0%
10%
20%
30%
40%
50%
60%
70%
Compression ratio
Figure 6: Comparison of gradient compression methods for the influence function approximation.
In Figure 6, we ablate different compression methods by computing the cetraininn-dattapoon influence scores with both a) compressed query (measurement) gradients, and b) the uncompressed gradients, and comparing the Pearson and rank correlations between the scores computed with (a) and (b). We hope to see a correlation of close to 100%, in which case the results for our method would unaffected by compression. We find that using 8-bit quantisation for compression results in almost change the ordering over training datapoints. This is in contrast to the SVD compression scheme used Grosse et al. (2023). This is likely because the per-ex per-example gradients naturally have a low-rank (Kronecker) structure in the classification, regression, settings, such as the one considered by Grosse et al. (2023). On the other hand, the diffusion training loss and other measurement functions considered in this work do not have this low-rank structure. This is because computing them requires multiple forward passes; for example, for the diffusion training loss, we need to average the mean square error loss in Equation (2) over multiple noise samples e(1) and multiple diffusion timesteps. We use 3 bit quantisation with query gradient batching (Grosse et al., 2023) for all K-FAC experiments throughout this work.
autoregressive language modelling
G DAMPING LDS ABLATIONS
We report an ablation over the LDS scores with the GGN approximated with different damping factors for TRAK/D TRAK and K-FAC influence in Figures 7 to 12. The reported damping factors for TRAK are normalised by the dataset size so that they correspond to the equivalent damping factors for method when viewing TRAK alternative approximation the GGN (see Section 3.1).
24

---

Published as a conference paper at ICLR 2025
CIFAR-2 K-FAC Influence Ablation
Rank correlation (LDS)
GGN-o5s
K. K--FC--o X-FAC-expand
0.25
GGN
K-FA F-Crueccee
GGNNodal
K-FAC-reduce
GGN
K- PrFFAeexpaad
0.15
CCN
EKK-FAC-xxpadd
reduce
0.10
AACespand
0.05
0.00
10-10 10
10
10-2
Damping factor
Figure 7: Effect of damping on the LDS scores for K-FAC influence on CIFAR-2. 250 samples were used for Monte Carlo estiamtion of all quantities (GGN and the training loss/measurement gradients). Ta Target/Mea sure re indicate what measurement was used for the ground-truth and the approximation respectively.
CIFAR-2 TRAK Influence
Rank correlation (LDS)
Target
Nessure Train, Losa
0.20
Loss
Sq.Horx. (D-TRAK)
0.15
ELBO
ELBO
Loss
EL20
Sq. Norm Sq. Horz. (D-TRAK)
0.10
0.05
0.00
10 10
10
10
10
10a
10a
Damping factor
Figure 8: Effect of damping on the LDS scores for TRAK (random projection) based influence on CIFAR-2. 250 samples were used for Monte Carlo estiamtion of all quantities (GGN and the training loss/measurement gradients). In the legend: Ta Target indicates what measurement we'r re trying to predict the change after retraining, Measure indicates what measurement function was substituted into the influence function approximation, and Tra in. Loss indicates what function was substituted for the training loss in the computation of the GGN and gradient of the training loss in the influence function approximation.
H EMPIRICAL ABLATIONS FOR CHALLENGES TO USE OF INFLUENCE FUNCTIONS FOR DIFFUSION MODELS
In this section, we describe the results for the observations discussed in Section 4.1.
Observation is based on Figures 13 and 14. Figure 13 shows the LDS scores on CIFAR-2 when attributing per-timestep diffusion losses & (see Equation (2)) using influence functions, whilst varying what (possibly wrong) per-timestep diffusion loss le is used as a measurement function in the influence function approximation (Equation (6)). Figure 14 is counter-equivalent to Figure 18 where instead using influence functions to approximate the change measurement, we actually retrain model the randomly subsampled subset of data and compute the measurement.
25

---

Published as a conference paper at ICLR 2025
CIFAR-10- K-FAC Influence
Rank correlation (LDS)
Target Messure
Losa
Loza
ELEC- Loss
ELBO
0.15
ELEC
ILBO+
0.10
0.05
0.00
10 10-10
10
10-2
Damping factor
Figure 9: Effect of damping on the LDS scores for K-FAC based influence on CIFAR-10. 250 samples were used for computing the EK-FAC GGN approximation (125 for the eigenbasis computa- tions, 125 for the eigenvalue computations), and 250 for computing a Monte Carlo estimate of the training loss/measurement gradients. Target/eeasure indicate what measurement used for the foound-truth and in the approximation respectively.
CIFAR-10 TRAK Influence
Rank correlation (LDS)
0.15
Target
Reasure Train, Loss
Loas
Less
Sq.Norn (D-TRAK)
Less
Korm Form Loss
0.10
BL80+
Lose
Lose
Komm Norn (D-TRAK)
0.05
0.00
10
10
10
10
10
10
10
Damping factor
Figure 10: Effect of damping on the LDS scores for TRAK (random projection) based influence on CIFAR-10. 250 samples were used for Monte Carlo estiamtion of all quantities (GGN and the training loss/measurement gradients). In the legend: Target indicates what measurement we'r re trying to predict the change after retraining, Measure indicates what measurement function was substituted into the influence function approximation, and Train Loss indicates what function was substituted for the training loss in the computation of the GGN and gradient of the training loss in the influence function approximation.
natural question to ask with regards Observation 1 is: does this effect go away in settings where the influence function approximation should more exact? Note that, bar the non-convexity. of the training loss function Lp. the influence function approximation in Equation (6) linearisation of the actual change in the measurement for the optimum of the training loss functions with some examples weigheee by & around & O. Hence, we might expect the approximation to be more exact when instead of fully removing some data points from the dataset (setting & = 1/N), we instead down-weight their contribution to the training loss by smaller non-zero factor. To investigate whether this is the case, we repeat the LDS analysis in Figures 13 and 14, but with e 1/2N; in other words, the training loss terms corresponding to the "removed" examples are simply down-weighted by factor of 1/2 in the retrained models. The results shown in Figures 15 and 16. Perhaps somewhat surprisingly, a contrasting effect can observed, where using per ter-timettep diffusion losses for larger times yields higher absolute rank correlation, but with the opposing sign. The negative
26

---

Published as a conference paper at ICLR 2025
CIFAR-2 K-FAC Influence Ablation
Rank correlation (LDS)
GGN-
K. K-FAC expand
0.25
GGNNodal
K-FAC reduce
GGN
K-FAC expand
GGN2os
EK-FAC expand
0.15
CGN
BK-FAC
0.10
nddel
EK- FAC expand
0.05
0.00
10 10-10
10
10-4
10-2
Damping factor
Figure 11: Effect of damping on the LDS scores for various GGN K-FAC approximations on CIFAR-2. 250 samples were used for Monte Carlo estiamtion of all quantities (GGN and the training moss/meaaur gradients); for eigenvalue-corrected K-FAC, the 250 samples are split so that 125 samples are used for the computation of the eigen-basis, and 125 samples are used for the computation of the eigen-values, order to roughly match the runtime between K-FAC and EK-FAC.
CIFAR-2 (deq.) K-FAC Influence
Rank correlation (LDS)
Log-likelibood (with loss) Log likelihood (with ELBO)
0.20
0.15
0.10
0.00
10-10
10
10-8
10
10-4
10
Damping factor
Figure 12: Effect of damping the LDS scores for K-FAC influence for approximating the marginal pog-rrbbybility measurement on dequantised CIFAR-2.
correlation between measurement t,, by for can also be observed for the true measurements in the retrained models in Figure 16. We also observe that in this setting, influence functions fail completely to predict changes in l, with the correct measurement function for t < 200.
Figure 17 shows the changes in losses after retraining the model on half the data removed against the predicted changes in losses using K-FAC Influence for two datasets: CIFAR-2 and CIFAR-1C. In both cases, for vast majority of retrained models, the loss measurement sample increases after retraining. On the other hand, the influence functions predict roughly evenly that the loss will increase and decrease. This trend is amplified if we instead look at influence predicted for per- tim-tipettep diffusion losses l (Equation (2)) for earlier timesteps t, which can be seen in Figure 18. On CIFAR-2, actual changes in &,, EgO, 00 measurements positive, which the influence functions approximation completely misses. For all plots, K-FAC Influence ran with damping factor of 10-8 and 250 samples for all gradient computations.
Observation
actually always
27

---

Published as a conference paper at ICLR 2025
Influence measurement
Square
E1o
Ezo
Eso
tso
E100
8200
Eooo
Eaun
Egsa
norm
14.6
13.2
10.5
9.1
8.2
7.6
7.3
12.9
5.8
5.3
-12.8
210
13.2
10.5
15.6
40
Rank correlation (%)
820
7.7
11.1
14.2
13.4
12.7
13.1
16.4
3.1
True mensurenen
Eso
0.1
2.4
1.6
12.0
16.2
17.4
18.5
16.0
2.9
2.6
2.0
20
Eso
-0.2
1.1
2.1
8.7
14.9
17.7
21.7
15.9
3.0
2.6
2.5
0.8
1.4
7.1
23.2
16.1
3.2
2.6
E200
0.5
U.3
0.5
3.1
8.2
12.0
25.2
18.6
3.5
3.2
1.0
20
Eso0
2.2
1.7
1.9
2.1
3.2
4.4
12.3
45.6
5.3
4.6
-1.6
ESD0
5.3
2.4
1.2
0.5
0.6
0.6
2.6
20.8
59.4
50.8
-3.0
-40
Egon
Figure 13: Rank correlation (LDS scores) between influence function estimates with different measurement functions and different measurements CIFAR-2. 2. The plot shows how well different per-timestep diffusion losses & work as measurement functions in the influence function approximation, when trying to approximate changes the actual measurements when retraining model.
True measurement
E10
820
Ego
Éao
E100
Egoo
6500
Espo
Eggo
32.3
20.3
13.9
7.9
6.4
6.1
6.1
7.1
4.0
0.6
60
010
19.9
30.3
28.2
19.4
15.5
12.2
3.6
-0.0
True mensttreeen
820
13.8
28.6
32.1
27.4
22.5
20.5
16.6
11.4
3.0
0.1
40
Rank correlation (%))
E00
8.0
20.4
25.3
36.3
34.9
33.1
26.2
13.4
2.5
1.3
20
E0
6.7
16.7
23.8
35.8
39.1
39.0
33.3
14.7
2.6
1.6
E200
6.9
14.2
18.9
28.8
39.3
48.6
23.6
2.9
1.6
-20
E500
9.4
13.6
15.1
17.2
18.5
19.3
27.8
74..9
7.2
2.8
-40
E500
4.3
4.0
3.3
2.9
2.9
3.0
51.2
14.7
60
Eyon
Figure 14: Rank correlation between true measurements for losses at different diffusion timesteps on CIFAR-2.
Figures 17 and 18 also shows that influence functions tend to overestimate the magnitude of the change in measurement after removing the training data points. This is in contrast to the observation in (Koh et al., 2019) in the supervised setting, where they found that influence functions tend to underestimate the magnitude of the change in the measurement. There are many plausible reasons for this, ranging from the choice the Hessian approximation ((Koh et al., 2019) compute exact inverse-Hessiane Hessian-vector products, whereas we approximate the GGN), to the possible "stability" of the learned distribution in diffusion models even when different subsets of data are used for training (Observation and (Kadkhodaie et al., 2024)).
Observation 3 Lastly, the observation that marginal log-probability remains essentially constant for models trained on different subsets of data is based Figures 19a to 19c.
In that figure, we reproduce a similar behaviour to that observed by Kadkhodaie et al. (2024) - the samples generated from the models are all different for models trained subsets below certain size, and they are nearly identical for models trained on subsets past critical size. This "collapse" also happens for the models' liglliiel (marginal probabilyyy density), well as the other
28

---

Published as a conference paper at ICLR 2025
Influence measurement
Square
810
Ego
fso
fao
E100
$200
Eooo
Eaon
Egsa
norm
-3.5
-133
-0.4
-0.3
-0.7
-0.9
-2.4
-8.5
-3.3
-2.5
0.1
10
010
-2.9
Rank correlation (%)
True mensurenen
820
-3.2
-3.0
2.77
-2.4
-266
-2.9
-8.3
-28.9
-11.7
-9.8
20
Es0
-2.1
-2.1
-2.3
-2.4
-2.8
-3.3
-9.4
-33.5
-11.9
-10.0
1.3
Eso
-2.5
-2.0
-1.9
-2.1
-2.6
-3.1
+9.3
-34.8
-11.4
-9.4
1.0
-2.0
-1.9
-2.1
-35.2
-11.3
E200
-2.8
-1.6
-1.6
-1.9
-2.0
-2.4
-6.6
-34.9
-10.7
-8.8
1.3
20
Eso0
0.4
0.1
0.5
0.8
1.7
2.4
6.4
16.8
-1.1
-0.7
-0.3
EsD0
2.7
2.2
1.1
0.5
0.3
1.6
15.5
45.7
38.7
-2.3
Egon
-40
Figure 15: Rank correlation (LDS scores) between influence function estimates with different measurement functions and different true measurements CIFAR-2 2, but with the retrained models trained on the full dataset with a random subset of examples having down-weighted contribution to training loss by a factor of x0.5.
True measurement
f1
F10
820
Ego
Fgo
1100
F200
8000
fano
Fggg
1
6.1
4.7
3.8
3.7
3.8
-0.1
210
4.8
True measurenent
820
4.3
13.4
16.6
18.2
18.2
18.3
17.0
-9.1
-5.8
-0.9
Rank correlation (%
150
4.4
14.7
19.4
23.4
24.2
24.4
22.6
-13.1
-6.7
-1.3
E30
4.6
14.9
19.8
24.6
25.9
26.3
24.8
-14.4
-6.8
-1.5
1100
14.9
19.9
24.9
$200
4.5
14.0
18.5
23.2
26
26.0
28.1
-12.8
-6.9
-1.5
-10
1600
-1.0
-7.2
-10.3
-14.0
-15.0
-15.2
-13.2
29.9
-1.7
0.5
-1.6
-6.1
-6.6
-6.6
-6.7
-6.6
-1.7
19.3
5.2
-20
Eg9
-0.2
-0.7
-0.7
-1.1
-13
-133
-1.1
0.1
3.8
-2.0
Figure 16: Rank correlation between true measurements for losses at different diffusion timesteps on CIFAR-2, but with the retrained models trained on the full dataset with random subset of examples having down-weighted contribution to a training loss by factor of x0.5.
measurements considered in this work (ELBO and training loss). Past a certain subset size threshold, ELBO, loss and the log-likelihood of models retrained different subsets become identical. Before that subset size threshold, there is some difference among the models. The threshold seems to be smaller for the ELBO compared the loss measurement. We have computed these quantities with models trained on uniformly dequantised CIFAR-I4 -10 Song et al. (2021b) so that we can meaningfully compute the marginal log- trg-prbbbbylity density, but we have also verified this observation for the ELBO and the training loss measurements regular STFAR-10 C. Each measurement was computed with Monte-Carlo estimate using 5000 samples, except for log-likelihood, where the measurement was averaged over 5 uniformly dequantised samples.
I LDS RESULTS FOR THE MARGINAL LOG-PROBABILLN MEASUREMENT
The results for the "marginal log Nog-probability" measurement shown in Figure 21. This measurement uses the interpretation of the diffusion model as a score estimator Song et al. (2021b), and computes the marginal hooprobbyyylity density assuming normalising flow reformulation of the model Song
29

---

Published as a conference paper at ICLR 2025
x10
CIFAR-2
x10
CIFAR-10
Frequency (# of samples bin)
Actual change in f(x., 0)
10a
1.0
10a
0.6
0.4
10
0.0
10a
-0.050
0.025 0.075
Estimated change in ((x,0)
Estimated change in f(x,, 0)
Figure 17: Change in diffusion loss in Section 2.1 when retraining with random subsets of 50% of the training data removed, as predicted by K-FAC influence (æ-axis), against the actual change in the measurement (y-axis). Results are plotted for measurements ((z, 0) for 50 samples generated from the diffusion model trained on all the data. The scatter color indicates the sample for which the change in measurement is plotted. The figure shows that influence functions tend to overestimate how often the loss will decrease when some training samples are removed; in reality, happens quite rarely.
et al. (2021b). Note that this will general be different from the marginal log-density under the stochastic differential equation (SDE) model that used for sampling Song al. (2021a). For these measurements, we trained models on the dequantised CIFAR-2 2 dataset as described in Song al. (2021b). This was done that the empirical data distribution which train the models has a density. Note that this different from the setting targeted by Ho et al. (2020) which the data is assumed live in discrete space.
EXPERIMENTAL DETAILS
In this section, we describe the implementation details for the methods and baselines, as well as the evaluations reported in Section 4.
J.1 DATASETS
We focus the following datasets in this paper:
CIFAR-10 CIFAR-10 is a dataset small RGB images size 32 x 32 Krizhevsky (2009). We use 50000 images (the train split) for training.
CIFAR-2 For CIFAR-2, we follow Zheng al. (2024) and create subset of CIFAR-10 with 5000 examples of images only corresponding to classes car and horse. 2500 examples of class car and 2500 examples class horse are randomly subsampled without replacement from among all CIFAR-10 images of that class.
ArtBench-10 The ArtBench-10 dataset (Liao et al., 2022) is a dataset 60000 artworks from 10 artistic styles. The RGB images of the artworks standardised 256 256 resolution. We use the full original train-split (50000 examples) from the original paper (Liao et al., 2022) for our experiments.
30

---

Published as a conference paper at ICLR 2025
x10
x10
Actual change in 4
Actual change in E50
0.4
-0.2
-0.1
Estimated change in &
Estimated change in Esa
2
Actual change 0100
Actual change in E000
-0.01
0.02
Estimated change in 0003
Estimated change in E000
10
10a
10a
10a
Frequency (# of samples in bin)
Figure 18: Change in terdiffusion-timeste losses l when retraining with random subsets 50% of the training data removed, predicted by K-FAC influence (æ-axiis), against the actual change in the measurement (g-axis). Results are plotted for the CIFAR 2 dataset, for measurements & (2, 0) for 50 samples æ generated from the diffusion model trained on all of the data. The scatter color indicates the sample x for which the change in measurement is plotted. The figure shows that: 1) influence functions predict that the losses will increase or decrease roughly equally frequently when some samples are removed, but, in reality, the losses almost always increase; 2) for sufficiently large time-steps (4500), this pattern seems subside. Losses l the 200 500 range seem to work well for predicting changes in other losses Figure 13.
J.2 MODELS
CIFAR For all CIFAR datasets, we train regular Denoising Diffusion Probabilistic Model using standard U-Net architecture described for CIFAR-1 0 in Ho et al. (2020). This U-Net architecture contains both convolutional and attention layers. We use the same noise schedule as described for the CIFAR dataset in Ho et al. (2020).
ArtBench For the ArtBench- 10 experiments, we use Latent Diffusion Model (Rombach et al., 2022) with the diffusion backbone trained from scratch. The architecture for the diffusion in the latent space is based on a U-Net with transformer layers, and fully described in the codebase at https: github. com/ Brund BrunoKM/ diffu fflfffsiin mensiinniii coon-nnflun ifluneee. For the latent-space encoder- decoder, we use the pretrained autoencoder from Stable Diffusion version 2 (Rombach et al., 2022), and fix its parameters throughout training.
31

---

Published as a conference paper at ICLR 2025
Suhbett size
10
50
500
5000
10000
25000
Correlation
1.0
0.323
01.33
0,98
0.97
1.00
1.00
1.00
1.00
1.00
1.00
0.82
0.84
0.5
E
Spearmann
Pearson
Logelikollhos
Log- Log-lllellhood for 1st model
Log- likelhhood for 1st model
Log- likelihood for 1st model
Log- likelhhood for 1st model
Log- hog-llkolthood for 1st model
Log Log-Ilkelihood for 1st model
Samples generated from same random noise
(a) Correlations in the marginal log- oog-rrobybility density for models trained on independent training sets of different sizes.
Subset size
10
500
5000
10000
25000
1.0
Correlation
0.91
0.99
0.79
0.82
0.299 (.80
0.5
E
Spearman
Pearson
0.0
model
Troes
Troes
Looe
Looes
Looe
Loee
for 1st. model
for 1st. model
for 1st. model
for Tat: model
For 1at: model
For 1at model
Correlations in measurement models trained on independent training sets of different sizes.
Subset size
Correlation
1.0
1.00
1.00
1.00
1.00
1.00
1.00
0.90
0.32
0.30 0.91
0.5
Spearman Pearson
ELBO for model
ELBO for 1at. model
ELBO for 1at. model
ELBO for 1at. model
ELBO for Tat. model
BLBBO for 1at model
BIBBO for 1at model
(c) Correlations in the ELBO measurement for models trained on independent training sets of different sizes.
Figure 19: Marginal probabbbayyyity density, the generated samples, the ELBO and the loss measure- ments eventually collapse to being near-identical for diffusion models trained on independent training sets of sufficient size.
32

---

Published as a conference paper at ICLR 2025
ELBO weighting
Training loss weighting
10a
10
200
100
600
800
1000
Diffusion timestep
Figure 20: The diffusion loss and diffusion ELBO formulated in (Ho et al., 2020) (ignoring the reconstruction term that accounts for the quantisation of images back to pixel space) are equal up to the weighting of the individual per-diffusion-timent parameters. This plot illustrates the relatives difference in the weighting for per-diffusion-timsste losses applied in the ELBO in the training loss.
loss terms and constant independent of the
CIFAR-2 (uniformly dequantised)
4. 19
CLIP Casine Similarity
1. 1.0%...
(1.49)
TRAK (m.. ELBO)
10..7
(0.49))
TRAK (m. loso)
14.9%zoa
(0.398)
D-TRAK
10.1% 12
(10.290)
K-FAC Influence (mo ELBO)
2o.
K-FAC Inffuence (m. loes)
35. 2%..%%...
Exact Retraining
Rank Correlation % (LDS)
Figure 21: Linear Data-modelling Score (LDS) for the marginal log-probability density
surement. The plot follows the same format as that of Figures 2a and 2b, with the exception that the colours instead only represent method family groupings. Overall, the method and the proxies proposed in this paper seem to work well on estimating the changes in the marginal log- probability.
J.3 SAMPLING
For the DDPM models (CIFAR), we follow the standard DDPM sampling procedure with a full 1000 timesteps create the generated samples as described by Ho (2020). DDPM sam- pling usually gives better samples (in terms of visual fidelity) than Denoising Diffusion Implicit Models (DDIM) sampling Song et (2022) when large number of sampling steps used. As described Section 2.1, when parameterising the conditionals pe (z)))))) with neural networks we have choice in how to set the variance hyperparame- ters (a?) - The of hyperparameters do not appear in the training loss; however, they do make difference when sampling. We use the "small" variance variant from Ho et al. (2020, $3.2), i.e. we
(1( 11111 (1) ()(=x)))),
set:
A (1 Ax)
of [T A
For ArtBench-11 experiments, we follow (Rombach et al., 2022) using the full 1000 timesteps for sampling in the latent space before decoding the sample.
33

---

Published as a conference paper at ICLR 2025
J.4 DETAILS ON DATA ATTRIBUTION METHODS
TRAK For TRAK baselines, we adapt the implementation of Park et al. (2023); Georgiev et al. (2023) to the diffusion modelling setting. When running TRAK, there are several settings the authors recommend to consider: 1) the projection dimension døro} for the random projections, 2) the damping factor X, and 3) the numerical precision used for storing the projected gradients. For (1), we use a relatively large projection dimension of 32768 as done in most experiments in Zheng al. (2024). We found that the projection dimension affected the best obtainable results significantly, and so we couldn't I get away with a smaller one. We also found that using the default float16 precision in the TRAK codebase for (3) results in significantly degraded results (see Figure 22, and so we recommend using Float32 precision for these methods for diffusion models. In all experiments, we use f Float 32 throughout. For the damping factor, we report the sweeps over LDS scores in Figures 8 and 10, and use the best result each benchmark, as these methods fail drastically if the damping factor is too small. The damping factor reported in the plots is normalised by the dataset size N, to match the definition of the GGN, and to make comparable with the damping reported for other influence functions methods introduced in this paper. For non-LDS experiments, we use the best damping value from the corresponding LDS benchmark.
CLIP cosine similarity One of the data attribution baselines used for the LDS experiments is CLIP cosine similarity (Radford et al., 2021). For this baseline, we compute the CLIP embeddings (Radford et al., 2021) of the generated sample and training datapoints, and consider the cosine similarity between the two as the "influence" of that training datapoint on that particular target sample. See (Park et al., 2023) for details of how this influence is aggregated for the LDS benchmark. Of course, this computation does not in any way depend on the diffusion model or the measurement function used, so is a pretty naive method for estimating influence.
-FAC We use curvl inpp (Dangel et al., 2025) package for our implementation of (E)K-FAC for diffusion models. Except where explicitly mentioned otherwise, we use the K-FAC (or EK-FAC) expand variant throughout. We compute (E)K-FAC for PyTorch nn. Conv2d and nr. Linear modules (including all linear maps in attention), ignoring the parameters in the normalisation layers.
Compression for all K-FAC influence functions results, we use int quantisation for the query gradients.
Monte Carlo computation of gradients and the GGN for influence functions Computing the per-example training loss É(O, x.)) in Section 1, the gradients of which are necessary for computing the influence function approximation (Equation (6)), includes multiple nested expectations over diffusion timestep i and noise added to the data e (t)) This is also the case for the GGN in Equation (9) and for the gradients Vol(O, in the computation of GGN in Equation (11), as well as for the computation of the measurement functions. Unless specified otherwise, we use the same number of samples for Monte Carlo estimation of the expectations for all quantities considered. For example, if we use K samples, that means that for the computation of the gradient of the per-example-loss VoE(O, xn) we'll sample tuples of (ł, e (0)) independently K times to form a Monte Carlo estimate, For GGN we explicitly iterate over all training data points, and draw K samples of (i, (1), In for each datapoint. For GGN]ess we explicitly iterate over all training
data points, and draw K samples of (i,)(0), am (7)
to compute the gradients Vol(O, xx,) before taking
an outer product. Note that, for GGN because we're averaging over the samples before taking the outer product of the gradients, the estimator of the GGN is no longer unbiased. Similarly, K samples are also used for computing the gradients of the measurement function.
For all CIFAR experiments, and GGN computations for K-FAC Influence, TRAK, D-TRAK), unless explicitly indicated in the caption otherwise.
250 samples throughout for all methods (including all gradient
J.5 DAMPING
For all influence function-like like methods (including TRAK and D-TRAK), use damping to improve the numerical stability of the Hessian inversion. Namely, for any method that computes the inverse of the approximation the Hessian H V2LD V/l/N Σ l(0, xn), we add a damping factor X to
34

---

Published as a conference paper at ICLR 2025
CIFAR-2 TRAK Influence
Rank correlation (LDS)
0.15
Targut
Measure Train, Losa
0.10
Loss
Nora (D-TRAZ)
KL30
BL80+
Loos
0.05
$1.00
Eq. Norm
Bq. Sorm (O-TRAE)
0.00
10 10-7
10-5
10
10
Damping factor
Figure 22: LDS scores on for TRAK (random projection) based influence on CIFAR-2 when using half-precision (float16) for influence computations. Compare with Figure 8. NaN results indicated with x.
the diagonal before inversion:
(H + X1)-1
where / is dparam dpa-am identity matrix. This is particularly important for methods where the Hessian approximation at high risk of being low-rank (for example, when using the empirical GGN in Equation (11), which is the default setting for TRAK and D-TRAK), For TRAK/D-TRAK, the approximate Hessian inverse is computed a smaller projected space, and so we add À to the diagonal directly in that projected space, done in Zheng et al. (2024)). In other words, if P E Ro Xdearan is the projection matrix (see (Park et al., 2023) for details), then damped Hessian- inverse preconditioned vector inner products between vectors U1, 02 E Rddassx (e.g. the gradients in Equation (6)) would be computed as:
(Pu)) (H + Pu.-
where H R PVLLDPT Rdoros Xdgeos is approximation to the Hessian in the projected space.
For the "default" values used for damping for TRAK, D-TRAK and K-FAC Influence, we primarily follow recommendations from prior work. For K-FAC Influence, the default is small damping value 10 throughout added for numerical stability of inversion, as done in prior work (Bae et al., 2024). For TRAK-based methods, Park al. (2023) recommend using damping: "[...] computing TRAK does not require the use of additional regularization (beyond the one implicitly induced by our use of random projections)" (Park et al., 2023, 6). Hence, 10-9 as the default value throughout.
the lowest numerically stable value of
Note that all damping values reported in this paper are reported as if being added the GGN for the Hessian of the loss normalised by dataset size This differs from the damping factor in the TRAK implementation (https: qithubb. www.o///oooooooooooo Soon/MadzyLaa trak), which added the GGN for the Hessian of an unnormalised loss (E,, E(O, )). Hence, the damping values reported in (Zheng al., 2024) are larger by factor of N (the dataset size) than the equivalent damping values reported in this paper.
J.6 LDS BENCHMARKS
For all CIFAR LDS benchmarks Park et al. (2023), we sample 100 sub-sampled datasets (M = 100 in Equation (13)), and we train 5 models with different random seeds (K 5 in Equation (13)), each with 50% of the examples in the full dataset, for total of 500 retrained models for each benchmark. We compute the LDS scores for 200 samples generated by the model trained the full dataset.
The only difference from the above for the Art Bench experiments is that we sample 50 sub-sampled datasets (M 50 in Equation (13)). This gives total of 250 retrained models for this benchmark.
Monte Carlo sampling of measurements For all computations of the "true" measurement functions for the retrained models in the LDS benchmarks use 5000 samples estimate the measurement.
35

---

Published as a conference paper at ICLR 2025
J.7 RETRAINING WITHOUT TOP INFLUENCES
For the retraining without top influences experiments (Figure 3), we pick 5 samples generated by the model trained on the full dataset, and, for each, train model with fixed percentage of most influential examples for that sample removed from the training dataset, using the same procedure as training on the full dataset (with the same number of training steps). We then report the change in the measurement on the sample for which top influences were removed.
Monte Carlo sampling of measurements Again, for all computations of the "true" measurement functions for the original and the retrained models used for calculating the difference in loss after retraining we use 5000 samples to estimate the measurement.
J.8 TRAINING DETAILS
For CIFAR- 10 and CIFAR 2 again follow the training procedure outlined in Ho et al. (2020), with the only difference being a shortened number of training iterations. For CIFAR-10, 10, we train for 160000 steps (compared to 800000 in Ho et al. (2020)) for the full model, and 80000 steps for the subsampled datasets (410 epochs in each case). On CIFAR- 2, we train for 32000 steps for the model trained on the full dataset, and 16000 steps for the subsampled datasets (~ 800 epochs in each case). We train for significantly longer than Zheng et al. (2024), as we noticed the models trained using their procedure were noticeably undertrained (some te-ddiffusion-iimmstte training losses & (0) have not converged). We also use a cosine learning-rai schedule for the CIFAR-2 2 models.
For Art we use the pretrained autoencoder from Stable Diffusion 2 (Rombach al., 2022), but we train the diffusion backbone from scratch (initialise randomly). We follow the training procedure in (Rombach et al., 2022) and train the full model for 200000 training iterations, and the models trained on the subsampled data for 60000 iterations. We use linear warm-up for the learning rate schedule for the first 5% of the training steps. We use the Adamm optimiser with learning rate of 10 weight-decay of 10-6 gradient norm clipping of 1, and exponential moving average (EMA) with maximum decay rate of 0. 0.9999 and EMA warm-up exponential factor of 0.75 (see the https: / /githbb. com/ hugg in face / dii parameters). We don't 't use cosine learning rate decay. We only train the diffusion backbone, leaving the original pretrained autoencoder unchanged.
ffusers library for details on the EMA
J.9 HANDLING OF DATA AUGMENTATIONS
In the presentation in Section 2, we ignore for the sake of clear presentation the reality that in most diffusion modelling applications we also apply data augmentations to the data. For example, the training loss Lp in Equation (3) in practice often takes the form:
1 N CD =
Σ n=1
E., [8(0, 2,),
where än is the data point In after applying a (random) data augmentation to it. This needs to be taken into account 1) when defining the GGN, as the expectation over the data augmentations IE,,,, can either be considered as part of the outer expectation F.,, or as part of the loss p (see Section 2.3), 2) when computing the per-example train loss gradients for influence functions, 3) when computing the loss measurement function.
When computing GGN in Equation (9), we treat data augmentations being part of the out "empirical data distribution". In other words, we would simply replace the expectation E,, in the definition of the GGN with nested expectation Ex, Ex.:
GGN
(0) Ez
Ex
(()
(0(
(21)
with 2 now being sampled from the diffusion process (xx (T.,) conditioned on the augmented sample ăn The terms changing from the original equation are indicated in yellow. The "Fisher" expression amenable MC sampling takes the form:
(0) Ez,
Ex..
E,
E_to
E
(0)
(
Erve (0)
Emod ~ N (2.0) 1).
36

---

Published as a conference paper at ICLR 2025
where, again, 9n. (0) = VollEmod (1. )12
When computing GGN in Equation (11), however, we treat the expectation over daea augmenta- tions as being part of the loss p, in order to be more compatible with the implementations of TRAK (Park et al., 2023) prior works that rely on an empirical GGN (Zheng et al., 2024; Georgiev et al., 2023). Hence, the GGN in Equation (11) takes the form:
G
Loss Exa
Vo (Ez, (((0,Tn)) VJ (Ez, #n),
((0,...)
EEn
Vol(o, kn) V
where ? is the per-example loss in expectation over data-augmentations. This is how the Hessian approximation is computed both when we're using K-FAC with GGN in presence data augmentations, or when we're using random projections (TRAK and D-TRAK).
When computing the training loss gradient in influence function approximation in equation Equa- tion (5), we again simply replace the per-example training loss ((0",z,) with the per-example training loss averaged over data augmentations {(0", Tr so that the training loss Lp still be written finite sum of per-example losses as required for the derivation of influence functions.
For the measurement function m Equation (6), we assume we are interested in the log probability of (or loss on) a particular query example in the particular variation in which it has appeared, so we do take data augmentations into account in the measurement function.
Lastly, since computing the training loss gradients for the influence function approximation for diffusion models usually requires drawing MC samples anyways (e.g. averaging per-diffusion timestep losses over the diffusion times and noise samples ((t)), we simply report the total number of MC samples per data point, where data augmentations, diffusion time etc. are all drawn independently for each sample.
1hee implementations of these methods store the (randomly projected) per-example training loss gradients for each example before computing the Hessian approximation. Hence, unless data augmentation is to be part of the per-example training loss, the number of gradients to be stored would he increased by the number of data augmentation samples taken.
considered
37


Reviews:
[{"review_id": "oBhPJRWtgz", "reviewer": "Reviewer_nmGh", "summary": "This paper extends influence functions for data attribution of diffusion models. With an exposition focusing on approximating a Hessian with a generalized Gauss-Newton matrix (GCN), this paper justifies specific design choices (linearizing the neural networks combined with K-FAC-expand approximation) that lead to better performing influence functions. The theoretical framework also subsumes previous work such as Journey-TRAK and D-TRAK. Empirically, influence functions perform similarly to D-TRAK on CIFAR-2 and CIFAR-10, with lower sensitivity to the choice of the damping parameter. Finally, empirical observations are made to show the challenges when applying influence functions to diffusion models.", "questions": "- In practice, how would we tune the damping parameter for K-FAC Influence? How would the tuning procedure differ from that for D-TRAK?\n- Observation 2 shows that F-FAC Influence overestimates the impact of data removal, whereas Koh et al. (2019) show that influence functions underestimates the impact of data removal. What are possible explanations for these differing observations?\n- Observation 3 seems to have been shown in Kadkhodaie et al. (2024; Figure 2 in particular). What additional insights are provided? For example, does Observation 3 only hold for ELBO, or does it also hold for the simple diffusion loss?\n\n\nReferences\n- Koh et al. (2019) - On the Accuracy of Influence Functions for Measuring Group Effects.\n- Kadkhodaie et al. (2024) - Generalization in Diffusion Models Arises from Geometry-Adaptive Harmonic Representations.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- A theoretical framework is shown to unify influence functions, Journey-TRAK, and D-TRAK for diffusion models. This provides some clarity to the field of data attribution for diffusion models, which currently seems rather empirical.\n- The theoretical framework motivates the design choices for approximating influence functions, which actually lead to better performance (as shown in Figure 4).\n- Empirical observations (Section 4.1) are made to better understand how to apply influence functions to diffusion models.", "weaknesses": "- The LDS and counterfactual evaluations are limited to DDPM on CIFAR. The paper would be strengthened if K-FAC Influence is also evaluated on LDM and LoRA-fine-tuned Stable Diffusion (as in the Journey-TRAK and D-TRAK papers). Furthermore, since the observations in Section 4.1 are empirical, not theoretical, they need to be validated in other model architectures and datasets to make the claims more general.\n- The proposed method K-FAC Influence is also sensitive to the choice of the damping parameter (LDS can range from ~15% to ~5% on CIFAR-10). This range of variation seems not much different from D-TRAK.\n- Observation 3 might be sensitive to the amount of training data removed. For example, if 90% of the data are removed, there could be more variation in the retrained diffusion models, as suggested by Kadkhodaie et al. (2024; Figure 2 in particular). In this sense, the claim in Observation 3 might need to be adjusted to be less general.\n- This is minor: It would be useful to include a brief summary of Grosse and Martens (e.g., the reduce vs. expand settings), perhaps in the Appendix.\n\nReferences\n - Kadkhodaie et al. (2024) - Generalization in Diffusion Models Arises from Geometry-Adaptive Harmonic Representations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "MjbByZd6m1", "reviewer": "Reviewer_cA9c", "summary": "In this work the authors adapt influence functions to work with diffusion models. Computing influence functions typically requires one to compute the inverse of the Hessian, which is intractable for deep learning models. Therefore the authors make use of a series of approximations. Firstly, they approximate the Hessian with the generalized Gauss-Newton matrix (GGN). Then they approximate the GGN with a block diagonal matrix, where the diagonal blocks are defined by layers. Finally, they approximate each block as a kronecker product.\n\nThe authors then evaluate their method on “linear data modelling score” and “retraining without top influencers”, and find their methods achieve state-of-the-art on both metrics.", "questions": "You explicitly mention that K-FAC is defined for linear and convolutional layers (page 6), but this doesn’t seem to include attention layers, which all diffusion architectures that I’m aware of use. Is K-FAC also defined for attention in prior work, or could you define it?\n\nCould you provide a runtime analysis, or some idea of what the runtime is (e.g. how much time would it take to compute a all influence values for a single sample on an A100)? How does this runtime (and possibly space requirement) scale with the complexity of the model? How much computational resources would I need if I wanted to run this on my favorite existing diffusion model?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "This paper combines interesting ideas to tractably approximate influence functions for diffusion models, and achieve state-of-the-art performance on training data attribution on diffusion models, which is a challenging problem.", "weaknesses": "While it is clear that the proposed method is much more scalable than computing the raw influence function, it is difficult to assess how scalable this method actually is based on the main body of this paper (e.g. things like what is the runtime, how does it scale with the model size, what are the space requirements). A reader ought to have a good idea of how feasible this would be to implement for themselves.\n\nIt is not clear whether the proposed method outperforms the existing methods given the overlaps in standard error. Perhaps statistical significance could be reported here to give a better idea?\n\nThe experimental evaluation is only performed on CIFARs 2 and 10, so it is unclear how well this method scales to larger image generators, both in the quality of attributions and in terms of compute.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "e3TylMO2dk", "reviewer": "Reviewer_VBdr", "summary": "The paper provides a systematic treatment of tractable influence function (IF) estimation for generative diffusion models. Several forms of IFs are considered: IF of the ELBO, IF of the loss, IF of the path probabilities. These approaches are compared both conceptually and experimentally. The IF of the marginal probabilities of the images is discussed but not implemented.\nA large fraction of the paper deals with the intractability of the Jacobian estimation involved in the calculation of the IF and offers and compares several approximate solutions.\n\nThe paper provides a thorough experimental analysis of the different IF approaches on several datasets and compares the results with the exact retraining method.\n\n[edit] I increased my score due to the new analysis and results provided during the rebuttal phase, which answered some of the opened questions and provided links with the recent literature.", "questions": "In the last paragraph before the discussion (starting from line (507), the authors argue that the ELBO could be a poor proxy of the marginal probability due to the fact that it is rather insensitive to the resampling of up 50% of the dataset, meaning that such a resampling do not meaningfully change the ELBO of the generated images.\n\nI think that this conclusion is incorrect. There is in fact evidence that diffusion training enters a 'generalization phase' after a finite dataset size, after which the generative probabilities remain approximately constant even under a non-overlapping split of the training set. This for example can be seen in (Zahra, 2023).\n\nReferences:\nKadkhodaie, Zahra, et al. \"Generalization in diffusion models arises from geometry-adaptive harmonic representation.\" arXiv preprint arXiv:2310.02557 (2023).", "limitations": "", "rating": 10, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "This work is very valuable for anyone interested in IF estimation in generative diffusion models. This is likely going to be a very important tool due to privacy and copyright concerns associated with image and video generation. \n\nBoth the theoretical discussion and the experimental analysis is systematic and well executed. The paper is very easy to follow and it manages to both provide the intuition and the theoretical justification for its design choices.\n\nI particularly appreciated the discussion of the results, which highlight both the strengths and weaknesses of each approach and also the limitation of the approximate IF estimation when compared with exact retraining.", "weaknesses": "Arguably, the marginal likelihood logp(x) is the ideal target for a IFa analysis. While it is somewhat challenging to compute using the deterministic integration formula with stochastic or deterministic  trace estimation, it is not intractable and in my opinion it should have been analyzed. For example, having the analysis of the IF of logp(x) would help to make sense of the counterintuitive behavior of the IF of the ELBO.\n\nIt would have also appreciated to see an analysis on more tractable low-dimensional models with small networks, where the Hessian can be computed exactly. This analysis can be used to evaluate the relative performance of the Hessian approximation techniques, at least in this simple regime.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)