Intriguing Properties of Data Attribution on Diffusion Models

Retraining-based methods offer better efficacy than gradient-based methods.

This claim follows the TRAK paper [2]. In terms of optimizing LDS scores, retraining-based methods like the datamodels approach [3] accurately predict the model output on a target test input by retraining at most 1.5 million models on subsets of CIFAR-10. This approach also outperforms various methods in terms of predicting data counterfactuals. In Appendix A.1 of the TRAK paper [2], the authors show that datamodels (300K models) are the best when estimating prediction brittleness. Thus, with unlimited computation, retraining-based methods offer better efficacy than gradient-based methods by simply retraining many models to obtain accurate estimation.

---

Is there a good reason why D-TRAK should be easier to estimate than TRAK?

We deduce that this is because TRAK suffers from gradient obfuscation or saturation as we mentioned above, and therefore TRAK needs to aggregate information from more timesteps to provide effective attribution. In contrast, D-TRAK can obtain sufficient information from much less timesteps.

---

Can the authors explain the consistent drops in LDS scores when shifting from validation examples to generations?

This is because training examples and validation examples are both sampled from the true data distribution, whereas generated examples are sampled from the model distribution. In practice, there is always a gap between the true data distribution and the learned model distribution, so there are consistent drops in LDS scores when attributing generated data to training data, compared to attributing validation data.

---

Do the authors have any other ideas for why their alternative $\\mathcal{L}$ formulations make sense?

As reported in Table 1, $\\mathcal{L}\_{\\textrm{Square}}$, $\\mathcal{L}\_{\\textrm{Avg}}$, \\mathcal{L}\_{\\textrm{2-norm}}, and \\mathcal{L}\_{\\textrm{1-norm}} consistently outperform $\\mathcal{L}\_{\\textrm{Simple}}$, while \\mathcal{L}\_{\\textrm{\infty-norm}} has worse performance. We deduce that this is because $\\mathcal{L}\_{\\textrm{Square}}$, $\\mathcal{L}\_{\\textrm{Avg}}$, \\mathcal{L}\_{\\textrm{2-norm}}, and \\mathcal{L}\_{\\textrm{1-norm}} can better retain the information from $\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}$, while $\\mathcal{L}\_{\\textrm{Simple}}=\\mathbb{E}\_{t,\\boldsymbol{\\epsilon}}[2\\cdot(\\boldsymbol{\\epsilon}\_{\\theta}-\\boldsymbol{\\epsilon})\^{\\top}\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}]$ becomes ill-behaved when model converges (i.e., $\\boldsymbol{\\epsilon}\_{\\theta}\\approx \\boldsymbol{\\epsilon}$ and thus $\\boldsymbol{\\epsilon}\_{\\theta}-\\boldsymbol{\\epsilon}$ becomes a noisy error term). As to \\mathcal{L}\_{\\textrm{\infty-norm}}, it has worse performance since it discards too much information and only retains the maximum value.

---

Can the authors justify their choice to focus on the LDS metric defined via $\\mathcal{L}\_{\\textrm{Simple}}$?

Indeed, the design choice of defining LDS via $\\mathcal{L}\_{\\textrm{Simple}}$ follows prior work [4], and it seems that $\\mathcal{L}\_{\\textrm{ELBO}}$ may be a more reasonable choice since it provides a lower bound on log-likelihood. Nevertheless, recent work [5] shows that $\\mathcal{L}\_{\\textrm{Simple}}$ can be regarded as a weighted integral of ELBOs over different noise levels, and can be equivalent to ELBO under monotonic conditions.

Besides, in Section 4.4, we empirically perform counterfactual evaluation by measuring the $\\ell_{2}$-distance and CLIP similarity between the images generated before/after the exclusion of top-1000 positive influencers identified by D-TRAK and TRAK. Both the quantitative (Figures 3 and 11) and visualized (Figure 4) results show that D-TRAK is better than TRAK, which is consistent with the trend concluded from the LDS metric.

---

References:
[1] Data Shapley: Equitable Valuation of Data for Machine Learning. ICML 2019
[2] Trak: Attributing Model Behavior at Scale. ICML 2023
[3] Datamodels: Predicting Predictions from Training Data. ICML 2022
[4] The Journey, Not the Destination: How Data Guides Diffusion Models. ICML Workshop 2023
[5] Understanding Diffusion Objectives as the ELBO with Simple Data Augmentation. NeurIPS 2023

