Fast Data Attribution for Text-to-Image Models

TL;DR

The paper addresses the data attribution bottleneck in text-to-image models by distilling a slow unlearning-based attribution method into a fast, embedding-based retrieval framework. It introduces a two-stage data curation workflow and a learning-to-rank objective that maps synthesized prompts to training-image embeddings, enabling rapid identification of influential training data with minimal overhead. Across MSCOCO and Stable Diffusion on LAION data, the approach delivers substantial speedups (up to thousands of times faster than prior methods) while achieving competitive attribution performance, with a combined DINO+CLIP-Text feature space performing best. The method scales to hundreds of millions of training samples and large-scale models, offering a practical path toward real-world data attribution and potential compensation mechanisms, albeit with limitations tied to the teacher-based derivation and the need for further validation on extremely large models.

Abstract

Data attribution for text-to-image models aims to identify the training images that most significantly influenced a generated output. Existing attribution methods involve considerable computational resources for each query, making them impractical for real-world applications. We propose a novel approach for scalable and efficient data attribution. Our key idea is to distill a slow, unlearning-based attribution method to a feature embedding space for efficient retrieval of highly influential training images. During deployment, combined with efficient indexing and search methods, our method successfully finds highly influential images without running expensive attribution algorithms. We show extensive results on both medium-scale models trained on MSCOCO and large-scale Stable Diffusion models trained on LAION, demonstrating that our method can achieve better or competitive performance in a few seconds, faster than existing methods by 2,500x - 400,000x. Our work represents a meaningful step towards the large-scale application of data attribution methods on real-world models such as Stable Diffusion.

Figures (17)

• Figure 1: Attribution performance vs. throughput. Previous methods (AbU wang2024attributebyunlearning, D-TRAK zheng2023intriguing) offer high attribution performance but are computationally expensive for deployment. Fast image similarity using off-the-shelf features (DINO) lacks attribution accuracy. We distill slower attribution methods into a feature space that retains attribution performance while enabling fast deployment.
• Figure 2: Our method. Given a synthesized sample, data attribution aims to find which elements in the training set are more influential. (Top) Attribution by Unlearning (AbU) is a slow but accurate method. It works by unlearning a synthesized example and evaluating the change in reconstruction loss on each training image, where each evaluation takes many forward passes. We generate AbU scores, using them to train an attribution-focused embedding (bottom), so that attribution can be performed by fast similarity search, while retaining the accuracy of the slower AbU method.
• Figure 3: (Left) Feature spaces. We compare different feature spaces, before and after tuning for attribution. We measure mAP to the ground truth ranking, generated by AbU+. While text-only embeddings perform well before tuning, image-only embeddings become stronger after tuning. Combining both performs best and is our final method. (Right) Ranking loss functions. Simple MSE regression does not converge well. Ordinal loss works well, but does not support fast similarity search at inference time. We use cross-entropy, which achieves performance similar to ordinal loss while supporting similarity search. We report 1-standard error in the plots.
• Figure 4: (Left) Data scaling. We investigate the impact of the number of synthesized queries. Note that each synthesized query contains attribution scores with 10k training points. We find that the performance quickly improves and saturates. (Right) Sampling outside the neighbor set. We vary the probability of selecting non-nearest neighbor images when building the attribution dataset. Using a few randomly sampled, unrelated images from the training set helps keep the learned attribution model, while having too many impedes the learning. We report 1-standard error in the plots.
• Figure 5: Sampling strategies of dataset construction. (Left) For each query, we randomly select from the 10k neighbors to learn from. Reliable rankings can be learned, even with relatively fewer training images per query are provided in the dataset. (Right) Given a fixed budget of 2.45M query-training attribution ranks, we test trading off between fewer training images per query and more synthesized queries. We find that at this budget, more query images with fewer training images are beneficial. We report 1-standard error in the plots.