Concept-TRAK: Understanding how diffusion models learn concepts through concept-level attribution
This work addresses transparency and copyright issues in generative AI for stakeholders like developers and regulators, though it is incremental as it builds on existing influence functions.
The paper tackles the problem of understanding how diffusion models learn specific concepts like styles or objects, which is crucial for addressing copyright and transparency concerns, by introducing Concept-TRAK, a concept-level attribution method that shows substantial improvements over prior methods on the AbC benchmark.
While diffusion models excel at image generation, their growing adoption raises critical concerns around 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 matter most to stakeholders. To bridge this gap, we introduce \emph{concept-level attribution} via a novel method called \emph{Concept-TRAK}. Concept-TRAK extends influence functions with two key innovations: (1) a reformulated diffusion training loss based on diffusion posterior sampling, enabling robust, sample-specific attribution; and (2) a concept-aware reward function that emphasizes semantic relevance. We evaluate Concept-TRAK on the AbC benchmark, showing substantial improvements over prior methods. Through diverse case studies--ranging from identifying IP-protected and unsafe content to analyzing prompt engineering and compositional learning--we demonstrate how concept-level attribution yields actionable insights for responsible generative AI development and governance.