Thank you for your supportive review and suggestions. Below we try to offer appropriate explanations for why D-TRAK works and respond to the detailed comments.

---

It would be helpful for the authors to clarify why they can't use ArtBench itself, for example. It seems like the cost of running TRAK/D-TRAK should be manageable, given that retraining is optional in these methods.

Indeed, running TRAK/D-TRAK is computationally efficient and scalable to larger datasets, given that retraining is optional in these methods. However, to compute LDS scores on a large-scale dataset, such as ArtBench, we must first build the corresponding LDS benchmark, which requires retraining a large number of models and accounts for the majority of the computational cost. We tried our hardest and scaled up to CIFAR10 and ArtBench-5 based on our maximum computation budget.

---

Why do the interpolation in Section 3.2?

As reported in Table 1, $\\mathcal{L}\_{\\textrm{Square}}$, $\\mathcal{L}\_{\\textrm{Avg}}$, \\mathcal{L}\_{\\textrm{2-norm}}, and \\mathcal{L}\_{\\textrm{1-norm}} consistently outperform $\\mathcal{L}\_{\\textrm{Simple}}$. The counterintuitive results in Table 1 motivate us to take a closer look at the difference between $\\nabla\_{\\theta}\\mathcal{L}\_{\\textrm{Simple}}$ and $\\nabla\_{\\theta}\\mathcal{L}\_{\\textrm{Square}}$. In Eq. (8), we observe that $\\nabla\_{\\theta}\\mathcal{L}\_{\\text{Simple}}$ consists of two terms: $\\mathbb{E}\_{t,\\boldsymbol{\\epsilon}}\\left[2\\cdot\\boldsymbol{\\epsilon}^{\\top}\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}\\right]$ and $\\mathbb{E}\_{t,\\boldsymbol{\\epsilon}}\\left[2\\cdot\\boldsymbol{\\epsilon}\_{\\theta}^{\\top}\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}\\right]$, where the latter term is exactly $\\nabla\_{\\theta}\\mathcal{L}\_{\\text{Square}}$. Thus, we propose to interpolate between these two terms to see if there is a combination ratio that leads to a higher LDS score.

Surprisingly, as shown in Figure 1, the most theoretically-justified $\\nabla\_{\\theta}\\mathcal{L}\_{\\textrm{Simple}}$ (i.e., $\\eta=0.5$) consistently has the lowest LDS scores. It seems that the poor performance of $\\mathcal{L}\_{\\textrm{Simple}}$ is a result of the cancel-out effect of the above two terms, which might be related to the gradient obfuscation/saturation problem (as discussed in Section 5).

---

The role of non-convexity.

Non-convexity is always a major factor leading to inferior performance of theoretically justified (in convex settings) attribution methods such as TRAK. However, non-convexity may not explain why D-TRAK empirically outperforms TRAK, because there is no evidence that D-TRAK should be superior in non-convex settings.

In addition to non-convexity, we notice another factor influencing TRAK performance, which may be associated with gradient obfuscation or saturation. Specifically, as seen in Figure 2, TRAK achieves its highest LDS score on intermediate checkpoints, whereas D-TRAK achieves its highest LDS score on final checkpoints. This indicates that the gradients of $\\nabla\_{\\theta}\\boldsymbol{\\mathcal{L}}$ used in TRAK tend to saturate as the model converges, which may not be a good design choice for attributing the training objective $\\boldsymbol{\\mathcal{L}}$ itself, as discussed in Section 5.

Given these observations, it is hypothesized that D-TRAK benefits from using loss functions that are not identical but are connected to the training objective. This would prevent gradient obfuscation or saturation while retaining sufficient information for attribution.

---

The role of Shapley values.

Thank you for pointing out, we have corrected the reference to Ghorbani & Zou [1] in the revision. Regarding the use of Shapley values in evaluation, for both feature attribution and data attribution, the exact Shapley value should be able to be viewed as the ground truth to evaluate different attribution methods that aim to efficiently approximate the exact Shapley values.

---

The lead-up to Definition 2 is hard to follow.

In the revision, we explicitly state that the LDS score considers the sum of attributions as an additive proxy for $\\boldsymbol{\\mathcal{F}}$ following your suggestion.

