Localizing Knowledge in Diffusion Transformers
Authors: Arman Zarei
Samyadeep Basu
Keivan Rezaei
Zihao Lin
Sayan Nag
Soheil Feizi
Venue: NeurIPS 2025 poster ()
Keywords: Localization
Interpretability
Diffusion Models
Diffusion Transformers
Primary area: deep_learning

Abstract: Understanding how knowledge is distributed across the layers of generative models is crucial for improving interpretability, controllability, and adaptation. While prior work has explored knowledge localization in UNet-based architectures, Diffusion Transformer (DiT)-based models remain underexplored in this context. In this paper, we propose a model- and knowledge-agnostic method to localize where specific types of knowledge are encoded within the DiT blocks. We evaluate our method on state-of-the-art DiT-based models, including PixArt-$\alpha$, FLUX, and SANA, across six diverse knowledge categories. We show that the identified blocks are both interpretable and causally linked to the expression of knowledge in generated outputs.
Building on these insights, we apply our localization framework to two key applications: *model personalization* and *knowledge unlearning*. In both settings, our localized fine-tuning approach enables efficient and targeted updates, reducing computational cost, improving task-specific performance, and better preserving general model behavior with minimal interference to unrelated or surrounding content.
Overall, our findings offer new insights into the internal structure of DiTs and introduce a practical pathway for more interpretable, efficient, and controllable model editing.

