Research map / Image generation / Text-to-image diffusion
Localizing and Editing Knowledge In Text-to-Image Generative Models
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.
Adapts causal mediation analysis to Stable Diffusion to locate where visual-attribute knowledge is stored, finds a single causal state in the CLIP text-encoder, and uses it for Diff-QuickFix, a data-free closed-form concept-editing method that runs in under a second.
Problem
Where knowledge about visual attributes (style, objects, color, action) is stored inside text-to-image diffusion models, and how this localization can be used for fast concept editing.
Approach
Adapts causal tracing to text-to-image diffusion by corrupting subject token embeddings and restoring single layers from a clean model, measuring CLIP-Score; finds knowledge is distributed across UNet layers but localized to the first text-encoder self-attention layer, and edits that layer's W_out with a closed-form least-squares update.
Key results
- The CLIP text-encoder of Stable Diffusion contains only one causal state across attributes: the first self-attention layer at the last subject token.
- Diff-QuickFix performs a closed-form edit in under one second, claimed as 1000x faster than fine-tuning-based concept ablation, with CLIP-Score comparable to Concept-Erase and Concept-Ablation.
- Removing 50 artist styles at once lowers CLIP-Score from 0.29 to 0.21 for the multi-concept ablated model.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use Diff-QuickFix to update stale knowledge | Stable-Diffusion v1.4 · Stale knowledge update (incorrect fact) | CLIP-Score: 0.28 to 0.23 | Helped |
| Use Diff-QuickFix to update stale knowledge | Stable-Diffusion v1.4 · Stale knowledge update (correct fact) | CLIP-Score: 0.22 to 0.29 | Helped |
| Ablate 50 artistic styles with Diff-QuickFix | Stable-Diffusion v1.4 · Removing 50 artistic styles | CLIP-Score: 0.29 to 0.21 | Helped |
| Edit the sole causal text-encoder layer (self-attn-0) rather than a non-causal layer (ablation) | Diff-QuickFix on Stable-Diffusion · Style, object, and fact edits | CLIP-Score with the original attribute | Helped |
| Edit non-causal text-encoder layers instead of the causal layer (ablation) | Diff-QuickFix on Stable-Diffusion · Style, object, and fact edits | CLIP-Score with the original attribute | No clear effect |
| Restore down-1-resnet-1 in the UNet (ablation) | Stable-Diffusion causal tracing · Objects | caption fidelity | Helped |
| Restore mid-block cross-attention in the UNet (ablation) | Stable-Diffusion causal tracing · Action | caption fidelity | Helped |
| Restore down-blocks.1.attentions.0.transformer-blocks.0.ff in the UNet (ablation) | Stable-Diffusion causal tracing · Color | caption fidelity | Helped |
| Restore the first self-attention layer in the UNet (ablation) | Stable-Diffusion causal tracing · Style | caption fidelity | Helped |
| Restore down-blocks.0.attentions.0.transformer-blocks.0.attn2 in the UNet (ablation) | Stable-Diffusion causal tracing · Objects | caption fidelity | No clear effect |
| Restore down-blocks.2.attentions.1.transformer-blocks.0.attn2 in the UNet (ablation) | Stable-Diffusion causal tracing · Action and style | caption fidelity | No clear effect |
| Restore the first self-attention layer for the last subject token in the text-encoder (ablation) | Stable-Diffusion causal tracing · Visual attributes | caption fidelity | Helped |
Surprising findings
- Cross-attention layers in the UNet are generally not causal states; knowledge is instead stored in ResNet and self-attention blocks.
- Unlike GPT-style language models where factual knowledge sits in mid-MLP layers, the text-encoder's causal state is the first self-attention layer.
Limitations
- Focused only on Stable Diffusion; other architectures left for future work.
- Concepts are not removed from the UNet, so prompt engineering or fine-tuning the text-encoder can re-introduce them.
- Edits do not generalize to neighboring concepts (e.g., Eiffel Tower still appears in Paris scenes).
Experimental setup
Datasets: Probe dataset of captions for objects, style, color, action (constructed by authors) Concept-Ablation prompt dataset
Benchmarks: Authors' probe dataset of attribute captions
Metrics: CLIP-Score
Baselines: Concept-Ablation Concept-Erasure TIME Original Stable-Diffusion
Builds on: Causal Tracing (ROME) Concept Ablation Concept Erasure
Closest papers by meaning
- On Mechanistic Knowledge Localization in Text-to-Image Generative Models (ICML 2024 · similarity 0.93)
- Unveiling Concept Attribution in Diffusion Models (NeurIPS 2025 · similarity 0.91)
- Progressive Text-to-Image Diffusion with Soft Latent Direction (AAAI 2024 · similarity 0.90)
- Continuous, Subject-Specific Attribute Control in T2I Models by Identifying Semantic Directions (CVPR 2025 · similarity 0.90)
- Magnet: We Never Know How Text-to-Image Diffusion Models Work, Until We Learn How Vision-Language Models Function (NeurIPS 2024 · similarity 0.90)
- The Hidden Language of Diffusion Models (ICLR 2024 · similarity 0.89)
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.