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

ICLR 2026 Conference Submission9959 Authors

Keywords: Diffusion models, Data attribution, Concept

TL;DR: Our method, Concept-TRAK, identifies which training examples influenced specific concepts within the diffusion model, not just entire images, enabling targeted attribution for copyright compliance and model interpretability.

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.

Primary Area: generative models

Submission Number: 9959