Paper: The Journey, Not the Destination: How Data Guides Diffusion Models (ICLR-2024)

Reviewer comment (Presentation):
The structure of the paper is a bit confusing and not easy to follow.

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

Reviewer comment (Evaluation):
The fine-grained analysis in section C.2 is intriguing. It appears that the model has a tendency to generate images by referencing specific semantic parts of the training images rather than generating the entire images. Therefore, does it still make sense to attribute these generated images to image-level data?

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

Reviewer comment (Theory):
In general, 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. This is due to a series of approximations that make the computation possible. The interesting question is how those approximations influence the final observations (e.g. using one-step approximation of $\hat{x}_0^t$ as an approximation of the distribution’s expectation $E[x_0|x_t]$). If I understand it correctly, this assumption indicates that there would be no change in the final distribution due to later diffusion steps. For some cases with higher t values (e.g. timesteps>600 in Figure 3) this approximation might be wrong.

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

Reviewer comment (Theory):
“However, the effect of a single time step on the final distribution may be small and hard to quantify. Hence, we assume that attributions change gradually over time and replace the denoising model for a small interval of time steps (i.e., between steps t and t − Δ).” - This assumption is counterintuitive to the previous derivations on what authors consider attributions for diffusion models in their framework.

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

Reviewer comment (Presentation):
The submission is hard to follow without in-depth understanding of the Trak method (Perk et al. 2023) it is heavily based on. It is also almost impossible to understand the experiments without reading the appendix.

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

Reviewer comment (Presentation):
The presentation of the submission seems ad hoc and sloppy. There are multiple different thoughts that lack cohesion. For example, Figure 4 appears on page 5 without any connection to the text, while it is referred to on page 8.

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

Reviewer comment (Reproducibility):
The provided code in its current form is far from being useful for the full reproduction of the results presented in this work. It is a total of ~200 useful lines of code (excluding imports and comments) with some definitions of functions that should be connected to something in an unclear way.

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

Reviewer comment (Presentation):
The submission does not follow ICLR guidelines - the margins are significantly smaller (1 in instead of required 1.5) and the font is different.

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

Reviewer comment (Experiments):
Are the CIFAR10 models evaluated in this submission conditioned on the class identity? Otherwise, how for example the plot in Fig 3 (left) calculated - using only 15 examples presented on the right?

Author response:
Our CIFAR-10 DDPM models are unconditional. To calculate the likelihood of a given image containing a horse in Figure 3, we apply a CIFAR-10-trained ResNet classifier to 200 samples from the distribution $p_{\theta}(\cdot|x_t)$. 15 of these 200 samples are visualized in the figure. Following the comment of reviewer KCCM, we ablate this result over three different choices of classifiers for identifying the horse features and find that the plot is consistent in each case (see the newly added Figure C.11 in Appendix C.5).

Reviewer comment (Evaluation):
“Thus, if we treat this likelihood as a function of t, the steps at which there is the largest increase in likelihood (i.e., the steepest slope) are most responsible for the presence of this feature in the final image. In fact, it turns out that this likelihood often increases rapidly” - The example presented in Fig 3 is rather extreme. Do you have any intuition why it is so?

Author response:
“The example presented in Fig 3 is rather extreme. Do you have any intuition why it is so?”
Indeed, we note that the transition in Figure 3 happens in a remarkably short interval of steps. We acknowledge that the nature of our finding is empirical and we currently lack a theoretical understanding. One possible (very crude) intuition for the rapid increase in likelihood is that before the transition, the noisy latent $x_t$ is approximately equidistant to multiple “clusters” of images (e.g., a cluster of “horse” images and a cluster of “car” images), and after the transition $x_t$ is closer to one of these clusters, thus making the diffusion model “commit” to that cluster. Importantly, this would happen if the latent space in which the diffusion process happens is aligned with human-aligned features---i.e., images closer in the latent space are also semantically more similar to us. As some additional motivation for the existence of this phenomenon, we would like to note that discrete events like bifurcations often appear in continuous systems.

Reviewer comment (Evaluation):
A limitation of the proposed approach is the fact that attribution scores are only provided for a single denoising step. This is unintuitive, as it requires multiple steps to be analyzed to understand how an image was generated. It would be good to obtain a single-shot attribution score for the entire diffusion trajectory. While simple heuristics can be employed to obtain this from the current approach, it's unclear if these are useful and interpretable.

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

Reviewer comment (Evaluation):
There is little analysis regarding how attributions change throughout the diffusion trajectory. It would be interesting to analyze more how ranking b/w attributions stay consistent b/w different timesteps. For example, what's the correlation coefficient b/w attributions of two timesteps close to each other v/s further away? How many of the +ve influence samples in the initial/middle timesteps stay positive throughout the diffusion trajectory?

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

Reviewer comment (Experiments):
The claim regarding conditioning likelihood increasing in small time intervals is a bit weak (i.e. features appear in specific timesteps). This should be more rigorously studied for multiple generated images on CIFAR-10 and MS-COCO.

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

Reviewer comment (Presentation):
The font, and margins of this submission violate ICLR guidelines. This should be corrected in the next version of the paper.

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

Reviewer comment (Reproducibility):
Several important details are missing regarding the attribution approach. TRAK uses a random projection matrix to compress gradients to low dimensional space, the dimension hasn't been mentioned at all. Is this random projection step not done for attributing diffusion models? Are multiple checkpoints used to estimate attribution scores? Are these trained on different subsets of training data? How much storage and compute is required for estimating the attribution scores?

Author response:
We appreciate the reviewer pointing out the missing details regarding the hyperparameters we use with TRAK. Below, we provide a summary, which we have also added to Appendix A of the paper:
- We do indeed use a random projection. In particular, we used a projection dimension of 4096 for CIFAR-10 and 16384 for MS COCO; it is worth noting that we achieved similar quantitative results by using a projection dimension of 2048.
- We do use multiple model checkpoints to estimate attribution scores. In particular, we use 100 model checkpoints for CIFAR-10, and 20 for MS COCO.
- We train the models that we use for estimating attributions on the entire training set
- The TRAK features for CIFAR-10 and MS COCO take up a total of ~20 GB of disk space. The compute time for obtaining a set of TRAK features is approximately proportional to training for one epoch, since we compute (and project) gradients for all samples in the training set.