/AttributeByUnlearning

Code for the paper "Data Attribution for Text-to-Image Models by Unlearning Synthesized Images."

Primary LanguagePython

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

Project | Paper

Sheng-Yu Wang1, Aaron Hertzmann2, Alexei A. Efros3, Jun-Yan Zhu1, Richard Zhang2.
Carnegie Mellon University1, Adobe Research2, UC Berkeley3
In NeurIPS, 2024.

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.

Code layout

Please refer to the subdirectories for different sets of experiments.

Directory Experiments
mscoco attribution for text-to-image diffusion models trained on MSCOCO
custom_diffusion attribution for customized text-to-image models from AbC benchmark

Citation

@inproceedings{wang2024attributebyunlearning,
  title={Data Attribution for Text-to-Image Models by Unlearning Synthesized Images},
  author={Wang, Sheng-Yu and Hertzmann, Aaron and Efros, Alexei A and Zhu, Jun-Yan and Zhang, Richard},
  booktitle={NeurIPS},
  year = {2024},
}

Acknowledgements

We thank Kristian Georgiev for answering all of our inquiries regarding JourneyTRAK implementation and evaluation, and providing us their models and an earlier version of JourneyTRAK code. We thank Nupur Kumari, Kangle Deng, Grace Su for feedback on the draft. This work is partly supported by the Packard Fellowship, JPMC Faculty Research Award, and NSF IIS-2239076.