NeurIPS 2025

Localizing Knowledge in Diffusion Transformers
Arman Zarei1∗, Samyadeep Basu3, Keivan Rezaei1, Zihao Lin2, Sayan Nag3, Soheil Feizi1
1University of Maryland
2University of California, Davis
3Adobe
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-α, 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. 1
1
Introduction
Diffusion and flow models [14, 26, 22, 13, 17, 8, 19] have rapidly become the leading paradigm for
a 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. A growing body of work [11, 32, 16, 37] has focused on localizing different types of
knowledge and capabilities within these models, enabling more targete...