Data Attribution for Text-to-Image Models by Unlearning Synthesized Images

TL;DR

This work tackles data attribution for text-to-image diffusion models by proposing an unlearning-based counterfactual framework to identify training images that most influence a given synthesis. The method unlearns the synthesized image via an elastic weight consolidation objective and a Newton update, then scores training samples by how much they are forgotten, with updates focused on cross-attention KV. It is evaluated with a gold-standard leave-K-out retraining on MSCOCO and the Customized Model Benchmark, where it outperforms influence-function and feature-matching baselines across K and model categories. The approach offers a practical, model-aware way to attribute training data, with implications for accountability, data provenance, and potential compensation in AI-generated content.

Abstract

The goal of data attribution for text-to-image models is to identify the training images that most influence the generation of a new image. Influence is defined such that, for a given output, if a model is retrained from scratch without the most influential images, the model would fail to reproduce the same output. Unfortunately, directly searching for these influential images is computationally infeasible, since it would require repeatedly retraining models from scratch. In our work, we propose an efficient data attribution method by simulating unlearning the synthesized image. We achieve this by increasing the training loss on the output image, without catastrophic forgetting of other, unrelated concepts. We then identify training images with significant loss deviations after the unlearning process and label these as influential. We evaluate our method with a computationally intensive but "gold-standard" retraining from scratch and demonstrate our method's advantages over previous methods.

Figures (17)

• Figure 1: (a) Our algorithm: We propose a new data attribution method using machine unlearning. By modifying the pretrained model $\theta$ to unlearn the synthesized result $\hat{\mathbf{z}}$, the model also forgets the influential training images crucial for generating that specific result. (b) Evaluation: We validate our method through counterfactual evaluation, where we retrain the model without the top $K$ influential images identified by our method. When these influential images are removed from the dataset, the model fails to generate the synthesized image.
• Figure 2: Attribution results on MSCOCO models. We show generated samples used as a query on the left, with training images being identified by different methods on the right. Qualitatively, our method retrieves images with more similar visual attributes. Notably, our method better matches the poses of the buses (considering random flips during training) and the poses and enumeration of skiers.
• Figure 3: Leave-$K$-out analysis for MSCOCO models. We compare images across our method and baselines generated by leave-$K$-out models, using different $K$ values, all under the same random noise and text prompt. A significant deviation in regeneration indicates that critical, influential images were identified by the attribution algorithm. Our method leads to image generation that deviate significantly, even with as few as 500 influential images removed ($\sim$0.42$\%$ of the dataset).
• Figure 4: Spatially-localized attribution. Given a synthesized image (left), we crop regions containing specific objects using GroundingDINO liu2023grounding. We attribute each object separately by only running forgetting on the pixels within the cropped region. Our method can attribute different synthesized regions to different training images.
• Figure 5: Qualitative examples on the Customized Model benchmark. The red boxes indicate ground truth exemplar images used for customizing the model. Both our method and AbC baselines successfully identify the exemplar images on object-centric models (left), while our method outperforms the baselines with artistic style models (right).