Research map / Image generation / Autoregressive and latent generation
Localizing Knowledge in Diffusion Transformers
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Localizes where specific knowledge (styles, celebrities, safety content, copyrighted characters, landmarks, animals) is encoded in Diffusion Transformer blocks using attention contribution, and uses the localized blocks for cheaper DreamBooth personalization and concept unlearning.
Problem
Knowledge localization has been studied for UNet diffusion models but remains underexplored for Diffusion Transformers (DiTs), which limits targeted fine-tuning for personalization and unlearning.
Approach
A model- and knowledge-agnostic localization that ranks DiT blocks by the aggregated attention contribution of knowledge-specific text tokens to image tokens, validated by swapping in knowledge-neutral prompts at the top-K blocks; the top-K blocks then restrict fine-tuning for personalization and unlearning.
Key results
- In the personalization setting, localized DreamBooth (K=9 of 28 blocks on PixArt-α) outperforms full-model fine-tuning on prompt alignment, identity preservation, and FID while using less memory and training time.
- For concept unlearning on PixArt-α with K=5 blocks, localized fine-tuning achieves erasure comparable to full fine-tuning with about 15-20% speedup and roughly 30% lower memory usage.
- On PixArt-α with K=6, the attention-contribution method's style removal (CLIP drop 0.0194, CSD drop 0.0700) is comparable to brute-force window search (CLIP drop 0.0232, CSD drop 0.0812), at roughly 28x lower search cost.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fine-tune only localized blocks instead of the full model | DreamBooth personalization on PixArt-α · Model personalization | prompt alignment | Helped |
| Fine-tune only localized blocks instead of the full model | DreamBooth personalization on PixArt-α · Model personalization | surrounding identity preservation | Helped |
| Fine-tune only localized blocks instead of the full model | DreamBooth personalization on PixArt-α · Model personalization | FID | Helped |
| Fine-tune only localized blocks instead of the full model | Concept unlearning on PixArt-α · Concept unlearning | target concept erasure | No clear effect |
| Fine-tune only localized blocks instead of the full model | Concept unlearning on PixArt-α · Concept unlearning | surrounding identity preservation | Helped |
| Fine-tune only localized blocks instead of the full model | Concept unlearning on PixArt-α · Concept unlearning | anchor alignment | Helped |
| Fine-tune only localized blocks instead of the full model | Concept unlearning on PixArt-α · Concept unlearning | FID | Helped |
| Use attention-contribution localization instead of exhaustive block-window search | PixArt-α style localization · Style localization | CLIP score after intervention: 0.2255 to 0.2061; brute force: 0.2255 to 0.2023 | No clear effect |
| Use attention-contribution localization instead of exhaustive block-window search | PixArt-α style localization · Style localization | CSD score after intervention: 0.8481 to 0.7781; brute force: 0.8481 to 0.7669 | No clear effect |
Surprising findings
- Knowledge distribution differs strongly across architectures: SANA concentrates knowledge in few blocks while PixArt-α is more diffuse.
- Animal knowledge in FLUX is encoded more diffusely than copyright, place, or celebrity knowledge.
Limitations
- The method ranks blocks but does not determine the exact K needed for full removal; evaluation relies on external metrics such as CSD and CLIP.
- No ground-truth benchmark exists for knowledge localization.
- Unlearning and personalization experiments are primarily on PixArt-α, with the style category focus.
Experimental setup
Datasets: LocKK WikiArt Artists COCO
Benchmarks: LocKK COCO
Metrics: CLIP score LLaVA VQA score CSD score FID identity accuracy
Baselines: full-model fine-tuning DreamBooth brute-force localization
Builds on: DreamBooth noise-based concept ablation attention contribution
Closest papers by meaning
- Unveiling Concept Attribution in Diffusion Models (NeurIPS 2025 · similarity 0.89)
- On Mechanistic Knowledge Localization in Text-to-Image Generative Models (ICML 2024 · similarity 0.89)
- Localizing and Editing Knowledge In Text-to-Image Generative Models (ICLR 2024 · similarity 0.87)
- Blink of an eye: a simple theory for feature localization in generative models (ICML 2025 · similarity 0.87)
- Precise Parameter Localization for Textual Generation in Diffusion Models (ICLR 2025 · similarity 0.87)
- Knowledge Circuits in Pretrained Transformers (NeurIPS 2024 · similarity 0.87)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.