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

2025-07-09

AI Summary

Diffusion-based image generation raises critical challenges regarding copyright attribution and transparency: existing attribution methods can only identify training samples influencing entire images, failing to trace fine-grained semantic concepts such as style or object categories. This paper introduces the first concept-level attribution paradigm for diffusion models. We reformulate the training loss under posterior sampling and design a concept-aware reward function to enable precise溯源 of specific semantic elements. Further, we integrate influence functions, differentiable rendering, and posterior sampling into a concept-sensitive attribution framework. Evaluated on the AbC benchmark, our method significantly outperforms state-of-the-art approaches. It enables practical applications including intellectual property identification, harmful content detection, and prompt engineering analysis. By enabling interpretable, semantically grounded attribution, our work provides a foundational technical advance toward responsible governance and explainability of generative AI systems.

Technology Category

• Computer Vision: Diffusion Models for Vision
• Natural Language Processing: Generation
• Cognitive Modeling & Cognitive Systems: Conceptual Inference and Reasoning

Application Category

• User Modeling, Personalization and Recommendation: Fairness-aware retrieval and ranking
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Web Mining and Content Analysis: Large pretrained models with web data

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.

Problem

• Understanding concept learning in diffusion models
• Addressing copyright and transparency concerns in diffusion models
• Isolating contributions to specific image elements like styles or objects

Innovation

• Concept-level attribution via Concept-TRAK method
• Reformulated diffusion training loss for robust attribution
• Concept-aware reward function emphasizing semantic relevance