Influence Functions for Scalable Data Attribution in Diffusion Models

Authors: Bruno Kacper Mlodozeniec, Runa Eschenhagen, Juhan Bae, Alexander Immer, David Krueger, Richard E. Turner
Primary area: interpretability and explainable AI
Keywords: diffusion models, influence functions, generalised gauss newton, ggn, data attribution, hessian approximation, interpretability, curvature, kronecker-factored approximate curvature, k-fac

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.

Decision: accept (oral)
Reviewer scores: 6, 8, 10

Reviews and discussion:
<thread>
creation date:
1730201660336

thread_type:
official_review

<thread_object>
creator: reviewer_vbdr

<content>
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.
A 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.

The paper provides a thorough experimental analysis of the different IF approaches on several datasets and compares the results with the exact retraining method.

[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.

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. 

Both 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.

I 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.

It 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.

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.

I 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).

References:
Kadkhodaie, Zahra, et al. "Generalization in diffusion models arises from geometry-adaptive harmonic representation." arXiv preprint arXiv:2310.02557 (2023).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730435223615

thread_type:
official_review

<thread_object>
creator: reviewer_ca9c

<content>
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.

The 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.

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.

It 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?

The 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.

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?

Could 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730777911522

thread_type:
official_review

<thread_object>
creator: reviewer_nmgh

<content>
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.

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.
- The theoretical framework motivates the design choices for approximating influence functions, which actually lead to better performance (as shown in Figure 4).
- 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.
- 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.
- 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.
- 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.

References
 - Kadkhodaie et al. (2024) - Generalization in Diffusion Models Arises from Geometry-Adaptive Harmonic Representations.

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?
- 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?
- 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?


References
- Koh et al. (2019) - On the Accuracy of Influence Functions for Measuring Group Effects.
- Kadkhodaie et al. (2024) - Generalization in Diffusion Models Arises from Geometry-Adaptive Harmonic Representations.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733313607644

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
ArtBench Results

comment:
As requested by the reviewers, **we ran the LDS scoring in a more challenging setting on the ArtBench dataset (`50000` images, `256x256` resolution) with unconditional Latent Diffusion Models (LDMs).** You can see the [new results figure for the loss measurement here](https://anonymous.4open.science/api/repo/iclr2025-rebuttal-supplementary-C798/file/lds_scores_loss.pdf). We also summarise these ArtBench LDS results below:

| Group | Method | Rank Correlation % (LDS) | Untuned Result |
|-------|--------|-----------------|----------------|
|  | CLIP Similarity | 5.2% ±1.0 | - |
| Correct measurement | TRAK | 9.8% ±1.0 | 7.5% |
| ↳ | **K-FAC Influence** | 10.0% ±0.9 | 8.7% |
| ↳ | **EK-FAC Influence** (eigenvalue-corrected K-FAC) | 16.9% ±0.8 | 15.5% |
| Incorrect measurement | D-TRAK | 13.8% ±1.0 | 1.0% |
| Baseline | Exact Retraining | 21.3% ±0.3 | - |

Firstly, we have updated our K-FAC results with the eigenvalue correction from George et al. (2018) [1], as recommended for influence functions in Grosse, Bae and Anil et al. (2023) [2].  Eigenvalue-corrected Kronecker-factored Approximate Curvature (EK-FAC) is a minor modification to the K-FAC method that adjusts the eigenvalues of the approximation to closer match the approximated GGN. In all experiments, we adjust the number of samples used for K-FAC computation to hold the computational cost approximately fixed (half as many samples for EK-FAC as for K-FAC). The eigenvalue correction appears to consistently and significantly improve the quality of the influence approximation in all our experiments.

Secondly, we found a bug in a public library on which we have been relying for our experiments, which adversely affected the previously reported K-FAC LDS scores. The bug led to the effective damping for bias parameters being higher than for weight parameters. With this fix, we get slightly improved performance on all LDS benchmarks (see the [updated results figure](https://anonymous.4open.science/api/repo/iclr2025-rebuttal-supplementary-C798/file/lds_scores_loss.pdf)). Morevoer, we observe even higher robustness to the choice of the damping factor for EK-FAC Influence (see [this CIFAR-2 figure](https://anonymous.4open.science/api/repo/iclr2025-rebuttal-supplementary-C798/file/cifar2_lds_damping_ablation_bias_fix.pdf) and [this CIFAR-10 figure](https://anonymous.4open.science/api/repo/iclr2025-rebuttal-supplementary-C798/file/cifar10_lds_damping_ablation_bias_fix.pdf)). The previously shared log-likelihood LDS benchmark figure already includes this fix. 
    
We have integrated these new results into the paper and have expanded the summary of prior work on K-FAC in Appendix C to cover EK-FAC. We will also update the new log-likelihood LDS results figure with the ArtBench experiments once those are finished.

[1] George et al. (2018), "Fast Approximate Natural Gradient Descent in a Kronecker-factored Eigenbasis"

[2] Grosse, Bae and Anil et al. (2023), “Studying Large Language Model Generalization with Influence Functions”
    
    
---
    
    
Lastly, we'd like to thank all the reviewers for the feedback on the paper. Interesting points that we haven't considered before have been raised in the reviews, and the paper has become significantly more well-rounded thanks to the discussion and experiments we have included thanks to the reviewers' suggestions.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper extends influence functions to diffusion models, addressing the computational challenges of Hessian inversions with scalable approximations, including Kronecker-factored Approximate Curvature (K-FAC) and its eigenvalue-corrected variant (EK-FAC). The proposed method demonstrates state-of-the-art performance on benchmarks like Linear Data Modeling Score (LDS). The authors provide theoretical insights that unify influence functions with prior work (e.g., TRAK and D-TRAK).