Localizing Knowledge in Diffusion Transformers
Arman Zarei', Samyadeep Basu', Keivan Rezaei', Zihao Lin², Sayan Nag', Soheil Feizi! Uniivrriiyy of Maryland
University of California, Davis AAdobe
Abstract
Understanding how knowledge is distributed across the layers of generative models is crucial for improving interpretability, controllability, and adaptation. While prior work has explored knowledge localization in UNet-based architectures, Diffusion Transformer (DiT)-based models remain underexplored in this context. In this paper. we propose a model- and Knowledge-agnostic method localize where specific types of knowledge are encoded within the DiT blocks. We evaluate our method state-of-the-art DiT-based models, including PixArt-a, FLUX, and SANA, across six diverse knowledge categories. We show that the identified blocks are both interpretable and causally linked the expression of knowledge generated outputs. Building on these insights, we apply our localization framework to two key applications: model personalization and knowledge unlearning. both settings, our localized fine-tuning approach enables efficient and targeted updates, reducing computational cost, improving task-specific performance, and better preserving general model behavior with minimal interference unrelated or surrounding content. Overall, our findings offer new insights into the internal structure of DiTs and introduce practical pathway for more interpretable, efficient, and controllable model editing.
Introduction
have rapidly become the leading paradigm for wide range of generative tasks, particularly text-to-image (T2I) synthesis. With access to such powerful pretrained models, it is crucial to explore their potential for applications beyond mere generation. growing body of work [ 32
Diffusion and flow models 14 26 22
13 17
19
has focused on localizing different types of knowledge and capabilities within these models, enabling more targeted usage. For example, [I1] showed that cross-attention layers are key to incorporating prompt compositional information, while 32 20 demonstrated that structural information is often concentrated in the self-attention modules of UNet-based architectures. These insights have been applied to tasks such as image editing and structure-preserviing generation [1] 20).
I6 37
Localizing where specific knowledge resides within models is essential for interpretability, targeted interventions, and understanding model behavior. In generative models, it plays a critical role in applications such as model unlearning and personalization. Several works [66 29 9 36 have shown that generative models often memorize unsafe or unwanted content (e.g., copyrighted or NSFW content) and proposed methods to selectively erase such concepts. model personalization, the goal is to generate novel renditions of subject using only few reference photos across diverse scenes and poses. In both cases, localizing knowledge within the model crucial for enabling targeted interventions that make fine-tuning more efficient and effective, while better preserving the model's prior capabilities and overall generation quality.
"Correspondence to: azarci @umd.cdu
Project page available at: http::///mmmaaaaeeegeeeottoolllu
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Copyright
Safety
Animal
blocks
FLUX
Style
Place
blocks
blacks
PixArt
SANA
Figure I: Localization across various DřT models and knowledge categories. For each model, heatmaps indicate the frequency of each block being selected as dominant carrier of different target knowledge. Green-bordered images are standard generations, while red-bordered images result from withholding wewledge-speciffe information in the localized blocks. Our method successfully localizes diverse knowledge types, with variation in localization patterns across models.
Recent advances text-to-image generation have marked notable evolution, transitioning from UNet-based architectures (27) to Transformer-based models (3], particularly the Diffusion Trans- former (DiT) 23 DiT architectures, with their purely attention- baseid structure, have achieved state-of-the-ar aarte-rttaga generation quality compared to UNet counterparts. While extensive research has ex- plored interpretability and localization in UNet-based architectures, DiT-based models have received comparatively little attention in this regard despite their recent emergence and strong performance.
In this paper, we thoroughly investigate the localization of different types of knowledge within the blocks of diffusion transformers across a range of state-of-the art models, including FLUX (I7), SANA 35 and Pix [5] We introduce a model-agnostic and knowledge-agnostie method that provides strong and reliable signal for identifying the blocks most responsible for generating specific types of knowledge. Our approach demonstrates strong performance and robustness across all evaluated models and diverse set of knowledge categories, such copyrighted content, NSFW material, and artistic styles (Figure D}. While the distribution of localized blocks varies from model to model, our method consistently identifies the key regions responsible for encoding each knowledge.
Building upon our knowledge localization technique, we further propose practical applications of our method for model personalization and unlearning in DiTs (Figure 2 Whether the objective to inject new knowledge or remove undesired content, our approach first localizes the relevant information within the blocks of the DiT and then enables targeted interventions to modify it. Through extensive experiments, we demonstrate that our method outperforms baseline approaches that modify all blocks, achieving superior preservation of generation quality and consistency on unrelated and surrounding prompts. Notably, in the personalization setting, our method achieves improved task- specific performance compared full- full-model fine-tuning, while also minimizing interference with surrounding knowledge. Additionally, our method is more efficient, offering faster training and lower memory usage compared to these baselines.
In summary, our contributions are: (1) We are the first to explore knowledge localization in DiTs by introducing large-scale probing dataset covering diverse categories, and by proposing an automatic, model- and knowledge-agnostic method for identifying where such information resides within the model's blocks. (2) We conduct extensive evaluations across multiple DiT architectures and diverse knowledge types to validate the generality and robustness of our approach. (3) Building on this localization, demonstrate practical applications for efficient model personalization and unlearning.

---

[V] Deg playing
[V] Dog Swimming
Localized Personalization
Finetuning
[V] Dog
Localization
Copyright
NSFW
Locolized Unlearning
"The Botmon inside snow globe"
"A naked woman"
Figure 2: Targeted fine-tuning via knowledge localization. Given concept personalize or remove, our method first identifies the most relevant blocks via knowledge localization and restricts fine-tuning those blocks. This enables efficient adaptation (top) and targeted suppression (bottom) with minimal impact on surrounding content, while better preserving the model's prior performance.
Our method enables targeted fine-tuning that is faster and more memory-efficient, while also achieving superior preservation of generation quality and consistency unrelated surrounding prompts.
2
Related Works
2.1 Diffusion and Flow Models
Diffusion models are a class of generative models based on stochastic differential equations (SDEs), where noise is progressively added to data through a stochastic forward process, ultimately transform- ing the data distribution into standard Gaussian. corresponding reverse process is then learned to reconstruct the original data from noise. Flow matching methods, closely related to diffusion models, instead define deterministic mapping from noise to data using ordinary differential equation (ODE). These methods learn a time-dependent vector field (1 t) that is trained to approximate a target field (x)), which directs samples from the noise distribution toward the data distribution, by minimizing a flow matching loss. See Appendix A.1 for further details.
2.2 Interpretability of Diffusion Transformers
The internal mechanisms of text-to-image diffusion models have been primarily explored in the context of UNet-based models 3 37]. These studies reveal that knowledge of various visual concepts-sude artistic style-is either localized or distributed across a small subset of layers within the UNet architecture. Beyond offering interpretability, these localization insights have been leveraged to address practical challenges, including the removal of copyrighted content, without the need for full model retraining. With the recent shift toward transformer-bast Flux (17) and PixArt-a: (5), understanding how and where concepts has become an emerging area of interest. Recent work has begun to uncover the interpretability of diffusion transformers. For instance, [10] show that attention maps in models like Flux can act as high-quality saliency maps, while I identify a subset of critical layers that are particularly effective for downstream tasks such as image editing. However, it remains unclear how knowledge of visual concepts- surchts-suchhaae as copyrighted objects, artistic styles, or safety-related content iss represented and localized within diffusion transformers. Gaining such insights could enable targeted removal of undesirable content and enhance model personalization for various downstream applications.
models such as encoded in these new models
3 Localizing Knowledge in Diffusion Transformers
3.1 Probe Dataset Description
To systematically evaluate the generalizability and robustness our localization method, we first introduce new dataset called LocK (Localization of Knowledge) designed around six distinct categories of knowledge and concepts: artistic styles (e.g., "style of Van Gogh "), celebrities (e.g.. "Albert Einstein" "), sensitive safety-related content (c.g.. "a naked woman". "a dead body covered
3

---

in blood"), copyrighted characters (e.g.. "the Batman") "), famous landmarks (e.g., "the Eiffel Tower"), and animals (e.g.. "a black panther"). These categories are selected to cover diverse range of visual and semantic information, while also being representative of key use cases in model unlearning (e.g.. removing copyrighted or harmful content) and personalization (e.g.. adding user-specific characters or styles). Compared to prior datasets used in localization and model editing literature, our probing set is significantly larger in both scale and semantic diversity, enabling a more comprehensive evaluation.
For each target knowledge we construct a set of knowledge-specific prompts {pr .'......., .. J. designed capture that knowledge in diverse contexts. For example, for バ: "the Batman" ", prompt p could be "the Batman walking through a desert". To isolate the contribution of the target knowledge e in each prompt p,, we also define knowledge-agnostic prompts (denoted as yf-mcutaa)) for every knowledge-speeiffe prompt p, where p is derived from p by replacing the target knowledge with semantically related but generic placeholder (e.g., "a character walking through a desert"). These paired prompt sets allow us to perform controlled interventions for evaluating knowledge localization and editing behavior in subsequent sections. For more details the dataset construction, statistics, and representative prompt examples, please refer to Appendix B.
3.2 Localization Method
Our goal is to identify which layers within a DiT-based text-to-image model are responsible for encoding specific semantic knowledge. Specifically, given a prompt p (c.g., "Albert Einstein walking in the street"), where x denotes the target knowledge (i.e.. "Albert Einstein" "), we aim to pinpoint which blocks in the model are primarily responsible for representing K. By localizing the internal representation of such knowledge, we can better understand how knowledge distributed across the model's 's architecture and enable targeted interventions such as editing, personalization, or unlearning.
We leverage atlention contribution 7 38 to identify the layers responsible for generating specific content in the image. At a given layer, the attention contribution of a text token to image tokens quantifies how much that token influences the embeddings of the image tokens. We localize the layers where a text token exhibits higher attention contribution, interpreting them as the stages where the corresponding style, object, or pattern is synthesized. We adopt attention contribution as our localization signal because it offers an intuitive and principled way to trace how textual information propagates through the model and influences the generated image. Moreover, it can be universally applied across wide range of DiTs, as it builds on the shared mechanism of attention computation.
More formally, consider layer l of diffusion transformer with L layers equipped with multi- head cross-attentiorf mechanism comprising H heads. For each head h E [H]. let the query, key, value, and output projection matrices be denoted by W g. W W .. and W respectively. Let X1, X2,..... ................................ XT represent the token embeddings of the input prompt, and y1, Y2,-.. ................................ yr denote the token embeddings of the image tokens at layer L. For head h, let the projection of the text token x; onto the key and value matrices be denoted by k// and vý. respectively, and let the projection of the image token y; onto the query matrix be denoted by q Then, the attention contribution of text token x, to image token y:, aggregated over all heads, can be expressed as:
H
Σ attm va Wh
cont, =
where attmf,, is the attention weight between image token y, and text token Xy, computed as:
(f
(af.kk)) k
dn r=1 }
T
attn SOFTMAX
where dn, is the head dimensionality, and the softmax is taken over all text tokens for a fixed image token y,. To compute the overall contribution of a text token X,, we average over all image tokens. Finally, for tokens of interest (X,, Xjz- Xj, } corresponding to the target knowledge x in the prompt pt, compute their attention contribution across all layers E [L]. and identify the layers
2 shared attention mechanism in the of MMDiTs

---

Attention Contrbbuton
LLM Prompt Seneration
Augeegghed
The Batman inside globe"
'The Batman sculpture'
The Batman
Block Index
arin(Q.
Knowledge Attribution Ronking
Q from image
attiiy X Pas Wy
"A choracter inside globe
"A characten
lego
sculpture"
"The Batman inside globe"
"The Batman as lego sculpture"
from
from
The Buttion
Attention Contribution
Prompt Intervention
Figure 3: Overview of our knowledge localization method. We first generate images from prompts {pt containing target knowledge K, and compute token-level attention contributions across layers. Aggregated scores identify the top-K blocks 3% most responsible for encoding k. Replacing their inputs with mhowledge-agnostee prompts {pr (pr-neiiaa)) suppresses the knowledge in the output.
with the highest aggregated contribution as those most responsible for generating the corresponding style, object, or pattern in the image
Figure 3 illustrates the overall pipeline of our knowledge localization method. Given a target first construct set of prompts {pr .p......... } that contain the knowledge, either manually or using an LLM. Using the DiT model, we generate images and compute the attention in each prompt p at each layer (step in Figure These values are averaged across seeds and prompts to obtain a per-layer score indicating how much each block contributes to injecting the knowledge into the image (step 2 in Figure 3}.
knowledge
contribution of the tokens {x1 x Xj. corresponding
We then select the top-K most dominant blocks (B%) as the most informative.
To verify the role of the localized blocks Bi prompts {p" .p...,........ p but replace the inputs to the 3% with knowledge-agnostic prompts py meuuuual)) which omit the knowledge (step Figure 3). In models like PixArt-a, this is done by swapping the cross-attention input, and for MMDiT-based models like FLUX, which use separate prompt branch, we perform two passes, one with {pe} and one with (pr-noural) ), and overwrite the text branch input in the 3% of the first pass with those from the second.
generate images using the original
(ey- ppgucuural
3.3 Experiments and Results
In this section, we present the results of our proposed knowledge localization method, evaluating its effectiveness across multiple model architectures and diverse knowledge categories.
Baselines and Architectures We evaluate knowledge localization method on three state-of- the-art models: PixArt-a, FLUX, and SANA, covering a range of DiT-hased architectural designs. PixArt-a: injects prompt information into the image (latent) space via cross-attention blocks using a pretrained T5 encoder (25), while SANA uses a lightweight LLM-based encoder [31] instead. In contrast, MMDiT-based models such FLUX maintain a separate prompt branch, parallel the image branch, which is updated throughout the model and merges with image representations through shared attention layers. This architectural diversity allows us to assess the generality of our method across different prompt injection mechanisms.
'We found that excluding A from the computation of cont,,y can enhance the localization signal in certain models; however, retained for consistency across experiments.
5

---

PicArt
FLUX
SANA
oppyright
capyrigh
copyright
caldonty
caleniyy
safety
coldanit
caleanty
k 0%
= 10%
K
5
Figure 4: Differences in how knowledge is localized across categories and models. LLVV--baed evaluations and generation samples as the number of intervened blocks increases, where K denotes the top-K most informative blocks identified by our localization method. Some knowledge types (e.g.. copyright) are highly concentrated in a few blocks, while others (e.g., animals) are more distributed across the model. Examples include outputs from the base models and their intervened counterparts.
Evaluation Metries To evaluate the presence of target knowledge in generated images, multiple complementary metrics. First, we use CLIP [24] to measure the semantic alignment between the target knowledge and the generated image. Second, answering capabilities by explicitly querying whether the knowledge appears in the image. Finally, for style-related concepts, where CLIP and LLaVA less reliable, we employ the CSD (Contrastive Style Descriptors) [30] metric, which is more robust for assessing stylistic consistency.
leverage LLaVA's (1] visual question
proposed dataset, LOcK,, spanning all six knowledge categories. The training split is used to perform knowledge localization for each target, and the evaluation split is used to assess the effectiveness of localization via prompt intervention and the metrics described above.
Dataset We
Results Figure presents the results our localization method across different model architectures and knowledge categories. For each model, the heatmap bars show how frequently each block is selected among the top-K most informative blocks (with aggregated across knowledge categories. We also include generation samples with and without prompt intervention to validate the effect of the localized blocks. Our method consistently identifies the blocks most responsible for encoding each knowledge type. Notably, observe that knowledge distributed quite differently across model architectures. In SANA, knowledge tends to be highly concentrated in narrow set of blocks, whereas in PixArt-a. the distribution is more diffuse- diffuse- though certain blocks still emerge as consistently dominant. This architectural disparity in how knowledge is stored underscores the importance of localization methods adaptable across architectures, approach is capabee of reliably identifying the relevant regions where knowledge is encoded.
40% of the model's total blocks),
Figure illlstratts how different knowledge categories are localized across models. In each column, for every target knowledge in our dataset, we first identify the top- K most informative blocks per model using localization method. We then evaluate the effect of prompt intervention on these blocks using the LLaVA-based evaluation metric, which is shown on 0 100% scale. As the plots show, both quantitative and qualitative results reveal that different types of knowledge localize differently. Some knowledge types are concentrated in just a few blocks, while others are more widely distributed. For instance, in FLUX, the drop in the LLaVA score is significantly larger for categories such as copyright, place, or celebrity, compared to the animal category- suggesting that animal-related knowledge is encoded more diffusely throughout the model's blocks.
To further examine how knowledge localization varies within single category, we analyze the differences across individual target knowledge-fo example, comparing "Pablo Picasso" and "Van

---

156
444
109
082
C5D Score
working style of Pablo Pieasse"
over mountain range style of Patrick Caulfield"
280
0.8
078
10
cityycape at dusk in the style of Vincent Van 6ogh"
"a village equare with people chatting in the style of Claude Monet"
K
Figure 5: Variation how artistic styles are localized within the model. We report CSD scores for various artists in the PixArt-a model as the number of intervened blocks K increases. The numbers indicate how many artist styles remain identifiable at each K. While styles like Patrick Caulfield are localized in fewer blocks, others like Van Gogh are distributed more.
Gogh" within the artistic style category. Using localization method, we identify the top-K most informative blocks and evaluate how well the target style is preserved in the generated images under prompt intervention, using the CSD metric and predefined threshold. Figure shows results for the PixArt-a model. The plot reports CSD scores across varying values of K, with annotations indicating how many artists can still be generated in their correct style. Notably, using as few as K 4 blocks, we can localize the stylistic identity for approximately 80% of the artists. However, some styles remain preserved even when their corresponding information is not present in the top-K blocks, indicating that additional blocks are needed for full localization. On the right side of Figure 5S] we show qualitative examples from different artists. As illustrated, styles like Picasso are more localized (typically requiring 6-8 blocks), whereas styles like Van Gogh are more distributed and require larger set of blocks (around 12) for effective representation. We further explore whether there is correlation between the nature of the artistic style (e.g., level of abstraction or detail) and the number of blocks needed for localization. Additional analysis can be found in Appendix [B.3.2
For more qualitative and quantitative results, see Appendix B.3.I Also, to highlight the efficiency and effectiveness of our method, we compare it with a brute-force localization approach in Appendix B.2
Applications
4.1 Model Personalization
Model personalization aims to synthesize high-fidelity images of a subject in novel scenes, poses, colors, and configurations using only a few reference images. We follow the DreamBooth setup (2x. where unique identifier token is assigned to the new subject, and the model is fine-tuned for a few epochs internalize the subject's visual identity and associate it with that token. For details on the DreamBooth setup and task formulation, please refer to Appendix C.I
Unlike conventional DreamBooth, leverage knowledge localization precisely guide which parts of the model to fine-tune. Given a new subject, we first infer its semantic class (e.g.. dog for specific dog instance), then use our method identify the blocks most responsible for encoding knowledge related to that class. Fine-tuning is then restricted to only those blocks. This targeted approach reduces computational cost for training, while also yielding better qualitative and quantitative results. Our method leads to stronger preservation of surrounding concepts and scene consistency, and exhibits superior prompt alignment novel scenarios compared to full-mode fine-tuning (see Section 4.3
4.2
Concept Unlearning
We define concept unlearning as the task of removing a specific target concept from a generative model's knowledge, such that the model can no longer synthesize images corresponding to that concept. Rather than retraining the model from seratch on a dataset with the concept manually excluded, goal is to achieve this effect through minimal and targeted intervention. To this end,
7

---

Input Images
Boseline Genercted Images
Golden
Husky
Torget Knowledge
Surrounding Knowledge
Dreombooth
Localized Dreambooth
Dreambooth
Lecalized Dreambooth
"A [V] dog in origami style"
"A golden netriever dog walking the jungle"
*A dog lego sculpture"
"A husky wolking in the jungle"
Figure 6: Improved prompt alignment and surrounding identity preservation via localized DreamBooth. Left: Localized fine-tuning better adheres to prompt specifications. Right: Surrounding class-level identities are better preserved, demonstrating reduced interference with other concepts.
Customized Knowledge
Surrounding Knowlodge
Promot Aligment Store
Prompt Aligneett Score (1)
Identity Score
Ftb Scone
0310
0,775
Stiee
0.17
FID
0.033
03000
0286
patr
0.390
25
M00 500 1980
1830
400
400
600
$50
12.0
400 Trainin Sheze
400
1000
40.0
810
1200
1350
Trinting Stope
Tricking Shops
Figure 7: Improved quantitative performance with localized DreamBooth. Localized fine-tuning outperforms full-model tuning across all metrics, achieving higher prompt alignment, better identity preservation, and improved FID, while being more efficient in memory usage and training time.
follow the setup proposed Kumari al. [6]. For given target concept (e.g., "The Batman"), assume access to an associated anchor concept (e.g., "a character") broader or semantically related category that serves as a neutral substitute. The objective the target concept with that of the anchor, effectively erasing the specific knowledge while preserving the model's general generative capabilities. For more details on the setup, see Appendix C.22
align the model's output distribution for
As with model personalization, we incorporate concept localization to identify which blocks in the model encode information related the target concept. Rather than updating the entire model, restrict fine-tuning to these localized regions. This targeted unlearning not only improves memory efficiency and speeds up training, but also leads comparable improved results in both qualitative and quantitative evaluations, as we will demonstrate in the following section.
4.3 Experiments and Results
In this section, we present the results of our proposed localized personalization and unlearning methods, evaluating their effectiveness across both qualitative and quantitative metrics.
Setup We base our experiments on the publicly available PixArt-a [5] model. For model personaliza- tion, we follow the setup introduced in DreamBooth (28), and apply localized fine-tuning by updating only 9 9 out of the model's 28 transformer blocks. For concept unlearning, we adopt the setup
proposed by Kumari et al. [16]. and apply our method by fine-tuning only K = - blocks. We focus primarily on the style category, selecting the 30 artists the model best producing- based on CSD scores-and apply unlcarning to each. Further experimental details provided in Appendix C.3

---

Baseline
Full Unlearning
Localized Unleanning
Baseline
Full Unlearning
Lacalized Unlearning
horse grazing in an open field in the style of Pablo Picasso"
mountain trail with hikers in the style Vincent Van Gogh"
Figure 8: Effective and efficient concept unlearning with localized fine-tuning. Our method
removes targeted styles as effectively full-model tuning, while requiring much less computation.
Identity Score (1 (114)
Identity Accuracy (1
FID (4)
0.35
fiaettring
35
Torget
Surnenning
Score
Accurdyy
0.25
Ancher
34
0.20
CLIP
FID
0.15
0.88
0.30
0.6
0.3
100
150
200
0
50
100
200
50
100
160
200
Steps
Steps
Steps
Figure 9: Quantitative results showing better performance with localized unlearning. Localized unlearning achieves comparable target erasure while better preserving surrounding identities, anchor alignment, and overall generation quality (FID), compared to full-model fine-tuning.
Evaluation Metrics For personalized models, we employ multiple metrics to comprehensively assess their performance. Prompt Alignment Score is measured via the CLIP Score between the generated image and the corresponding prompt, capturing how well the image reflects intended scene, style, and semantics described by the prompt. Identity Score, also based on CLIP similarity, measures how well the generated image preserves the subject's 's visual identity, e.g., retaining the distinctive appearance of Husky Golden Retriever when fine-tuming surrounding class-level concepts are selected from LLM-generated prompts that resemble the subject's broader category. Finally, to assess overall image quality, compute the Fréchet Inception Distance (FID) 12 on a 10k sample subset of the COCO dataset
specifie dog. These
[8 throughout the fine-tuning process.
For concept unlearning, we follow the evaluation protocol of Kumari et al. [6]. As in model personalization, we report the Identity Score, which uses CLIP similarity to measure how well the target concept (e.g.. artist's style) is removed from the generated image. We also report Identity Accuracy, a binary metric that checks whether the CLIP similarity to the target (e.g., "Van Gogh style") falls below that of the anchor (c.g., "a painting"). Additionally, compute FID to evaluate the preservation of overall generation quality.
Results As for model personalization, Figure 7 shows that our targeted fine-tuning consistently outperforms full fine-tuning across all metrics. In terms of prompt adherence, qualitative results (Figure 6 left side) show that our method more faithfully reflects user prompts such as "a [V] dog in origami style". ". On the right, we observe that the identities of surrounding concepts (e.g., Husky, Golden Retriever) are better preserved, demonstrating our method's ability to preserve broader scene integrity and maintain surrounding class-level concepts while adapting to a new subject.
As for concept unlcarning, Figure shows shows that our localized unlearming approach better preserves the identity surrounding concepts and maintains alignment with the anchor prompts, while achieving comparable erasure performance on the target concept (see results at 200 steps). In terms of FID, our method demonstrates superior ability to retain the model's prior generation quality compared to full fine-tuning. Moreover, as illustrated in Figure 9 our method effectively removes the targeted styles with performance on par with full-model fine fine-tuning- yet with significantly lower computational cost (15- 50% speedup and approximately 30% reduction in memory usage).

---

5 Conclusion
In this paper, we presented model- and chowledge-agnossie method for localizing where specific knowledge resides within the blocks of Diffusion Transformers. Through extensive experiments across multiple DiT architectures and diverse knowledge categories, we demonstrated the generalizability and robustness of our method. We further introduced a new comprehensive localization dataset designed to support future research in this area. Building on our localization, we applied our method to practical downstream tasks, showing that localized fine-tuning improves task-specific performance while being less disruptive unrelated model behavior and being more efficient. We hope this work serves as foundation for more interpretable, controllable, and efficient adaptation of DiTs.
Acknowledgement
This project supported in part by grant from an NSF CAREER AWARD 1942230, the ONR PECASE grant N00014-25-1-2378. ARO's Early Career Program Award 310902-00001, Army Grant No. W91 INF2 1 12076,, the NSF award CCF2212458 NSF Award No. 2229885 (NSF Institute for Trustworthy AI in Law and Society, TRAILS), a MURI grant 14262683, DARPA AIQ DARPA AIQ grant HR001 12590066 and an award from meta 314593-00001.
References
[1] O. Avrahami, O. Patashnik, O. Fried, E. Nemchinov, K. Aberman, D. Lischinski, and D. Cohen- Or. Stable flow: Vital layers for training-free image editing, 2025. httttss///axxi org/
abb/2411. 14430
[2] S. Basu, K. Rezaei, P. Kattakinda, V. Morariu, N. Zhao, R. A. Rossi, V. Manjunatha, and S. Feizi. On mechanistic knowledge localization in text-to-image generative models. In Forty- first International Conference on Machine Learning, 2024. URL forum?id=faVEax RER
net/
[3] S. Basu, N. Zhao, V. I. Morariu, S. Feizi, and V. Manjunatha. Localizing and editing knowledge in text-to-image generative models. In The Twelfth International Conference on Learning Representations, 2024. URL htttps://openrevier: wtwpeeerreiieeeeee//rrrmmmmm
[4] M. Cao, X. Wang, Z. Qi, Y. Shan, X. Qie, and Y. Zheng. Masactrl: Tuning-free mutual self- attention control for consistent image synthesis and editing. In Proceedings of the IEEF/CVF international conference on computer vision, pages 22560-22570, 2023.
[5] J. Chen, J. Yu, C. Ge, L. Yao, E. Xie, Y. Wu, Z. Wang. J. Kwok, P. Luo, H. Lu, et al. Pixart-o: Fast training of diffusion transformer for photorealistic text-to-image synthesis. arXiv preprint arXiv:2310.mm 00426, 2023.
[6] G. Dar, M. Geva, A. Gupta, and J. Berant. Analyzing transformers in embedding space. arXiv preprint arXiv:2209.02553 2022.
[7] N. Elhage, N. Nanda, C. Olsson, T. Henighan, N. Joseph, B. Mann, A. Askell, Y. Bai, A. Chen, T. Conerly, N. DasSarma. D. Drain, D. Ganguli, Z. Hatfield-Dodds, D. Hernandez, A. Jones, J. Kernion, L. Lovitt, K. Ndousse, D. Amodei, T. Brown, J. Clark, J. Kaplan, S. McCandlish, and C. Olah. A mathematical framework for transformer circuits. Transformer Circuits Thread, 2021. https//transformer-crrcitss pubpsssuuppuul/lu wfta2//////aeaeo///nnooaoa /ffr///wwrr/innxx..iie html.
[8] Esser, S. Kulal, Blattmann, Entezari, J. Müller, H. Saini, Y. Levi, D. Lorenz, A. Sauer, F. Boesel, et al. Scaling rectified flow transformers for high-resolution image synthesis. In Forty-first international conference on machine learning, 2024.
[9] R. Gandikota, H. Orgad, Y. Belinkov, J. Materzyáska, and D. Bau. Unified concept editing in diffusion models. In Proceedings of the IEEÉ/CVF Winter Conference on Applications of Computer Vision, pages 5111-5120, 2024.
[10] A. Helbling, T. H. S. Meral, B. Hoover, P. Yanardag, and D. H. Chau. Conceptattention: Diffusion transformers learn highly interpretable features, 2025. URL abs/ hos//250.nn 04320
10

---

[11] A. Hertz, R. Mokady. J. Tenenbaum, K. Aberman, Y. Pritch, and D. Cohen -r. Prompt-to- prompt image editing with cross attention control. arXiv preprint arXiv:2208.01626. 2022.
[12] M. Heusel, H. Ramsauer, T. Unterthiner, B. Nessler, and S. Hochreiter. Gans trained by a two time-scale update rule converge to local nash equilibrium. Advances in neural information processing systems, 30, 2017.
[13] J. Ho and T. Salimans. Classifier-free diffusion guidance. arXiv preprint 22X. 2022.
[14] J. Ho, A. Jain, and P. Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
[15] A. Hurst, A. Lerer, A. P. Goucher, A. Perelman, A. Ramesh, A. Clark, A. Ostrow, A. Welihinda, A. Hayes, A. Radford, et al. Gpt-4o system card. arXiv preprint arXiv:2410.21276. 2024.
[16] N. Kumari, B. Zhang, S.-Y. Wang, E. Shechtman, R. Zhang, and J.-Y. Zhu. Ablating concepts in text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 22691 2023.
[17] B. F. Labs. Flux.
hwt.:/bbaakk-orstt.o.ooo labsllaaaaafu 2024.
[18] T.-Y. Lin, M. Maire, Belongie, J. Hays, P. Perona, D. Ramanan, Dollár, and C. Zitnick. Microsoft coco: Common objects in context. In Computer vision-ECCV 2014: 13th European conference, zurich, Switzerland, September 6-12, 2014, proceedings, part 13, pages 740-755. Springer, 2014.
[19] Y. Lipman, R. T. Chen, H. Ben-Hamu, M. Nickel, and M. Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
[20] B. Liu, C. Wang. T. Cao, K. Jia, and J. Huang. Towards understanding cross and self-attention in stable diffusion for text -exidg image editing. Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 7817-7826, 2024.
[21] H. Liu, C. Li, Q. Wu, and J. Lee. Visual instruction tuning. Advances in neural information processing systems, 36:34822-34916, 2023.
[22] A. Q. Nichol and Dhariwal. Improved denoising diffusion probabilistic models. In Interna- tional conference on machine learning, pages 8162-8171. PMLR, 2021.
[23] W. Peebles and S. Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 4195-4205 2023.
[24] A. Radford, J. W. Kim, C. Hallacy, A. Ramesh, G. Goh, S. Agarwal, G. Sastry, A. Askell, Mishkin, J. Clark, Learning transferable visual models from natural language supervision. In International conference on machine learning, pages 8748-877 8763. PmL.R, 2021.
[25] C. Raffel, N. Shazeer, A. Roberts, K. Lee, S. Narang, M. Matena, Zhou, W. Li, and J. Liu. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21(140):1-67, 2020.
[26] R. Rombach, A. Blattmann, D. Lorenz, P. Esser, and B. Ommer. High-resolution image synthesis with latent diffusion models. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 10684-10695, 2022.
[27] O. Ronneberger, P. Fischer, and T. Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and computer-assise assisted Intervention-MICC 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part III 18, pages 234-241. Springer, 2015.
[28] N. Ruiz, Y. Li, V. Jampani, Y. Pritch, M. Rubinstein, and K. Aberman. Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEFF/CVF conference computer vision and pattern recognition, pages 22500-225100 2023.
11

---

[29] P. Schramowski, M. Brack, B. Deiseroth, and K. Kersting. inappropriate degeneration in diffusion models. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages
Safe latent diffusion: Mitigating
2023.
[30] G. Somepalli, A. Gupta, K. Gupta, S. Palta, M. Goldblum, J. Geiping, A. Shrivastava, and Goldstein. Measuring style similarity in diffusion models. arXiv preprint arXiv:2400.01222. 2024.
[31] G. Team, M. Riviere, S. Pathak, P. G. Sessa, C. Hardin, S. Bhupatiraju. L. Hussenot, T. Mesnard, B. Shahriari, A. Ramé, et al. Gemma 2: Improving open language models at practical size. arXiv preprini arXiv:2408.00/ 18, 2024.
[32] N. Tumanyan, M. Geyer, Bagon, and T. Dekel. Plug-and-play diffusion features for text- driven image- tongge translation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 1921-1930, 2023.
[33] A. Vaswani, Shazeer, N. Parmar, Uszkoreit, L. Jones, N. Gomez, Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[34] WikiArt. WikiArt: Visual Art Encyclopedia, n.d. URL wwtphttt://// bttbpp::/www... www.wwwwww..tt......www.k... org/
[35] E. Xie, J. Chen, J. Chen, H. Cai, H. Tang, Y. Lin, Z. Zhang, M. Li, L. Zhu, Y. Lu, et al. Sana: Efficient high-resolution image synthesis with linear diffusion transformers. arXiv preprini arXiv:2410. 10629, 2024.
[36] J. Yoon, S. Yu, V. Patil, H. Yao, and M. Bansal. Safree: Training-free and adaptive guard for safe text-to-image and video generation. arXiv preprint arXiv:2410.12761, 2024.
[37] A. Zarei, K. Rezaci, S. Basu, M. Saberi, M. Moayeri, P. Kattakinda, and S. Feizi. Understand- ing and mitigating compositional issues in text-to-image generative models. arXiv preprint arXiv:2406.0 07844, 2024.
[38] A. Zarei, K. Rezaei, S. Basu, M. Saberi, M. Moayeri, P. Kattakinda, and S. Feizi. Improving com- positional attribute binding in text-to-image generative models via enhanced text embeddings, 2025. URL https: // wawiiv.org//////eaa..........lla abs//406...... 07844
12

---

A Related Works
A.1 Diffusion and Flow Models
Diffusion models belong to a class of generative models based on stochastic differential equations (SDE). The central idea is progressively add noise to the original data through stochastic forward process, eventually transforming the data distribution into a simple Gaussian distribution. This forward process is mathematically expressed as
dx f(z, t)dt + g(t)dW,,
where fax t) denotes the drift term, gt) is the diffusion coefficient, and dW, represents the Wiener process (the infinitesimal increment of standard Brownian motion at time f, intuitively understood as an instantaneous Gaussian random perturbation). The reverse process, which aims at reconstructing the original data distribution from noise, is formulated as
dx (f t) 9² (t) log + g(t)dW,
Here, the term V a log (z))), called the score function, describes the gradient of the data distribution at time t. The model is trained to approximate this score function by minimizing the score matching loss, formulated as:
Ee- (a.t [ a log - f)13)
where s0 is parameterized neural network and A(t) is time-dependent weighting function.
Closely related to diffusion models are flow matching methods, designed for training Continuous Normalizing Flows. Flow matching aims to deterministically map an initial noise distribution to target data distribution via an ordinary differential equation (ODE). The trajectory is determined by learned vector field described as:
dr di
Ug (v, (),
where vo (ar, i) is the parameterized vector field to be trained. The training objective involves minimizing the loss function:
(x()))
where (z) represents the target vector field, and (x( denotes intermediate distributions along the path from the initial to the final data distribution. Compared to diffusion models, flow matching meth- ods employ deterministic ODE paths instead of stochastic SDE paths, making them computationally more efficient. Hence, flow matching can be viewed as an efficient alternative to diffusion models.
B Localizing Knowledge in Diffusion Transformers
B.1 Probe Dataset Description
In this section, we describe the construction of our proposed dataset, LocK (Cocalization of Knowledge), which is organized around six distinct categories of knowledge and concepts: artistic styles (e.g., "style of Van Gogh"), "), celebrities (e.g.. "Albert Einstein "), sensitive or safety-related content (e.g., "a naked woman", "a dead body covered in blood"), copyrighted characters (e.g.. "the Batman") "), famous landmarks (e.g., "the Eiffel Tower"), and animals (e.g.. "a black panther") These categories were selected span diverse range of visual and semantic information while reflecting key use cases model unlearning (e.g., removing copyrighted or harmful content) and personalization (e.g., adding user-specific characters or styles).
To construct the target knowledge samples: for the artistic style category, selected 1,108 samples from the WikiArt Artists dataset [34]. For the remaining categories, we used ChatGPT-40 [15) to generate a list of representative examples, initialized through few-shot prompting setup. Prompt augmentation was similarly performed using ChatGPT-4o. For each category. we provided several examples and asked the model to generate diverse, semantically meaningful prompts corresponding to target knowledge instances. Table [aolvedessv statistics for each category, including the number of target knowledge entries, number of augmentation prompts, and total dataset size. Table P also presents examples of prompts across the six categories. Compared to prior datasets used in localization and model editing, LocK is substantially larger in both scale and semantic diversity, facilitating a more comprehensive and rigorous evaluation of knowledge localization methods.
13

---

Table 1: Dataset statistics across six knowledge categories in Lock
Style
Copyright
Safety
Celebrity
Place
Animal
# Target Knowledge
1108
31
50
30
20
40
Train Prompts
20
20
10
20
10
20
# Eval Prompts
30
30
20
25
20
30
Dataset Size
55400
1550
1500
1350
600
2000
B.2 Comparison with Brute-Force Localization
To demonstrate the robustness, efficiency, and effectiveness of our localization method, we compared it against a brute-force baseline. Specifically, we implemented brute- force approach that exhaustively evaluates all possible contiguous block windows of size within the model. For each candidate window, we applied prompt intervention and evaluated the results using the CLIP score and the CSI) score, as described in Section 3.3 In the case of the PixArt model, where the number of blocks is 28, the brute-force method explores 28 possible windows (including circular windows), making it computationally expensive.
We set K = 6 and computed both CLIP and CSD scores for each window. For comparison, also ran our proposed localization method under the same settings (K' 6). The brute-force method, when selecting the optimal window for removing style information, resulted in CLIP score drop of 0.0232 (from 0.2255 to 0.2023) and a CSD score drop of 0.0812 (from 0.84 0.8481 to 0.7669). In contrast, method achieved CLIP score drop of only 0.0194 (from 0.2255 to 0.2061) and CSD score drop of 0.0700 (from 0.8481 to 0.7781).
These results indicate that our localization method performs comparably to the brute-force approach while being approximately 28 times faster on the PixArt model. More generally, for DiT model with B blocks, our method offers a Bx speedup. This highlights both the efficiency and reliability of our approach.
B.3 Experiments and Results
B.3.1 Qualitative and Quantitative Results
In this section, we present additional quantitative and qualitative results from localization experiments. As described in Section 3.3] we evaluate localization performance using multiple metrics, including LLaVA score, CLIP score, and the CSD distance.
Figure 12 provides a comprehensive overview of the CLIP score across varying values of K-thee number of blocks selected by localization method for different models and knowledge categories. For reference, PixArt-a has 28 total blocks, FLUX has 57, and SANA has 10. In each case, we evaluate localization performance blocks. As shown in the figure and discussed in Section 3.3 the localization trends vary significantly across both models and knowledge types. This highlights that different types of knowledge distributed differently within each architecture. For example, in FLUX, using only 2 localized blocks leads to noticeable drop in CLIP score inciiaaiiiiiinnngnnec successful removal of the target knowledge for categories such as copyright. However, this pattern does not consistently appear in other models, underscoring the architectural differences in how knowledge is represented.
K ranges from 0% to approximately 50% of the model's 's total
As discussed in Section 3.3 different artistic styles exhibit varying degrees localization. Figure I0 presents additional qualitative results showing generations with and without prompt intervention across different values of K. Each subfigure includes a colored label, ranging from red to green, representing the average CSD distance for the corresponding artist. A redder label indicates that the style remains strongly present (i.e., less removed), while greener labels indicate more effective style removal. As shown, styles from artists like James Turrell and Parrick Caulfield can be removed with very few blocks, while more detailed or textured styles, such as those of Monet or Van Gogh, require intervention larger number of blocks to achieve comparable removal.
Finally, Figure [3 presents qualitative examples of knowledge localization in the FLUU model across different values of K for various knowledge categories.
14

---

vincent van gogh
pablo picasso
cloude monet
patrick coulfield
Jomes turnell
K=O
07792
08023
0.7440
k: 1:14
0:7528
0.7479
0.7733
0J344
"a Serene countryside"
"a row of housesalong canal"
woman reading by the window"
"a lighthouse overlooloing the
"a farmer working in the field"
Figure 10: Qualitative results of artistic style localization across difffrentt values of K. Each column shows generations from the PixArt-a model for a specific artist, comparing outputs with and without prompt intervention on the top-K localized blocks. The colored labels indicate the average CSD distance using a red-to-green spectrum: more reddish colors signify higher similarity the original style (i.e., style is still preserved), while more greenish colors indicate greater deviation (i.e., the style has been more effectively removed).
B.3.2 Impact of Artistic Style Characteristics on Localization Block Distribution
In this section, we investigate whether there is correlation between the nature of artistic such as its level of abstraction or detail and the number of blocks required for localization. Specifically. for each artist, we determine the minimum number of blocks, denoted by K, that must be intervened upon (via prompt intervention) suppress the presence of that artist's style in the generated image. We quantify the number of blocks needed to remove an artist's style using CSD metric. For each artist, gradually increase the number of intervened blocks K. and compute the CSD distance between the resulting generations and bascline images generated without the style. We define style "removed" when this distance exceeds threshold of 0.82. This threshold represents the point at which the intervention removes the style to a degree comparable to omitting from the prompt entirely. The corresponding value of K is then recorded as the number of blocks that encode the artist's 's style.
The resulting K value for each artist reflects how distributed or localized their stylistic features are across model's layers. We then group artists into m clusters based on these K values (e.g., cluster, {K = 2, 4). cluster, {K = 6, 8}. clusters {K 12, 14)), and explore whether these groupings align with stylistic characteristics such as abstraction, simplicity, texture richness, or level of detail.
15

---

James Turrell
Patrick Caulfield
Abstraction Simplicity
Pablo Picasso
Malevich
6=
Detail Texture Complexity
Vincent van Gogh
Claude Monet
K 12, 14
1
Figure 11: Relationship between artistic style complexity and the number of blocks required for localization. For each artist, we identify the minimum number of blocks K needed to localize their style. Artists with more abstract and minimalist styles tend have lower K values, indicating their styles are encoded in fewer blocks. In contrast, artists with more detailed and textured styles require higher K values, suggesting more distributed representation across the model.
Ideally, this analysis would involve structured dataset containing detailed annotations of each artist's style. However, this is beyond the primary scope of our paper, we adopt lighter-weits alternative: we use GPT-4o to analyze the artist clusters. Given the list of artists each cluster, we prompted GPT-4o to assess whether the groupings aligned with known characteristics of their artistic styles. Interestingly, GPT-4o identified clear pattern: styles characterized by higher abstraction and simplicity tended to correspond to lower K values, whereas styles with greater detail and texture complexity were associated with higher K. We further validate this observation through qualitative examples presented in Figure I1 which visually illustrate the relationship between stylistic complexity and block localization.
B.4 Localization at the Attention Head Level
To further investigate the granularity of our localization approach, we analyze whether knowledge is concentrated within specific attention heads distributed across multiple heads within each attention block. Specifically, we extend our experiments to perform localization at the attention head level within each attention block.
apply our localization framework to identify the top K' important heads (out of 16) within the top K = 6 attention blocks of the Pix model, based on their attention contribution. We then evaluate the effectiveness of these heads using prompt intervention at the head level. The results are summarized below.
16

---

Prempt Interertion
No Dntervertion
Full Ditereettion
style
place
copyright
aninal
celebrity
safety
FixArt-a
27
3.7%
977
0.50
0250
3.20
0200
6250
642
0110
011
copyright
celebrity
2.15
0.00
1.00
0300
0.00
0.00
6.00
1.45
0.20
FLUX
230
0700
1750
0000
1.95
017
417
style
place
copyright
aninal
celebrity
safety
2343
SANA
270
3.50
079
0005
0.25
0225
3.22
0225
6225
3.00
Figure 12: CLIP score vs. number localized blocks (K) across models and knowledge categories. Localization trends vary notably across architectures and knowledge types. For instance, FLUX. just blocks suffice reduce the CLIP score the copyright category, while other models require more blocks-highttt highlighttg differences in how knowledge is encoded.
K'
CLIP Score (L)
LLaVA ())
0.2855
0,9478
4
0.2579
0,8063
8
0,2506
0.7174
12
0.2487
0.6797
16
0.2435
0.6378
Here, K" represents the baseline (no intervention) and K' 16 represents the block-level localization and intervention. The results indicate that knowledge is not localized in individual heads, but rather distributed across broader As result, isolating small subset of heads does not yield substantial benefit, and meaningful intervention often requires modifying many heads.
This supports focus bloek-level localization, which provides more practical and effective abstraction for identifying and intervening on knowledge within diffusion transformers.
B.5 Ablation on Transformer Block Components
To further examine the roles of different modules within transformer blocks, analyze the contribu- tion of components beyond cross attenniotniion models like PixArt, the cross-attention module the primary-and and the only-mechanism through which image tokens receive information from the text tokens, making it natural and impactful target for intervention. Nonetheless, to more thoroughly evaluate this assumption, conducted ablation study examining the individual roles of other modules within the transformer block.
Each block in PixArt consists of self-atte attentieenon normalization, cross-attention, and feedforward (MLP) components. For this analysis, we first localized the top K most important blocks based on our localization metrie. Then, for each module within these blocks, we performed isolated interventions
17

---

by running two forward passes at each denoising step step-one one with knowledge-specing prompt (containing the target information) and one with knowledge-agnostic prompt. We then replaced the output of the module under study in the knoclicge-ppecific forward pass with the corresponding output from the knowledge-agnostic pass.
Intervention Block Type
CLIP Score (.)
LLaVA (1)
Baseline (No Intervention)
0.2854
0.9588
Normalization
0,2830
0.9537
FeedForward
0.2802
0.9437
Self-Attention Cross-Attention
0.2770
0.9301
0.2628
0.8636
As shown in the results, interventions the cross-attention module yield the strongest localization signal, reinforcing our original design choice. These results confirm that cross-attention plays the most significant role in transmitting knowledge-related information between text and image representations.
C
Applications
C.1 Model Personalization
Given only few (typically 3-5) casually captured images of specific subject-withour any accompanying textual descriptions-
goal, following the setup in Ruiz et al. (28), is to synthesize high-fidelity images of that subject in novel scenes and configurations guided solely by text prompts. These prompt-driven variations may involve changes in location, appearance (e.g., color or shape), pose, viewpoint, and other semantic attributes.
The objective is to implant new (identifier, subject) pair into the model's vocabulary in a way that preserves the subject's visual identity while enabling compositional generation. To avoid the overhead of manually writing detailed descriptions for each reference image, we adopt the labeling scheme introduced in Ruiz et al. (28), where each input image is annotated with the phrase "a [identifier] [class noun]". Here, [identifier] is a unique token assigned to the subject, and [class noun] is coarse semantic category (c.g.. dog, cat). The class noun can either be manually specified or inferred using a classifier. This setup allows us to leverage the model's prior for the specified class while learning a new embedding for the subject identifier. During fine-tuning, DreamBooth adjusts the model backbone over few epochs, enabling it to entangle the subject's identity with the learned identifier and synthesize novel views, articulations, and contexts consistent with the reference images.
C.2 Concept Unlearning
We define concept unlearning as the task of removing a specific target concept from a generative model's knowledge, such that the model can no longer synthesize images corresponding to that concept. Unlike retraining from scratch on a dataset with the concept manually excluded-an approach that is both impractical and computationally expensive- we aim to directly modify the model's 's behavior through minimal, targeted intervention. A key challenge in this process ensuring that unlcarning a concept does not degrade the model's performance on semantically related concepts or compromise its general prior capabilities.
To achieve this, we follow the setup proposed by 16 For given target concept (e.g., "The Batman") anchor concept (e.g.. "a character") more general semantically related category that serves as a neutral replacement for the target. The anchor concept should preserve the contextual meaning of the original prompt while abstracting away the target identity. In this setting, the goal is to align the model's ss output distribution for the target concept with that of the anchor, thereby erasing the specific concept while maintaining broader generative capabilities.
assume access
Formally, given a set of target prompts {c*} containing the target concept, and semantically related anchor prompt c, we minimize the KL divergence between the model's 's conditional distributions:
arg min DxL. (p( c) po ( c²))
18

---

where p(ro:T c)) is the reference distribution conditioned on the anchor concept, and pe (oor c*) the model's distribution when prompted with the target concept. Intuitively, we encourage the model to treat prompts containing the target concept c" as if they referred to the anchor concept c.
We apply noise-based concept ablation from 16 by fine-tuning the model on these image-prompt pairs using a standard diffusion loss:
c(x,c,c") = Be ot [lle é (,
where z; is the noisy version of image z at timestep t, and e is the noise to be predicted. As baseline, we fine-tune all model weights, which [16] report to be the most effective among standard unlcarning techniques.
To construct training data, we use the dataset described in Section B.1 and form triplets (x, c, c²). where x is an image generated from prompt c, and c* is derived by replacing the anchor concept in with the target concept. For example, if c = "aa photo of a character ranning", " then ="a photo of the Batman rumning", and x is the image generated from
C.3 Experiments and Results
C.3.1 Setup
Model Personalization We adopt the dataset and experimental setup proposed by [28), and base our experiments on PixArt-a [5). using their publicly available DreamBooth fine-tuning scripts. model with batch size of 1, using the AdamW optimizer with a learning rate of 5 x 10-6 and a weight decay of 3 x 10-2, All input images are resized fixed resolution of 512 x 512, maintaining consistent aspect ratio throughout training. For our localized fine-tuning approach, we update only 9 blocks out of the model's 28 total blocks.
Specifically, we fine-tune the SitArr-XL-2-512--0
Concept Unlearning We adopt the experimental setup proposed by Kumari et al. [16] and, consistent with our model personalization experiments, base our work PixArt-a (5), using their publicly released fine-tuning scripts. Specifically, we fine-tune the PixArt-XL-22-1 XX--51222222-2 model with a batch size of 16, using the AdamW optimizer with a learning rate of 1 x 10 and weight decay of 3 x 10 10-2. To enable memory-efficient training. we clip the gradients to a maximum norm of 0.01. All images resized to fixed resolution of 5 11xx11,, ensuring consistent aspect ratio across training samples. In our localized fine-tuning approach, we restrict updates to only K = 5 blocks out of the model's 28 total transformer blocks. For the style category, we select the top 30 artists whose styles are most easily reproduced by the model, based on CSD scores, and apply unlearning to each. For the copyright category, we use all samples from dataset LocK. All experiments conducted using an RTX A6000 GPU.
C.3.2 Comparison with Random and Low-Importance Block Selection
In order to evaluate the effectiveness of our localization strategy and its impact downstream tasks such as unlearning, we compare our Top- K block selection method against two additional baselines: (1) randomly selected K blocks, and (2) the K least important blocks (referred to Bottom- K). This comparison helps assess both the specificity of our localization metric and its influence on model performance when applied to different objectives.
We first evaluate the localization performance of these selection strategies. The table below reports prompt intervention results for each method.
19

---

Category
Block Selection Policy
CLIP (L)
LLaVA (I)
Baseline (No Intervention)
0.2854
0.9588
Bottom-K K
0.2845
0,9585
Copyright
Random
0.2756
0.9117
Top-K (Ours)
0.2337
0.5463
Baseline (No Intervention)
0.2808
0.7607
Bottom-K
0.2801
0.7602
Celebrity
Random
0,2753
0.6857
Top-K (Ours)
0.2420
0.2743
As shown, the Top-K K selection consistently outperforms both the random and bottom-K baselines by large margin, demonstrating the effectiveness of our localization method in identifying knowledge- bearing blocks.
To further assess the impact of these selection strategies finetuning based unlearning, applicd localized unlearning using the same three block selection methods. Specifically, we fine-tuned only the selected blocks using each of these strategies. The results are shown below.
Block Selection Policy
CLIP Score ())
CLIP Acc. (L)
Baseline (No Finetuning)
0.2735
0.91
Bottom-K K
0.2636
0.82
Random
0.2526
0.73
Top-K (Ours)
0.2293
0.55
These results show that finetuning the Top-K most important blocks leads to significantly better unlearning performance compared random bottom-K selections, reinforcing the relevance of our localization strategy.
D Limitations
Our work introduces framework for localizing knowledge within the blocks of diffusion transformers by ranking blocks based on their relative importance- ffromffrmm the most to the least significant. This ordered localization already highly effective and allows selective focus on the most relevant blocks for analysis, editing, unlearning tasks. However, our approach does not determine the exact value of required for fully removing or representing a given piece of knowledge without relying on prompt intervention and external evaluation metrics such as CSD CLIP scores. A promising future direction is to estimate K automatically using internal model signals- as patterns in our attention contribution metric (e.g., entropy, peak sharpness) other structural the need for external feedback. Additionally, while our evaluations are based on carefully designed prompts and validated metrics, the lack of ground-truth benchmarks for knowledge localization presents another challenge. Developing benchmarks or synthetic datasets with known localization properties could strengthen the validation of future methods.
20

---

K=0%
K=10%
K=30%
K=50%
"the R2-D2 in pixel art form"
"the Spongebob made of clay"
"a naked woman, in vibrant and dreamy style"
person with blood on their hands, minimalist line art style"
"Lionel Messi looking surprised"
"Queen Elizabeth II at an art gallery"
"the Pyramids of Giza as a miniature model"
"the Colosseum watercolor painting"
Figure 13: Qualitative examples of knowledge localization FLUX across different values of K
21

---

Table 2: Example prompts from each knowledge category in LOCK
Category: Style
Anchor = "a painting"
"a bustling harbor with boats in the style of vincent van gogh
"a river flowing through a forest the style of claude monet"
"a cozy cafe by the sidewalk in the style of pablo picasso" "a horse-drawe carriage on dirt road the style of salvador dali"
Category: Copyright
Anchor "a character"
"the deadpool watching something from a distance" "the harry potter on hot summer day" "the r2-d2 in watercolor on textured paper"
"the spongebob walking the jungle"
Category: Safety
Anchor = "a person"
"a topless woman, in minimalist line art style"
"a person exposing genitalia in public, mimicking real-world photograhhhh "a person being beaten with stick, sharp and finely detailed" "a child being punished with belt, with soft natural lighting"
Category: Celebrity
Anchor "a person"
"taylor swift in Times Square at night"
"barack obama in a casual streetwear outfit"
"cristiano ronaldo sitting on the steps of a building"
"queen elizabeth with neon lights reflecting on their face" "
Category: Place
Anchor "a place" "
"the eiffel tower covered in fog"
"the statue of liberty in van gogh starry night style"
"the pyramids of giza lit by fireworks" "
"the acropolis of athens inside snow globe"
Category: Animal
Anchor "an animal"
"a buffalo standing next to a person"
"a penguin looking surprised" "a giraffe in origami style" panther standing on mountain cliff"'
22

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification: The abstract and introduction clearly state the core contributions: proposing a model- and nnowledge-agnostie method localize knowledge in DiT', evaluating across diverse models and knowledge categories, and applying the method to personalization and unlearning. These claims substantiated by extensive empirical results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The paper includes dedicated "Limitations" section (Appendix D
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "I "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on a few datasets or with few runs. In general, empirical results often depend implicit assumptions, which should be articulated.
The authors should reflect the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full of assumptions and complete (and correct) proof?
23

---

Answer: [NA]
Justification: The paper does not present formal theorems or mathematical proofs.
introduces a methodology and provides detailed mathematical formulations (e.g., attention contribution), but no formal theoretical results requiring assumptions or proofs.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions the paper (regardless whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: The paper includes implementation details: model types, dataset construction. hyperparameters, prompt templates, evaluation metrics, and procedures (Sections 3.3 4.3
C.3
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (c.g., to registered users), but it should be possible for other researchers to have some path to reproducing verifying the results.
24

---

5. Open access to data and code
Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [Yes]
Justification: The paper introduces dataset (Lock)) and along with the code will be included in the supplemental material
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co)
setpibimsssooooooccycy... for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
7 //nips. www.c/ppwwppccccccecocoeoooooece gtte://wiissseeeee wwwweeewwisssoonnooooooooooooeee
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
The paper includes details data splits, training setups (e.g.. batch size, optimizer, learning rate, resolution), model configurations, and evaluation procedures, both in the main text and Appendix (Sections
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
The proposed method has been thoroughly evaluated under diverse conditions and model setups, and been compared with baselines such brute-force localization
Justification:
(Sections 3.3
Guidelines:
The answer NA means that the paper does not include experiments.
25

---

The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
show in tables or
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Experimental details are included in Appendix C,3
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental rums as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (c.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer:
[Yes]
Justification:
The research conforms to the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: The paper discusses societal implications (for example, safer and more con- trollable generative models) in Section
26

---

Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied partieular applications, let alone deployments. However, if there is direct path any negative applications, the authors should point it out, For example, it is legitimate to point that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of MI...
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: Instructions will be provided along with the dataset
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: Existing assets are properly cited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
27

---

For scraped data from a particular source (e.g., website). the copyright and terms of service of that souree should be provided.
If assets are released, the license, copyright information, and terms use in the package should be provided. For popular datasets, paperswipphaeeoooccaaappp thcode. com/ com/dataseta has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided. If this information is not available online, the authors are encouraged to reach out to the asset's ' creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: The LoCK dataset is introduced as a new asset. The paper provides thorough documentation in the Appendix B.1 including the construction process, categories, example prompts, and statistics.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details of the dataset/codeemmodll as part their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and sereenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
28

---

Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if' applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [Yes]
B.1 and B.3.2 discuss how LLM was used in the localization method, dataset construction, and some experiments.
Justification: Sections 3.2
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to our LLM policy http::///eerrrpp.cc/oooofereerrr cc///oofeeeennes/2225//.... for what should or should not be described.
29


Reviews:
[{"review_id": "4e92G6uaZI", "reviewer": "Reviewer_vyfy", "summary": "This paper tackles the problem of locating transformer blocks that are responsible for / critical for generating certain knowledge — certain object or certain style. They analyzed cross attention strength at each layer and use prompt ablation for given layer to evaluate their importance for the goal. Then they applied this knowledge to efficient finetuning — DreamBooth personalization, and Unlearning. — where they focused on tuning the most critical blocks this leads to more efficient fine tuning and smaller off target effect.", "questions": "- Conceptually, how do the authors think about how the “knowledge” locates inside each block?\n- Maybe it’s out of scope, but as authors discussed in limitation, have you observed any structural / intrinsic property of attention that correlate with the importance of the block to the final perturbation outcome?\n- Even though fewer blocks were finetuned, the speed up is not that dramatic, is it because it still needs the full backward pass to compute gradient to each target block?\n- Seems the attention contribution method is not limited to DIT models no? UNet based stable diffusion 1 also used cross attention to fetch textural information to guide generation. So the method should be applicable there too?", "limitations": "Authors clearly discussed limitation in supplementary.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "### Strength\n\n- They tackle an open question of interpretability in diffusion models esp. transformers, and they showed the generality of their method in a wide range of recent open source models.\n- The method was used to locate the most relevant blocks, and the knowledge was used to finetuning locally, which is effective, efficient and easy to use.\n- The method is simple to use and principled.\n\n### Weakness\n\n- Even the paper is talking about knowledge, what does authors mean by knowledge is not totally clear. Maybe a definition could be added somewhere. Do you mean knowledge by word to visual content mapping? Or knowledge just mean a certain word / meaning?\n- Maybe a deeper understanding of where knowledge is inside each block was not explored (c.f. to localizing knowledge in language models). So from mechanistic interpretability perspective it’s relatively weak.", "was_revised": true, "final_justification": "The new experiments address my concern about the depth of the paper. With the clearer definition of “knowledge,” the within block fine-grained localization experiments, and the strong correlation between output norm and block importance, the paper’s findings are both deeper and more impactful.", "rebuttal": "We sincerely thank the reviewer for their thoughtful and constructive feedback. Below, we address the reviewer’s concerns in detail:\n\n> Even the paper is talking about knowledge, what does authors mean by knowledge is not totally clear. Maybe a definition could be added somewhere. Do you mean knowledge by word to visual content mapping? Or knowledge just mean a certain word / meaning?\n\nIn our paper, we use the term knowledge to refer to various forms of visual concepts, such as stylistic elements, safety considerations, and popular locations. In particular, these are represented as textual tokens that guide the generation of specific visual content—for example, a particular set of tokens may correspond to the generation of Van Gogh’s artistic style, as noted by the reviewer. We provide detailed descriptions of these categories in Appendix Sec. B.1. Notably, such definitions of visual concepts have been explored in prior work on diffusion model interpretability [1, 2, 3]. In the final version of the paper, we will include a formal definition of *knowledge* and explicitly relate it to these prior works, as highlighted in this rebuttal.\n\n\n> Maybe a deeper understanding of where knowledge is inside each block was not explored (c.f. to localizing knowledge in language models). So from mechanistic interpretability perspective it’s relatively weak. \n\n> Conceptually, how do the authors think about how the “knowledge” locates inside each block?\n\nThe reviewer raises an insightful question regarding the localization of *knowledge* within individual blocks. In response, we extended our framework to explore whether knowledge can be further localized to finer-grained components, such as individual attention heads. Specifically, we applied our localization method to identify the top $K'$ most influential heads (out of 16) within the top $K=6$ attention blocks of the PixArt model. Heads were ranked based on their attention contribution scores, and their impact was assessed through intervening on attention heads.\n\n| $K'$ | CLIP Score $(\\downarrow)$   | LLaVA $(\\downarrow)$  |\n|:----|:--------:|:--------:|\n| $\\mathcal{0}$  | $0.2855$ | $0.9478$ |\n| $\\mathcal{4}$  | $0.2579$ | $0.8063$ |\n| $\\mathcal{8}$  | $0.2506$ | $0.7174$ |\n| $\\mathcal{12}$ | $0.2487$ | $0.6797$ |\n| $\\mathcal{16}$ | $0.2435$ | $0.6378$ |\n\nwhere $K'=0$ represents the baseline (no intervention) and $K'=16$ represents the block-level localization and intervention. The results indicate that **knowledge is not localized in individual heads, but rather distributed across a broader set**. As a result, isolating a small subset of heads does not yield substantial benefit, and meaningful intervention often requires modifying many heads. This supports our focus on block-level localization, which provides a more practical and effective abstraction for identifying and intervening on knowledge in diffusion transformers. \n\nFurthermore, we conducted an ablation study examining the individual roles of other modules within the transformer block. Each block in PixArt consists of self-attention, normalization, cross-attention, and feedforward (MLP) components. For this analysis, we first localized the top $K$ most important blocks. Then, for each module within these blocks, We performed isolated interventions by running two forward passes at each denoising step—one with a knowledge-specific prompt (containing the target information) and one with a knowledge-agnostic prompt. We then replaced the output of the module under study in the knowledge-specific forward pass with the corresponding output from the knowledge-agnostic pass.\n\n| Intervention Block Type    | CLIP Score $(\\downarrow)$ | LLaVA $(\\downarrow)$ |\n|:----------------------------|:------------:|:--------:|\n| Baseline (No Intervention) | $0.2854$     | $0.9588$ |\n| Normalization                       | $0.2830$     | $0.9537$ |\n| FeedForward                | $0.2802$     | $0.9437$ |\n| Self-Attention             | $0.2770$     | $0.9301$ |\n| Cross-Attention            | $0.2628$     | $0.8636$ |\n\nAs shown in the results, interventions on the cross-attention module yield the strongest localization signal, reinforcing our original design choice. We thank the reviewer for this valuable suggestion, which has helped us strengthen the empirical foundation of our work. We will include the full results and experimental details in the final camera-ready version of our paper.\n\n\n> Maybe it’s out of scope, but as authors discussed in limitation, have you observed any structural / intrinsic property of attention that correlate with the importance of the block to the final perturbation outcome?\n\nWe thank the reviewer for the insightful question. While our main focus was not on characterizing the intrinsic properties of attention that correlate with block importance, we agree this is a valuable direction. To further investigate and more directly answer the reviewer’s concern, we conducted a series of additional ablation experiments during the rebuttal phase to explore this connection more systematically.\n\nWe analyzed several structural properties of attention blocks and measured their correlation with the attention contribution metric (used in our paper as a proxy for block importance). Specifically, we considered the following metrics:\n\n- *Attention Map Aggregated Norm*: For each block, we extracted the attention maps corresponding to the text tokens of interest attending to image tokens, across all heads, and computed the average $\\ell_2$ norm of these attention weights.\n- *Attention Entropy*: Using the attention scores (after softmax), we calculated the entropy of each head’s distribution via $-\\sum_j M_{i,j} \\log M_{i,j}$ (computed for each image token $i$), then aggregated the mean and variance of these entropies across heads and image tokens.\n- *Output Norm*: We computed the $\\ell_2$ norm of the attention block’s output for image tokens and aggregated the result across tokens.\n- *Output/Input Norm Ratio*: To assess the extent of transformation in each block, we measured the ratio between the output norm and the input norm for image tokens, giving an estimate of how much the representation is altered.\n- *Attention Contribution* (As discussed in the main paper)\n\nWe performed these analyses over $M$ generations using diverse knowledge-specific prompts, and for a model with $N$ blocks (below results are for the PixArt model with 28 blocks), we obtained $M \\times N$ values per metric. We then computed Pearson and Spearman correlations between each metric and the attention contribution scores across all generations.\n\n| Property of Attention    | Pearson Correlation | Spearman Correlation |\n|--------------------------|:---------------------:|:----------------------:|\n| *Attention Map Agg. Norm*  | $0.1477$              | $0.3444$               |\n| *Mean Entropy*             | $-0.1690$             | $-0.3489$              |\n| *Variance Entropy*         | $0.0162$              | $0.2237$               |\n| *Output Norm*              | $0.9722$              | $0.9245$               |\n| *Output/Input Norm Ratio*  | $0.8365$              | $0.8256$               |\n\nAs shown in the table, some metrics such as the *output norm* exhibit meaningful correlation with block importance. These findings provide preliminary evidence that other certain internal signals in the attention mechanism do reflect the role of the block in knowledge localization.\n\nWe are grateful to the reviewer for this suggestion, which helped us further deepen our analysis. We will include the results and detailed methodology of these experiments in the final camera-ready version of the paper.\n\n> Even though fewer blocks were finetuned, the speed up is not that dramatic, is it because it still needs the full backward pass to compute gradient to each target block?\n\nYes, that’s exactly right. To update the earlier, localized blocks, gradients must still be computed through the later, unlocalized blocks so that the backward signal can propagate all the way back. As a result, while we do observe a noticeable speed-up by restricting updates to fewer blocks, the improvement is not dramatic due to the full backward pass still being required.\n\n> Seems the attention contribution method is not limited to DIT models no? UNet based stable diffusion 1 also used cross attention to fetch textural information to guide generation. So the method should be applicable there too?\n\nWhile in principle the attention contribution method could be applied to UNet-based architectures, in practice it's more challenging due to the architectural differences. In UNet, the spatial resolution and feature dimensions change across layers (e.g., in the encoder, features become deeper and the spatial size shrinks). This variability affects the structure and semantics of the attention, making the contribution signals less interpretable and harder to align across blocks. As a result, we don’t obtain a straightforward or consistent signal as we do in DiT-style models.\n\n-----\nWe hope these clarifications address the reviewer’s concerns, and we thank them again for their valuable insights.\n\n----\n\n[1] Gandikota, Rohit, et al. \"Erasing concepts from diffusion models.\" Proceedings of the IEEE/CVF international conference on computer vision. 2023.\n\n[2] Basu, Samyadeep, et al. \"Localizing and editing knowledge in text-to-image generative models.\" (2024).\n\n[3] Basu, Samyadeep, et al. \"On mechanistic knowledge localization in text-to-image generative models.\" ICML, 2024."}, {"review_id": "E12R63sq5y", "reviewer": "Reviewer_AmLP", "summary": "Prior research has studied the knowledge localization problem in UNet-based diffusion models; however, that understanding might not be transferable to recent transformer-based models. This paper aims to diminish that gap by proposing a method to locate where knowledge is encoded in transformer blocks. They create a probing dataset, LOCK, covering a wide range of concepts and knowledge. Their localization method tries to find the attention contribution of concept text tokens in the prompt to image tokens, and identify layers that yield the highest contribution. They validate the role of important layers by performing causal tracing with the knowledge-agnostic prompts. This paper also proposes to utilize important layers to edit knowledge more effectively. In particular, when performing personalization or unlearning, we can fine-tune those layers only. Experimental results show that this approach improves the performance of knowledge editing methods, likely due to the knowledge isolation in the model.", "questions": "- Can we apply this framework to identify important attention heads? What if we intervene on those heads only, for example, causal tracing or finetuning?\n- What is the performance when we fine-tune random or unimportant layers?", "limitations": "yes", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- This paper explores the knowledge localization problem in transformer-based models, which hasn't been studied before.\n- The causal tracing experiments in Sec. 3 validate the role of important layers.\n- The experiments show that finetuning important layers could be more effective than full finetuning, while requiring less resource.\n\n## Weaknesses\n- The term \"block\" in the paper is not clearly defined. It could be layers, modules, or attention heads, while this paper only studies layers. From the formulation, it seems we can also find important attention heads, which could help localize knowledge better since they are subcomponents of layers. Furthermore, this approach identifies important layers from attention contribution, while a transformer layer contains other modules as well, such as MLP and normalization layers. However, this approach does not take into account the effect of these modules.\n- The motivation for a separate work for transformer-based models is not clear. Previous works [1,2] have already localized important layers in diffusion models; this paper lacks a discussion on the challenge of those methods on transformer-based models. It'd be more convincing if you could provide the difference between the important layers identified by this framework and other methods.\n- This work fine-tunes only important layers and observes an improvement over full finetuning. However, it's not conclusive enough to connect the role of knowledge localization to finetuning methods. The experiments should also include other settings where we fine-tune random layers or the least important layers and show that the performance in those cases are worse than full-finetuning.\n\n\n[1] Basu, Samyadeep, et al. \"Localizing and editing knowledge in text-to-image generative models.\" The Twelfth International Conference on Learning Representations. 2023.\n\n[2] Basu, Samyadeep, et al. \"On mechanistic knowledge localization in text-to-image generative models.\" Forty-first International Conference on Machine Learning. 2024.", "was_revised": true, "final_justification": "The paper proposes a new approach to localizing knowledge by using attention contribution; however, it could be improved by explaining why this method shows a stronger signal than causal intervention in prior works. The practicality of this paper is also questionable, it's unclear that the proposed method indeed offers more efficiency than standard finetuning methods. Also the experiments could be strengthened by comparing with more recent and efficient methods.", "rebuttal": "We genuinely appreciate the reviewer’s insightful and helpful feedback. Below, we provide detailed responses to each of the concerns raised.\n\n> The term \"block\" in the paper is not clearly defined. It could be layers, modules, or attention heads, while this paper only studies layers. From the formulation, it seems we can also find important attention heads, which could help localize knowledge better since they are subcomponents of layers…\n\n> Can we apply this framework to identify important attention heads? What if we intervene on those heads only, for example, causal tracing or finetuning?\n\nWe appreciate the reviewer’s insightful comments regarding the potential benefits of analyzing attention heads individually. To address this concern and further strengthen our work, we extended our experiments to perform localization at the attention head level within each attention block. Specifically, we applied our localization framework to identify the top $K'$ important heads (out of 16) within the top $K=6$ attention blocks of the PixArt model. We used attention contribution to rank heads and evaluated the effectiveness of these heads using prompt intervention on the head level.\n\n| $K'$ | CLIP Score $(\\downarrow)$   | LLaVA $(\\downarrow)$  |\n|:----|:--------:|:--------:|\n| $\\mathcal{0}$  | $0.2855$ | $0.9478$ |\n| $\\mathcal{4}$  | $0.2579$ | $0.8063$ |\n| $\\mathcal{8}$  | $0.2506$ | $0.7174$ |\n| $\\mathcal{12}$ | $0.2487$ | $0.6797$ |\n| $\\mathcal{16}$ | $0.2435$ | $0.6378$ |\n\nwhere $K'=0$ represents the baseline (no intervention) and $K'=16$ represents the block-level localization and intervention. The results indicate that **knowledge is not localized in individual heads, but rather distributed across a broader set**. As a result, isolating a small subset of heads does not yield substantial benefit, and meaningful intervention often requires modifying many heads. This supports our focus on block-level localization, which provides a more practical and effective abstraction for identifying and intervening on knowledge in diffusion transformers.\n\n> … Furthermore, this approach identifies important layers from attention contribution, while a transformer layer contains other modules as well, such as MLP and normalization layers. However, this approach does not take into account the effect of these modules.\n\nWe also appreciate the reviewer’s suggestion to investigate the role of other modules within the transformer block. We fully understand this concern and agree that it is important to evaluate the contribution of components beyond cross-attention.\nWe would first like to emphasize that in models like PixArt, the cross-attention module is the primary and *the only* mechanism through which image tokens receive information from the text tokens. This makes it a natural and impactful target for intervention and localization in these models. \nNonetheless, to more thoroughly address the reviewer’s point, we conducted an ablation study examining the individual roles of other modules within the transformer block. Each block in PixArt consists of self-attention, normalization, cross-attention, and feedforward (MLP) components. For this analysis, we first localized the top $K$ most important blocks. Then, for each module within these blocks, We performed isolated interventions by running two forward passes at each denoising step—one with a knowledge-specific prompt (containing the target information) and one with a knowledge-agnostic prompt. We then replaced the output of the module under study in the knowledge-specific forward pass with the corresponding output from the knowledge-agnostic pass.\n\n| Intervention Block Type    | CLIP Score $(\\downarrow)$ | LLaVA $(\\downarrow)$  |\n|:----------------------------|:------------:|:--------:|\n| Baseline (No Intervention) | $0.2854$     | $0.9588$ |\n| Normalization                       | $0.2830$     | $0.9537$ |\n| FeedForward                | $0.2802$     | $0.9437$ |\n| Self-Attention             | $0.2770$     | $0.9301$ |\n| Cross-Attention            | $0.2628$     | $0.8636$ |\n\nAs shown in the results, interventions on the cross-attention module yield the strongest localization signal, reinforcing our original design choice. We thank the reviewer for this valuable suggestion, which has helped us strengthen the empirical foundation of our work. We will include the full results and experimental details in the final version of the paper.\n\n\n\n> The motivation for a separate work for transformer-based models is not clear. Previous works [1,2] have already localized important layers in diffusion models; this paper lacks a discussion on the challenge of those methods on transformer-based models. It'd be more convincing if you could provide the difference between the important layers identified by this framework and other methods.\n\nThe reviewer raises a valid and important point. Both [1] and [2] employ “causal tracing” to identify localized layers useful for downstream tasks such as unlearning.  While such a technique can theoretically be applied to DiTs -  [2] explicitly notes in Section (M) of the Appendix the challenges of applying these methods to diffusion models with a T5-based text encoder. The authors were unable to modify/edit layers amenable to such analysis in T5-based models which can be useful for downstream applications for unlearning. Since the modern DiT architecture predominantly relies on T5, we chose to explore alternative approaches as in our paper that scale better to this setting. In fact, our initial attempts to apply causal tracing to DiT did not yield promising results and we weren’t able to identify localized layers corresponding to different visual concepts. This observation further motivated our proposed method, which—unlike previous approaches that rely on computationally expensive brute-force searches for localization—leverages efficient attention-based attribution signals. Our method scales effectively to large architectures and avoids the need for exhaustive layer-wise searches. We will add this note in the final camera-ready version of the paper to provide a more comprehensive picture. \n\n\n> This work fine-tunes only important layers and observes an improvement over full finetuning. However, it's not conclusive enough to connect the role of knowledge localization to finetuning methods. The experiments should also include other settings where we fine-tune random layers or the least important layers and show that the performance in those cases are worse than full-finetuning.\n\n> What is the performance when we fine-tune random or unimportant layers?\n\nWe greatly appreciate the reviewer’s suggestion to incorporate additional baselines such as random block selection and selection of the least important blocks. This helps evaluate both the specificity of our localization method and its impact on downstream applications like unlearning.\n\nTo begin, we conducted experiments to assess the localization performance of these alternative block selection strategies. While our main method selects the top-$K$ most important blocks based on our localization metric, we compared it against two baselines: (1) randomly selected $K$ blocks, and (2) the $K$ least important blocks (referred to as Bottom-K in the tables below). The table below reports prompt intervention results for each of these selection strategies.\n\n| Category  | Block Selection Policy     | CLIP $(\\downarrow)$ | LLaVA $(\\downarrow)$  |\n|-----------|----------------------------|:------------:|:--------:|\n| $\\mathcal{Copyright}$ | Baseline (No Intervention) | $0.2854$     | $0.9588$ |\n|           | Bottom-K                   | $0.2845$     | $0.9585$ |\n|           | Random                     | $0.2756$     | $0.9117$ |\n|           | **Top-K (Ours)**               | $\\mathbf{0.2337}$     | $\\mathbf{0.5463}$ |\n| $\\mathcal{Celebrity}$ | Baseline (No Intervention) | $0.2808$     | $0.7607$ |\n|           | Bottom-K                   | $0.2801$     | $0.7602$ |\n|           | Random                     | $0.2753$     | $0.6857$ |\n|           | **Top-K (Ours)**               | $\\mathbf{0.2420}$     | $\\mathbf{0.2743}$ |\n\nAs shown, the Top-$K$ selection consistently outperforms both the random and bottom-$K$ baselines by a large margin, demonstrating the effectiveness of our localization method in identifying knowledge-bearing blocks. \n\nFurther, to evaluate the effect of these selection strategies on finetuning-based unlearning, we applied localized unlearning using the same three block selection methods. Specifically, we fine-tuned only the selected blocks using each of the three selection strategies. The results are shown below.\n\n| Block Selection Policy   | CLIP Score $(\\downarrow)$ | CLIP Acc. $(\\downarrow)$ |\n|--------------------------|:------------:|:-----------:|\n| Baseline (No Finetuning) | $0.2735$     | $0.91$      |\n| Bottom-K                 | $0.2636$     | $0.82$      |\n| Random                   | $0.2526$     | $0.73$      |\n| **Top-K (Ours)**                    | $\\mathbf{0.2293}$     | $\\mathbf{0.55}$      |\n\nThese results show that finetuning the Top-$K$ most important blocks leads to significantly better unlearning performance compared to random or bottom-$K$ selections, reinforcing the relevance of our localization strategy. We thank the reviewer for this valuable suggestion, which has helped us strengthen the empirical rigor of our work. We will include the full set of results and implementation details in the final camera-ready version of the paper.\n\n\n-----\nWe hope these clarifications address the reviewer’s concerns, and we thank them again for their valuable insights."}, {"review_id": "6tismnUMTn", "reviewer": "Reviewer_gcMX", "summary": "The paper proposes a model-agnostic and knowledge-agnostic method to localize where specific types of knowledge are encoded within Diffusion Transformer blocks, enabling efficient and targeted model editing.", "questions": "1.\tDoes the final performance in the experiment using the llava model get affected by the preferences of that model?\n2.\tThe paper refers to knowledge-agnostic prompts in Section 3.1. Could you clarify their specific formulation and provide the template used in implementation?\n3.\tFigure 4 shows that FLUX experiences the steepest performance decline across multiple metrics as K increases. What specific characteristics of FLUX might explain its heightened sensitivity to larger K values?", "limitations": "Yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \n\n1.\tThe proposed method in this paper enables targeted fine-tuning that is both faster and more memory-efficient, while maintaining high generation quality and consistency.\n2.\tThe authors conduct comprehensive experiments to validate the effectiveness of the proposed method.\n\nWeaknesses: \n\n1.\tFigures 1 and 4 demonstrate the differences in the knowledge distribution patterns among various models. What causes these differences, and are they due to variations in model architectures? \n2.\tWhile the paper adopts K=9 for model personalization and K=5 for concept unlearning, providing more details on the selection process would enhance the understanding of the method. \n3.\tWhy does an increase in K lead to a decrease in the CLIP score, and what is the reason behind this phenomenon?", "was_revised": true, "final_justification": "The authors addressed most of the concerns.", "rebuttal": "We sincerely thank the reviewer for their constructive feedback. We are pleased that the reviewer found our experimental results comprehensive. Below, we address the reviewer’s concerns in detail:\n\n> Figures 1 and 4 demonstrate the differences in the knowledge distribution patterns among various models. What causes these differences, and are they due to variations in model architectures?\n\nThe observed differences in knowledge distribution patterns across models can be attributed to several factors, including variations in model architecture, differences in training procedures, the diversity and composition of training datasets, and the inherent complexity of the high-dimensional optimization landscape. These elements jointly influence how and where knowledge is encoded within each model, leading to the distinct patterns observed. Prior work [1] also highlights this phenomenon, showing that knowledge is distributed differently across architectures\n\n> While the paper adopts K=9 for model personalization and K=5 for concept unlearning, providing more details on the selection process would enhance the understanding of the method.\n\nWe conducted a small-scale hyperparameter sweep over values of $K \\\\in \\\\{3, 5, 7, 9\\\\}$ and empirically found that $K = 5$ yields the best results for concept unlearning, while $K=9$ performs best for model personalization. These values provided a good balance between edit effectiveness and minimal side effects. We will clarify this selection process in the final version of the paper.\n\n> Why does an increase in K lead to a decrease in the CLIP score, and what is the reason behind this phenomenon?\n\nThe CLIP score used in our localization setup measures how strongly the generated images reflect the presence of the target knowledge. As $K$ increases, a greater number of blocks identified as encoding this knowledge are intervened. Specifically, instead of providing these blocks with a knowledge-specific prompt (one containing the target knowledge), we supply a knowledge-agnostic prompt, so that the target information is no longer fetched by these blocks during generation. As a result, the model becomes increasingly unable to express the target knowledge in its outputs, leading to a corresponding drop in the CLIP score.\n\n\n> Does the final performance in the experiment using the llava model get affected by the preferences of that model?\n\nWe understand the reviewer's concern about potential biases in using LLaVA as an evaluator. While LLaVA is a commonly used VLM in the community for such evaluations [2, 3], we complemented it with additional evaluation metrics, such as CLIP and CSD score, to improve the robustness of our results. To further address the reviewer's concern directly, **we extended our evaluation by incorporating two additional VLMs**—**Qwen and Instruct-BLIP**—to reduce model-specific biases. Below, we include a partial view of these results for the *copyright*, *celebrity*, and *place* categories using the FLUX model across different values of $K$. The complete results will be included in the final version of the paper.\n\n\n| Category   | $K$    | CLIP  | LLaVA | Qwen  | BLIP  |\n|:------------|:------|:-------:|:-------:|:-------:|:-------:|\n| $\\mathcal{Copyright}$  | $\\mathtt{0\\\\%}$  | $0.2947$ | $0.9530$ | $0.7105$ | $0.9097$ |\n|            | $\\mathtt{20\\\\%}$   | $0.2240$ | $0.4703$ | $0.1598$ | $0.3068$ |\n|            | $\\mathtt{40\\\\%}$   | $0.1934$ | $0.1827$ | $0.0009$ | $0.0490$ |\n| $\\mathcal{Celebrity}$  | $\\mathtt{0\\\\%}$  | $0.2778$ | $0.5225$ | $0.3435$ | $0.7240$ |\n|            | $\\mathtt{20\\\\%}$   | $0.2223$ | $0.1635$ | $0.0575$ | $0.1894$ |\n|            | $\\mathtt{40\\\\%}$   | $0.1956$ | $0.0464$ | $0.0017$ | $0.0180$ |\n| $\\mathcal{Place}$      | $\\mathtt{0\\\\%}$  | $0.3046$ | $0.9850$ | $0.8778$ | $0.9742$ |\n|            | $\\mathtt{20\\\\%}$   | $0.2341$ | $0.4506$ | $0.2833$ | $0.4225$ |\n|            | $\\mathtt{40\\\\%}$   | $0.1776$ | $0.0322$ | $0.0100$ | $0.0172$ |\n\n\n\n> The paper refers to knowledge-agnostic prompts in Section 3.1. Could you clarify their specific formulation and provide the template used in implementation?\n\nAs discussed in Section 3.1, knowledge-agnostic prompts are constructed by taking a knowledge-specific prompt and replacing the target knowledge with a semantically related but generic placeholder. For example, for the knowledge “the Batman”, the specific prompt “the Batman walking through a desert” becomes “a character walking through a desert” in the agnostic version. Table 2 in the appendix includes several examples of knowledge-specific prompts along with their anchors. The knowledge-agnostic prompts are created by replacing the knowledge with the corresponding anchor in those examples.\n\n> Figure 4 shows that FLUX experiences the steepest performance decline across multiple metrics as K increases. What specific characteristics of FLUX might explain its heightened sensitivity to larger K values?\n\nAs noted earlier, differences in model architecture, training data, and optimization dynamics lead to varying localization patterns across models. In the case of FLUX, our intuition for its steeper drop in localization metrics as $K$ increases lies in its architecture. Specifically, FLUX uses the MMDiT architecture, which concatenates text and image tokens and processes them jointly. Rather than injecting textual information only via cross‑attention (as seen in models like PixArt), FLUX’s design allows bidirectional, layer‑wise interaction between text and visual streams.This architectural design means that intervening on some blocks can change the flow of information more broadly, potentially affecting downstream layers and leading to a sharper drop as more blocks are intervened.\n\n----\n\nWe hope these additional clarifications address the reviewer’s concerns, and we thank them again for their valuable insights.\n\n----\n\n[1] Basu, Samyadeep, et al. \"On mechanistic knowledge localization in text-to-image generative models.\" ICML, 2024.\n\n[2] Han, Evans Xu, et al. \"Progressive compositionality in text-to-image generative models.\" arXiv preprint arXiv:2410.16719 (2024).\n\n[3] Sun, Kaiyue, et al. \"T2v-compbench: A comprehensive benchmark for compositional text-to-video generation.\" Proceedings of the Computer Vision and Pattern Recognition Conference. 2025."}]

Author rebuttal:

Meta-review: This paper introduces a method to localize knowledge within the blocks of Diffusion Transformer (DiT) architectures and applies it to tasks such as personalization and unlearning. The reviewers generally agree on the technical soundness and relevance of the proposed framework. Reviewer vyfy supports acceptance, citing the generality and utility of the method, and notes that concerns about clarity and depth were addressed through additional experiments and discussion. Reviewer gcMX is moderately supportive (borderline accept), appreciating the thorough experiments but requesting better justification for some empirical choices and clearer explanation for architectural differences across models. Reviewer AmLP, however, remains more skeptical (borderline reject), acknowledging that many initial concerns were addressed, but still questioning the theoretical grounding of attention-based localization versus causal tracing, the practicality of the proposed efficiency improvements, and the limited use of only older baseline methods. Although AmLP raises late-stage concerns regarding baseline choice, the authors correctly note these were not raised during the initial review and should not weigh heavily at this stage. Given the novelty of studying knowledge localization specifically in DiT architectures, the empirical effectiveness of the proposed approach, and the extensive rebuttal efforts that clarified methodology and addressed concerns, the AC recommends acceptance and wishes the author can adopt the suggestions from reviewers in the final version.
Decision: Accept (poster)