Paper: Intriguing Properties of Data Attribution on Diffusion Models (ICLR-2024)

Reviewer comment (Theory):
The point above is the reviewer’s main concern, and they wonder whether this paper could be more valuable given time to revise it and offer an appropriate explanation for why D-TRAK works.

Author response:
Thank you for your timely feedback and suggestions. Though we have made every effort to provide an understanding of D-TRAK through empirical ablation studies, we admit that there is still room for further improvement, both empirically and theoretically. We will do our best to design new experiments to gain a better understanding of how D-TRAK works.

Reviewer comment (Theory):
There is an emphasis in the introduction on the role of non-convexity in this setting, and its potential impact on TRAK not working as expected. After reading the paper, the reviewer is not sure it ultimately addresses this point, or offers any explanation for how D-TRAK circumvents the issue. The reviewer suggests that the authors expand on that subject in the paper.

Author response:
***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.

Reviewer comment (Writing):
In Section 2.2, the lead-up to Definition 2 is hard to follow. It might be helpful to explicitly state that the LDS score considers the sum of attributions as an additive proxy for $\mathcal{F}$, as it's explained in the original datamodels work.

Author response:
***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.

Reviewer comment (Evaluation):
The reviewer asks the authors to explain the consistent drops in LDS scores when shifting from validation examples to generations.

Author response:
***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.

Reviewer comment (Experiments):
Out of the listed baselines, to the best of my knowledge only Journey TRAK [1] has been explicitly used for diffusion models in previous work. As the authors note, Journey TRAK is not meant to be used to attribute the *final* image $x$ (i.e., the entire sampling trajectory). Rather, it is meant to attribute noisy images $x_t$ (i.e., specific denoising steps along the sampling trajectory). Thus, the direct comparison with Journey TRAK in the evaluation section is not on equal grounds.

Author response:
***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.

Reviewer comment (Experiments):
For the counterfactual experiments, I would have liked to see a comparison against Journey TRAK [1] used at a particular step of the the sampling trajectory. In particular, [1, Figure 2] shows a much larger effect of removing high-scoring images according to Journey TRAK, in comparison with CLIP cosine similarity.

Author response:
***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.

Reviewer comment (Evaluation):
Why is the rank correlation close to $0$ when $\mathcal{L}\_{square}$ is used for both $\phi^s$ and $\mathcal{F}$ (Table 4)? The authors have observed that $\mathcal{L}\_{square}$ is useful as a model output function when predicting $\mathcal{L}\_{simple}$ or $\mathcal{L}\_{ELBO}$. It is particularly odd then that it does not do a good job at predicting *itself*. Is there any particular reason why that might be the case?

Author response:
***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.