Influence Functions for Scalable Data Attribution in Diffusion Models

Thank you for a detailed review and the feedback! We very much appreciate that you think the work is really valuable for anyone interested in influence functions in generative diffusion models! Here, we will try to address your outstanding points and describe what we have done, or are planning on doing, to address them.

Weaknesses

Arguably, the marginal likelihood logp(x) is the ideal target for a IF analysis.

Yes, you're right. Although deriving an influence function approximation for the $\log p(x)$ measurement would have been difficult (since computing $\nabla_\theta \log p(x)$ is challenging), we could have still considered using $\log p(x)$ as a ground-truth measurement.

The primary reason we haven't done it is that, in the standard DDPM training setup, the data distribution doesn't have a density (the pixel values are discrete). This makes us worried about then interpreting the model trained on that data as a density model (the density values can be made arbitrarily high by just concentrating the density around the discrete pixel values everywhere). The DDPM ELBO acknowledges the fact that the data is discrete by including a reconstruction term (which, in all fairness, is also numerically quite poorly behaved). Hence, the DDPM ELBO targets a different “marginal log-probability” than that computed by the continuous normalising flow reformulation of the model (the latter is a marginal log-density). The works that compute the exact log-likelihoods in the diffusion settings usually train on uniformly dequantised data to get around this issue.

We have CIFAR-2 models trained on uniformly dequantised data that we investigated in an unrelated ablation. We could look at how well influence functions/TRAK with ELBO/loss proxies approximate the exact log-likelihood in that setting?

It would have also [been] appreciated to see an analysis [...] with small networks, where the Hessian can be computed exactly.

We're considering an analysis with a smaller network on CIFAR-2 ($<1$ million parameters), for which we can compute exact inverse Hessian vector products (and exact inverse GGN vector products) with iterative solvers. We're unsure if this will be computationally tractable for trained models that still learn something meaningful from the data, but we will report our results.

Questions:

[...] the authors argue that the ELBO could be a poor proxy of the marginal probability [...] 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.

We were not aware of the results from Zahra Kadkhodaie et al., thank you for pointing these out! The results in Kadkhodaie et al. seem to corroborate “Observation 3”, but we agree they would imply that the ELBO is not necessarily a bad proxy for marginal log-probability of generating a sample. We revised the discussion around “Observation 3” to reference the results from Kadkhodaie et al., and removed sentences conjecturing that observation 3 could imply that ELBO is a poor proxy.

Furthermore, the results of computation of exact marginal probabilities on dequantised data might shed some extra light on this, since we'll be able to look at the correlation between the ground-truth ELBO and log-likelihood measurements (albeit in a slightly different setting).

Given the results from Kadkhodaie et al., it seems likely that observation 3 will be dependent on subset size used for retraining. Therefore, we will also try to check what happens to the correlation of the ELBO measurements between retrained models as we vary the amount of training data. Currently, we are running retraining of models with random subsets of 10, 50, 500, 5000, and 25000 examples from CIFAR-10. These will hopefully allow us to make a closer connection between the generalisation results observed by Kadkhodaie et al. and our results with the ELBO. We will report these results and further revise Observation 3 discussion once they are available.

Thank you for a thorough review and the feedback! The overall verdict on the paper appears positive, which we very much appreciate. Here, we will try to address your outstanding points and describe what we have done, or are doing, to address them.

Weaknessses

[...] it is difficult to assess how scalable this method actually is [...]. A reader ought to have a good idea of [what is the runtime, how does it scale with the model size, what are the space requirements].

That's a good point. We added an appendix section with the goal of giving the reader an impression of the runtime and memory requirements (new Appendix E) of K-FAC Influence. We discuss both asymptotic compute and memory scaling of K-FAC Influence against TRAK, and report empirical runtimes for a single NVIDIA A100 for our experiments. The experiments for which we report the runtimes include ones on a more complex ArtBench dataset which we are currently running an LDS benchmark on.

The new appendix section will, hopefully, also answer your question: “Could you provide a runtime analysis, or some idea of what the runtime is [...]?”.

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 standard error bars are indeed overlapping in the “retraining without top influences” experiments in Figure 3. In response, we ran more experiments to double the number of retraining runs for that experiment. We updated Figure 3 in the paper.

