GUDA: Counterfactual Group-wise Training Data Attribution for Diffusion Models via Unlearning

AI summary

Overview

Research area: Machine learning — training data attribution for diffusion generative models, combining machine unlearning with counterfactual (leave-one-group-out) reasoning.

Technical level: Advanced. The paper assumes familiarity with diffusion models, ELBO-based likelihood surrogates, machine unlearning objectives, and attribution evaluation metrics.

Scope in one sentence: The paper introduces GUDA, a framework that approximates costly Leave-One-Group-Out retraining with machine unlearning in order to attribute a generated image back to the groups of training data that most influenced it.

What This Paper Is About

Practitioners often need to know which groups of training data — an artistic style, an object class — shaped a generated image, not just which individual examples did. The gold-standard answer is counterfactual: retrain the model with that group removed and see how its behavior on the sample changes, but this Leave-One-Group-Out (LOGO) procedure requires training N+1 models from scratch, which becomes impractical as the number of groups grows. GUDA instead starts from a single full-data model and applies machine unlearning to produce an approximate counterfactual model per group, then scores group influence by the difference in a likelihood-based measure (ELBO) between the full model and each unlearned model.

Key Contributions

1. A formal estimand for group attribution. The authors define LOGOA (LOGO Attribution), a per-generation group-removal counterfactual estimand based on the ELBO difference between the full-data model and a model retrained on the retain set, with LOGO retraining serving as the oracle.

2. An unlearning-based approximation framework. GUDA approximates each counterfactual model by applying machine unlearning to a shared full-data model rather than retraining from scratch, using a loss of the form L_unlearn = L_forget + λ_pres L_preserve, where the preservation term is a score-matching loss against the frozen full model on the retain set.

3. Two setting-specific instantiations. For unconditional generation, GUDA-U uses a ReTrack-based redirection that trains toward importance-weighted targets from the retain set. For conditional text-to-image attribution, GUDA-C redirects forget-condition responses toward CLIP-weighted style-selection anchors that keep prompt content fixed and swap only the style descriptors.

4. Empirical validation against LOGO oracles. Experiments on CIFAR-10 and on artistic style attribution with Stable Diffusion 1.5 show GUDA identifies primary contributing groups more reliably than semantic similarity (CLIPA), gradient-based attribution (DAS, D-TRAK, TRAK), and instance-level unlearning (Wang et al., 2024), while reaching about 100× speedup over LOGO retraining on CIFAR-10.

Main Findings

• CIFAR-10 head identification: GUDA (ReTrack) achieves 0.727 Top-1, 0.798 MRR, 0.677 NDCG@3, 0.475 Top-3, 0.423 RBO, and 0.265 Spearman over 10 classes and 2,048 queries — the best or tied-best head-focused metrics among compared methods.

• Unlearning quality drives attribution quality: GUDA with the ReTrack forget loss reaches 0.727 Top-1 versus 0.619 for GUDA with the ESD forget loss, which the authors attribute to ReTrack's theoretical connection to LOGO retraining through importance sampling.

• Similarity is not sufficient: CLIPA reaches 0.662 Top-1 on CIFAR-10 and 0.338 Top-1 on UnlearnCanvas — competitive on some rank-based metrics but below GUDA on head-identification metrics.

• Gradient-based attribution struggles: On CIFAR-10, DAS reaches 0.716 Top-1, D-TRAK 0.609, and TRAK 0.118 Top-1 (versus 10% chance), which the authors connect to limitations in adapting influence functions to unconditional generative models.

• Wall-clock efficiency: GUDA totals 2h 02m on CIFAR-10 versus 207h 47m for the LOGOA oracle, a ~100× speedup. The paper states unlearning uses roughly 1/120 of the optimization steps of full retraining (20 epochs versus 2,400 epochs), and that GUDA's query-time evaluation is faster than DAS (1.6s versus 11.6s per image).

• UnlearnCanvas style attribution: Over 16 evaluated styles and 320 queries, GUDA reaches 0.456 Top-1, 0.582 MRR, 0.734 NDCG@3, 0.405 Top-3, 0.446 RBO, and 0.239 Spearman, outperforming CLIPA (0.338 Top-1) and Wang et al. (2024) (0.047 Top-1). GUDA also achieves the highest Spearman correlation in this setting.

• Full-ranking metrics can mislead: The authors note Wang et al. (2024) obtain a higher Spearman correlation than CLIPA (0.147 versus 0.117) despite near-random Top-1 agreement (4.7%), and therefore emphasize Top-k metrics as the primary evaluation criterion.

• ELBO as a proxy: The paper reports that ΔELBO and Δlog p obtained via probability-flow ODE estimation on CIFAR-10 are strongly correlated, with head identification remaining stable despite potential KL-gap asymmetry.

Methodology in Plain English

