The Journey, Not the Destination: How Data Guides Diffusion Models

We thank the reviewer for the helpful feedback.

Structure of the Presentation

To address the reviewer’s concerns over the structure of the presentation, we have made a number of changes to the structure of the main text: We made revisions to Section 3 (which introduces our attribution framework) to simplify the language and make the presentation more modular. Additionally, we added an in-depth section on TRAK [1] to the background section in order to better contextualize our method. We shortened Section 4 (which describes our method) and made more explicit connections to the preliminaries in Section 2.

Simplicity of the proposed method

As the reviewer acknowledges, our paper formalizes and studies a new research problem. We introduce a method building on top of state-of-the-art work for data attribution in the supervised setting. While our method is relatively simple, we view the simplicity of our method as a feature, not a bug. In particular, we demonstrate through our evaluations that, despite its “simplicity,” our method surfaces attributions that are counterfactually significant to the behavior of diffusion models. Due to its simplicity, our method is easier to implement (by extending upon the open source implementation of TRAK) and easier for other works to build upon. Conceptually, the success of our approach indicates that despite the complexity underlying modern NNs, we can reason about them by approximating them with kernel models.

Attribution beyond image-level data

We agree that it would be interesting to extend the data attribution task to attribute behavior not only to individual training examples, but rather to parts of images (our understanding is that is what the reviewer is suggesting). For example, the diffusion model might be using the background of some training image at an earlier step, but might then use an object in the foreground at a later point (we allude exactly to this in Appendix C.2). More broadly, such a perspective on data attribution would also be useful in other settings, including supervised tasks. We view this direction as a great potential future extension of the work.

However, we believe that it is still very valuable to attribute generated samples back to image-level data. First, for many possible applications of data attribution, it is most convenient and natural to intervene on image-level data. For example: Copyright violations and data leakage are generally only discussed at the level of individual images (or sets of images, e.g., an artist’s body of work). Data curation is almost always formalized on a per-image level. Data poisoning is usually performed on individual examples [2].

Additionally, attributing with respect to patches might limit the scalability of the method. For example, if we divide images into 16 patches, we would need to compute and store 16x more gradients.

[1] Park, S. M., Georgiev, K., Ilyas, A., Leclerc, G., & Madry, A. (2023). TRAK: Attributing model behavior at scale. arXiv preprint arXiv:2303.14186.

[2] Khaddaj, Alaa, et al. "Rethinking Backdoor Attacks." (2023).

We thank the reviewer for the helpful feedback.

Mismatch between goal and final experiments

The reviewer raises a concern that “there is a significant mismatch between the goal of the method as expressed at the beginning of the submission and the final experiments being evaluated […] due to a series of approximations”. We would like to respectfully disagree with this statement.

On one hand, we acknowledge that our method indeed involves a series of approximations, which enable us to run large scale experiments. On the other hand, importantly, we do not make these simplifications in our evaluation. In particular, note that in Figure 5 (Retraining without top influencers), we simply retrain the diffusion models (DDPMs and LDMs) without the training examples we flag as influential, and sample directly from $p_\theta(\cdot | x_t)$ for each of the resulting models. Then, we report FID between those samples, and the ones generated by models trained on the entire training set. Note that neither of the approximations we employ in our method (using $\hat{x}\_0^t$ in lieu of samples from $p_\theta(\cdot|x_t)$, and using the training loss instead of the likelihood) is present in this evaluation. Hence, despite the approximations used in computing attribution scores, we show that our method significantly outperforms the similarity-based baselines in attributing the full distribution.

The reviewer also raises the question whether replacing $\mathbb{E}\_{x\_0 \sim p\_{\theta}(\cdot|x\_t)}[x\_0]$ with $\hat{x}\_0^t$ implicitly creates an additional assumption that the later diffusion steps do not contribute to the final generated image. This is not the case: the use of $\hat{x}\_0^t$ to approximate $\mathbb{E}\_{x\_0 \sim p\_{\theta}(\cdot|x\_t)}[x\_0]$ simply provides us with an computationally-efficient estimate of $\mathbb{E}\_{x\_0 \sim p\_{\theta}(\cdot|x\_t)}[x\_0]$. As we note in the paper, the additional assumption we introduce is that the diffusion model is approximately consistent, as defined in [2].

Overall, we found that despite the simplicity of the approximations used, our approach enables meaningful attributions. That said, using more sophisticated approximations is an interesting direction for future work.

Effect of a single step

