PaperlayerView exact revision on arXiv

Source-linked AI summary

Influence Functions for Scalable Data Attribution in Diffusion Models

Bruno Mlodozeniec, Runa Eschenhagen, Juhan Bae, Alexander Immer, David Krueger, Richard Turner

arXiv:2410.13850v5cs.LGcs.AI

TL;DR

Diffusion-model data attribution lacks a scalable influence-function approach, despite interest in how training examples affect generated samples. The paper develops a framework with diffusion-specific GGN and (E)K-FAC approximations, and reports stronger performance than prior attribution methods on common evaluations. Its observations also show that influential examples can matter while numerical changes remain difficult to predict.

  • Problem

    Influence functions had not been scalably adapted to general diffusion modeling, where attribution aims to estimate how training data affect the probability of generating a particular example.

    Source: §1 · p. 2

  • Method

    The paper formulates influence functions for diffusion-model proxy measurements and develops scalable Hessian approximations using diffusion-specific GGNs and (E)K-FAC.

    Source: §1 · p. 2, §2.3 · p. 4, §3.1.2 · p. 6

  • Results

    The recommended method outperforms existing diffusion-model data-attribution methods on common evaluations, including LDS and retraining without top influences.

    Source: §1 · pp. 2–3

  • Takeaways & Limitations

    Influence functions can identify examples that affect training loss and are perceptually relevant to generated samples, but do not accurately predict numerical changes in measurements.

    Source: §5 · p. 10

    Influence functions in diffusion models fall short of accurately predicting numerical changes in loss measurements, despite identifying influential examples.

    Source: §5 · p. 10

Abstract

from arXiv · show

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 functions framework. 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 systematically develop K-FAC approximations based on generalised Gauss-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 Linear Data-modelling Score (LDS) or retraining without top influences, without the need for method-specific hyperparameter tuning.

1 INTRODUCTION

The paper develops a scalable influence-function framework for diffusion-model data attribution, addressing the lack of scalable adaptation to general diffusion modeling. It formulates diffusion-specific GGN approximations and reports better attribution performance than existing methods on common evaluations.

1 INTRODUCTION

  • Influence functions estimate how model outputs would change if training examples were down-weighted or removed, but had not been scalably adapted to general diffusion modeling.For diffusion models, the target can be the probability of generating a particular example.Source: §1 · p. 2
  • The framework formulates diffusion-specific GGN matrices and uses (E)K-FAC approximations to make influence-function Hessian computations scalable.The authors discuss the assumptions behind alternative GGN formulations and empirically ablate design choices.Source: §1 · p. 2
  • The proposed method outperforms existing diffusion-model attribution methods on common metrics, including Linear Datamodeling Score and retraining without top influences.The framework also unifies prior methods as design choices and evaluates variations of its GGN approximation and other choices.Source: §1 · pp. 2–3

2 BACKGROUND

The background introduces diffusion-model training, influence functions for estimating effects of data removal, and generalized Gauss-Newton matrices as Hessian approximations. It explains the modeling and approximation concepts used in the later framework.

2.1 DIFFUSION MODELS

  • Diffusion models learn reverse conditionals that approximate the forward process’s reverse conditionals, and train by minimizing expected per-timestep losses over the dataset.The model predicts the noise added to data at a timestep; the training loss is a sum of timestep-specific terms.Source: §2.1 · p. 3

2.2 INFLUENCE FUNCTIONS

  • Influence functions approximate the change in optimal parameters after training examples are down-weighted, then estimate the resulting change in a measurement at a test input.Removing examples is approximated by setting the down-weighting parameter to 1/N and dropping higher-order terms, under suitable assumptions.Source: §2.2 · p. 4