Thanks to the authors for their response. The clarifications are helpful, but the degree of insight on 1) why these alternative versions of TRAK are easier to estimate and 2) why they may be intrinsically more useful (i.e., even if all versions were perfectly estimated) remains somewhat unsatisfactory. I wonder if the authors could do a better job with these points given more time to revise the work, and perhaps design new experiments.

Furthermore, regarding the last point about the choice of evaluation via LDS with ${\mathcal{L}}_{\text{simple}}$, my question was less about L_simple vs. L_ELBO - I was curious whether either have a clear relationship with the likelihood of a generation under the model's probability distribution. It seems like that's what we should try to assess with data attribution methods for diffusion models: for any generation, we care which training examples made them more likely. The current formulation seems convenient, and I see that it follows from prior work, but this paper doesn't seem to argue that it represents what we really care about.

Overall, I'm not ready to raise my score. In terms of soundness and completeness I'm tempted to lower my score to marginally below acceptance, but I recognize that this is a timely topic that would be of interest to ICLR attendees.

Thank you for your valuable review and suggestions. Below we try to provide some insights about the success of D-TRAK over TRAK, and respond to the comments in Weaknesses.

---

Is LDS not an appropriate metric?

We deem LDS as an appropriate metric based on two considerations:

• First, both D-TRAK and TRAK satisfy the additive assumption of LDS [1,2], so it is reasonable to follow LDS in evaluating (from the aspect of Spearman rank correlation) the counterfactual predictors constructed from D-TRAK and TRAK.

• Second, in Section 4.4, we empirically perform counterfactual evaluation by measuring the $\\ell_{2}$-distance and CLIP similarity between the images generated before/after the exclusion of top-1000 positive influencers identified by D-TRAK and TRAK. Both the quantitative (Figures 3 and 11) and visualized (Figure 4) results show that D-TRAK is better than TRAK, which is consistent with the trend concluded from the LDS metric.

---

Is the non-convexity at fault?

Partially yes, non-convexity is always a major factor leading to inferior performance of theoretically justified (in convex settings) attribution methods such as TRAK. However, non-convexity may not explain why D-TRAK empirically outperforms TRAK, because there is no evidence that D-TRAK should be superior in non-convex settings.

In addition to non-convexity, we notice another factor influencing TRAK performance, which may be associated with gradient obfuscation or saturation. Specifically, as seen in Figure 2, TRAK achieves its highest LDS score on intermediate checkpoints, whereas D-TRAK achieves its highest LDS score on final checkpoints. This indicates that the gradients of $\\nabla\_{\\theta}\\boldsymbol{\\mathcal{L}}$ used in TRAK tend to saturate as the model converges, which may not be a good design choice for attributing the training objective $\\boldsymbol{\\mathcal{L}}$ itself, as discussed in Section 5.

Given these observations, it is hypothesized that D-TRAK benefits from using loss functions that are not identical but are connected to the training objective. This would prevent gradient obfuscation or saturation while retaining sufficient information for attribution.

---

Is there something special about diffusion models that TRAK fails to capture?

As discussed in Section 5, our main insight is that when the model output function of interest $\\boldsymbol{\\mathcal{F}}$ is the same as the training objective $\\boldsymbol{\\mathcal{L}}$, the gradients of $\\nabla\_{\\theta}\\boldsymbol{\\mathcal{L}}$ may not be a good design choice for attributing $\\boldsymbol{\\mathcal{L}}$ itself. Although in this paper we focus on diffusion models, we deduce that this insight also holds on classification models. For example, [2] has reported that using $\\nabla\_{\\theta}\\boldsymbol{\\mathcal{L}}$ in TRAK leads to worse performance than using $\\nabla\_{\\theta}\\log(\\exp(\\boldsymbol{\\mathcal{L}}(\\boldsymbol{x};\\theta))-1)$ on classification problems, even if the LDS score is computed by setting $\\boldsymbol{\\mathcal{F}}=\\boldsymbol{\\mathcal{L}}$. Besides, selecting intermediate checkpoints has also been discussed when attributing classifiers [3].

---

How do different norm losses lead to better or worse LDS scores?

