Localizing Knowledge in Diffusion Transformers

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:

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?

The 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

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.

We 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.

Why does an increase in K lead to a decrease in the CLIP score, and what is the reason behind this phenomenon?

The 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.

Does the final performance in the experiment using the llava model get affected by the preferences of that model?

We 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.

The paper refers to knowledge-agnostic prompts in Section 3.1. Could you clarify their specific formulation and provide the template used in implementation?

As 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.

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?

As 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.

We hope these additional clarifications address the reviewer’s concerns, and we thank them again for their valuable insights.

[1] Basu, Samyadeep, et al. "On mechanistic knowledge localization in text-to-image generative models." ICML, 2024.

[2] Han, Evans Xu, et al. "Progressive compositionality in text-to-image generative models." arXiv preprint arXiv:2410.16719 (2024).

[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.

We genuinely appreciate the reviewer’s insightful and helpful feedback. Below, we provide detailed responses to each of the concerns raised.

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…

Can we apply this framework to identify important attention heads? What if we intervene on those heads only, for example, causal tracing or finetuning?

We 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.

where $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.

… 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.

We 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. We 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. Nonetheless, 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.

As 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.

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.

The 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.

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.

What is the performance when we fine-tune random or unimportant layers?

We 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.

To 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.

As 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.

Further, 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.

These 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.

We hope these clarifications address the reviewer’s concerns, and we thank them again for their valuable insights.

Thank you for your detailed comment. I believe that the analysis of attention heads would complement the paper.

Regarding the motivation, it's still unclear to me how the attention contribution could localize knowledge in diffusion models that use T5 encoder while [2] couldn't. Theoretically, [2] performs intervention on the model and therefore should pinpoint the causal effect and locate important layers. From the current state, it seems that the paper only empirically observes that attention contribution is a stronger signal; a deeper explanation in Section 3.2 would strengthen the paper. However, I agree with the author that the brute-force approach in [2] could be a challenge, and the proposed method helps mitigate that.

Regarding the experimental results of random-, full-, and top-k finetuning, it'd be better if the author could use a consistent setup on diverse knowledge, i.e., different animals, styles, and objects.

I am also concerned about the practicality of this work. It seems that full-finetuning is only slower by less than 3 minutes, which is not too significant. Also, we need to take into account the computation to localize knowledge and the hyperparameter search process to find $K$. The results in the paper also give minimal practical insights; I am not sure which points in Figure 7 and 9 I should take when personalizing or unlearning diffusion. Additionally, the experiments only use one personalization method (DreamBooth) and unlearning method (Concept Ablating), which could be outdated compared to more recent and efficient methods such as InstantBooth or RECE.

For those reasons, I decided to keep the original rating.

Shi, Jing, et al. "Instantbooth: Personalized text-to-image generation without test-time finetuning." CVPR 2024

Gong, Chao, et al. "Reliable and efficient concept erasure of text-to-image diffusion models." ECCV 2024

We sincerely thank the reviewer for their thoughtful and constructive feedback. Below, we address the reviewer’s concerns in detail:

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?

In 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.

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.

Conceptually, how do the authors think about how the “knowledge” locates inside each block?

The 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.

where $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.

Furthermore, 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.

As 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.

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?

We 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.

We 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:

• 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.
• 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.
• Output Norm: We computed the $\ell_2$ norm of the attention block’s output for image tokens and aggregated the result across tokens.
• 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.
• Attention Contribution (As discussed in the main paper)

We 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.

As 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.

We 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.

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?

Yes, 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.

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?

While 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.

We hope these clarifications address the reviewer’s concerns, and we thank them again for their valuable insights.

[1] Gandikota, Rohit, et al. "Erasing concepts from diffusion models." Proceedings of the IEEE/CVF international conference on computer vision. 2023.

[2] Basu, Samyadeep, et al. "Localizing and editing knowledge in text-to-image generative models." (2024).

[3] Basu, Samyadeep, et al. "On mechanistic knowledge localization in text-to-image generative models." ICML, 2024.