We find that the (one-sided) p-values for the paired difference test of the score across the 4 settings in Figure 3 being greater for K-FAC Influence are: 0.0003% when comparing against TRAK, and 1.7% when comparing against D-TRAK respectively. If you think it'd be informative, we can add the previous sentence to the experiments section.

The p-values for the paired difference tests on individual benchmarks are:

• CIFAR-2 — 2% removed | K-FAC Influence against TRAK: $p=0.00048$
• CIFAR-2 — 10% removed | K-FAC Influence against TRAK: $p=0.0024$
• CIFAR-10 — 2% removed | K-FAC Influence against TRAK: $p=0.093$
• CIFAR-10 — 10% removed | K-FAC Influence against TRAK: $p=0.031$
• CIFAR-2 — 2% removed | K-FAC Influence against DTRAK: $p=0.2$
• CIFAR-2 — 10% removed | K-FAC Influence against DTRAK: $p=0.043$
• CIFAR-10 — 2% removed | K-FAC Influence against DTRAK: $p=0.29$
• CIFAR-10 — 10% removed | K-FAC Influence against DTRAK: $p=0.17$

The experimental evaluation is only performed on CIFARs 2 and 10

You're right, it would be great to demonstrate our method on more complex datasets with larger models. To address this, we are currently running the LDS experiments with a larger latent diffusion model (LDM) on a higher resolution dataset. We are running the benchmarks on the full ArtBench dataset ($256\times 256$ resolution images). We will hopefully update the results in the paper within the discussion period.

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?

Yes, K-FAC can be defined for all linear layers with weight sharing, including attention (Eschenhagen et al. (2023)). Example 1 in Section 2.2 of Eschenhagen et al. (2023) is a good explanation of how the attention mechanism can be thought of as a linear layer with weight sharing. Previously, in the main text body we said that we implement "K-FAC-expand whenever weight sharing is used" (line 292/293), but we have adjusted the text to explicitly mention attention.

We have also added an additional appendix section which summarises the general K-FAC formulation for linear layers with weight sharing from previous work (new Appendix C).

---

Thanks for the feedback. If you have any further thoughts on how to better incorporate the changes mentioned above, please let us know!

Thank you for taking the time to write a thorough review and for all the suggestions! In the following, we will address each listed weakness, describe what we will aim to do about it, and answer the questions that were raised.

Weaknesses

The LDS and counterfactual evaluations are limited to DDPM on CIFAR.

You're right, we agree that our conclusions could be strengthened by considering more complex datasets and latent diffusion models (LDM).

We will do the LDS experiments with an LDM model on a higher resolution dataset, and, hopefully, update the results within the discussion period. We are running the benchmarks on the full ArtBench dataset with a larger, unconditional LDM model trained from scratch. We were thinking of first focusing on unconditional LDMs to limit the number of changes to the previous settings.

The proposed method K-FAC Influence is also sensitive to the choice of the damping parameter [...].

That is a fair point. However, while both methods might benefit from tuning, K-FAC Influence seems to work more consistently well with a default value. Choosing the smallest possible damping value that is numerically stable appears to be a good heuristic, and the value of 1e-8 used in prior non-diffusion literature works well, so it appears like a good universal default (see Figure 7 & 9). For TRAK these two points do not hold (see Figure 8 & 10). Lastly, TRAK-based methods do not work at all if the damping is too small, and what “too small” means seems to vary between the datasets considered.

We decided to tone down the claim that “TRAK and D-TRAK appear to be significantly more sensitive to tuning the damping factor” (line 374) by dropping “significantly”. We would also welcome other suggestions for how to appropriately adjust the wording to describe these differences. Lastly, we updated Figure 2 to make the comparison of what one can expect with tuned and untuned damping factors a little bit clearer.

Observation 3 might be sensitive to the amount of training data removed [...] as suggested by Kadkhodaie et al. (2024) [...]

We were not aware of the results from Kadkhodaie et al. (2024), thank you for pointing these out! These seem to corroborate “Observation 3”, but we agree they would imply that it might be dependent on the subset size. We revised the discussion around “Observation 3” to reference the results from Kadkhodaie et al., and adjusted the wording to make it clear the observation might depend on the dataset size.

