PaperlayerView exact revision on arXiv

Source-linked AI summary

Distilling Diffusion Score Discrepancy for Efficient Training Data Attribution

Shixuan Liu, Joan Serrà, Kin Wai Cheuk, Jinju Kim, Woosung Choi, Yukara Ikemiya, Wei-Hsiang Liao, Jiaqi W. Ma, Yuki Mitsufuji

arXiv:2609.38776v1cs.LGcs.AI

TL;DR

Existing diffusion attribution can rely on proxy losses and costly per-sample gradients or query-specific optimization. This paper defines influence through local score discrepancy, estimates it with TID, and distills its rankings into forward-only TIDE. Across three datasets, TID matches or surpasses unlearning-based attribution, while TIDE retains most of its accuracy at far lower per-query cost.

  • Problem

    Diffusion attribution often measures effects on proxy losses rather than generated behavior, while gradient-based methods can require costly per-sample gradients and unlearning methods query-specific fine-tuning.

    Source: §1 · pp. 1–2

  • Method

    TID estimates a training sample’s expected local score discrepancy as a preconditioned gradient similarity, and TIDE distills TID rankings online into a forward-only student using the diffusion model’s activations.

    Source: §1 · p. 2, §2 · p. 2

  • Results

    Across three datasets, TID matches or surpasses unlearning-based attribution, while TIDE retains most of TID’s accuracy at four-to-five orders of magnitude lower per-query cost.

    Source: §1 · p. 2, §4 · p. 10

  • Takeaways & Limitations

    TIDE attributes generated samples in milliseconds, faster than generating the sample itself.

    Source: §4 · p. 10

    The evaluation covers a single diffusion variant, although the formulation extends to DDPM, EDM, and flow matching.

    Source: §4 · p. 10

Abstract

from arXiv · show

Training data attribution for diffusion models aims to identify the training samples that influence a generated instance, but existing methods either require costly per-sample gradient computation or query-specific model optimization. Moreover, most methods attribute changes in a proxy loss rather than changes in the actual model's generative behavior. We address these limitations by formulating attribution directly with a local score discrepancy measure, which applies to any diffusion variant (including DDPM, EDM, and flow matching), and by showing that such measure can be estimated without retraining, as a preconditioned gradient similarity. We instantiate this estimator as Training-data Influence via score Discrepancy (TID), which uses Kronecker-factored curvature to avoid random projections and per-sample gradient storage. We then distill TID into TIDE, a forward-only student trained online to reproduce the teacher's rankings from the diffusion model's internal activations. Under counterfactual evaluation on CIFAR-10, ArtBench-10, and MS-COCO, TID matches or outperforms state-of-the-art approaches, while TIDE retains most of TID's accuracy at four to five orders of magnitude lower per-query cost, attributing generated samples in milliseconds and faster than the generation itself.

1 INTRODUCTION

Diffusion attribution often targets proxy losses rather than generated behavior, while existing approaches can be costly. The paper introduces score-discrepancy attribution with TID and a forward-only distilled student, TIDE, and reports strong counterfactual results at substantially reduced query cost.

1 INTRODUCTION

  • Existing diffusion attribution methods often measure training-sample effects on proxy utilities such as training loss, while gradient-based methods can require costly per-sample gradients and unlearning methods query-specific fine-tuning.A loss shift need not reflect a change in generated output.Source: §1 · pp. 1–2
  • TIDE distills TID into a forward-only student using the diffusion model’s internal representations and online teacher supervision, producing compact embeddings for nearest-neighbor retrieval.Online supervision is generated within each training batch rather than from a fixed offline dataset.Source: §1 · p. 2
  • TID defines influence through the expected local score discrepancy caused by removing a training sample, then estimates it as a preconditioned gradient similarity using Kronecker-factored curvature.This target applies across diffusion variants and avoids random projections and per-sample gradient storage.Source: §1 · p. 2
  • On CIFAR-10, ArtBench-10, and MS-COCO, TID outperforms state-of-the-art unlearning-based TDA on two datasets, while TIDE retains most accuracy at four-to-five orders of magnitude lower per-query cost.The evaluation also reports that ablations validate the core design choices.Source: §1 · p. 2