2.3 GENERALISED GAUSS-NEWTON MATRIX

  • The generalized Gauss-Newton matrix is a Hessian approximation formed by linearizing a component of the loss; for locally linear mappings, it equals the Hessian.Different decompositions of the loss yield different valid GGNs when the outer function is convex.Source: §2.3 · p. 4
  • The model-output GGN is exact when the network is linear, while the loss-based GGN is an empirical Fisher and requires a more restrictive linearity condition to equal the Hessian.For nonlinear diffusion losses, the text therefore motivates using the model-output GGN for Hessian approximation.Source: §2.3 · p. 5

3 SCALABLE INFLUENCE FUNCTIONS FOR DIFFUSION MODELS

The framework adapts influence functions to diffusion models by formulating diffusion-specific GGN approximations and making them scalable with (E)K-FAC. It also describes practical attribution computations and proxy measurements for generation probability.

3 SCALABLE INFLUENCE FUNCTIONS FOR DIFFUSION MODELS

  • The framework adapts influence functions to diffusion models, interprets prior attribution methods as design choices, and motivates its own alternatives.Its design space includes different GGN matrices for the diffusion objective and their implicit assumptions.Source: §3 · p. 5

3.1 APPROXIMATING THE HESSIAN

  • The diffusion training loss’s output dimensionality makes explicit Jacobian computation intractable, motivating Monte Carlo formulations of its GGN approximations.The section considers both model-output and loss-based decompositions; the authors choose the model-output GGN as a better-motivated Hessian approximation than the loss-based alternative.Source: §3.1.1 · pp. 5–6

3.1.1 GGN FOR DIFFUSION MODELS

  • For the model-output GGN, sampled model outputs provide auxiliary targets, and the resulting matrix can be interpreted as a Fisher-information matrix but not the Fisher of the marginal model distribution.The formulation avoids having the diffusion noise targets affect the Hessian approximation, unlike the loss-based empirical-Fisher approach discussed for prior work.Source: §3.1.1 · p. 6

3.1.2 (E)K-FAC FOR DIFFUSION MODELS

  • K-FAC makes GGN computation scalable by approximating it with layerwise blocks built from Kronecker products, enabling efficient inversion and vector multiplication.The recommended method uses Monte Carlo expectations, K-FAC-expand for weight sharing, EK-FAC, and damping; K-FAC covers common diffusion layers including convolutions and attention.Source: §3.1.2 · p. 6

3.2 GRADIENT COMPRESSION AND QUERY BATCHING

  • Query batching avoids recomputing training-example gradients for multiple attributed samples, while quantization-based gradient compression works better than SVD-based compression in this setting.The authors suggest the difference may reflect that diffusion-model gradients are not low-rank.Source: §3.2 · p. 7

3.3 WHAT TO MEASURE

  • Because exact marginal log-probability evaluation and its gradients are challenging in most settings, the framework considers loss, sampling-trajectory probability, and ELBO as proxy measurements.The marginal log-probability is a natural target for asking how training examples affect the probability of generating a sample.Source: §3.3 · p. 7

4 EXPERIMENTS

K-FAC Influence outperforms prior attribution methods on the reported LDS and counterfactual retraining evaluations. The experiments also reveal limits: proxy measurements can outperform the matching measurement, and influence predictions can diverge from observed loss changes.

4 EXPERIMENTS

  • LDS measures Spearman rank correlation between predicted measurement changes and changes observed after retraining on independently subsampled 50% training sets.The other evaluation removes a fixed percentage of top-ranked examples and measures the resulting change after retraining from scratch.Source: §4 · pp. 7–8
  • TRAK and D-TRAK are more sensitive to damping choices than K-FAC Influence, for which generally any sufficiently small damping value works reasonably well.Tuning damping through an LDS benchmark would require repeated model retraining.Source: §4 · p. 8

4.1 POTENTIAL LIMITATIONS OF INFLUENCE FUNCTIONS FOR DIFFUSION MODELS

  • For predicting ELBO or marginal log-probability changes, substituting the training-loss measurement can yield higher LDS than using the matching measurement.The paper reports this pattern as a peculiarity of the LDS results.Source: §4.1 · p. 8