The researchers frame group attribution as a counterfactual question: what would the model's behavior on a generated sample look like if an entire group had never been in the training data? The honest way to answer is to retrain the model without that group, but doing this for every group is far too expensive.

Their substitute is machine unlearning. They train one full model, then for each group they fine-tune that model for a small number of steps to remove the group's influence, producing an approximate "counterfactual" model. To check how much a group mattered, they compare how well the full model and the unlearned model explain a generated sample using the ELBO, a tractable lower bound on log-likelihood. A large drop in ELBO when a group is unlearned means that group contributed to the sample.

The unlearning step balances two goals: suppressing the target group (the forget term) and preserving behavior on everything else (a preservation term that matches the frozen original model's score predictions on retained data). How the forget term works depends on the setting. For unconditional models, they redirect denoising targets toward nearby retained samples weighted by how plausibly those samples could have generated the current noisy latent (ReTrack-style). For text-to-image models, they face an extra problem: removing a style also removes the prompts describing it, so the forget prompts fall outside the retain distribution. They solve this by building anchor prompts that keep the original content but swap in a retained style chosen by CLIP-weighted sampling, and train the unlearned model to match the frozen full model's prediction under that anchor condition.

Once all the unlearned models are computed in advance, attributing any new generated sample only requires evaluating the ELBO under each one — no re-running of unlearning per query.

Why This Matters

Impact on research. The paper reconceives unlearning not as a deletion mechanism but as a computational approximation of a counterfactual oracle, and it targets group-level effects directly rather than summing instance-level scores, which prior work (Koh et al., 2019; Basu et al., 2020) shows fails to capture nonlinear group interactions. It also provides an explicit LOGO-validated estimand (LOGOA), addressing a gap the authors identify in both the attribution literature and the unlearning literature.

Real-world applications:

• Copyright assessment for AI-generated artwork that blends multiple artistic styles.
• Fair compensation for data providers whose work contributed to a model's outputs.
• Debugging generated outputs by tracing them to specific data groups.
• Dataset curation and licensing decisions informed by which groups actually drive generations.

Industry relevance. The reported ~100× wall-clock speedup on CIFAR-10 (2h 02m versus 207h 47m) and the ability to precompute counterfactual models once and reuse them across queries make group attribution feasible at scales where per-group retraining is not. The code is released at https://github.com/sony/guda, and the work is a collaboration between Sony AI, Sony Group Corporation, and Stanford University.

Future Directions

• Overlapping groups. The current formulation assumes a partition into N disjoint groups. The authors explicitly flag extending to overlapping groups, where a sample may belong to multiple groups, as an important practical direction left for future work.

• Better unlearning operators. The gap between GUDA-ReTrack (0.727 Top-1) and GUDA-ESD (0.619 Top-1) on CIFAR-10 suggests attribution fidelity is bounded by how well the chosen unlearning method approximates the LOGO counterfactual, so future unlearning improvements should translate into better attribution.

• Scaling to full style sets and larger benchmarks. UnlearnCanvas experiments evaluate attribution on the first 16 styles in alphabetical order out of 60 trained styles for computational feasibility; scaling the evaluation to the full group set is a natural next step.

• Beyond vision and beyond group-level attribution. The paper positions GUDA as complementary to unlearning-based attribution for text-to-music models (Choi et al., 2025), and notes that using unlearning for attribution is not a plug-in choice — deletion-oriented objectives do not inherently target an explicit LOGO counterfactual.

Target Audience

Researchers and practitioners working on training data attribution, machine unlearning, and diffusion generative models; engineers building provenance, licensing, or auditing tooling for generative image systems; and readers already comfortable with diffusion model objectives and likelihood-based evaluation who want to understand how counterfactual reasoning can be made computationally tractable at the group level.

Authors’ abstract

Training-data attribution for vision generative models aims to identify which training data influenced a given output. While most methods score individual examples, practitioners often need group-level answers (e.g., artistic styles or object classes). Group-wise attribution is counterfactual: how would a model's behavior on a generated sample change if a group were absent from training? A natural realization of this counterfactual is Leave-One-Group-Out (LOGO) retraining, which retrains the model with each group removed; however, it becomes computationally prohibitive as the number of groups grows. We propose GUDA (Group Unlearning-based Data Attribution) for diffusion models, which approximates each counterfactual model by applying machine unlearning to a shared full-data model instead of training from scratch. GUDA quantifies group influence using differences in a likelihood-based scoring rule (ELBO) between the full model and each unlearned counterfactual. Experiments on CIFAR-10 and artistic style attribution with Stable Diffusion show that GUDA identifies primary contributing groups more reliably than semantic similarity, gradient-based attribution, and instance-level unlearning approaches, while achieving ~100x speedup on CIFAR-10 over LOGO retraining.