We need to clarify that we did not adversarially search for loss functions to optimize the LDS metric. Actually, we happened to find that $\\mathcal{L}\_{\\textrm{Square}}$ counterintuitively works much better than $\\mathcal{L}\_{\\textrm{Simple}}$ and $\\mathcal{L}\_{\\textrm{ELBO}}$, so we also tried other common $\\ell_{p}$-norm and average pooling losses.

As reported in Table 1, $\\mathcal{L}\_{\\textrm{Square}}$, $\\mathcal{L}\_{\\textrm{Avg}}$, \\mathcal{L}\_{\\textrm{2-norm}}, and \\mathcal{L}\_{\\textrm{1-norm}} consistently outperform $\\mathcal{L}\_{\\textrm{Simple}}$, while \\mathcal{L}\_{\\textrm{\infty-norm}} has worse performance. We deduce that this is because $\\mathcal{L}\_{\\textrm{Square}}$, $\\mathcal{L}\_{\\textrm{Avg}}$, \\mathcal{L}\_{\\textrm{2-norm}}, and \\mathcal{L}\_{\\textrm{1-norm}} can better retain the information from $\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}$, while $\\mathcal{L}\_{\\textrm{Simple}}=\\mathbb{E}\_{t,\\boldsymbol{\\epsilon}}[2\\cdot(\\boldsymbol{\\epsilon}\_{\\theta}-\\boldsymbol{\\epsilon})\^{\\top}\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}]$ becomes ill-behaved when model converges (i.e., $\\boldsymbol{\\epsilon}\_{\\theta}\\approx \\boldsymbol{\\epsilon}$). As to \\mathcal{L}\_{\\textrm{\infty-norm}}, it has worse performance since it discards too much information and only retains the maximum value.

---

References:
[1] Datamodels: Predicting Predictions from Training Data. ICML 2022
[2] Trak: Attributing Model Behavior at Scale. ICML 2023
[3] Estimating Training Data Influence by Tracing Gradient Descent. NeurIPS 2020

Thank you for your valuable review and suggestions, we have uploaded a paper revision including additional experiment results. Below we respond to the comments in Weaknesses (W) and Questions (Q).

---

W1: The direct comparison with Journey TRAK in the evaluation section is not on equal grounds.

As you mentioned, only Journey TRAK (J-TRAK) [1] has been explicitly used for diffusion models in previous work, but still not meant to be used to attribute the final image. This means that there is no off-the-shelf baseline for D-TRAK, so we have to actively adapt previous methods (including TRAK and J-TRAK), in order to construct baselines on attributing final images generated by diffusion models.

Although we have made every effort to adapt previous methods as described in Appendix A.3, it is possible that stronger baselines could still be developed. We will make more rigorous clarification on the selection of baselines in the final reversion.

---

W2: A comparison against Journey TRAK on counterfactual experiments.

Following your suggestions, we perform counterfactual experiments to compare D-TRAK and J-TRAK at specific steps of the sampling trajectory (please see Appendix E for details). As shown in Figures 21, 22, 23, and 24, J-TRAK and its ensembling version J-TRAK(8) can indeed induce a large effect on the generation results via removing high-scoring training images specific to a timestep. Nevertheless, our findings indicate that D-TRAK remains more effective at identifying images that have a larger impact on the generation results at a specific timestep (we evaluate at $t=300,400$ following [1]). Even more counterintuitive are these results in light of the fact that D-TRAK was not originally intended to attribute a specific sampling timestep.

---

W3: A more thorough attempt at explaining/justifying the changes.