4 EXPERIMENTS

  • K-FAC Influence outperforms TRAK in all LDS settings, beats benchmark-tuned D-TRAK with the loss measurement, and consistently identifies more influential examples than baselines in counterfactual retraining.The LDS comparison is shown in Figure 2; counterfactual retraining results are shown in Figure 3.Source: §4 · p. 8, Figure 2 · p. 9, Figure 3 · p. 9

4.1 POTENTIAL LIMITATIONS OF INFLUENCE FUNCTIONS FOR DIFFUSION MODELS

  • Higher-timestep losses act as better proxies for lower-timestep losses, and in some cases the higher, incorrect timestep yields better LDS than the matching timestep.The reported comparison varies the per-timestep loss used as the influence approximation’s measurement function.Source: §4.1 · p. 10
  • Influence functions predict both positive and negative loss influence, whereas removing data points predominantly increases loss in practice.Source: §4.1 · p. 10
  • Influence functions tend to overestimate how often removing groups of data points improves loss on generated samples.Source: §4.1 · p. 10
  • With sufficiently large training sets, marginal log-probability is close to constant across removed 50% subsets; ELBO and training-loss measurements show the same pattern.The reported pattern concerns a given sample generated by the model.Source: §4.1 · p. 10

5 DISCUSSION

Influence functions identify examples that matter to training loss and generated samples, but do not accurately predict numerical changes after retraining, especially for marginal-probability proxies.

5 DISCUSSION

  • Influential examples significantly affect training loss when removed and appear perceptually relevant to generated samples.However, influence functions fall short of accurately predicting numerical measurement changes after retraining.Source: §5 · p. 10
  • Numerical prediction is especially weak for measurement functions that proxy the marginal probability of sampling an example.Source: §5 · p. 10
  • Influence functions remain useful for exploratory model understanding and data curation, while copyright-style numerical attribution needs better probability proxies and approximations.Source: §5 · p. 10

A DERIVATION OF INFLUENCE FUNCTIONS

The derivation treats down-weighting training examples as a continuous change in the loss, then uses the implicit function theorem to approximate how the optimum moves.

A DERIVATION OF INFLUENCE FUNCTIONS

A.1 IMPLICIT FUNCTION THEOREM

  • The implicit function theorem guarantees a differentiable local solution when its stated differentiability and invertibility conditions hold.Source: §A.1 · p. 14

A.2 APPLYING THE IMPLICIT FUNCTION THEOREM TO QUANTIFY THE CHANGE IN THE OPTIMUM OF A LOSS

  • For the loss optimum to vary differentiably, the derivation assumes a twice continuously differentiable loss and an invertible parameter Hessian.Source: §A.2 · p. 15
  • Strict convexity ensures the stationary solution is the unique minimum, and a loss linear in ε simplifies the mixed derivative used in Equation (5).Source: §A.2 · p. 15

B DERIVATION OF THE FISHER “GGN” FORMULATION FOR DIFFUSION MODELS

The appendix develops curvature approximations for diffusion models and explains how K-FAC variants handle shared-weight layers and diffusion-specific expectations.

B DERIVATION OF THE FISHER “GGN” FORMULATION FOR DIFFUSION MODELS

C.1 BACKGROUND: KRONECKER-FACTORED APPROXIMATE CURVATURE

  • K-FAC approximates a neural network’s Fisher or GGN with layer-wise block-diagonal curvature, representing each block as a Kronecker product.Source: §C.1 · p. 16
  • Weight-sharing layers apply shared weights across inputs such as attention tokens or convolutional spatial dimensions, motivating expand and reduce loss settings.Source: §C.1 · p. 17
  • K-FAC-expand matches the layer-wise block-diagonal GGN for deep linear networks with mean-square-error loss and corresponds to established convolutional and attention approximations.Source: §C.1 · p. 17
  • K-FAC-reduce is exact for deep linear networks with mean-square-error loss and scaled-sum reduction; the reduction function aggregates the weight-sharing dimension.Source: §C.1 · p. 17
  • Either K-FAC variant can be applied in either setting, with an independent choice for each weight-sharing dimension.Source: §C.1 · p. 17