2 METHOD

Section 2 introduces TID, which estimates training-sample influence through local score discrepancy, then distills its rankings into TIDE embeddings from the diffusion model’s own activations.

2.1 LOCAL SCORE DISCREPANCY

  • TID measures a removed sample’s influence by expected score discrepancy around a query’s noisy states, a formulation applicable to DDPM, EDM, and flow-matching models.A larger discrepancy indicates a stronger shift in the learned score field around the query.Source: §2.1 · p. 3
  • A first-order output linearization and one-step Newton update estimate the leave-one-out effect using the gradient of the summed model output.The method drops the residual and leverage factor, which the authors report is negligible to ranking in their ablations.Source: §2.1 · pp. 3–4
  • Layerwise K-FAC estimates the gradient-similarity kernels without random projections; L2-normalized gradient factors limit domination by large norms across noise scales.After normalization, TID drops the schedule-dependent noise-scale weight so sampled scales contribute standardized influence values.Source: §2.1 · p. 4

2.2 DISTILLING ATTRIBUTION INTO FORWARD EMBEDDINGS

  • TIDE learns low-dimensional embeddings from the diffusion model’s internal activations, using cosine similarity to reproduce TID rankings with a single forward pass.Training-side embeddings can be cached, while each query uses M forward passes at inference.Source: §2.2 · pp. 5–6
  • Online batchwise teacher supervision avoids a fixed offline pair set, while a top-heavy pairwise ranking objective concentrates learning on candidates near the top of TID’s rankings.Teacher scores and ranks are detached, so only the embedding model receives gradients.Source: §2.2 · pp. 5–6

3 EXPERIMENTS AND RESULTS

The experiments evaluate attribution through counterfactual retraining across three datasets, comparing accuracy, per-query cost, and TID design choices. TID leads the accuracy comparison, while TIDE offers a much lower-cost alternative.

3.2 COMPARISON WITH EXISTING METHODS

  • TID surpasses MUCS on ArtBench-10 and MS-COCO while performing almost on par on CIFAR-10; other gradient-based methods perform worse across all three datasets.TIDE closely approaches TID's accuracy while requiring only forward passes and lower dimensionality.Source: Table 1 · p. 7, §3.2 · p. 8

3.1 METHODOLOGY

  • The counterfactual evaluation removes the union of each query’s top-2% training samples, retrains from scratch, and compares regenerated images with originals and random-removal controls.Experiments use unconditional CIFAR-10, class-conditional ArtBench-10, and text-conditioned MS-COCO; attribution quality is summarized by AUC across four image-similarity metrics.Source: §3.1 · pp. 6–7

3.3 EFFICIENCY TRADE-OFFS BETWEEN TDA METHODS

  • TIDE combines embedding-based methods’ low per-query cost and storage profile with the accuracy of expensive gradient- or unlearning-based methods, and can attribute faster than generation.Figure 2 compares counterfactual AUC with warm-cache per-query cost on CIFAR-10, ArtBench, and MS-COCO using a logarithmic cost axis.Source: Figure 2 · p. 8, §3.3 · p. 8

3.4 DISSECTING THE DESIGN OF TID

  • Removing the alternatives to TID’s summed model output substantially reduces accuracy, with the denoising-loss target causing the largest drop.The target comparison fixes TID’s curvature model.Source: §3.4 · p. 9
  • Independent corruption draws decrease µ by 9.2%, while restoring gradient magnitude decreases µ by 6.2%.The authors attribute these drops to unmatched noise-dependent gradient components and to sample difficulty and noise level affecting gradient norms.Source: §3.4 · p. 9
  • Adding either residual or leverage corrections changes µ by less than 1%, while sampled GGN preconditioning decreases µ by 5.1%.The authors favor the simpler kernel for its lower computational cost and report that changing the target affects performance more than changing curvature approximation.Source: §3.4 · p. 9

