Research map / Image generation / Fast diffusion sampling
Intriguing Properties of Data Attribution on Diffusion Models
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Shows that theoretically unjustified surrogate gradient functions (e.g., L_Square instead of L_Simple) in a TRAK-style attribution, called D-TRAK, substantially outperform TRAK for attributing DDPMs and LoRA-finetuned Stable Diffusion on LDS and counterfactual evaluation.
Problem
Data attribution for diffusion models: tracing generated images back to training samples efficiently and effectively in non-convex settings, where theoretically motivated choices (TRAK with the training loss gradient) may not perform best.
Approach
Generalize TRAK to D-TRAK by building the projected gradient features from alternative per-sample functions (e.g., squared noise-prediction norm L_Square, average-pooled and p-norm variants) instead of the training loss, dropping the weighting term; empirically these outperform the loss-gradient choice by large margins, and interpolation analysis suggests norm-based losses retain more information
Key results
- On CIFAR-2 with 10 timesteps, D-TRAK reaches 26.79% validation LDS vs 11.42% for TRAK.
- On ArtBench-2 generation LDS with 10 timesteps, D-TRAK reaches 24.16% vs 7.78% for TRAK.
- Median counterfactual l2 distance on CIFAR-2 is 8.97 for D-TRAK vs 5.90 for TRAK, with lower CLIP similarity (0.881 vs 0.943).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Simple-loss feature gradients with Square-loss feature gradients | TRAK attribution on DDPM · CIFAR-2, CIFAR-10, and ArtBench-2/5 validation and generation sets | LDS: +15.37%/+10.15%, +11.55%/+8.67% and +15.35%/+5.1% validation; +13.04%/+9.8%, +9.82%/+6.76%, and +16.38%/+6.51% generatio | Helped |
| Use Square-loss gradients instead of Simple-loss gradients for features (ablation) | D-TRAK on CIFAR-2 · CIFAR-2 validation and generation | LDS: 19.50 -> 30.81; 12.05 -> 22.62 | Helped |
| Use Avg-loss gradients instead of Simple-loss gradients for features (ablation) | D-TRAK on CIFAR-2 · CIFAR-2 validation and generation | LDS: 19.50 -> 24.91; 12.05 -> 16.76 | Helped |
| Use 1-norm gradients instead of Simple-loss gradients for features (ablation) | D-TRAK on CIFAR-2 · CIFAR-2 validation and generation | LDS: 19.50 -> 23.44; 12.05 -> 15.10 | Helped |
| Use 2-norm gradients instead of Simple-loss gradients for features (ablation) | D-TRAK on CIFAR-2 · CIFAR-2 validation and generation | LDS: 19.50 -> 24.72; 12.05 -> 15.75 | Helped |
| Use ELBO-loss gradients instead of Simple-loss gradients for features (ablation) | D-TRAK on CIFAR-2 · CIFAR-2 validation and generation | LDS: 19.50 -> 8.46; 12.05 -> 3.49 | Hurt |
| Use infinity-norm gradients instead of Simple-loss gradients for features (ablation) | D-TRAK on CIFAR-2 · CIFAR-2 validation and generation | LDS: 19.50 -> 5.22; 12.05 -> 3.99 | Mixed |
| Set interpolation weight eta away from 0.5 toward 0 or 1 (ablation) | D-TRAK interpolation between Simple and Square losses · CIFAR-2 | LDS | Helped |
| Use uniform rather than cumulative timestep sampling (ablation) | TRAK · CIFAR-2 validation and generation under 100-computation budget | LDS: 13.45 -> 16.57; 4.58 -> 9.13 | Helped |
| Use uniform rather than cumulative timestep sampling (ablation) | D-TRAK · CIFAR-2 validation and generation under 100-computation budget | LDS: 29.35 -> 30.71; 19.86 -> 22.07 | Mixed |
| Set the LDS output function to Square, 1-norm, or infinity-norm loss instead of Simple loss (ablation) | Data attribution evaluation · CIFAR-2 | LDS | Hurt |
| Use Square-loss feature gradients rather than Simple-loss feature gradients | D-TRAK on CIFAR-2 counterfactual evaluation · CIFAR-2 | pixel-wise L2 distance: 5.90 -> 8.97 | Helped |
Surprising findings
- Using the squared noise-prediction norm (L_Square) or other norm-based surrogates instead of the training loss gradient consistently yields higher LDS, contrary to the theoretical design.
- Setting the model output function to norm-based losses makes LDS near zero, suggesting those outputs are not well approximated linearly.
Limitations
- Theoretical explanation for why unjustified design choices outperform remains missing.
- Retraining-based LDS evaluation is computationally expensive, limiting scale of ground-truth benchmarks.
- Experiments limited to small-resolution DDPMs and one LoRA-finetuned Stable Diffusion setting.
Experimental setup
Datasets: CIFAR-10 CelebA ArtBench
Benchmarks: CIFAR-2 CIFAR-10 CelebA ArtBench-2 ArtBench-5
Metrics: linear datamodeling score (Spearman) pixel-wise l2 distance CLIP cosine similarity
Baselines: TRAK Gradient (dot prod.) TracInCP GAS Relative Influence Renormalized Influence Journey TRAK Raw pixel CLIP similarity Empirical Influence Datamodel
Builds on: TRAK DDPM LoRA TracInCP
Closest papers by meaning
- Diffusion Attribution Score: Evaluating Training Data Influence in Diffusion Models (ICLR 2025 · similarity 0.94)
- Influence Functions for Scalable Data Attribution in Diffusion Models (ICLR 2025 · similarity 0.90)
- Training Data Attribution in Diffusion Models via Mirrored Unlearning and Noise-Consistent Skew (NeurIPS 2026 · similarity 0.90)
- An Efficient Framework for Crediting Data Contributors of Diffusion Models (ICLR 2025 · similarity 0.88)
- The Emergence of Reproducibility and Consistency in Diffusion Models (ICML 2024 · similarity 0.88)
- Diffusion Models for Attribution (AAAI 2025 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.