C.2 DERIVATION OF K-FAC FOR DIFFUSION MODELS

  • The diffusion derivation constructs both K-FAC variants for the model GGN and permits an eigenvalue correction to obtain EK-FAC.Source: §C.2 · p. 18
  • For K-FAC-reduce, the diffusion derivation changes how expectations and sums are handled; K-FAC-expand ignores the target distribution in the approximation.Source: §C.2 · p. 18

C.3 EMPIRICAL ABLATION

The ablations favor MC-Fisher, EK-FAC, and K-FAC-expand, while runtime and scaling analyses describe the compute–memory trade-offs of K-FAC influence and TRAK. The section also reviews related diffusion-attribution methods and notes that Hessian approximation quality was not directly evaluated.

C.3 EMPIRICAL ABLATION

  • MC-Fisher modestly outperforms empirical Fisher; EK-FAC improves every configuration, K-FAC-expand beats reduce, and empirical-Fisher reduce performs worst.The authors note that expand is exact for deep linear networks, while the difference between expand and reduce is difficult to attribute across the differing evaluation settings.Source: §C.3 · p. 19
  • The experiments do not directly evaluate the approximation quality of the Hessian estimates.The authors say results seem to imply that a better Hessian approximation leads to better downstream attribution performance.Source: §C.3 · p. 19

D RELATED WORK

  • Prior diffusion-attribution work extends TRAK, while this paper explores measurement and loss choices and GGN matrices for diffusion models.Kwon et al. (2023) use an empirical-Fisher approximation for LoRA-finetuned diffusion models; Lin et al. (2024) explore alternative measurement functions with TRAK.Source: D RELATED WORK · p. 20

E RUNTIME MEMORY AND COMPUTE

  • Influence implementations trade memory against compute; for periodic queries, the authors propose storing compressed preconditioned training gradients, and report runtimes on one NVIDIA A100 GPU.The runtime discussion covers 200 query points; training-gradient computation can be amortized by storing the compressed gradients.Source: E RUNTIME MEMORY AND COMPUTE · p. 20, §E.2 · p. 23

E.1 ASYMPTOTIC COMPLEXITY

  • For a feed-forward network, K-FAC influence runtime is O((N + Q)SW^2L + NQW^2L + W^3L), with storage O(QW^2L + NQ).The analysis uses N training examples, Q queries, S Monte Carlo samples, width W, and depth L; the feed-forward assumption is illustrative and extends to other common architectures.Source: §E.1 · p. 21

E.1.1 K-FAC AND K-FAC INVERSION

  • Computing K-FAC over the training set costs O(NSW^2L), inversion adds O(LW^3), and the inverse requires O(LW^2) storage.The inversion cost is described as negligible relative to K-FAC computation for larger datasets.Source: §E.1.1 · p. 22

E.1.2 PRECONDITIONED QUERY GRADIENTS COMPUTATION

  • Query-gradient computation costs O(QSW^2L + QW^3L), while training-gradient scoring costs O(NSW^2L + NQW^2L).Storing preconditioned query gradients requires O(QW^2L); storing final scores requires O(NQ).Source: §E.1.2 · p. 22, §E.1.3 · p. 22

E.1.4 COMPARISON WITH TRAK

  • TRAK runtime and memory depend on projection dimension R, whose scaling with network size is unclear and may need to grow proportionally to preserve accuracy in the worst case.The stated costs are O((N + Q)SW^2L + NQR + R^3) runtime and O((N + Q)R + R^2) memory.Source: §E.1.4 · pp. 22–23

E.2.1 RUNTIME RESULTS

The runtime tables cover K-FAC influence and TRAK across datasets, while the memory tables describe their respective implementation overheads.

