Concept-TRAK: Understanding How Diffusion Models Learn Concepts through Concept-Level Attribution

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

ICLR-26

2026

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.