Concept-TRAK: Understanding How Diffusion Models Learn Concepts through Concept-Level Attribution
Abstract
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.
Authors
- Yonghyun Park*
- Chieh-Hsin Lai
- Satoshi Hayakawa*
- Yuhta Takida
- Naoki Murata
- Wei-Hsiang Liao
- Woosung Choi
- Kin Wai Cheuk
- Junghyun Koo
- Yuki Mitsufuji
*External Authors
Venue
ICLR-26
Date
2026