4 CONCLUSION

The paper presents TID as a score-discrepancy approach to diffusion attribution and distills its rankings into TIDE embeddings. It reports that TID matches or surpasses unlearning-based attribution across three datasets, while identifying limitations in the student’s scope and the evaluation.

4 CONCLUSION

  • Across three datasets, TID matches or surpasses unlearning-based attribution, while TIDE retains most of TID’s accuracy and attributes queries in milliseconds, faster than generation.The conclusion describes TID’s target as the expected local score discrepancy between the trained model and its leave-one-out counterfactual.Source: §4 · pp. 9–10
  • A general theory of when forward activations suffice for gradient-based influence remains open; TIDE is specific to its trained model and dataset, and evaluation covers one diffusion variant.The formulation extends to DDPM, EDM, and flow matching.Source: §4 · p. 10

A RELATED WORK

  • Diffusion attribution methods include subset-retraining, gradient-based, and mirrored-unlearning approaches; the latter require query-specific optimization, while gradient-based approaches incur gradient, curvature, and storage costs.The paper distinguishes TID by defining influence through score discrepancy rather than proxies for generation probability.Source: A RELATED WORK · p. 15
  • TIDE differs from Fast Data Attribution by distilling per-query rankings online from the diffusion model’s activations rather than using an offline set of annotated pairs and a standalone visual encoder.Its compact embeddings support nearest-neighbor retrieval while retaining model-specific attribution information.Source: A RELATED WORK · p. 15

A.1 DETAILS OF GRADIENT-BASED METHODS

Gradient-based diffusion attribution methods share a common template but differ in their scalar target, curvature approximation, and gradient compression and normalization.

A.1 DETAILS OF GRADIENT-BASED METHODS

  • Gradient-based diffusion attribution differentiates a scalar function on a corrupted training sample, applies a curvature approximation, then compares training samples with the query using curvature-aware gradient inner products.The paper organizes methods and TID ablations by scalar target, curvature, kernel representation, and normalization.Source: §A.1 · p. 15

B AN OVERVIEW OF TID AND TIDE

TID’s teacher and TIDE’s student share a pass through the frozen diffusion model during distillation. Teacher scores come from Kronecker-factored gradient factors, while student embeddings come from activations and are trained to reproduce teacher rankings.

B AN OVERVIEW OF TID AND TIDE

  • One shared noise draw and a single model pass supply gradient factors for TID’s teacher matrix and residual-stream activations for TIDE’s embeddings; a ranking loss compares their per-row rankings.The teacher matrix is detached, so the ranking-loss update changes only the embedding model.Source: B AN OVERVIEW OF TID AND TIDE · pp. 15–16
  • At inference, TIDE embeds a generated query from forward activations and scores it against cached training embeddings by cosine similarity, with no backward pass.The query is noised using aligned draws.Source: B AN OVERVIEW OF TID AND TIDE · p. 16
  • During distillation, the Kronecker-factor inverses are estimated once and held fixed; each batch recomputes per-sample gradient factors and embeddings, at a cost of one backward pass per sample plus O(B^2) inner products.Table 3 organizes methods by curvature, gradient kernel, and normalization; its lower block’s TID variants are evaluated in Table 2.Source: B AN OVERVIEW OF TID AND TIDE · p. 16, Table 3 · p. 16

C.1 THEORETICAL BACKGROUND

The background connects forward noising and score matching to the score estimates used across diffusion variants. It then derives variant-specific weights and extends factorized gradient computation to common diffusion-network layers.

