Data Attribution for Diffusion Models: Timestep-induced Bias in Influence Estimation

TL;DR

This work extends data attribution to diffusion models by introducing Diffusion-TracIn, which operates over the denoising timesteps, and Diffusion-ReTrac, a re-normalized variant to counteract a timestep-induced gradient-norm bias that inflates influence estimates. The authors show that gradient norms vary with timesteps, causing generally influential training samples and unstable attributions; they address this with a two-pronged approach: (i) sampling a sparse set of timesteps to approximate the diffusion trajectory and (ii) renormalizing gradient information to produce more localized, test-targeted attributions. Across image tracing, targeted attribution, and outlier-detection tasks, Diffusion-ReTrac yields substantially more diverse and intuitive influential training samples, reducing the share of generally influential samples and improving precision in attribution (e.g., higher unique-sample counts in top-k proponents). The proposed methods offer a practical, scalable pathway to fairer and more interpretable diffusion-model attributions, with some limitations related to replay-based estimation and the need for deeper theory on large-norm timesteps.

Abstract

Data attribution methods trace model behavior back to its training dataset, offering an effective approach to better understand ''black-box'' neural networks. While prior research has established quantifiable links between model output and training data in diverse settings, interpreting diffusion model outputs in relation to training samples remains underexplored. In particular, diffusion models operate over a sequence of timesteps instead of instantaneous input-output relationships in previous contexts, posing a significant challenge to extend existing frameworks to diffusion models directly. Notably, we present Diffusion-TracIn that incorporates this temporal dynamics and observe that samples' loss gradient norms are highly dependent on timestep. This trend leads to a prominent bias in influence estimation, and is particularly noticeable for samples trained on large-norm-inducing timesteps, causing them to be generally influential. To mitigate this effect, we introduce Diffusion-ReTrac as a re-normalized adaptation that enables the retrieval of training samples more targeted to the test sample of interest, facilitating a localized measurement of influence and considerably more intuitive visualization. We demonstrate the efficacy of our approach through various evaluation metrics and auxiliary tasks, reducing the amount of generally influential samples to $\frac{1}{3}$ of its original quantity.

Figures (16)

• Figure 1: Correlation of Timestep and Norm. We plot the norm ranking and distance between training timestep to $t_{\text{max}}$ for 50 randomly selected samples. The resulting correlation is $0.7$ and the linear regressor (red) has a slope of 6.038.
• Figure 2: Samples' Norm vs. Training Timestep. We plot the norm and timestep of 2,000 randomly selected training samples. We observe that loss gradient norms tend to increase when the training timestep falls in the later range (towards noise). This upward trend is consistent at other checkpoints tested. The sample with the largest norm (red) and smallest norm (green) are shown; no exceptional visual patterns are noticed. Additional distributions for others as well as open-sourced models are included in Appendix \ref{['append_norm_timestep']} and \ref{['append_additional_norm']}.
• Figure 3: One Sample's Norms Varying Timestep. Example of norm distributions for one randomly selected sample is shown. On each checkpoint, the norm distribution is skewed to later timesteps for each individual sample. This trend exists regardless of where the actual training timestep (red) is sampled at.
• Figure 4: Image Tracing Accuracy. We measure precision by evaluating the ratio of correctly attributed samples among top $k$ proponents. While (a) shows both methods successfully traced MNIST test samples, (b) shows that Diffusion-TracIn fails to attribute CIFAR-plane test samples since the outlier MNIST samples with large norms are characterized as strong proponents, giving low precision for small $k$. Note that random attribution is a strong baseline in (b) since the 5,000 CIFAR-planes are more likely to be proponents compared to the 200 MNIST samples.
• Figure 5: Image Tracing on Outlier Model. We show the attribution results on 4 test samples, using both training (leftmost top 2) and generated samples (leftmost bottom 2); influence scores are stated below each image. While both methods successfully attribute MNIST zero test samples, Diffusion-TracIn incorrectly characterizes MNIST training samples as influential to a CIFAR-plane test sample. This is notably mitigated with Diffusion-ReTrac.

Theorems & Definitions (1)

• Definition 1