E.2.1 RUNTIME RESULTS

  • Table 2 reports K-FAC influence-score runtime across datasets, with starred parameters belonging to pretrained model components such as a VAE.The tables report runtimes on a single NVIDIA A100 GPU.Source: Table 2 · p. 23, §E.2.1 · p. 23
  • Table 3 reports TRAK score-computation runtime across datasets, with starred parameters belonging to pretrained components such as a VAE.The tables report runtimes on a single NVIDIA A100 GPU.Source: Table 3 · p. 23, §E.2.1 · p. 23

E.2.2 MEMORY USAGE

F GRADIENT COMPRESSION ABLATION

The compression ablation compares influence scores from compressed and uncompressed query gradients and finds that 8-bit quantisation preserves their ordering almost unchanged.

F GRADIENT COMPRESSION ABLATION

  • 8-bit quantisation causes almost no change to the ordering of training datapoints by influence score.The comparison uses Pearson and rank correlations between scores computed with compressed and uncompressed query gradients.Source: F GRADIENT COMPRESSION ABLATION · p. 24
  • The authors contrast this result with SVD compression, noting that diffusion-loss and other measurement gradients lack the natural low-rank Kronecker structure of some other settings.Computing diffusion-loss gradients requires averaging across noise samples and diffusion timesteps.Source: F GRADIENT COMPRESSION ABLATION · p. 24
  • The paper uses 8-bit quantisation with query-gradient batching for all K-FAC experiments.Source: F GRADIENT COMPRESSION ABLATION · p. 24

G DAMPING LDS ABLATIONS

The damping ablations compare LDS scores for TRAK/D-TRAK and K-FAC influence; the supplied figure captions describe the measurements and estimation setup.

G DAMPING LDS ABLATIONS

  • The ablation reports LDS scores across damping factors for TRAK/D-TRAK and K-FAC influence.TRAK damping factors are normalized by dataset size to correspond to equivalent factors for the paper’s method.Source: G DAMPING LDS ABLATIONS · p. 24
  • For K-FAC influence on CIFAR-2, Figure 7 uses 250 Monte Carlo samples and identifies the target and measure used for ground truth and approximation.Source: Figure 7 · p. 25
  • For TRAK random-projection influence on CIFAR-2, Figure 8 uses 250 Monte Carlo samples and distinguishes target, measure, and training-loss functions in its legend.Source: Figure 8 · p. 25

H EMPIRICAL ABLATIONS FOR CHALLENGES TO USE OF INFLUENCE

The ablations expose systematic mismatches between influence-function predictions and retrained diffusion models, especially for loss changes after data removal. Proxy performance also depends on timestep and retraining conditions, while several measurements converge beyond a subset-size threshold.

Damping ablations

  • Damping sweeps evaluated LDS for TRAK and K-FAC on CIFAR-10 and CIFAR-2, including marginal log-probability measurement on dequantised CIFAR-2.The captions report 250 Monte Carlo samples for the relevant estimates; no sweep outcome is stated in these passages.Source: Figure 9 · p. 26, Figure 10 · p. 26, Figure 11 · p. 27

Observation 1

  • With 0.5-weighted retraining, larger-timestep proxies had higher absolute rank correlation but the opposite sign, while influence functions failed to predict ℓt changes correctly for t ≤ 200.Different-timestep losses were also negatively correlated in the retrained models.Source: H EMPIRICAL ABLATIONS FOR CHALLENGES… · pp. 26–27

Observation 2

  • After removing half the data, losses usually increased in retrained models, although influence functions predicted increases and decreases about equally often and overestimated change magnitudes.This pattern held across CIFAR-2 and CIFAR-10; on CIFAR-2, actual changes in ℓ1, ℓ50, and ℓ100 were always positive.Source: H EMPIRICAL ABLATIONS FOR CHALLENGES… · pp. 27–28

Observation 3

  • Across sufficiently large training subsets, generated samples and marginal log-probability, ELBO, and training-loss measurements became nearly identical across models.The subset-size threshold appeared smaller for ELBO than for training loss.Source: H EMPIRICAL ABLATIONS FOR CHALLENGES… · pp. 28–29

