🤖 machine learning

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

The paper proposes GUDA, a method that leverages machine unlearning to efficiently approximate counterfactual group-wise training data attribution for diffusion models, achieving reliable identification of influential data groups with a significant speedup over traditional retraining approaches.

Original authors: Naoki Murata, Yuhta Takida, Chieh-Hsin Lai, Toshimitsu Uesaka, Bac Nguyen, Stefano Ermon, Yuki Mitsufuji

Published 2026-06-02

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Naoki Murata, Yuhta Takida, Chieh-Hsin Lai, Toshimitsu Uesaka, Bac Nguyen, Stefano Ermon, Yuki Mitsufuji

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are looking at a beautiful, complex painting created by an AI. You might wonder: "Which specific group of training pictures taught this AI to paint in this particular style?" Was it the group of "Impressionist" photos? The "Sketch" photos? Or maybe the "Watercolor" photos?

This paper, titled GUDA, introduces a new way to answer that question for AI image generators (called Diffusion Models). Here is the breakdown using simple analogies.

The Problem: The "Too Many Cooks" Dilemma

AI models are trained on massive datasets containing millions of images, often grouped by style or category (like "dogs," "cars," or "Van Gogh style").

To find out which group influenced a specific AI-generated image, the "gold standard" method is called LOGO (Leave-One-Group-Out).

The Analogy: Imagine a choir singing a song. To find out which section (sopranos, altos, tenors, basses) contributed most to a specific note, you would have to record the choir, then record it again without the sopranos, then again without the altos, and so on.
The Issue: If you have 100 groups, you have to re-record the song 100 times. For an AI, this means retraining the model from scratch 100 times. This takes years of computing time and is practically impossible.

The Solution: The "Memory Eraser" (Unlearning)

The authors propose GUDA (Group Unlearning-based Data Attribution). Instead of retraining the AI from scratch every time, they use a technique called Machine Unlearning.

The Analogy: Instead of firing the whole choir and hiring a new one without the sopranos, you take the original choir and gently "teach them to forget" the sopranos' part. You keep the rest of the choir exactly the same but tweak their memory so they no longer know how to sing the soprano notes.
How it works:
1. Start with the fully trained AI (the "Full Model").
2. Apply a special "unlearning" process to remove the influence of one specific group (e.g., "Impressionism").
3. This creates a "Counterfactual Model"—a version of the AI that acts as if it never saw Impressionist photos.
4. Compare the original AI's output with this "forgotten" AI's output. If the "forgotten" AI can't make the image anymore, you know Impressionism was the key ingredient.

The Secret Sauce: How to "Forget" Correctly

Simply deleting data isn't enough; the AI needs to be guided on what to replace the forgotten data with so it doesn't get confused. The paper describes two ways to do this:

For General Images (Unconditional): They use a method called ReTrack.
- Analogy: If you tell the AI to "forget" a picture of a red car, ReTrack doesn't just leave a blank space. It looks at all the other cars in the database and says, "Okay, since you can't remember the red car, imagine a blue car or a truck instead, weighted by how similar they are." It redirects the AI's attention to the remaining data.
For Text-to-Image (Conditional): This is trickier because the AI is also listening to text prompts (e.g., "A dog in the style of Van Gogh"). If you remove "Van Gogh," the text prompt itself becomes weird.
- Analogy: The authors use Anchors. If the prompt says "Dog, Van Gogh style," and you want to forget Van Gogh, they swap "Van Gogh" with a different style (like "Sketch") in the prompt while keeping the "Dog" part the same. This teaches the AI to handle the "Dog" concept without the "Van Gogh" influence, without breaking the text instructions.

The Results: Fast and Accurate

The paper tested this on two things:

CIFAR-10: A simple dataset of 10 object classes (like cats, airplanes, trucks).
Artistic Styles: Using Stable Diffusion to generate images in different art styles.

The Findings:

Speed: GUDA is roughly 100 times faster than the "gold standard" method of retraining from scratch. It's like finishing a marathon in 2 hours instead of 200 hours.
Accuracy: It is much better at identifying the most important groups than other methods that just look for "visual similarity" (like asking, "Does this image look like a Van Gogh painting?"). GUDA asks, "If we removed Van Gogh from the AI's brain, would this image still exist?"
Comparison: It beat other popular methods that try to guess influence by looking at gradients (mathy slopes) or by unlearning individual pictures one by one.

Summary

GUDA is a shortcut. It allows researchers to ask, "What part of the training data made this AI image?" by gently "erasing" groups of data from the AI's memory and seeing what changes, rather than painfully rebuilding the AI's brain from scratch every single time. It makes understanding AI generation faster and more practical.

Technical Summary: GUDA – Counterfactual Group-wise Data Attribution for Diffusion Models via Unlearning

1. Problem Definition

Training data attribution for vision generative models aims to identify which training data influenced a specific output. While existing methods focus on instance-level attribution (identifying specific training samples), practitioners often require group-level attribution (e.g., identifying which artistic styles or object classes contributed most to a generation).

The core challenge is that group influence is non-additive and nonlinear; simply summing instance-level scores fails to capture interactions within shared representations. The theoretical "gold standard" for group attribution is the Leave-One-Group-Out (LOGO) counterfactual: measuring how a model's behavior on a generated sample changes if a specific group $k$ were entirely absent from the training data. However, realizing this counterfactual requires retraining $N+1$ models (one full model plus $N$ models with each group removed), which is computationally prohibitive for large-scale diffusion models.