We will also try to check what happens to the correlation of the ELBO measurements between retrained models as we vary the amount of training data. Currently, we are running retraining of models with random subsets of 10, 50, 500, 5000, and 25000 examples from CIFAR-10. We will report these results and further revise Observation 3 discussion once they are available.

It would be useful to include a brief summary of Grosse and Martens (e.g., the reduce vs. expand settings) [...].

We will add a more detailed explanation of K-FAC for the general case of linear layers with weight sharing to the appendix before the end of the discussion phase.

Observation 3 might be sensitive to the amount of training data removed.

Observation 3 seems to have been shown in Kadkhodaie et al. (2024). [...] does Observation 3 only hold for ELBO, or does it also hold for the simple diffusion loss?

Given the results by Kadkhodaie et al. (2024), we further investigated “Observation 3” and the behaviour of the ELBO as the training data size varies. Whereas Kadkhodaie et al. (2024) looked at the “collapse” to identical generated samples across different retrained models, we added new experiments that look at the exact log-likelihood, ELBO, and diffusion loss measurements most relevant to the influence setting. Indeed, we do observe that “Observation 3” is subset size dependent, and holds for other measurements like exact log-likelihood or diffusion loss (although the transition occurs marginally ‘earlier’ for the ELBO than for the diffusion loss).

For these experiments, we retrained multiple models on subsets of varying size of (de-quantised) CIFAR-10.

We do reproduce the same behaviour as Kadkhodaie et al. (2024) — the samples generated from the models are all different for models trained on subsets below a certain size, and they are nearly identical for models trained on subsets past a critical size (see figure). This “collapse” also happens for the models' log likelihoods, as well as the ELBO and the diffusion loss. Past a certain subset size threshold, both the ELBO, the loss and the log-likelihood of models retrained on different subsets become identical (see the ELBO results figure, and the loss results figure). Before that subset size threshold, our “Observation 3” does not hold. The threshold seems to be smaller for the ELBO compared to the loss measurement, and therefore we didn't see perfect correlation among the loss measurements on in our original experiments.

That being said, the ELBO and the log-likelihood measurements are not identical, and ELBO is not a perfect proxy for the log-likelihood (as seen from the LDS results). We do observe that the ELBO starts becoming a better proxy for the log-likelihood once the training data sizes are sufficient for the model to enter the generalisation phase (see figure).

---

We will revise section 4.1 to include a discussion of these new results. In particular, we have revised “Observation 3” to say:

For sufficiently large training set sizes, ELBO, log-likelihood and the diffusion loss are close to constant on generated samples, irrespective of which examples were removed from the training data.

We will also start off the discussion of “Observation 3” by pointing out Kadkhodaie et al. (2024) were the first to notice this kind of “collapse”. We merely verify that it holds for the ELBO, log-likelihood and the diffusion loss, which has implications for using and interpreting influence functions in this setting.

Overall, I raised my rating from a 5 to 6.

Thank you for reading and considering our response carefully. We're glad that the response has addressed many of your concerns!

The score is not higher because currently all the empirical results are about the CIFAR datasets. [...] it still remains possible that, on a different dataset and model architecture, the damping parameter needs to be set to a large value as for D-TRAK.

We agree, and we've been working on new experiments to address this. We successfully ran the LDS benchmark 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. We describe the results in a global response to all the reviewers. In short: K-FAC matches TRAK with a matching measurement function, and, after introducing eigenvalue correction for K-FAC recommended in prior work, we observe that eigenvalue-corrected K-FAC improves upon TRAK/D-TRAK in all settings.

In this new ArtBench LDM setting with a new dataset and architecture, we observe similar sensitivity to the damping factor selection: K-FAC/EK-FAC influence functions can have good LDS performance with the default damping parameter $\lambda=1e^{-8}$. You can see the ArtBench damping ablation for TRAK and D-TRAK in this figure and for K-FAC Influence in this figure. We agree this is a strong argument for the use of K-FAC Influence in practice, and we will highlight this point in the paper as you recommended.

Overall, we would like to again thank the reviewer for the feedback and comments on the paper, and suggestions for this and other experiments that we believe have made the paper more complete.