Concept-TRAK: Understanding how diffusion models learn concepts through concept-level attribution

TL;DR

Concept-TRAK is a novel method that isolates the influence of specific concepts in diffusion models, improving interpretability and addressing copyright concerns in image generation.

Contribution

We introduce Concept-TRAK, a new influence attribution method with specialized loss functions to isolate concept-specific influences in diffusion models.

Findings

Significant improvement over prior methods in concept attribution benchmarks.

Effective on synthetic, CelebA-HQ, and AbC datasets.

Versatile in real-world text-to-image generation scenarios.

Abstract

While diffusion models excel at image generation, their growing adoption raises critical concerns about copyright issues and model transparency. Existing attribution methods identify training examples influencing an entire image, but fall short in isolating contributions to specific elements, such as styles or objects, that are of primary concern to stakeholders. To address this gap, we introduce concept-level attribution through a novel method called Concept-TRAK, which extends influence functions with a key innovation: specialized training and utility loss functions designed to isolate concept-specific influences rather than overall reconstruction quality. We evaluate Concept-TRAK on novel concept attribution benchmarks using Synthetic and CelebA-HQ datasets, as well as the established AbC benchmark, showing substantial improvements over prior methods in concept-level attribution scenarios. We further demonstrate its versatility on real-world text-to-image generation with compositional and multi-concept prompts.