2. Methodology: GUDA Framework

The authors propose GUDA (Group Unlearning-based Data Attribution), a framework that approximates the expensive LOGO counterfactual models using machine unlearning rather than retraining from scratch.

Core Concept:
Instead of training $\theta_{logo}^{-k}$ (the model trained without group $k$ ) from scratch, GUDA starts with the full-data model $\theta_{full}$ and applies a machine unlearning operator to selectively remove the influence of group $k$ , yielding an approximate counterfactual model $\theta_{ul}^{-k}$ .

Attribution Scoring (LOGOA):
The attribution score is defined as the difference in a likelihood-based scoring rule (specifically the Evidence Lower Bound, ELBO) between the full model and the unlearned counterfactual model:
$\text{GUDA}_k(x_0, c) = \text{ELBO}(x_0|c; \theta_{full}) - \text{ELBO}(x_0|c; \theta_{ul}^{-k})$
A positive score indicates that group $k$ contributed significantly to the generation of sample $(x_0, c)$ .

Unlearning Implementation:
The unlearning objective consists of two terms:

Preservation Term ( $L_{preserve}$ ): A score-matching loss that ensures the unlearned model maintains performance on the retain set ( $D_{-k}$ ), preventing catastrophic forgetting.
Forget Term ( $L_{forget}$ ): Varies based on the generation setting:
- GUDA-U (Unconditional Generation): Utilizes ReTrack-based redirection. It trains the model to denoise samples from the forget set ( $D_k$ ) toward importance-weighted targets derived from the retain set ( $D_{-k}$ ), effectively approximating the posterior mean of the retain-only distribution.
- GUDA-C (Conditional Text-to-Image): Addresses the challenge where removing a style group shifts both image and prompt distributions. It employs Weighted Style-Selection Anchors. For a forget pair (image, prompt), the method constructs an "anchor" prompt by keeping the content description but swapping the style descriptor with one from the retain set (sampled via CLIP-weighted similarity). The model is then trained to match the full model's prediction under this anchor condition, resolving the condition-distribution mismatch.

3. Key Contributions

Formalization of Group Attribution: The paper formalizes group-wise attribution as a counterfactual estimand (LOGOA) based on the difference in ELBO between a full model and a group-removed model, establishing LOGO retraining as the oracle.
Unlearning-Based Approximation: It introduces a practical framework that approximates these counterfactuals via machine unlearning, achieving a $\sim 100\times$ speedup over LOGO retraining on CIFAR-10.
Two Specific Instantiations:
- GUDA-U: For unconditional diffusion, using ReTrack-based importance-weighted redirection.
- GUDA-C: For conditional text-to-image, using CLIP-weighted style-selection anchors to handle prompt distribution shifts.
Empirical Validation: Demonstrates that directly targeting the group-removal counterfactual via unlearning outperforms semantic similarity baselines (CLIPA) and instance-level aggregation methods (DAS, TRAK, D-TRAK) in identifying primary contributing groups.

4. Experimental Results

The authors evaluated GUDA on two benchmarks:

CIFAR-10 (Unconditional): With 10 classes and 2,048 queries, GUDA achieved 72.7% Top-1 agreement with the LOGO oracle, outperforming CLIPA (66.2%) and gradient-based methods (DAS: 71.6%, but significantly slower). GUDA achieved a $\sim 100\times$ wall-clock speedup over LOGO retraining (2 hours vs. 207 hours).
UnlearnCanvas (Conditional Text-to-Image): Using Stable Diffusion 1.5 and 16 artistic styles, GUDA achieved 45.6% Top-1 agreement, significantly outperforming CLIPA (33.8%) and a prior unlearning-based attribution method (Wang et al., 2024) which scored only 4.7%. GUDA also showed superior performance in NDCG@3 and RBO metrics, indicating better alignment with the most influential styles.

Key Observations:

Unlearning Quality Matters: GUDA using ReTrack-based unlearning significantly outperformed GUDA using ESD (Erasing Stable Diffusion) unlearning, validating the theoretical connection between ReTrack and LOGO retraining.
Limitations of Alternatives: Semantic similarity (CLIPA) failed to capture counterfactual group effects, and instance-level methods struggled to scale to group-level attribution, often performing near-randomly on head-identification metrics.

5. Significance and Claims

The paper claims that GUDA provides a scalable and practical solution for understanding the provenance of generative model outputs at a group level. By bridging the gap between the theoretical ideal of counterfactual retraining and the computational reality of large-scale models, GUDA enables:

Fair Compensation: More accurate identification of which data groups (e.g., artists) contributed to outputs.
Debugging: Identifying problematic data sources or styles that negatively impact generation.
Transparency: Offering a method to audit generative models without the prohibitive cost of full retraining.

The authors emphasize that GUDA is not a certified deletion mechanism but a computational approximation for attribution. They note that as more LOGO-aligned unlearning methods emerge, they can be integrated into the GUDA framework to further improve attribution quality. The work is positioned as a step toward scalable, group-level understanding of generative models, moving beyond instance-level approximations that fail to capture nonlinear group interactions.