Actually, we happened to find that $\\mathcal{L}\_{\\textrm{Square}}$ counterintuitively works much better than $\\mathcal{L}\_{\\textrm{Simple}}$ and $\\mathcal{L}\_{\\textrm{ELBO}}$ on attributing diffusion models, so we also tried other common $\\ell_{p}$-norm and average pooling losses. As reported in Table 1, $\\mathcal{L}\_{\\textrm{Square}}$, $\\mathcal{L}\_{\\textrm{Avg}}$, \\mathcal{L}\_{\\textrm{2-norm}}, and \\mathcal{L}\_{\\textrm{1-norm}} consistently outperform $\\mathcal{L}\_{\\textrm{Simple}}$, while \\mathcal{L}\_{\\textrm{\infty-norm}} has worse performance. We deduce that this is because $\\mathcal{L}\_{\\textrm{Square}}$, $\\mathcal{L}\_{\\textrm{Avg}}$, \\mathcal{L}\_{\\textrm{2-norm}}, and \\mathcal{L}\_{\\textrm{1-norm}} can better retain the information from $\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}$, while $\\mathcal{L}\_{\\textrm{Simple}}=\\mathbb{E}\_{t,\\boldsymbol{\\epsilon}}[2\\cdot(\\boldsymbol{\\epsilon}\_{\\theta}-\\boldsymbol{\\epsilon})\^{\\top}\\nabla\_{\\theta}\\boldsymbol{\\epsilon}\_{\\theta}]$ becomes ill-behaved when model converges (i.e., $\\boldsymbol{\\epsilon}\_{\\theta}\\approx \\boldsymbol{\\epsilon}$ and thus $\\boldsymbol{\\epsilon}\_{\\theta}-\\boldsymbol{\\epsilon}$ becomes a noisy error term). As to \\mathcal{L}\_{\\textrm{\infty-norm}}, it has worse performance since it discards too much information and only retains the maximum value.

Furthermore, non-convexity is always a major factor leading to inferior performance of theoretically justified (in convex settings) attribution methods such as TRAK. However, non-convexity may not explain why D-TRAK empirically outperforms TRAK, because there is no evidence that D-TRAK should be superior in non-convex settings. In addition to non-convexity, we notice another factor influencing TRAK performance, which may be associated with gradient obfuscation or saturation. Specifically, as seen in Figure 2, TRAK achieves its highest LDS score on intermediate checkpoints, whereas D-TRAK achieves its highest LDS score on final checkpoints. This indicates that the gradients of $\\nabla\_{\\theta}\\boldsymbol{\\mathcal{L}}$ used in TRAK tend to saturate as the model converges, which may not be a good design choice for attributing the training objective $\\boldsymbol{\\mathcal{L}}$ itself, as discussed in Section 5.

Given these observations, it is hypothesized that D-TRAK benefits from using loss functions that are not identical but are connected to the training objective. This would prevent gradient obfuscation or saturation while retaining sufficient information for attribution.

Q1: Why is the rank correlation close to 0 when $\\mathcal{L}\_{\\textrm{Square}}$ is used for both $\\phi\^{s}$ and $\\boldsymbol{\\mathcal{F}}$ (Table 4)?

We have discussed this phenomenon in Appendix B (the paragraph of $\\boldsymbol{\\mathcal{F}}$ selection). Specifically, as mentioned in Appendices E.1 and E.6 of the TRAK paper [2], there is a latent assumption of linearity on the model output function $\\boldsymbol{\\mathcal{F}}$ in linear datamodeling prediction tasks. If the chosen $\\boldsymbol{\\mathcal{F}}$ is not well approximated by a linear function of training examples, then that puts an upper bound on the predictive performance of any attribution method.

Therefore, note that $\\boldsymbol{\\mathcal{F}}=\\mathcal{L}\_{\\textrm{Simple}}$ can be linearly approximated since its value is small when the model converges (the model is trained by $\\boldsymbol{\\mathcal{L}}=\\mathcal{L}\_{\\textrm{Simple}}$). In contrast, the value of $\\boldsymbol{\\mathcal{F}}=\\mathcal{L}\_{\\textrm{Square}}$ remains large and cannot be well approximated by a linear function.

---

Q2: Why set $\\mathcal{Q}$ to be the identity matrix even when $\\boldsymbol{\\mathcal{F}}\\neq\\boldsymbol{\\mathcal{L}}$?

In our paper, we primarily focus on the settings where $\\boldsymbol{\\mathcal{F}}=\\boldsymbol{\\mathcal{L}}$ and notice that the weighting term $\\mathcal{Q}$ becomes an identity matrix in TRAK. When we generalize TRAK to D-TRAK, we simply drop the weighting term $\\mathcal{Q}$ and discover that the performance is already satisfactory. As a result, there is no weighting term $\\mathcal{Q}$ in D-TRAK when we conduct ablation studies on D-TRAK in Table 4.

---

References:
[1] The Journey, Not the Destination: How Data Guides Diffusion Models. ICML Workshop 2023
[2] Trak: Attributing Model Behavior at Scale. ICML 2023

 While this consistency of D-TRAK across timesteps is somewhat expected given the choice of $\mathcal{L}$, it is intriguing to see it holding up in practice.