C.1 THEORETICAL BACKGROUND

  • Affine-drift diffusion variants use Gaussian transition kernels, and their common parameterizations recover score estimates as affine functions of network outputs.For flow matching on the linear path, the conditional kernel is also Gaussian, allowing the same conditional-score expression.Source: §C.1 · pp. 18–19

C.2 VARIANT WEIGHTS

  • The score-difference conversion weight depends on the noise schedule and remains bounded above and below by positive constants on the stated truncated window.The weight is obtained from the schedule-dependent affine output-to-score map; flow matching additionally requires σmax < 1.Source: §C.2 · p. 19

D FACTORIZED GRADIENTS BEYOND LINEAR LAYERS

  • Factorized gradient formulas extend beyond linear layers to tokenized linear layers, convolutions, transposed convolutions, and normalization parameters.The constructions preserve factorized computation where an outer-product structure is available, evaluate inner products without materializing gradients, and use diagonal curvature for normalization parameters.Source: D FACTORIZED GRADIENTS BEYOND LINEAR… · pp. 20–21

E EXPERIMENTAL DETAILS

The experiments use EDM-trained DiT models and paired counterfactual retraining, with a fixed-query evaluation protocol. Attribution settings include aligned noise draws for TID and a shared, compact student architecture for TIDE.

E EXPERIMENTAL DETAILS

  • All diffusion models use EDM-trained DiT backbones, and counterfactual models share the full-data model’s training configuration and random seed.The paired design excludes only selected influential samples to reduce optimization variation and isolate removal effects.Source: E EXPERIMENTAL DETAILS · p. 21
  • Each seed uses 20 held-out generated queries, with every method selecting the top 2% per query before retraining without the union of selected samples.Because query-specific selections overlap, effective removal fractions range from 21% to 33%; Table 6 reports the effective removal sizes.Source: E EXPERIMENTAL DETAILS · p. 21, Table 6 · p. 23
  • TID uses 32 Karras-schedule timesteps, while TIDE uses the same 3.2M-parameter architecture on all three datasets.TIDE's four-head attention pooling and two-layer MLP account for about 3% of the roughly 100M-parameter backbone; Table 1 uses M=32 and Figure 2 tests M∈{1, 32, 100}.Source: E EXPERIMENTAL DETAILS · pp. 21–22, Table 5 · p. 23
  • TIDE's within-batch training transfers to generated queries, with student–teacher Spearman correlations of 0.87, 0.89, and 0.86 on CIFAR-10, ArtBench-10, and MS-COCO.The authors suggest forward noising may reduce differences between training and generated distributions as a potential explanation.Source: E EXPERIMENTAL DETAILS · p. 22

F.1 QUALITATIVE ANALYSIS OF ATTRIBUTION RANKINGS

The qualitative rankings show that semantic similarity at the top does not by itself establish faithful attribution, while the lowest-ranked examples distinguish TID and TIDE from encoder baselines. These single-query examples illustrate ranking content but do not independently establish quantitative results.

F.1 QUALITATIVE ANALYSIS OF ATTRIBUTION RANKINGS

  • Figures 4–6 show the eight highest- and lowest-ranked training samples for one generated query per dataset, and those examples are illustrative rather than independent quantitative evidence.The rankings are taken from deployed attribution outputs without recomputation or retuning.Source: §F.1 · pp. 23–24
  • Top-ranked examples are often semantically plausible across methods, but this agreement alone does not establish faithful attribution.For the skateboarding query, all methods retrieve skateboarding scenes; for the automobile query, all but D-TRAK primarily retrieve automobiles.Source: §F.1 · p. 23
  • TID and TIDE often place weak-visual-signal examples at the ranking tail, unlike encoder baselines that retrieve semantically dissimilar examples.Examples include sparse line drawings, isolated objects on white backgrounds, and dim nighttime scenes; this pattern is compatible with gradient-geometric influence rather than a simple semantic-opposition axis.Source: §F.1 · p. 24