I LDS RESULTS FOR THE MARGINAL LOG-PROBABILITY MEASUREMENT

The section evaluates marginal log-probability attribution using a density measurement derived from a score-estimator interpretation of diffusion models. Its measurement setup uses dequantised CIFAR-2, and the accompanying loss-removal result shows a mismatch between predicted and actual changes.

I LDS RESULTS FOR THE MARGINAL LOG-PROBABILITY MEASUREMENT

  • The marginal log-probability measurement uses a score-estimator interpretation and a normalising-flow reformulation of the diffusion model.The resulting density can differ from the marginal log-density of the SDE model used for sampling.Source: I LDS RESULTS FOR THE MARGINAL… · pp. 29–30
  • When half the training data was removed, influence functions overestimated how often diffusion loss would decrease; actual decreases were rare.Figure 17 compares predicted changes with measured changes for generated CIFAR samples.Source: Figure 17 · p. 30

J EXPERIMENTAL DETAILS

J.1 DATASETS

The appendix specifies the datasets, model and attribution implementations, evaluation procedures, training settings, and treatment of data augmentation used in the experiments.

J.1 DATASETS

  • Experiments use CIFAR-10, a 5,000-example car-and-horse CIFAR-2 subset, and 50,000 ArtBench-10 training images.CIFAR-2 contains 2,500 randomly subsampled images from each class; ArtBench images are 256 × 256 RGB artworks across 10 styles.Source: §J.1 · p. 30

J.2 MODELS

  • CIFAR models use DDPM U-Nets, while ArtBench uses a latent diffusion model; figure captions report measurement collapse across sufficiently large subsets and good LDS performance for marginal log-probability.The diffusion loss and ELBO differ in per-timestep weighting and a parameter-independent constant, ignoring the reconstruction term.Source: §J.2 · p. 31, Figure 19 · p. 32, Figure 20 · p. 33

J.3 SAMPLING

  • CIFAR sampling uses the full 1,000 DDPM timesteps with the “small” variance variant, and ArtBench sampling uses 1,000 latent-space timesteps before decoding.Source: §J.3 · p. 33

J.4 DETAILS ON DATA ATTRIBUTION METHODS

  • TRAK uses 32,768-dimensional projections and float32; (E)K-FAC covers convolutional and linear layers, while CLIP cosine similarity serves as a model-independent baseline.CIFAR influence computations generally use 250 samples; Monte Carlo estimates also apply to gradients, GGN terms, and measurements, with a noted biased GGN estimator when averaging precedes the outer product.Source: §J.4 · p. 34

J.5 DAMPING

  • Damping stabilises approximate Hessian inversion, with default values of 10^-8 for K-FAC Influence and 10^-9 for TRAK-based methods.TRAK/D-TRAK add damping in projected space; reported damping is normalised for the dataset-size-normalised GGN, and float16 TRAK influence results are also examined.Source: §J.5 · p. 35, Figure 22 · p. 35

J.6 LDS BENCHMARKS

  • CIFAR LDS benchmarks use 100 subsampled datasets and five retrained seeds per dataset, while ArtBench uses 50 subsampled datasets.Each benchmark computes LDS for 200 full-data generated samples; true measurements use 5,000 Monte Carlo samples.Source: §J.6 · p. 35

J.7 RETRAINING WITHOUT TOP INFLUENCES

  • Retraining without top influences removes a fixed percentage of the most influential examples for each of five generated query samples and reports the resulting measurement change.Original and retrained measurements use 5,000 Monte Carlo samples.Source: §J.7 · p. 36

J.8 TRAINING DETAILS

  • CIFAR full-data and subsampled models use shorter training schedules than the cited Ho et al. procedure, while ArtBench trains its diffusion backbone from scratch with a fixed pretrained autoencoder.CIFAR-2 uses a cosine learning-rate schedule; ArtBench uses warm-up and AdamW settings described in the passage.Source: §J.8 · p. 36
Loading 2410.13850v5…