The reviewer raises a concern about the assumption that our attributions change gradually over the course of the diffusion process. First, we note that we empirically observe this phenomenon. In particular, see Figure C.5---we show that there is a very high rank correlation between our attributions for steps in close proximity to each other. Second, we would like to point out that in practice we only use this assumption for small intervals, e.g., <=100 steps within a 1000-step diffusion trajectory.

Presentation of the submission

To address the reviewer’s concerns over the quality of the presentation, we have made a number of improvements to the structure of the main text:

• “The submission is hard to follow without in-depth understanding of the Trak method”: We added an in-depth section on TRAK [1] to the background section in order to better contextualize our method.
• “Figure 4 appears on page 5 without any connection to the text, while it is referred to on page 8”: Some of the figures were indeed poorly placed - we have adjusted the figure placements in the rebuttal, so that they appear closer to where they are referenced.
• “lack of cohesion”: we made revisions to Sections 3 and 4 that introduce our attribution framework and describe our method. In particular, we simplified the language and made the presentation more modular.

Code

Indeed, we do provide a concise (~200 loc) API built on top of the TRAK API (https://github.com/madrylab/trak/). But we view the simplicity of our implementation as one of the strengths of our method rather than a shortcoming. To address the ease of use of our API, we have added:

1. Two demo notebooks (examples/demo_MSCOCO.ipynb and examples/demo_CIFAR10.ipynb) that provides a simple way for users to score images synthesized by an LDM model trained on MS COCO, and a DDPM model trained on CIFAR-10, respectively.
2. A set of scripts (together with a README with an explicit set of commands) to reproduce the computation of our TRAK scores on both DDPM models trained on CIFAR-10 and LDMs trained on MS COCO.

Adhering to margin and font

We thank the reviewer for pointing out the incorrect margins and font size in our submission. This was an unintentional accident due to a git merge, and we fixed the revised version to adhere to the ICLR guidelines. The margins we used were wider than the ICLR guidelines, however it turns out our font was also wider. After fixing both the font and margins, we found that there was very little difference in the resulting paper length (about one paragraph). We apologize for this error.

We thank the reviewer for the helpful feedback.

Clarity of the submission

To address the reviewer’s concerns over the clarity of the presentation, we have made a number of changes to the structure of the main text:

• “The flow can be more straightforward”: We made revisions to Section 3 (which introduces our attribution framework) to simplify the language and make the presentation more modular. Additionally, we added an in-depth section on TRAK to the background section in order to better contextualize our method.
• “the paper is too wordy”: we shortened Section 4 (which describes our method) and made more explicit connections to the preliminaries in Section 2.

Use of Classifiers to Identify Features

In our view, using a classifier is the only reasonable (and scalable) approach to identifying the emergence of features, as any other approach would likely involve manually inspecting images for their presence. However, the reviewer raises a valid concern about the robustness of our approach.

To address the reviewer’s concerns, we have added an additional figure, Figure C.11, to Appendix C.5. In this figure, we visualize the emergence of a feature in the diffusion trajectory for three generated samples from the CIFAR-10 dataset using three different classifiers (ResNet trained on CIFAR-10, adversarially robust ResNet trained on CIFAR-10, and CLIP zero-shot initialized with the CIFAR-10 class names). We find that in each case, the outputs of all three classifiers are aligned.

Additionally, we would like to point the reviewer to Figure C.2 in Appendix C.4. Here, we visualize the emergence of features in a sample generated by Stable Diffusion which has images with more details, and find that features still appear often in a short interval of steps.

That said, we would like to highlight that neither our method nor our evaluation strategy depends on this empirical phenomenon. We instead consider this finding as a way to relate features of the synthesized images to specific steps of the diffusion process, thus providing a more intuitive interpretation of our per-step attributions.

Interpretation of FID evaluation

What the results of Figure 5 illustrate is that (as intended) the top-k most important training examples at step $t$ (as identified with our method) have a significant impact on the conditional distribution $p\_\theta(\cdot|x\_t)$ of generated images. Intuitively, we expect that the top-k (positively influential) images will share features with the original conditional distribution, but less so with the new conditional distribution (resulting after removing those top-k images).

The increase in FID is not simply the effect of training on a smaller dataset, as we compare our method to the baseline of removing the top k training images according to either CLIP and pixel similarity. In this case, the model is indeed trained on a smaller dataset (with the same number of examples removed), yet the increase in FID is significantly less than the increase resulting from using our attribution scores.

Evaluations of our Method

We are not sure what aspects of Figure 7 the reviewer is referring to as low quality. In this figure, we show that even though our attribution scores are targeted at a specific timestep, we are still able to leverage them to attribute the full diffusion process, and do so much better than popular baselines of using image similarity for data attribution.

Our patch-based attributions results (Figure 6) are only intended to convey qualitative results. In particular, this subsection of the experiments is intended to showcase a natural extension of our method, rather than a comprehensive evaluation of our framework.

Importantly, we do not present these results (Figures 6, 7) as the primary evaluation in our work. Our main evaluation is presented In Section 5.2, where we report quantitative results for two metrics for evaluating attribution methods: linear datamodeling score (LDS) and distributional changes (as measured by FID) after removing top influencers. In both cases, our attributions significantly outperform the similarity baselines. Overall, our method is able to identify training examples most important to the generation of a given image, and moreover, precisely quantify the counterfactual impact of those training examples. We believe that these results sufficiently convey the efficacy of our approach.

We thank the reviewer for the helpful feedback.

Attributing individual steps rather than the full trajectory

Indeed, within our framework we attribute at individual steps. In our view, attributing at either individual steps or the full trajectory have complementary benefits, and studying data attribution with respect to full diffusion trajectory is an important direction for future work. In particular, full-trajectory attribution might be more relevant towards matters of copyright violation and data valuation (as it might better assess how much each training image contributes to the final image). However, as we demonstrate in our work, step-specific attributions help us to gain a more fine-grained understanding of the way in which diffusion models use data, and can help to isolate attributions to individual features. For example, in Figure C.3 in Appendix C.2, we highlight cases in which the same image is highly influential at one step but negatively influential at another step. This shows that influences actually evolve over steps, which we could not have identified if we attributed the full diffusion trajectory as a whole.

That said, we agree that attributing the full trajectory is an interesting and important future direction to extend our current work.

How attributions change throughout the trajectory

We recap some existing analyses in our paper to study how attributions change throughout the trajectory:

1. We identify a pattern for how attributions relate to the formation of features over steps. Specifically, we find that positive influencers begin to reflect a given feature as soon as the likelihood of a feature being in the final image begins to increase, while negative influencers begin to reflect a given feature when the likelihood of the feature reaches near 100% probability.
2. In Appendix C.2, as we mentioned above, we study the overlap between top and bottom influencers. In fact, we find that whether a given training image was one of the top most influencers does not correlate to whether it will be a negative influencer as well at another step.

Following the reviewer’s suggestion, in Appendix C.2 we have added a plot of the correlation between attribution scores as a function of their proximity (Figure C.5). As expected, we find that the correlation between steps decreases as the distance between steps increases. Interestingly, we find that once attribution scores are computed at around 500 steps apart, their correlation reaches near zero. These findings again highlight the local nature of attributions.

Likelihood increasing in small time intervals

To address the reviewer’s concerns, we have added a new figure, Figure C.11, to Appendix C.5 that includes additional examples of the appearance of features over the steps of the diffusion trajectory. We would also like to point the reviewer to Figure C.2 in Appendix C.4. Here, we visualize the emergence of features in a sample generated by Stable Diffusion, and find that in this setting features appear often in a short interval of steps.

Additionally, we would like to highlight that neither our method nor our evaluation strategy depend on this empirical phenomenon. We instead consider this finding as an avenue to bridge features of the synthesized images with intervals of the diffusion process, thus providing a more intuitive explanation of our attributions.

Adhering to margin and font

We thank the reviewer for pointing out the incorrect margins and font size in our submission. This was an unintentional accident that happened due to a git merge, and we fixed the revised version to adhere to the ICLR guidelines. The margins we used were wider than the ICLR guidelines, however it turns out our font was also wider. After fixing both the font and margins, we found that there was very little difference in the resulting paper length (about one paragraph). We apologize for this error.

Fine-tuning text-to-image models

We thank the reviewer for the interesting suggestion. We agree that it would be valuable to compare our method to the data attribution framework in [1], which identifies “ground truth” attributions by fine-tuning a text-to-image model on a given training image, guaranteeing that it is highly influential. Since the provided code for this work focused on attributions for Stable Diffusion, to compare to this method we would need to either adapt it to MS COCO models, as suggested, or implement our approach for subsets of LAION. Both of these tasks were not feasible during the rebuttal period, but we will consider implementing this comparison for the camera-ready version of the paper.