Concept-TRAK: Understanding how diffusion models learn concepts through concept attribution
Authors: Yong-Hyun Park
Chieh-Hsin Lai
Satoshi Hayakawa
Yuhta Takida
Naoki Murata
Wei-Hsiang Liao
Woosung Choi
Kin Wai Cheuk
Junghyun Koo
Yuki Mitsufuji
Venue: ICLR 2026 Poster (2026)
Keywords: Diffusion models
Data attribution
Concept
Primary area: generative models
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. We further demonstrate its versatility on real-world text-to-image generation with compositional and multi-concept prompts.

Published as a conference paper at ICLR 2026
CONCEPT-TRAK: UNDERSTANDING HOW DIFFUSION MODELS LEARN CONCEPTS THROUGH CONCEPT-LEVEL ATTRIBUTION
Yonghyun Park 1*
Chich-Hsiin Lai²
Satoshi Hayakawa²
Yuhta Takida²
Naoki Murata²
Wei-Hsiang Liao?
Woosung Choi
Kin Wai Cheuk?
Junghyun Koo?
Yuki Mitsufuji2.8
University of Pennsylvania
SONY AI
3 Sony Group Corporation
ABSTRACT
While diffusion models excel at image generation, their growing adoption raises critical concerns about copyright issues and model transparency. Existing attribu- tion methods identify training examples influencing an entire image, but fall short in isolating contributions to specific elements, such styles or objects, that are of primary concern to stakeholders. To address this gap, we introduce concept-level attribution through novel method called Concept- TRAK, which extends influence functions with a key innovation: specialized training and utility loss functions designed to isolate concept-specte quality. We evaluate Concept- TRAK on novel concept attribution benchmarks us- ing Synthetic and CelebA- HHQ datasets, as well as the established AbC benchmark, showing substantial improvements over prior methods in concept-level attribu- tion scenarios. We further demonstrate its versatility on real-world text-to-image generation with compositional and multi multi-concept prompts.
influences rather than overall reconstruction
INTRODUCTION
Diffusion models (Lai et al., 2025; Ho et al., 2020; Song et al., 2021a;b; Rombach et al., 2022; Ramesh et al., 2022; Saharia al., 2022) have achieved remarkable success in image generation, not merely through generating high-fidelity images, but through their ability to learn and flexibly compose concepts from training data (Rombach et al., 2022; Okawa et al., 2023).
This capability raises important questions about accountability and transparency. When models learn and exploit specific concepts from training data, stakeholders need understand which training samples contributed to those concepts: whether for recognizing artistic contributions, ensuring fair compensation, safety auditing, model debugging, or copyright compliance (Wen et al., 2024; Somepalli al., 2023b; Carlini et al., 2023; Somepalli et al., 2023a; GenLaw2024; Brittain, 2023).
To address these diverse needs, data attribution methods (Koh and Liang, 2017; Park et al., 2023a; Ghorbani and Zou, 2019) have emerged promising tools, estimating how much each training example contributes to a generated output (Deng et al., 2023). These methods have proven valuable for tasks such as data valuation (Jia al., 2019a), curation (Min et al., 2025), and understanding model behavior (Ruis et al., 2024). While recent work has begun exploring attribution methods tailored for diffusion models (Zheng et al., 2024; Lin al., 2025; Mlodozeniec et al., 2024; Wang et al., 2024b), these approaches generally estimate contributions at the level of entire images. This broad perspective poses critical limitation: in many practical scenarios, stakeholders care about specific concepts within an image, rather than the whole composition.
For example, consider an Al-generated image depicting an IP-protected character (c.g., Pikachu) rendered in pencil drawing style as shown in Figure (a)) In such cases, copyright concerns from IP holders (e.g.. The Pokémon Company) would primarily focus on the character itself, not
Work done during an intersship SONY AI. Email: park 9eseas www.n..ee..... .ddu

---

