Title: Localizing Knowledge in Diffusion Transformers
Venue: NeurIPS 2025
Accepted: yes
Keywords: Localization

Interpretability

Diffusion Models

Diffusion Transformers
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.


Peer reviews:
comments: 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.

STRENGTHS AND WEAKNESSES:
Strengths: 

1.	The proposed method in this paper enables targeted fine-tuning that is both faster and more memory-efficient, while maintaining high generation quality and consistency.
2.	The authors conduct comprehensive experiments to validate the effectiveness of the proposed method.

Weaknesses: 

1.	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? 
2.	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. 
3.	Why does an increase in K lead to a decrease in the CLIP score, and what is the reason behind this phenomenon?

QUESTIONS:
1.	Does the final performance in the experiment using the llava model get affected by the preferences of that model?
2.	The paper refers to knowledge-agnostic prompts in Section 3.1. Could you clarify their specific formulation and provide the template used in implementation?
3.	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?

LIMITATIONS:
Yes
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-01

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nPrior 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\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- What is the performance when we fine-tune random or unimportant layers?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-03"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n### 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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nAuthors clearly discussed limitation in supplementary.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-03"}