F.2 SENSITIVITY TO THE REMOVAL BUDGET

On CIFAR-10, TIDE's counterfactual performance improves as the per-query removal budget rises, with a possible plateau beyond 3%.

F.2 SENSITIVITY TO THE REMOVAL BUDGET

  • TIDE performance improves consistently on CIFAR-10 as the per-query removal budget increases from 1% through 2%, 3%, and 5%.The authors note a possible plateau beyond 3%, suggesting influence may not extend significantly over that percentage of training items on CIFAR-10.Source: §F.2 · p. 24

F.3 RELATIVE MAGNITUDE OF THE REMOVAL EFFECT

The similarity ratio measures the magnitude of query–regeneration change after removing attributed samples, complementing AUC’s measure of separation from random removal. Its rankings broadly agree with AUC, while results vary by dataset.

F.3 RELATIVE MAGNITUDE OF THE REMOVAL EFFECT

  • The similarity ratio measures the percentage change in query–regeneration similarity after removal, while AUC measures separation from the random-removal control.The ratio captures the magnitude of induced change; more negative values indicate larger drift.Source: §F.3 · p. 24
  • The random-removal control yields approximately −2 to −3% similarity change.Source: §F.3 · p. 24
  • TID induces the largest average drift on ArtBench-10 and MS-COCO, while TID, TIDE, and MUCS perform similarly on CIFAR-10.The agreement with AUC is a consistency check because the metrics capture different properties.Source: §F.3 · p. 24

F.4 ADDITIONAL EFFICIENCY DETAILS

This section separates cached one-off computation from per-query work, explains how costs differ across attribution methods, and reports qualitative rankings and removal-effect measures. Preprocessing can be substantial, particularly for higher-resolution data, even when subsequent query costs are lower.

F.4 ADDITIONAL EFFICIENCY DETAILS

  • The complexity accounting distinguishes one-off cached work from warm per-query work, using N for training samples, Q for queries, and M for the timestep or draw budget.It also separates final score-product flops from model evaluations.Source: §F.4 · pp. 24, 26, Table 9 · p. 27
  • Table 8 reports mean relative similarity change and standard error across seeds; more negative values indicate larger removal-induced drift.Its random-removal control is approximately −2 to −3%.Source: Table 8 · p. 26
  • CLIP and DINO cache one embedding per training sample, whereas D-TRAK and DAS cache projected, timestep-averaged gradients.D-TRAK and DAS still require query gradients; DAS also computes normalization statistics during preprocessing.Source: §F.4 · p. 26
  • TIDE shifts gradient computation into distillation and training-set embedding, leaving per-query forward passes and an embedding comparison.The distillation phase costs O(EN)(cfb + cenc,fb), and embedding the training set costs NM(cfwd + cenc).Source: §F.4 · p. 26
  • MUCS has no query-independent computation: each query fine-tunes the model and then sweeps the training set to measure loss changes.Source: §F.4 · p. 27
  • DAS preprocessing takes 12.3–38.1 h, while TIDE takes 3.8–10.0 h; the embedding-bank construction itself requires only minutes.Model-agnostic encoder preprocessing takes less than five minutes; these are one-off costs on a single H100.Source: §F.4 · p. 27, Table 10 · p. 27
  • DAS’s per-image–timestep cost rises from 13.0 ms on CIFAR-10 to 23.2 ms on MS-COCO, so size-and-timestep scaling alone underpredicts MS-COCO preprocessing time.Scaling the CIFAR-10 measurement predicts 21.4 h on MS-COCO, versus 38.1 h measured; the passage attributes the difference to 64×64 inputs with CLIP conditioning.Source: §F.4 · p. 27
  • TID caches factor estimates and eigendecompositions, but recomputes training gradients for every attribution pass because full-scale storage is infeasible.Its per-query cost includes training-side evaluations as well as the query’s own evaluations.Source: §F.4 · p. 27
Loading 2609.38776v1…