Published as a conference paper at ICLR 2026
(a) Traditional attribution
(b) Concept attribution Concept-TRAK
Which tre training sample
training sample
Influenced to generated image?
Influenced the ( Pikeacuu concept?
P
TRAK
Concept-TRAK
6
WOLDAR'S
Al- generated Image
Top Influences
Al- generated images
Top Influences
Figure I: (a) Traditional attribution methods like TRAK identify training samples that influenced an entire generated image, often yielding influences unrelated to specific concepts of interest. (b) Our Concept-TRan -RAcK identifies training samples that specifically influenced a targeted concept (e.g.. "Pikachu"), " enabling precise attribution for features of interest.
the stylistic pencil rendering. Yet, traditional attribution methods, e.g., TRAK, identify training samples that influenced the generation of the full image, failing to isolate those tied to particular concepts. As Figure (a) demonstrates, these methods tend to retrieve stylistically similar images (e.g., pencil-drawn objects) but miss the character that is actually subject to copyright protection.
To address this gap, we introduce concept-level attribution, which estimates each training example's contribution to specific semantic features such as styles, objects, concepts. Building this, propose Concept-TRALE TRAK,, an extension of influence functions (Koh and Liang, 2017) that quantifies how training data affects the model's ability to generate individual concepts. Our key insight is that effective concept attribution requires designing loss functions that capture relevatt directions rather than optimizing for overall reconstruction quality. Concept- TRAK achieves this through reward-based formulations that explicitly target concept-relevant influences. As shown in Figure (b)) Concept-TRAL correctly identifies training samples responsible for the concept of Pikachu, rather than irrelevant stylistic cues.
To rigorously evaluate our method, we introduce novel concept-level attribution benchmarks on Synthetic and CelebA-HQ datasets. Concept-TRAK substantially outperforms baselines, especially for -aitr--distribution samples with unseen concept combinations. Additionally, we evaluate on the established AbC benchmark (Wang et al., 2023b), retrieval-based framework for text-to-imane model data attribution, where Concept- TRAK significantly outperforms prior methods by accurately retrieving training examples that influence specific concepts. Finally, case studies on IP-protected content, unsafe concept detection, model debugging, and relational concept learning highlight Concept-TRAK3 practical utility for real-world applications and understanding concept learning in diffusion models.
BACKGROUND
2.1 DIFFUSION MODEL.
Diffusion models (Sohl-Dickstein al., 2015; Song et al., 2021b; Ho al., 2020) are class of generative models that synthesize images through an iterative denoising process. Starting from clean image Zo, the forward process adds Gaussian noise produce a sequence increasingly noisy images 2,, following, q(r, æo) = ((/or (1 T)) where a; is a noise schedule controlling the level corruption at timestep
A neural network en (zt, 1) is trained to predict the added noise e, enabling reconstruction of xo from Fe at noise level The training objective is called the denoising score matching (DSM) loss: CDsM(#0; 0)) Ex.... [lle eø (z,, )) which encourages the model to approximate the gradient of the log-density (i.c., score function): eø (Z t) x V (x,e). For simplicity, we omit the timestep t in co(x,,1) t) when the context is clear.
Diffusion models can be extended to conditional generation by incorporating additional information such as a text prompt (Ho and Salimans, 2022; Nichol et al., 2021). In this setting, the model lears 6o (z c) & V log (z, c), allowing it to generate images that are not only realistic but also aligned with the conditioning input.
2

---

Published as a conference paper at ICLR 2026
2.2 DATA ATTRIBUTION
The goal of data attribution is to estimate the contribution of a training sample x' to a model's utility loss V, e-g., performance metric or objective function that quantifies how well the model performs (e.g., test loss) (Koh and Liang, 2017; Park et al., 2023a; Ghorbani and Zou, 2019). While Leave-One- Out retraining provides exact attribution, is computationally prohibitive for modern large-scale models.
To address this limitation, influence functions (Koh and Liang, 2017) efficiently approximate the effect of removing a training example x² using gradient- based estimates. Given a model with parameters e Idd and training loss ((:;),, the influence function is defined as:
I( V) A gy H g..
Here, 9r Va[(=j; 0) represents the gradient of the loss with respect parameters 0 for sample zo- gv = VoV the gradient of utility loss V, and = (D; denotes the Hessian matrix the training loss computed over the entire training dataset D f However, computing influence functions remains computationally challenging, each attribution query requires recomputing training gradients for the entire training set in addition to the expensive Hessian computation (Choe al., 2024).
To address this, TRAK (Park et al., 2023a) proposes projecting gradients into a lower-dimensional space using a random projection matrix P Ra with k < d:
I( V) (PT gv) P g:.
where Hp pT HP k is the projected Hessian. This enables efficient storage and reuse of gradient for multiple attribution queries.
2.3 DATA ATTRIBUTION FOR DIFFUSION MODELS
Prior work on diffusion models (Xie al., 2024; Zheng et al., 2024; Lin et al., 2025) has primarily focused on whole- attribution, typically using the objective for both training and utility losses. These studies reveal that attribution performance highly sensitive to the choice of loss function. In particular, the standard DSM loss introduces stochasticity via both the noise term and the perturbed input It, resulting in noisy gradient estimates that require extensive averaging to be reliable, making it suboptimal for attribution. To mitigate this, D-TRAK (Zheng et al., 2024) employs the squared f2-norm ||ea [13 and DAS (Lin et al., 2025) employs the squared fz-norm ||Ea | compute influence scores, achieving improved stability and accuracy.
These findings highlight that choosing a robust loss function for gradient computation is essential for reliable attribution in diffusion models. We extend this insight the concept-level setting, which demands loss functions specifically designed to capture concept-specific influence.
3 METHOD
We present Concept-TRAK, a framework for quantifying the contribution of individual training samples specific concepts learned by diffusion models. Unlike prior work that measures influence on entire generated images, our approach targets how training data affects the model's ability to represent particular semantic concepts.
3.1 DEFINITION: CONCEPT-LEVEL ATTRIBUTION
We define concept-level attribution as measuring how training sample æi influences the model's ability to generate concept Cuaget- We quantify this through the expected concept presence:
po (Caargg Eeg Pego--pumpe(-le) (p(Ctarea(>z0)) [20)]
where P(Curges (o) represents the probability that concept Curget is present in image Fp, and Psampte (-Ic)
represents the sampling distribution used for generation.
More specifically, consider two attribution scenarios (Figure 2):
3

---

Published as a conference paper at ICLR 2026
(a) Global concept attribution
Q: Which training sample Influenced to the dog)
( (log) )
Concept-TRAK
Gan
concept?
General concept
Top Influences
(b) Local concept attribution
Which training sample Influenced dag) ) concept in this image?
Concept-TRAK
Generated image
Top Influences
Figure 2: (a) Global concept attribution identifies training samples that influenced the learning of general concepts across all generations. (b) Local concept attribution identifies training samples that influenced the learning of specific concept manifestations appearing in particular generated image. For example, when applying local concept attribution the "dog" concept a generated image of a bulldog-Iig like dog. can observe that it retrieves images similar to bulldogs, demonstrating more targeted attribution.
Global attribution: Psample represents the model's generative distribution (e.g., conditional sam- pling with CFG), measuring concept probability across all generations
Local attribution: Psample 5(r0 xr(²),, measuring how training data influences the specific manifestation of a concept in a particular generated image 2x
To approximate how training samples contribute this concept probability, we employ influence function frameworks:
I( Currget) V Comeepp (Cuaree: 0 H V (=) 0),
(1)
where Cconcopu( CCarget: 0) measures model performance for specific concept generation (i.e., our utility loss V), and Luain (z); ()) captures the training sample's contribution.
In this work, we focus concepts Otagett that can be specified conditioning inputs (e.g., text prompts of "Pikachu", "Mario", or class index), enabling us to leverage existing conditional generation mechanisms for precise attribution. For general concept attribution including visual concepts, please refer to the Appendix C.
3.2 LOSS FUNCTION DESIGN
While the theoretical setup would be pe (Cuaget the utility loss and the standard DSM loss for training loss, prior work has shown that attribution performance is highly sensitive to loss function design (Section 2.3). The key challenge lies designing loss functions Leonceps and Caain that provide robust, concept-relevant gradients rather than generic denoising signals.
Geometric Motivation Our approach motivated by the hypothesis that meaningful concept directions correspond to tangent vectors of the diffusion model's latent space, which we leverage to design concept-aware loss functions.
The latent variables of diffusion models Zt, lie on a lower-dimensional manifold (Chung al., 2022b; Liu et al., 2022). Prior work has identified semantically rich structure within this manifold's tangent space, enabling concept-based editing approaches (Park et al., 2023b). Additionally, classifier-free guidance vector e ( c) - e cx(x)). which contain rich concept information (Brack et al., 2023; Wang al., 2023c). have been shown operate effectively in the tangent space of the data manifold (Chung et al., 2024; Kwon et al., 2025).
Motivated by these findings, hypothesize that concept-relevant directions can be more effectively captured by operating in the tangent space (Chung et al., 2022a; Park et al., 2023b; Wang al., 2023c) rather than through standard denoising objectives.
4

---

Published as a conference paper at ICLR 2026
Solution: Reward Optimization Our geometric motivation raises a practical question: how do identify concept-relevant directions within the tangent space and incorporate them into loss functions? We propose that reward optimization provides this capabillity: reward gradients V R(x,), serve as concept-specific guidance directions that point toward concept-enhancing regions in the tangent space.
Consider the reward optimization objective (Jaques et al., 2017; Rafailov et al., 2023):
max Ex [2))) - 3DxL (pe (+1c) Psample(-|c)
Pe
where R(xo) is reward function, pa pa(-(c) is our target model for reward optimization, and controls regularization strength. While the sampling distribution Psumple can vary practice, we use Po(-Ic) in our derivation for theoretical clarity. (For general case, please refers to Appendix B.1)
The optimal solution is (Rafailov et al., 2023; Korbak al., 2022):
p" (zolc) x Po ()) exp( R(za)/3).
For diffusion models, we can extend this to intermediate timesteps by defining rewards on noisy latents R(2;) (Wallace et al., 2024).
To analyze the gradient direction toward this reward-shaped distribution, we define loss based on Explicit Score Matching (ESM) (Vincent, 2011; Huang et al., 2021):
LESM (To;; 0) = [e,-o(2.120)) [||
log p(ge (c) - Vz, (x,lc))))
For all subsequent loss functions, we assume 2, ~ q(x (r)) unless otherwise specified.
The score function of p" (æz(c) decomposes V.,, log p" (xxle) v., log (zzccc) v R(,,).. Converting to diffusion model notation, this leads to our reward-based loss function:
Lreward (zo; 0) E.,, [||sgle((<ric) 1/B V R(x:)) (0( c)
(2)
where sg[-] is stop-gradient operation, and 99 is hyperparameter whose specific choice becomes irrelevant due to gradient normalization (Section 3.4).
Intuitively, Eq. (2) steers the model's ' output in the direction of the reward gradient V R(x,). We
now instantiate this framework with concrete reward designs that ensure the gradients operate in the tangent space.
3.3 CONC
We now instantiate our reward-based framework by designing specific reward functions for concept attribution. Following Eq. (2), we replace the general reward R(x,)) with two concept-specifer rewards: one that increases the probability of generating training sample s/. and another that increases the probability of concept c. These become our training and utility losses, respectively.
Training Loss
To capture how training sample zá influences the model's generation, define:
Raain (x:) A log p(z) (2))
where &o = E[ is the posterior mean predicted by diffusion model. This reward encourages the model to generate samples likely to have originated from zq-
Following the approach from Diffusion Posterior Sampling (DPS) (Chung et al., 2022a), we assume Gaussian distributions for the training data, i.e., p(2)(1z0) xx exp(-I|ag o (m giving us (xx, x 11/Oaas || & 2 While DPS uses this for posterior sampling, we propose to use it for attribution by constructing a training loss. This gradient V.. -l 1/ operates as tangent vectors on the data manifold (Chung al., 2022a), aligning with our geometric framework.
Substituting this into our framework (Eq. (2)) gives us the training loss:
Larii (Io; 0) = Eg,, [[sg[< (F)) + X Vz, 1l o 2
( 1)
Interestingly, this tangent space motivated formulation yields equivalent loss to V-DB from the GFlowNet framework (Liu al., 2025) under specific assumptions.
5

---

Published as a conference paper at ICLR 2026
where X, 13 is hyperparameter whose specific choice becomes irrelevant due to gradient normalization (Section 3.4).
While DSM loss has similar goal of capturing how training samples influence generation, train loss differs in how the learning signal is constructed. DSM provides reconstruction-drve signal, whereas DPS-based reward explicitly yields tangent-space guidance vectors, which empirically find more stable for concept-level attribution.
Utility Loss
To measure concept presence, we define:
Reuncept (3z)) _ log (Ctargea
Maximizing this reward corresponds to maximizing lower bound of our target concept probability Po (oCCurga) = Exeerppoee [Po Curgea [2)] (Appendix B.2). Here, Psample determines the attribution scope:
for global attribution, represents the model's 's generative distribution; for local attribution, Psample 5(zp ).
When Curget conditioning input, the reward gradient reduces classifier-free guidance vectors (ate; Clurget) - 6o (3re) (Ho and Salimans, 2022), which operate as concept-relevant tangent vec- tors (Chung al., 2024; Brack et al., 2023). For concepts embedded within condition c, we use concept slider guidance ea (zz; c) e0 ( c_) (Gandikota et al., 2024) to measure the target concept's contribution within the context, where e_ is the condition that removes the target concept (e.g.. "pencil drawing of Pikachu", c.: "pencil drawing").
Substituting this into our framework (Eq. (2)), the corresponding utility loss is:
Loonoope CCragetti = Ex [[sg]e + A (Ea (atri c) Eo (Zr; e- ))] c c)
where Ae is scaling constant whose specific value does not impact final attribution scores due to gradient normalization (Section 3.4).
Concept-Level Influence Function Having designed both training and utility losses to operate through reward gradients in the tangent space, we can now apply the influence function framework:
I I(ais Otarget) V Cooncepp (Curgett 0) H V Cinin (zij 0)
This measures the alignment between the guidance direction induced by training sample z and the guidance direction representing target concept Curgat High alignment indicates that the training sample significantly contributed to the model's 's ability to generate the concept.
3.4 ADDITIONAL TECHNIQUES
Deterministic Sampling via DDIM Inversion eliminate stochasticity from the forward diffusion process 2, ~ q=>//2)) for more stable attribution, we employ deterministic DDIM inversion derive deterministic noisy latents o;- DDIMiny (z). 0 + 1) from training samples xij. Combined with our loss functions, this approach removes all sources of randomness from gradient computation, resulting in more stable influence estimates through improved gradient fidelity.
Global vs. Local Concept Attribution The choice of sampling distribution Psample in our utility loss (Eq. (4)) determines the attribution scope. We implement this distinction follows: (1) Global attribution uses the full conditional distribution, sampling via DDIM from noise. Local attribution constrains this to Psample (2o) = 5(x0 - s, requiring tes0o = 2o We enforce this constraint by sampling via DDIM from noise that used for generating nrtsst.
Gradient Normalization Varying loss magnitudes across timesteps can cause certain gradients to dominate attribution results. To address this, normalize each timestep gradient 9: to unit norm, 9e = 9 ensuring that no single timestep exerts disproportionate influence on the final attribution score. This normalization also makes our method invariant to hyperparameters such as and data in our framework, providing additional robustness.
Gradient Projection Following TRAK (Park et al., 2023a), we project gradients to lower- dimensional space (k & d) for computational efficiency. Moreover, we approximate the Hessian

---

Published as a conference paper at ICLR 2026
using the Fisher Information Matrix, which requires only negligible overhead given pre-computed training gradients.
We refer to the complete method incorporating reward optimization based loss function, deterministic sampling, and gradient normalization within the influence function framework Eq. (5) as Concept- TRAK. Further implementation details provided in Appendix E.
4 EXPERIMENTS
In this section, evaluate Concept-TRAK across multiple concept attribution scenarios, comparing against TRAK (Park et al., 2023a), D-TRAK (Zheng et al., 2024), DAS (Lin et al., 2025). For text-to- image (T2I) model, additionally compare against unlearning-basee based attribution method (Wang et al., 2024b). To evaluate concept-level attribution, we conduct two controlled evaluations on class- conditional diffusion models and evaluate on an established real-world T2I model data attribution benchmark (AbC, Wang et al. (2023b)). Note that standard Linear Datamodeling Score (LDS, Park al. (2023a)) used traditional data attribution is inapplicable to concept-level attribution evaluation (see Appendix A for detailed discussion).
Scope of Experiments Since baseline methods were developed with concept-level attribution in mind, is natural that they struggle under our evaluation. Our comparisons should therefore be viewed not as criti- cisms of prior methods, but evidence that this newly defined task demands specialized approaches, highlight- ing the need for further research in this direction.
(a)
In-distribution
Out-of-distribttio
(b) sample
Concept attribution
Top Influences
In the main text, we focused on controlled evaluations on synthetic and CelcbA-HQ datasets, and evaluation on the established AbC benchmark for T2I models. In the appendix, provide extended analyses including: (1) set-level attribution as an alternative baseline (Ap- pendix D.I), (2) challenging real-world scenarios with semantically similar concepts (Appendix 1.2.1) and complex compositional prompts (Appendix D.2.2), (3)
Shape
Figure 3: Experimental setup. (a) Train dif- fusion models on image-uple pairs (shape, color), excluding all red-Iriangle combina- tions. (b) Generate ID/OOD samples and perform concept-level attribution; the pre- diction is correct if the top influential train- ing samples contain the target concept.
qualitative case studies on diverse applications (Ap- pendix D.2.3, Appendix D.4).
4.1 CONTROLLED EVALUATION: SYNTHETIC DATASET
Evaluating concept attribution requires knowing the ground truth source of each concept, which is unavailable in real datasets. To address this, we construct a controlled synthetic dataset with two binary concepts: color E (red, blue} and shape E |ttriangle circle). We train conditional diffusion model where each condition is encoded as concatenated one-hot vectors.
To comprehensively evaluate concept attribution methods, we design our dataset enable testing in both in-distribution (ID) and duitt--diittiuution (OOD) scenarios (Figure 3(a)). We exclude {red, triangle) combinations from training, creating ID cases where concept combinations were seen during training (e.g., blue circle) and OOD cases requiring novel concept combinations through generalization (c.g., red triangle). This setup allows us to understand how attribution methods behave when training data directly supports the generated output versus when the model must combine concepts in novel ways.
Evaluation Protocol We use Precision @ 10 to evaluate concept attribution. illustrated in Figure 3(b), we generate test image and perform concept attribution for specific concept (e.g.. "shape"). We then check whether the top-ranked training samples contain the same target concept as the generated image. In this example, training sample with triangle is correct (O) while one with a circle is incorrect (x). We generate 16 test images for each concept combination and report Precision @ 10 averaged across all test cases. Note that baseline methods perform standard data attribution on the generated image, while Soncept-TRAK performs -pecctt-peccfic attribution targeting individual concepts within that image, i.e., local concept attribution.
7

---

Published as a conference paper at ICLR 2026
(a) In-distribution
Top Influences
(b) Out-of-distriaution Generated sample
Top Influences
Generated sample
Ours (Shap))
Ours
(Shapa)
Ours
Ours
(Colod
(Cole)
DAS
DAS
D-TRAK
D-TRAK
Figure 4: The target concept for each method indicated in parentheses (Shape/Color). A data attribution method succeeds when the top influential training samples contain the concept as the generated sample. (a) In-distribution case: Both baseline methods and our approach successfully retrieve relevant training samples. (b) Out-of-distribution: Our method accurately retrieves training samples for each individual concept (triangle for shape, red for color), while baselines can only retrieve samples related to one concept due to image-level attribution limitations.
Table I: Precision @ 10 on synthetic dataset.
Table 2: Precision @ 10 on CelebA-HQ dataset.
In-distribution
Out-of-distribation
In-distribution
Out-of-distribun
Method
Shape
Color Avg.
Shape Color
Avg.
Method
Eyeglasses
Male Smiling Avg. Byeglasses Male Smiling Avg.
Ours
1.00
1.00
1.00
0.80
0.90
0.85
Ours
0,97
0.93
0.87
0.92
1.00
1.00
0.90
0.97
1.00
1.00
0.50
DAS
0.99
0.90
0.70
0.70
0.67
D-TRAK
1.00
1.00
1.00
1.00
0.00
0.50
D-TRAK
0,56
0.44
0.51
0.50
0.30
0.60
0.00
0.30
TRAK
0.67
0.93
0.80
0.60
0.30
0.45
TRAK
0.86
0.96
0.71
0.84
0.60
0.70
0.50
0.60
Results As shown in Table 1, Concept- TRAK maintains strong performance in both ID (1.00) and OOD (0.85) scenarios, while baseline methods exhibit a significant performance drop in the OOD setting (< 0.50). This gap underscores fundamental distinction between attribution settings. In ID cases, image-based attribution methods can succeed for concept retrieval only by leveraging visual similarity, since there exists training samples that includes the same concept combinations those in the generated output (Figure 4(a)). In contrast, OOD cases contain no training samples with the exact target concept combination, requiring methods to isolate the contribution of individual concepts from compositionally novel outputs (Figure 4(b)).
4.2 CONTROLLED EVALUATION: CELEBA-HO
We extend our evaluation to real images using CelebA-HQ with three binary concepts: eyeglasses, male, and smiling. We deliberately exclude all samples containing the combination {cyeglasses, male, smiling} from the training dataset, creating a more challenging OOD scenario where the model should compositionally combine three concepts. We follow the same Precision @ 10 evaluation protocol, generating 16 test images per available combination.
Results As shown in Table 2, Concept-TRee TRAK achieves consistently strong performance both ID (0.92) and OOD (0.97) scenarios. In contrast, DAS performs well in ID (0.96) but drops substantially in OOD settings (0.67). This difference reflects the distinct challenges of each setting: ID scenarios often benefit from image-level similarity, as retrieved samples visually resembling the generated image typically contain the target concept. However, OOD scenarios require isolating individual concepts from compositionally novel combinations, where image-level similarity insufficient for accurate concept attribution (see qualitative results in Appendix D.3).
4.3 ATTRIBUTION BY CUSTOMIZATION (ABC)
We use the Attribution by Customization Benchmark (AbC) (Wang et al., 2023b), an established benchmark for T21 model data attribution. AbC evaluates attribution methods on models fine-tuned with exemplar images to learn new concepts via special tokens (V), measuring whether attribution method successfully retrieves the exemplars from generated images. This setup offers a rare source of ground truth: generated outputs known to be directly influenced by the exemplars. While this setting lacks generality for large-scale training regimes, remains the most reliable way to evaluate concept-level attribution in current T2I models.
8

---

Published as a conference paper at ICLR 2026
AbC-style
AbC con-eettett
Table 3: Ablation study.
0.8
0.8
Recallgg 10
Config
Recall @ 10 (1)
0.6
Dura (Chobal
0.6
Durs Bacal
TRAK (Base: LUSM)
0.04
D-TRKX
+ (Config A: CRewaddDDS)
0.261
0.2
liua
0.2
DVS
+ (Config Cors)
0.335
300 1000
30
100
300
3000
+ (Config C: DDIMInv)
0.564
GPU
GPU hours
+ (Contig D: Normalize)
0.955
Figure 5: Recall@ 10 performance on AbC benchmark.
Generated Sample
Outo
Generated Sample
BeoA TATTOO
GOOCET
TROK
charoost drawdwng of share.
FIP
(v) swalis
Figure 6: Qualitative results on the AbC benchmark. Correctly retrieved samples are highlighted with red boxes. Previous methods (Unlearn, D-TRAK) struggle with interference from style elements and retrieve unrelated images, while Ours (Global) successfully isolates target concepts (V).
Evaluation Protocol Following the setup in Wang et al. (2023b), we report Recall 10, i.e., the proportion of times the exemplar images are successfully retrieved from pool containing the exemplars and 100K LAION images. While the original benchmark involves not only learning exemplar using special tokens but also fine-tume finettunin the model's parameter on the exemplar dataset, real-world use cases more commonly involve investigating the concepts generated or utilized by single pretrained model. To better reflect this, we adopt textual inversion (Gal et al., 2022) with a frozen base model (SD1.4v) (Rombach et al., 2022), which only trains with a special token (V), without parameter updates. Further implementation details are provided in the Appendix
Results Our method achieves significantly higher Recall@ 10 while maintaining computational efficiency comparable TRAK-based method, shown in Figure 5. As illustrated in Figure 6, prior methods often fail to isolate the concept of interest (V) due to interference from style or other visual elements in the generated image. In contrast, Concept-TRAK effectively isolates and attributes the target concept (V), demonstrating superior performance in concept-level attribution. These results can be explained by the inherently compositional nature of T2I generation. In this AbC benchmark, model is required to combine a learned concept (V) with diverse styles or objects to generate test samples. As we demonstrated from controlled evaluation, such compositional scenario makes precise concept attribution substantially more difficult for image-based attribution methods, thereby amplifying the performance gap observed with Concept-TRAL.
We conduct an ablation study using 48 samples from the AbC dataset to assess the impact of each design choice. Starting from the baseline TRAK with LOSM, adding concept-aware util- ity gradients (A), DPS- opS-based training gradients (B), DDIM inversion (C), and gradient normalization (D) progressively improves performance, with our full method achieving 0.955 Recall @ 10.
Ablation Study
4.4 REAL-WORLLE SCENARIOS
While prior evaluations rely on controlled settings with known ground-truth concept labels, we further demonstrate the versatility of CoRCept-TRAK on real-world text-to-image generation using SD1.4v Rombach et al. (2022) with a 100K subset of LAION Schuhmann et al. (2022). We present representative qualitative results below; quantitative evaluation with automated concept verification and additional qualitative examples are provided in Appendix D.2.
Compositional Prompts Real-world prompts often combine multiple concepts such objects and styles. key challenge is attributing training influence to each individual concept within a compositionally generated image. As shown in Figure 7, for prompt "cat in the style of graffiti art,"

---

Published as a conference paper at ICLR 2026
Prompt
Generated Image
Top Influence
(object)
Cat in the style of graffiti art
style)
Figure Qualitative results for compositional concept attribution using Concept-TRAK, For generated image from the prompt "cat in the style of graffiti art," Concept-TRAK separately retrieves training samples for the <object> and <style> concepts separately.
Prompt
Generated Image
Top Influence
Ituidyy bear
teddy bear
skateboard in Times Square
( sleatzboarrl
(Times Square
Figure 8: Qualitative results for complex multi-concept attribution using Concept-TRAPE For gen- erated image from the prompt "a teddy bear on a skateboard in Times Square," Concept-TRAK sepa- rately retrieves training samples for each constituent concept <teddy bear>, <skateboard>, and <T <Times Square>.
Concept-TRAK successfully retrieves cat images for the object concept and graffiti art images for the style concept.
Complex Multi-Concept Prompts We further evaluate on prompts that require composing three more distinct concepts. As shown in Figure 8, for "a teddy bear skateboard in Times Square," Concept-TRAie TRAK decomposes the prompt into its constituent concepts, separately retrieving teddy bear images, skateboard images, and Times Square images from the training data. These results demonstrate Concept-TRAKE 's practical utility for understanding how diffusion models compose multiple concepts from training data in real-world generation scenarios.
RELATED WORK
Data Attribution Established data attribution methods include influence functions (Pruthi et al., 2020), which approximate leave-one-our retraining via gradients. TRAK and LoGra (Park al., 2023a; Choe et al., 2024) improve scalability through random projections. Game-theoretic approaches like Data Shapley (Jia al., 2019b; Chorbani and Zou, 2019), based on Shapley values (Shapley et al. 1953), were initially limited by retraining costs, but recent work (Wang et al., 2023a; 2024a) improves efficiency by removing this requirement. Unlcarning-base methods (Wang et al., 2024b) offer alternative trade-offs between efficiency and theoretical rigor.
Data Attribution for Diffusion Models Early diffusion attribution methods adapted influence functions (Pruthi et al., 2020; Park et al., 2023a), but were biased by timestep-des -eeestep-deeendent gradient norms. Xie et al. (2024) addressed this via re- re-normalized formulation. Zheng et al. (2024) extended TRAK to diffusion models, exploiting its scalability. Lin al. (2025) later proposed the Diffusion Attribution Score, which quantifies per-sample influence by directly comparing predicted distributions, yielding more precise attributions than loss-based approaches.
CONCLUSION
In this work, we offer an initial investigation into concept-level attribution, introducing Concept- TRAK as a foundational framework. introduces specialized reward-based training and utility loss functions designed to isolate concept-specific influences. Concept-TRAK outperforms existing methods on novel concept attribution benchmarks using the Synthetic and CelebA-HQ datasets, as well as the AbC benchmark. We expect these contributions to inspire further research toward more robust concept-level attribution benchmarks and methods for increasingly sophisticated generative models.
10

---

Published as a conference paper at ICLR 2026
STATEMENT ON THE USE OF LARGE LANGUAGE MODELS
This work made use of large language models to assist with proofreading and improving the clarity of the writing. All research ideas, theoretical development, experiments, and coding were carried out solely by the authors.

---

Published as a conference paper at ICLR 2026
REFERENCES
Manuel Brack, Felix Friedrich, Dominik Hintersdorf, Lukas Struppek, Patrick Schramowski, and Kristian Kersting. Sega: Instructing text-to-image models using semantic guidance. Advances in Neural Information Processing Systems, 36:25365-25389, 2023.
Blake
Brittain. Getty images lawsuit says stability misused photos to train Renters, February 2023. URL
ai. getty-images att.://aaaaaaaa--aaaaaaaaaaaaaaa
Nicolas Carlini, Jamie Hayes, Milad Nasr, Matthew Jagielski, Vikash Sehwag, Florian Tramer, Borja Balle, Daphne Ippolito, and Eric Wallace. Extracting training data from diffusion models. In 32nd USENIX Security Symposium (USENIX Security 23), pages 5253-5270, 5270, 2023.
Sang Keun Choe, Hwijeen Ahn, Juhan Bae, Kewen Zhao, Minsoo Kang, Youngseog Chung, Adithya Pratapa, Willie Neiswanger, Emma Strubell, Teruko Mitamura, al. What is your data worth to gpt? Ilm-scale data valuation with influence functions. arXiv preprint arXiv:2405.. (tt.::2005..355, 2024.
Hyungjin Chung, Jeongsol Kim, Michael T Mccann, Marc L Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. arXiv preprint arXiv:2209.146% 2022a.
Hyungjin Chung, Byeongsu Sim, Dohoon Ryu, and Jong Chul Ye. Improving diffusion models for inverse problems using manifold constraints. Advances in Neural Information Processing Systems, 2022b.
Hyungjin Chung. Jeongsol Kim, Geon Yeong Park, Hyelin Nam, and Jong Chul Ye. Cfg++: Manifold- constrained classifier free guidance for diffusion models. arXiv preprint arXiv:2406. 08070, 2024.
Junwei Deng, Shiyuan Zhang, and Jiaqi Ma. Computational copyright: Towards royalty model for music generative ai. arXiv preprint arXiv:2312.066646, 2023.
Rinon Gal, Yuval Alaluf, Yuval Atzmon, Or Patashnik, Amit H Bermano, Gal Chechik, and Daniel Cohen-Or. An image is worth one word: Personalizing text-to-image generation using textual 2022.
inversion. arXiv preprint arXiv:2208.0166
Rohit Gandikota, Joanna Materzyúska, Tingrui Zhou, Antonio Torralba, and David Bau. Concept sliders: Lora adaptors for precise control in diffusion models. In European Conference Computer Vision, pages 172-188. Springer, 2024.
Genl.aw2024. 2nd workshop on generative ai and law (genlaw '24). https: /tops::/wwww.. genl aw. org/2024-1 cml,, July 2024. Held at the International Conference on Machine Learning (ICML) 2024, Vienna, Austria.
Amirata Ghorbani and James Zou. Data shapley: Equitable valuation of data for machine learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov, editors, Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 2242-2251. PMLR, 09-15 Jun 2019. URL https::/proceeddiie 2057/ghrrbe
mlr press/
v97 gho rbani inlaanis 19c. html.
Roger Grosse, Juhan Bae, Cem Anil, Nelson Elhage, Alex Tamkin, Amirhossein Tajdini, Benoit Steiner, Dustin Li, Esin Durmus, Ethan Perez, et al. Studying large language model generalization with influence functions. arXiv preprint arXiv:2308 03299,,, 2023.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12598 2022.
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in Neural Information Processing Systems, 33:6840-688 6851, 2020.
Yuzheng Hu, Pingbang Hu, Han Zhao, and Jiaqi Ma. Most influential subset selection: Challenges, promises, and beyond. Advances in Neural Information Processing Systems, 37:119778-11 119810, 2024.
12

---

Published as a conference paper at ICLR 2026
Chin-Wer Wei Huang, Jae Hyun Lim, and Aaron Courville. A variational perspective diffusion- based generative models and score matching. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. URL https:: http///eeeeeeeeeeeeeeeeeeee net / forum? d=bXehDYUjjx..
Natasha Jaques, Shixiang Gu, Dzmitry Bahdanau, José Miguel Hermandez- Richard Turner, and Douglas Eck. Sequence tutor: Conservative fine-tuning of sequence generation models with KL-control. In Doina Precup and Yee Whye Tch, editors, Proceedings of the 34th International Conference on Machine Learning. volume 70 of Proceedings of Machine Learning Research, pages 1645-1654, PMLR, 06-11 Aug 2017. URL https: http:///oooceeeess.e..e.. / v7 /70/jaques17a html.
mlr press/
Ruoxi Jia, David Dao, Boxin Wang, Frances Ann Hubis, Nezihe Merve Gurel, Bo Li, Ce Zhang, Costas J Spanos, and Dawn Song. Efficient task-specific data valuation for nearest neighbor algorithms. arXiv preprint 1908:v9908008199 2019a.
Ruoxi Jia, David Dao, Boxin Wang, Frances Ann Hubis, Nick Hynes, Nezihe Merve Gürel, Bo Li, Ce Zhang, Dawn Song, and Costas J. Spanos. Towards efficient data valuation based on the shapley value. In Kamalika Chaudhuri and Masashi Sugiyana. editors, Proceedings of the Sventy-Second International Conference on Artificial Intelligence and Statistics, volume 89 of Proceedings of Machine Learning Research, pages 1167-1176. PMLR, 16-18 Apr 2019b. URL https:/ htppppocedinng igg...ngg.. mml.. press / v8 9/ a .1t..tml...
Keller Jordan, Jeremy Bernstein, Brendan Rappazzo, @fernbear.ssky.occilll Boza Vlado, You Jiacheng, Franz Cesista, Braden Koszarsky, and @Grad62304977. modded-nanogpl: Speedrunning the nanogpt baseline, 2024. URL https: httpt:::///ittub..ooo. modded-nanogpt.
Pang Wei Koh and Percy Liang. Understanding black-box predictions via influence functions. In International conference on machine learning, pages 1885-1894. PMLR, 2017.
Tomasz Korbak, Hady Elsahar, Germán Kruszewski, and Mare Dymetman. On reinforcement learning and distribution matching for fine-tuning language models with no catastrophic for- getting. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, Advances in Neural Information Processing Systems, volume 35, pages 16203-16220. Curran Asso- ciates, Inc., 2022. URL httpp:: neurips wwppapprrppppee/ee/peeeeeee.. wellee/sspp/////// 1tt22//Leee/77290cfaaaaaaaa////3 oorpeaaccccccccaaaaaaaaooeoeeeee
Conferencee .aacceeppp.. .pf..
Mingi Kwon, Shin seong Kim, Jaeseok Jeong, Yi Ting Hsiao, and Youngjung Uh. Tefg: Tangential damping classifier-free guidance. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 2620-2629, June 2025.
Chieh-Hsin Lai, Yang Song, Dongjun Kim, Yuki Mitsufuji, and Stefano Ermon. The principles of diffusion models. arXiv preprint arXiv:2510.2189) 2025.
Jinxu Lin, Linwei Tao, Minjing Dong, and Chang Xu. Diffusion attribution score: Evaluating training data influence in diffusion model. In The Thirteenth International Conference on Learning Representations, 2025. URL https: wpppreveeeeeeeee////// net forum ? idd=kuuuiiieee
Luping Liu, Yi Ren, Zhijie Lin, and Zhou Zhao. Pseudo numerical methods for diffusion models on manifolds. In International Conference on Learning Representations, 2022.
Zhen Liu, Tim Z Xiao, Weiyang Liu, Yoshua Bengio, and Dinghuai Zhang. Efficient diversity- preserving diffusion alignment via gradient-infomin gflownets. In ICLR, 2025.
Taywon Min, Haeone Lee, Hanho Ryu, Yongchan Kwon, and Kimin Lee. Understanding impact of human feedback via influence functions. arXiv preprint arXiv:2501.05790 2025.
Bruno Mlodozeniec, Runa Eschenhagen, Juhan Bae, Alexander Immer, David Krueger, and Richard Turner. Influence functions for scalable data attribution in diffusion models. arXiv preprint arXiv:2410. 13850, 2024.
13

---

Published as a conference paper at ICLR 2026
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. Glide: Towards photorealistic image generation and editing with text-guided diffusion models. arXiv preprint arXh::21 2021.
Maya Okawa, Ekdeep S Lubana, Robert Dick, and Hidenori Tanaka. Compositional abilities emerge multiplicatively: Exploring diffusion models on synthetic task. Advances in Neural Information 2023.
Processing Systems,
Sung Min Park, Kristian Georgiev, Andrew Ilyas, Guillaume Leclere, and Aleksander Madry. Trak: attributing model behavior at scale. In Proceedings of the 40th International Conference on Machine Learning, pages 27074-27113 2023a.
Kiog--Hyu Park. Mingi Kwon, Jaewoong Choi, Junghyo Jo. and Youngjung Uh. Understanding the latent space of diffusion models through the lens of riemannian geometry. Advances in Neural Information Processing Systems, 36:24129-24111 2023b.
Garima Pruthi, Frederick Liu, Satyen Kale, and Mukund Sundararajan. Estimating training data influence by tracing gradient descent. Advances in Neural Information Processing Systems, 33: 19920-19930, 2020.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly reward model. Advances in Neural Informatinn Processing Systems, 36:53728-533 -33741, 2023.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text- conditional image generation with clip latents. arXiv preprint 4rXiv:2204.06125, 1(2):3 2022.
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Björn Ommer. High- resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF confer- on computer vision and pattern recognition, pages 10684-10695, 2022.
Laura Ruis, Maximilian Mozes, Juhan Bac, Siddhartha Rao Kamalakara, Dwarak Talupuru, Acyr Locatelli, Robert Kirk, Tim Rocktäschel, Edward Grefenstette, and Max Bartolo. Procedural knowl- edge in pretraining drives reasoning in large language models. arXiv preprint arXiv:2411. 12580, 2024.
Seyedmorteza Sadat, Manuel Kansy, Otmar Hilliges, and Romann M Weber. No training, no problem: Rethinking classifier-fre free guidance for diffusion models. arXiv preprint arXiv:2407.02683 2024.
Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Kamyar Ghasemipour, Raphael Gontijo Lopes, Burcu Karagol Ayan, Tim Salimans, et al. Photorealistic text-to-image diffusion models with deep language understanding. Advances neural information processing systems, 35:36479-36494, 2022.
Christoph Schuhmann, Romain Beaumont, Richard Vencu, Cade Gordon, Ross Wightman, Mehdi Cherti, Theo Coombes, Aarush Katta, Clayton Mullis, Mitchell Wortsman, et al. Laion-5b: An open large-scale dataset for training next generation image-text models. Advances in neural information processing systems, 15:25278--5294. 2022.
Lloyd S Shapley et al. A value for n-person games. 1953.
Jascha Sohl-Don Saal-DDckstein... Eric Weiss, Niru Maheswaranathan. and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. International Conference on Machine Learning, pages 2256-2265. PMLR, 2015.
Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Diffusion art digital forgery? investigating data replication in diffusion models. In Proceedings of the IEEF/CVF conference on computer vision and patlern recognition, pages 6048-6058, 2023a.
Gowthami Somepalli, Vasu Singla, Micah Goldblum, Jonas Geiping, and Tom Goldstein. Under- standing and mitigating copying in diffusion models. Advances in Neural Information Processing Systems, 36:47783-4788 47803, 2023b.
14

---

Published as a conference paper at ICLR 2026
Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In Interna- tional Conference on Learning Representations, 2021a.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021b.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 2317):16111167, 2011.
Bram Wallace, Meihua Dang, Rafael Rafailov, Linqi Zhou, Aaron Lou, Senthil Purushwalkam, Stefano Ermon, Caiming Xiong, Shafiq Joty, and Nikhil Naik. Diffusion model alignment using direct preference optimization. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8228-8223 2024.
Jiachen T Wang, Yuqing Zhu, Yu-Xiang Wang, Ruoxi Jia, and Prateek Mittal. Threshold knn-shapley: A linear-time and privacy-friendly approach to data valuation. arXiv preprint 17Xv:2330..15709 2023a.
Jiachen T Wang, Prateek Mittal, Dawn Song, and Ruoxi Jia. Data shapley in one training run. arXiv preprint arXiv:2406.1100 11011, 2024a.
Sheng- Yu Wang, Alexei A Efros, Jun- Yan Zhu, and Richard Zhang. Evaluating data attribution for text-to-image models. In Proceedings of the IEEF/CVF International Conference on Computer Vision, pages 7192-7203, 2023b.
Sheng- Yu Wang, Aaron Hertzmann, Alexei Efros, Jun-Yan Zhu, and Richard Zhang. Data attribution for text-to-image models by unlearning synthesized images. Advances in Neural Information Processing Systems, 37:4235-422 4266, 2024b.
Zihao Wang, Lin Gui, Jeffrey Negrea, and Victor Veitch. Concept algebra for (score-based) text- controlled generative models. Advances in Neural Information Processing Systems, 36:35331- 35349, 2023c.
Yuxin Wen, Yuchen Liu, Chen Chen, and Lingjuan Lyu. Detecting, explaining, and mitigating memo- rization in diffusion models. In The Twelfth International Conference on Learning Representations, 2024. URL https: /ttp::////e/eeeeeeeeeeeee n 3o
Tong Xie, Haoyu Li, Andrew Bai, and Cho-Jui Hsieh. Data attribution for diffusion models: Timestep- induced bias in influence estimation. Transactions on Machine Learning Research, 2024.
An Yang, Anfeng Li, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chang Gao, Chengen Huang, Chenxu Lv, et al. Qwen3 technical report. arXiv preprint arXiv:2505.0930 2025.
Jingfeng Yao, Cheng Wang, Wenyu Liu, and Xinggang Wang. Fasterdit: Towards faster diffu- sion transformers training without architecture modification. Advances in Neural Information Processing Systems, 37:56166 66-56189, 2024.
Xiaosen Zheng, Tianyu Pang, Chao Du, Jing Jiang, and Min Lin. Intriguing properties of data attribu- tion on diffusion models. The Twelfth International Conference on Learning Representations, 2024.
15

---

Published as a conference paper at ICLR 2026
A DISCUSSIONS
A.1 LDS BENCHMARK
Linear Datamodeling Score (LDS, Park et al. (2023a)) is the most widely used benchmark in data attribution methods. LDS measures influence on specific utility losses through leave- evaluation. More specifically, it linearly approximates each data point's impact utility loss by measuring the difference between models trained on randomly subsampled data versus the full dataset.
While LDS is sometimes treated ground-truth proxy, it is unsuitable for our setting due to both theoretical and practical limitations:
Structural limitation in clas-condiitonal setting: In datasets like CelebA-HQ. many attributes (e.g., "eyeglasses") appear in over half the data. Under LDS's standard protocol (removing the top 50% most influential samples), many positive examples remain. Since conditional models can reproduce a concept from even a small number of positives, the measured utility drop can be minimal, systematically underestimating influence. This makes measuring each concept's true impact nearly impossible, arising from the violation LDS's linearity assumption (Hu et al., 2024). The influence on class generation probability exhibits highly non-linear behavior, such as sharp increases after sufficient number of sample removals.
Computational infeasibility for text-to-image setting: Running LDS requires retraining dozens of models (often 64+) after removing large portions of the training set. For instance, training a single T2I model on MS-COCO to reach qualitatively minimally meaningful generation already requires ~ GPU-days, making LDS prohibitively expensive.
Retrieval-based evaluation as an alternative: Our retrieval- based evaluation avoids this prevalence issue by directly checking whether the top-ranked training samples indeed contain the target concept, without being confounded by leftover positives after partial data removal.
A.2 LIMITATIONS AND FUTURE WORK
Our research advances beyond traditional data attribution by identifying training samples that specifi- cally contribute to particular concepts. While Concept-TRAK demonstrates superior performance compared to existing methods concept-level attribution tasks, our approach has limitations. As illustrated in Figure 14(a) with the Simpsons example, our method occasionally retrieves stylistically similar but conceptually distinct images (e.g., other cartoon characters rather than Sippssns-speciffic content).
We hypothesize that this limitation stems from the fundamental challenge of gradient estimation in diffusion models. While our reward-opiin wwwaurd-ttiiizztionbboo based loss formulation and ddim inversion successfully eliminate stochasticity from the standard diffusion loss computation and provide more stable gradient estimates, perfect gradient estimation would theoretically require the true DSM loss computed with infinite Monte Carlo samples over both the noise term and noisy latents æt- Our deterministic approximation, though significantly improved, cannot fully capture this infinite sampling complexity. Consequently, some attribution errors persist.
For concept-level attribution to serve reliable tool for addressing copyright concerns and en- abling model debugging. further development is required in two key areas: (I) establishing more sophisticated benchmarks that measure concept-level attribution performance across diverse concept types and contexts, and (2) enhancing the precision and theoretical guarantees of concept-level attribution methods. This work represents initial investigation that introduces the concept-level attribution problem and proposes Concept-TRAK as a foundational framework, and we anticipate that these contributions will catalyze further research into more robust concept-level attribution methods suitable for increasingly sophisticated generative models.
A.3 IMPACT STATEMENTS
While our work provides tools for analyzing training data and understanding diffusion models for image generation without direct safety concerns, there exists potential for misuse by model developers who might exploit our tools to learn unsafe problematic concepts. We emphasize that our method
16

---

Published as a conference paper at ICLR 2026
intended for responsible model development and governance, including the identification and mitigation of harmful content training datasets, and understanding the model's ' behavior.
B THEORETICAL DETAILS
B. CLASSIFIER- GUIDANCE EXTENSION
We derive the reward optimization loss for the case where Paamle corresponds to classifier-free guidance sampling. Under certain assumptions, this yields loss gradients equivalent Eq. (2) derived from Po in the main manuscript.
Consider the reward optimization objective with classifier- free guidance sampling:
max Exx- [[[[))]] [B(2)))
] - 3DKL (pe (pp(-lc)llPaamppppl)) Psample
where Psample (-lc) represents the classifier-free guidance distribution used in practice. The optimal solution follows the same form:
p" (xule) x (aample(2u)) (e) exp((((z)/3)
Following the same ESM derivation as the main text, obtain the CFG-based reward loss:
CCRG (#o: 0) = Ex, [[|sgle) (x:: - 1/3 Vx,R(z)) - Eg CFG (1);))) CFG (xx;e) CFF(C c) where Eq
cx(#)) ((E0(3);)) Eo (2x)) is the CFG noise prediction.
Expanding this expression:
DewanaZoo 0) Es, [[ssg(e((z)) (eo(EEc)) cs(x,) - 1/3 V., 2,, R(x)) (co( + ((e(()))) <x(x))))
The gradient of this loss is:
1 V
V CCC (zoz 0) Ex,
+ -(((())))))c((==))) c((xx))))
Under the assumption that Voco(x)) contributes minimally to concept-specific attribution directions, this simplifies to:
Crewand (z0; 0) RJ Ex, V., VoE
This is equivalent the gradient of Eq. (2) up a constant factor, which becomes irrelevant under gradient normalization. While this assumption is strong, it may be justified since co ca(x)) captures general denoising pattemss across the dataset, which could be largely independent of specific concept directions. The effectiveness of Eq. (2) in our experiments provides some empirical support for this approximation.
B.2 LOWER BOUND JUSTIFICATION FOR UTILITY Loss
We provide the mathematical justification for why maximizing reward function (concop((12)) log pl (CCagga (2)) corresponds to optimizing lower bound of our target concept probability pe (Garget) =
E Exgepenee Po (Curgen [..
Since Pa (Cunger must be computed through the diffusion process, we have:
log (Crarget) log Exy "Paasix [po ((argge(/z)))
(6)
log Err Prompip [Ez, (eg(ar (ro) (po (targen [x:)]]
(7)
E Pooaptt (o ( [log po (Carged [2:]]
(8)
where the inequality follows from applying Jensen's inequality twice (due to the concavity of log). Therefore, maximizing Ex.... [log Pu (Carget probability.
[2,]] optimizes lower bound of our target concept
17

---

Published as a conference paper at ICLR 2026
C OTHER TYPES OF REWARD
In this section, we show how to apply Concept-TRAK beyond textual concepts. We cover two scenarios: explicit differentiable reward models and implicit reward models defined by preference datasets.
C.I EXTERNAL DIFFERENTIABLE REWARD MODELS
Suppose we have access to an explicit, differentiable classifier that can predict the probability of specific concept: log p( (rarger zo). Here, Curget can be any concept of interest- for example, visual features, image aesthetics, etc. If this concept classifier is trained only on clean images 2o. then similar to our earlier approach, we define the reward model as R(xa) = log P(Cungel (oo), where t. = [z [ra] is the posterior mean predicted by the diffusion model (Chung et al., 2022a).
This yields our utility loss based on external reward model:
DRewadd-Des (Iz; 8) = |sg (co(zr) - 1/3 Vx, log pl (runged [0)] co c0(2))
C.2 PREFERENCE DATASETS
For scenarios where concepts are defined through preference data rather than explicit classifiers, can adapt our framework to work with preference pairs. This is particularly useful when the desired concept is subjective (e.g., aesthetic quality, safety) difficult to define through explicit labels.
Given preference pairs (x 2n where zi" > To indicates that zo is preferred over zy we need define how noisy latent z; sampled from our diffusion process relates to these preferences. Following the intuition that we want to steer the denoising process toward preferred outcomes and away from non-preferred ones, define the reward function as:
R(I 2 _o = log p( p (p)) log p( (0);
where áu El is the posterior mean predicted from the current noisy latent Zt.-
This formulation captures likelihood that the current denoising trajectory will lead to the preferred sample z versus the non -reeerefede sample Zo Under our Gaussian assumption that p(r" (2o) oc sap(-1[āg - x² (12). we obtain:
R( at 25) x -1200 - x 112
|
25
const
(10)
Taking the gradient with respect to dt:
V., R(x;; at _o 1io 2ō
11 V
1
1
(11)
This gradient naturally encourages the denoising process to move toward preferred samples zo (negative gradient term) and away from non- non- preferefeedd samples in (positive gradient term), making it suitable for integration into our Reward-DPS framework.
The resulting utility loss becomes:
Cyreferene-Des ((r)) xn 0) Ez [llsg(((((r)) 1/B V R(I 20 2n 66 (xx)
E. [l)gg(e + 1/B. (V., 120 - zo 112 1l - 2 [1:) ea (x))
enabling concept attribution with preference data without explicitly training reward model.
D ADDITIONAL RESULTS
D.I ADDITIONAL BASELINE: SET-LEVEL ATTRIBUTION
In the main paper, we focus our evaluation on local concept attribution, which measures the influence of training samples on a specific concept within a particular generated image. However, an alternative setting is global concept attribution, which measures the influence of training samples the model's learned distribution over a concept in general. Unlike local concept attribution, global concept attribution be naturally addressed by baseline methods such DAS and D-TRAK through set-level attribution, where utility gradients are computed by averaging across multiple generated images.
18

---

Published as a conference paper at ICLR 2026
Set-level Attribution For target concept c, define set-level utility gradient as:
Lyst Poo (zoo c)]
(12)
This formulation estimates the expected loss over the model's generative distribution conditioned on concept Assuming L approximates p(z|c), this objective closely aligns with our concept attribution utility definition Section 3.1:
(o(c) = Ezo-=p(ale))p(clar)))
(13)
By averaging gradients over multiple samples from p(z(c), we obtain an estimate how training data influenced the model's 's learned representation of concept overall, rather than its manifestation single image.
Experimental Setup We evaluate set-level attribution across all three benchmarks (Toy, CelebA- HQ, AbC) by adapting both Concept- TRAK and baseline methods to the global attribution setting:
Toy & CelebA-HQ: For each target concept c, we fix that concept and randomize all other attributes. We generate 256 images from p(z|c) using the trained diffusson model. For each method, we compute the utility gradient as the average of individual gradients across all 256 samples.
AbC: We use the benchmark-prowee provided prompts (containing the special token) to generate 256 images. We then compute the average utility gradient across these samples.
We then rank training samples by their influence scores and evaluate whether the top-ranked samples contain the target concept using Precision 10.
Results Tables 4, 5, and present the results for global concept attribution across all benchmarks.
Table 4: Set-level attribution results Toy dataset. All methods achieve strong performance this controlled dataset.
Concept Concept-TRAK
DAS
D-TRAK
Shape Color
1.00
1.00
1.00
0.90
0.70
1.00
Average 0.95
0.85
1.00
Table Set-level attribution results on CelebA-HQ. All methods achieve perfect performance when attributing concepts across multiple generated images.
Concept
Concept-
DAS
D-TRAK
Eyeglasses
1.00
1.00
1.00
Male
1.00
1.00
1.00
Smile
1.00
1.00
1.00
Average
1.00
1.00
1.00
The results demonstrate that Concept-TRAY achieves comparable superior performance across all benchmarks, confirming that our design choices transfer well the global setting. Notably, while baseline methods can perform reasonably in the global setting by averaging over many samples, they fundamentally lack the capability perform local concept attribution for individual images-a critical limitation that Concept- TRAK addresses.
D.2 REAL-WORLD SCENARIOS
While our controlled benchmarks (Toy, CelebA-HQ. AbC) provide rigorous evaluation with ground- truth labels, real-world applications of concept attribution often involve more challenging scenarios,
19

---

Published as a conference paper at ICLR 2026
Table Set-level attribution results on AbC benchmark. Concept-TRAKK achieves the best average performance, with particular strength in style attribution.
Concept
Concept-TRAK
DAS
D-TRAK
Object Style
0.89
0.98
0.95
0.93
0.81
0.83
Average 0.91
0.895
0.89
including semantically similar concepts and complex compositional prompts. In this section, we evaluate Concept-TRAK's performance on such scenarios using large-scale text-to-image model (SD1,4v, (v, Rombach et al. (2022)). Following AbC (Wang et al., 2023b), we use 100k subset of the LAION dataset (Schuhmann et al., 2022) for data attribution.
Evaluation Protocol For real-world scenarios, establishing ground-truit truth attribution labels is infea- sible, as we cannot definitively know which training samples influenced specific concepts in generated images. However, we can perform a sanity-check evaluation by verifying whether retrieved training samples actually contain the target concept. a method retrieves training samples that do not contain the concept being attributed, this indicates clear failure of concept attribution.
We employ Owwn3-VL--8 (Yang al., 2025), an open-soures state-of-the-an rision-langual model, to automatically assess whether retrieved images contain the target concept. For each retrieved training sample, we query the model: "Does this image contain (concept value)? ? Please answer with inaccurate attribution. We report
only 'yes' or "no"' If the model responds "no", count Precision @ 10, measuring the fraction of top- 10 retrieved samples that contain the target concept.
For certain generated images, the actual number of training samples that contributed to specific concept may be fewer than 10. Therefore, the upper bound for precision not necessarily 1.0. These metrics should be interpreted as relative performance indicators comparing methods rather than absolute measures of attribution quality.
D.2.1 SIMILAR CONCEPTS
critical challenge in concept attribution distinguishing between semantically similar concepts that share visual features. We evaluate whether methods can correctly attribute training samples for one concept without incorrectly retrieving samples from visually similar but distinet concepts.
Experimental setup We select five semantically similar big cat species: cat, tiger, jaguar, leopard, and cheetah. For local attribution, we generate images using prompts containing only the target concept (c.g.. "a photo of a cat") with different random seeds. For each generated image, we perform concept attribution and measure whether the top-10 retrieved samples contain the target concept rather than similar concepts. For global attribution, generate 256 images per concept and compute set-level attribution by averaging utility gradients across all samples.
Results Tables and present the results for local and global attribution on similar concepts.
Table 7: Local attribution semantically similar concepts (averaged across 8 seeds per concept).
Concept
ConCept-TRAK
DAS
D-TRAK
TRAK
cat
0.925
0.000
0.000
0.000
tiger
0.662
0.000
0.000
0.000
jaguar
0.188
0.000
0.000
0.000
leopard
0.300
0.000
0.000
0.087
cheetah
0.325
0.000
0.000
0.000
Average
0.480
0.000
0.000
0.017
20

---

Published as a conference paper at ICLR 2026
Table 8: Global attribution on semantically similar concepts (averaged over 256 generated images per concept).
Concept
Concept-TRaL TRAK
DAS
D-TRAK
TRAK
cat
1.000
1.000
1.000
0.900
tiger
1.000
0.800
0.800
0.800
jaguar
0.400
0.300
0.200
0.200
leopard
0.600
0.600
0.600
0.300
cheetah
0.500
0.300
0.300
0.300
Average
0.700
0.600
0.580
0,500
Generated Image
Top Influence
whiing
oat)
VO
( tiger)
(jaguar)
( (eopara)
cheetah) )
Figure 9:
Qualitative results for semantically similar concepts using Mancept-TRAK.
For local attribution, baseline methods achieve near-zero precision, indicating complete failure. In contrast, Concept-TRAK achieves 0.480 average precision, successfully retrieving concept- -ppccifccc training samples even for visually similar classes. The performance gap particularly pronounced for common concepts (cat: 0.925, tiger: 0.662) compared to rarer concepts (jaguar: 0.188). As shown in Figure 9, for rare concepts like jaguar, Concept-TRe
sometimes retrieves training samples of visually similar animals such as tigers with similar spotted patterns. This suggests that the model may have learned the rare concept partially through transfer from similar visual features.
For global attribution, baseline methods show improved performance by averaging over many samples, but Concept-TRRE TRAK still achieves the best results. This demonstrates that our method's design choices benefit both local and global attribution settings.
D.2.2 COMPOSITIONAL CONCEPTS
Real-world text-to-image generation often involves compositional prompts that combine multiple concepts such as objects and styles. A key challenge is attributing training influence to each individual concept within compositionally generated image. We evaluate this capability at two difficulty levels.
Common Object and Style We generate images combining common objects with artistic styles using prompts of the form "[object) in the style of (style)". We use two objects (cat, dog) and two styles (graffiti art, stained glass), generating 8 images per object-style combination for total of 32 images. For each image. global attribution baselines, we generate 256 images from each full prompt and perform set-level attribution, then separately evaluate whether retrieved samples contain the object style.
perform separate concept attribution for the object and the style. For
21

---

Published as a conference paper at ICLR 2026
Prompt
Generated Image
Top Influence
( object )
Cat the style of graffiti
(style)
object
Cat in the style of stained glass
( stylee >
( objject
Dog in style graffiti art
( stryle >
( obbjec
Dog in the style stained glass
stryle
Figure 10: Qualitative results for compositional concepts using Concept-TRAK
As Table shows, Concept-TRAK substantially outperforms baselines. For local attribution, baseline methods completely fail, while Concept-TRANn TRAK achieves over 90% accuracy. For global attribution, Concept-TRAi also performs better overall. Interestingly, baseline methods appear biased toward style-based attribution, performing relatively well on global style attribution but completely failing on object attribution. For qualitative results, please refer to Figure 10.
Table 9: Compositional attribution results for natural objects with artistic styles.
Attribution Type
Concept-TRAK
DAS
D-TRAK
TRAK
Local (Object)
0.934
0.025
0.025
0.034
Local (Style)
0.919
0.047
0.047
0.013
Global (Object)
1.000
0.000
0.000
0.125
Global (Style)
0.950
0.925
0.925
0.850
Unique Objects and Artist Styles. To increase difficulty, we use unique objects (Pikachu, Simpson) and famous artist styles (Vincent van Gogh, Pablo Picasso) that require specific training data. We generate 8 seeds per combination for total of 32 images.
Table 10: Compositional attribution results for unique objects with artist styles.
Attribution Type
Concept-TRAK
DAS
D-TRAK
TRAK
Local (Object)
0.581
0.000
0.000
0.000
Local (Style)
0.581
0.000
0.000
0.000
Global (Object)
0.725
0.100
0.125
0.450
Global (Style)
0.775
0.475
0.475
0.075
As shown in Table 10, Concept-TRAK successfully attributes each concept independently compared to the baseline method, even when multiple concepts are intertwined in single image. Same the previous experiment, baseline methods show style bias in global attribution. This suggests that using traditional data attribution methods conflates the overall visual aesthetic with style, while missing object-specific contributions. Concept- TRAK achieves balanced performance on both concepts. For qualitative results, please refer to Figure 11.
22

---

Published as a conference paper at ICLR 2026
Prompt
Generated Image
Top Influence
deject
Pikachu in the style of Pablo Picasso
< (style)
Yost Soo
object
Cat in the style of Van Gogh
( (style) )
dbject
Simpson in the style of Pablo Picasso
< (style) >
object)
Simpson in the style of Van Gogh
( (style)
Figure 11: Qualitative results for compositional concept attribution using Concept-TRAK,
D.2.3 COMPLEX CONCEPTS (QUALITATIVE)
Beyond quantitative evaluation, showcase Concept- TRAK's capability complex, multi-concen prompts commonly used text-to-image model benchmarks. We select three challenging prompts: "An astronaut riding horse Mars", "A teddy bear skateboard in Times Square", and "Avocado
chair"
Prompt
G Generated Image
Top Influence
teddy bear)
teddy bear on skateboard
slentzborrdl
Times Sqaare
(TTmmes Square
astronaut
An astronaut riding on horse Mars
horse)
(Mass) )
(An avocedo
An avocado chair
( (chaúr)
Figure 12: Qualitative results for complex multi-concept prompts using Concept-TRAK.
Figure 12 shows qualitative results for these prompts. For "astronaut riding horse on Mars," Concept-TRAK successfully retrieves astronaut images, horse images, and Mars landscape images separately, demonstrating its ability to decompose spatially composed scenes. For "avocado chair," " the method retrieves not just any chair images, but specifically round-shaped chairs that visually
23

---

Published as a conference paper at ICLR 2026
(a) In-distribution
(b) Out-of-distriuution
Ours
Generated sample
(+)
Generated sample
jerl
Ours IF iFemaiel
(Nala)
DAS
DAS
D-TRAK
D-TRAK
(00, F )
( M,
Figure 13: The target concept for each method indicated in parentheses (Byeglassest-Barc eyes/Femalee-Malee A data attribution method succeeds when the top influential training samples contain the same concept as the generated sample. (a) In-distribution case: Both baseline methods and our approach successfully retrieve relevant training samples. (b) -iitri--dittibttionn: Our method accurately retrieves training samples for each individual concept (eyeglasses and male), while baselines only retrieve samples related to concept due to image-level attribution limitations.
match the avocado-like form in the generated image. This demonstrates our local concept attribution capturing fine-grained visual features. For novel concept combinations like "avocado chair."
These qualitative results demonstrate Concept-TRAK's practical utility for understanding how diffu- sion models compose multiple concepts from training data, providing insights valuable for copyright analysis, model debugging, and interpretability research.
D.3 CONTROLLED EVALCATION: CELEBA-HQ (QUALITATIVE)
In Figure 13, we present qualitative results for concept-level attribution on the CelebA-HQ dataset. This replicates the trends observed in the synthetic dataset: in ID scenarios, images with the same concept as the generated sample can be found through visual similarity alone, but OOD scenarios require isolating individual concepts from compositionally novel outputs, where visual similarity alone fails.
D.4 APPLICATIONS OF CONCEPT-LEVER AITRIBUTION
Our concept-level attribution method provides valuable insights across multiple domains, as shown in Figure 14. For copyright protection, we trace training samples that influenced IP-protected concepts like Mario and Mickey Mouse, addressing provenance concerns. In the realm of safety, our method identifies training samples contributing sensitive concepts, enabling targeted data curation for responsible AI development. For model debugging, Concept-TR. TRAK pinpoints sources of both desirable features and problematic outputs, enhancing our understanding of prompt engineering. Finally, for concept learning, approach reveals how models acquire complex relational concepts like "hug" and "shake hands." " These applications demonstrate how concept-level attribution provides practical tools for addressing key challenges in generative AI development and governance. Note that these experiments use global concept attribution.
E IMPLEMENTATION DETAILS
E. COMPUTATIONAL RESOURCES
All experiments were conducted on NVIDIA H100 GPUs with 80GB memory. To reduce computa- tional costs, all experiments were performed using fp16 precision.
Influence function cunetion-based attribution methods consist of two computational stages: (1) a one-time preprocessing cost of computing training gradients for all training samples, and (2) lightweight per-query step that computes utility gradients for each concept or query. For SD v1.4 (LAION-100)) concrete measurements follows:
24

---

Published as a conference paper at ICLR 2026
(a) Identify JP-protected contents
(Mario)
(usspioos)
(Micesy mouse)
(b) Detect Unsafe contents
(Nullty)
(iinlease) >
(c) Model debugging
realictic hanals, fingera, hyper handls)
high Rit, high reseluttion, high qualityy hyper replistio)
distorted Lands, Pwsputed fungers)
Blurryy bering, frecay, Law dettil, low resouutioon low quality
(d) Understand concept learning
Vuug
(shalie kawwtss )
Applications of concept-level attribution across diverse tasks, (a) Identifying training sources of IP-protected characters. (b) Detecting origins sensitive content for safety governance. of (c) Tracing sources of desirable and problematic features for model debugging. (d) Revealing how models acquire relational concept understanding.
Figure 14:
Training gradient computation The TRAK baseline requires approximately GPU hours, while Concept-TRAM requires 32 GPU hours, due to additional operations including DDIM inversion, and prompt-level guidance.
Per-concept attribution Once gradients are cached, the TRAK baseline requires minute per query, while Concept-TRAN TRAK requires minutes per query.
E.2 SYNTHETIC DATASET
Data Generation We randomly sample shape and color attributes and place them random positions within the image canvas. Image resolution is 64 64. We generate a total of 10,000 synthetic images with this procedure.
Model Training We observe that dropping class conditions for classifier-free guidance severely harms compositional generalization ability, so we avoid this practice. The model trained using the Muon optimizer. While Adam optimizer produces qualitatively similar results, we choose Muon for its stability and significantly faster convergence. We train separate ResNet-based classifiers for each concept and continue training until out-of-distribution (OOD) sample generation accuracy reaches 99%. Training hyperparameters follow Jordan et al. (2024): Muon learning rate of 1e-3 for 1000 epochs, with identical momentum and other hyperparameters. For non mntriiii parameters, Adam optimizer with learning rate 1c-4, We used LightningDiT, state-of of-the-art modernized DiT architecture (Yao al., 2024) for diffusion model.
25

---

Published as a conference paper at ICLR 2026
Test Sample Generation We generate images starting from random noise. Images regenerated until a separate classifier confirms they match the conditioned concept, with a maximum of regeneration attempts per sample.
Baseline methods: For each training sample, we sample 10 different Iy and compute gradients using each method's respective loss function. Ours: For the training loss computation, we apply DDIM inversion with guidance scale 2, which we find beneficial for performance. Since we do not train null tokens, apply CFG by sampling random conditions each step following Sadat et al. (2024). We hypothesize that applying guidance during DDIM inversion removes concept from zo and learns tangent vectors that restore this concept. thus positively affecting concept attribution. More detailed analysis and improvements remain interesting future work. For utility loss computation, we do not use CFG.
Gradient Computation
E.3 CELEBA-HQ DATASET
Data Preparation We use 30,000 images from CelebA-HQ dataset, excluding all samples con- taining the combination [eyeglasses male + smiling). We resize all images to 64 64, to reduce computation.
Model Training We follow the identical training recipe as the synthetic dataset. We train separate ResNet-based classifiers for each concept and continue training until OOD sample generation accuracy reaches 95%.
Test Sample Generation We generate images starting from random noise. Images regenerated until separate classifier confirms they match the conditioned concept, with maximum of regeneration attempts per sample.
Gradient Computation We follow the identical gradient computation recipe as described for the synthetic dataset.
E.4 ABC BENCHMARK
This subsection presents the detailed experimental setup for evaluation of the AbC benchmark.
Benchmark Construction To address more realistic data attribution scenarios. we modify the original AbC benchmark setup. Rather than fine-tuning model parameters on customization data, we freeze the base model parameters and train only special tokens through textual inversion (Gal et al., 2022). Following Wang al. (2023b), we create 20 special tokens corresponding to 20 customization concepts. For each special token, we generate 20 images, resulting in 400 total generated images for data attribution evaluation.
We perform textual inversion using the default hyperparameters provided by the diffusers library: AdamW optimizer with learning rate 5.0 x 10-1, batch size 4, and training epochs 3000.
Baseline Methods Both TRAK, D-TRAK, and DAS need specify regularization hyperparam- eter X. To be more specific, in TRAK (Park et al., 2023a), we approximate the inverted projected GEG G and Gt Vo, L(x;; 0). The regular- Hp Fr AI is invertible in practice. On the other hand,
Hessian as Hp' (F, + A/)-1, where Fp = + E ization À is applied make this regularization makes TRAK-based data attribution effectively ignore components with small cigenvalues, significantly impacting attribution performance (Choe al., 2024).
Previous work recommends X* = 0.1 x mean (eranneeevvll (FP)) (Grosse et al., 2023). For a fair comparison, we perform hyperparameter sweep for TRAK, D-TRAK, DAS and ours across XE [ 10 10-4, 4, X x 10*] and report the best performance achieved.
Concept-TRAK Configuration Our method focuses on measuring contributions to specific con- cepts We use the synthetic image generation prompt as et (e.g., "a tattoo of (V) snake") and remove only the special token to create c (e.g.. : "a tattoo of snake"). We used cm as the base prompt We set the regularization parameter to the recommended value X = 0.1 x mean eigenvalues(F (PP) (Grosse et al., 2023) and use 1024 n, from random trajectories sampled from DDIM sampler for gradient computation.
26

---

Published as a conference paper at ICLR 2026
E.5 QUALITATIVE RESULTS
For the qualitative case studies presented in Appendix D.4, we maintain consistent experimental settings with the AbC benchmark evaluation. The regularization parameter is set À 0.1 mean(eigenvalues|| across all methods.
Below the prompts used for each concept-level attribution task:
IP-protected content:
Mario:
c ""character, Mario"
e_ "character"
Mickey Mouse:
C "character, Mickey Mouse"
c_ = "character"
- Simpsons:
c= "character, Simpsons"
C. "character"
Unsafe content:
Nudity:
"person, nude, naked"
C_ = "person"
Violence:
= "zombie, bloody, violence, gore, horror, grotesque, skull, fear, death, creature"
c_ ""
Model debugging:
Distorted hands:
C c "hands, distorted hands, amputated fingers"
e_ "realistic hands, five fingers, hyper realistic
hands"
High quality images:
"high detail, 8k, intricate, detailed, high resolution, high res, high quality, hyper realistic"
e.
"blurry, boring, fuzzy, low detail, low resolution, low res, low quality"
Concept learning:
Hug:
c "people hug each other"
e_ "people"
Shake hands:
C "people shake their hands"
C_ "people"
Note on Model debugging For the model debugging study, we perform bidirectional attribution by swapping and c. identify both positive and negative influences. This allows training samples that contribute to both problematic and desirable generation.
F ALGORITHM
In this section, we provide detailed algorithms for computing training gradients using the train loss Eq. (3) (Algorithm 1) and utility gradients using utility loss Eq. (4) (Algorithm 2) used in Concept- TRAK. The key computational steps highlighted in red show the guidance terms that distinguish our approach from standard methods.
27

---

Published as a conference paper at ICLR 2026
Algorithm 1 Train loss Luain
Algorithm Utility loss Leonspp
Require: z/. N. {āct==
Require: N, (āt ] {me (mo) for = 1 to N do if local attribution then
1: for to do
2: r DDIMinv
aT), ) aT
2:
3: fo Vou (of v1 04 (2()))
N
ST Noise used to generate jestt
4: else
4:
ôpPs -o, 1ig =1/1B
FT (0,
5:
(o(x)) sg(co(r)) ôprs]
6:
end if
6: Lops + (( (2
Uniform T)
7: In VoLDPS
8:
DDIM (TT - 1)
8: end for
9:
to f 1
Vāl ( vI- āz ca(xz))
9: g + N Em 9n/ 1191111
10: Reward-DPS V,, R(xx)
10: return
11: EG(Iz) f sg(Eo( &Roward-DPS
CReward (1²) co (xx))
13:
f- Vo DPS
14: end for
15: 9 g NV En 9n 3n/119//12
16: return
28


Reviews:
[{"review_id": "9PipBfOuCH", "reviewer": "Reviewer_vod1", "summary": "This paper introduces a novel data attribution method for diffusion models designed to identify which training samples influenced a specific concept within a generated image. The key innovation is the use of specialized, reward-based loss functions for both the training sample and the target concept, which are designed to capture concept-relevant directions in the model's tangent space. Experiments show that Concept-TRAK substantially outperforms prior methods in these concept-level attribution scenarios.", "questions": "Please see Weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- The paper addresses a significant and overlooked limitation in existing data attribution for diffusion models, which is attributing influence for a *specific* concept.\n- The paper conducts controlled experiments and shows that Concept-TRAK performs well in both ID and OOD settings.", "weaknesses": "- If I understand correctly, other methods only allow attribution to a specific generated image. How do you evaluate other methods in the experiments? A naive way is to find influential training data to a set of generated images containing the target concept. Is this how you conduct the experiment? If not, the paper should include that as well.\n- The paper does not specify which model version is evaluated in the paper. A thorough experiment with many model types, such as SD 1.4, SD v2, SD XL, Flux, DeepFloyd, etc. would strengthen the claim", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "E8d7bmEdBE", "reviewer": "Reviewer_KaYm", "summary": "This paper introduces Concept-TRAK, a novel framework for concept-level attribution in diffusion models. The method reformulates influence functions through a reward optimization framework, where a concept-aware loss is defined using gradients along the tangent space of the diffusion manifold. The authors leverage Explicit Score Matching to approximate concept-relevant gradients and compute per-sample attributions without re-training the diffusion model. Empirical studies on synthetic datasets, CelebA-HQ, and the AbC benchmark demonstrate that Concept-TRAK outperforms baselines (TRAK, D-TRAK, DAS) both in in-distribution and out-of-distribution concept attribution, achieving substantially higher performance.", "questions": "1. What is the time and memory cost of running Concept-TRAK on a large model like Stable Diffusion XL or SD 2.1?\n2. How does attribution behave when concepts overlap semantically?\n3. In prompts with multiple intertwined concepts (e.g., “a cat in Van Gogh style”), how is the negative concept constructed, and how does Concept-TRAK avoid misattributing correlated concepts (cat vs. tiger, style vs. object)?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The paper moves beyond image-level attribution toward concept-level interpretability, a direction of genuine novelty and relevance for auditing large-scale generative models.\n2. The integration of reward optimization and influence functions is conceptually elegant.\n3. The method has potential applications in copyright auditing and transparency for generative AI, aligning with broader research interests in responsible model development.", "weaknesses": "1. The approach likely has high computational complexity due to Hessian approximations and multiple gradient evaluations per sample.\n2. The synthetic dataset uses only a 2×2 factor combination; CelebA-HQ omits a single triple attribute; AbC relies on textual inversion with a frozen base model. These experiments demonstrate the idea but fall short of realistic large-scale diffusion scenarios involving diverse prompts and complex concept interactions.\n3. The method is described as “training a diffusion model on a reward-shaped distribution,” but in reality, Concept-TRAK operates post-hoc on pretrained models without updating parameters. This wording could mislead readers regarding optimization vs. analysis.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "12nRfynug4", "reviewer": "Reviewer_Xw2f", "summary": "The paper introduces Concept-TRAK, which attributes training samples that taught a specific concept in diffusion models by aligning a per-sample training gradient with a concept gradient. They first introduce losses that operate through reward gradients in the tangent space and construct an influence function framework based on these losses. Evaluations are performed on controlled synthetic datasets as well as CelebA dataset and AbC benchmark for data attribution evaluations.", "questions": "1. I did not see any experiments that show the computational effectiveness of the approach. How much time does the approach take for a single concept?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper proposes a novel problem where the focus is to pinpoint which training images taught a specific concept rather than just influencing an entire image.", "weaknesses": "1. The evaluations are very narrow and limited with respect to the concepts used. In real-world settings, the concepts can be highly compositional and overlapping, whereas the concepts used in the paper are simple and clean, which doesn’t reflect the real-world setting.\n\n2. The paper lacks a human evaluation for the applications shown in Fig. 8. It is important to analyse if the retrieved samples actually reflect the target concept in these cases. Adding a human judgment study would strengthen the claims.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)