🤖 machine learning

Concept Unlearning via Cross-Attention Activation Projection for Diffusion Models

The paper proposes PURE, a closed-form concept unlearning method that projects cross-attention key and value weights based on activation space representations to effectively erase target concepts from diffusion models under paraphrased and adversarial prompts while preserving the generation of retained concepts.

Original authors: Saemi Moon, Suhyeon Jun, Seoyeon Lee, Dongwoo Kim

Published 2026-05-26

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Saemi Moon, Suhyeon Jun, Seoyeon Lee, Dongwoo Kim

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a super-talented artist (the AI) who has learned to draw everything in the world. But, for safety or legal reasons, you need to teach this artist to forget how to draw specific things, like a famous cartoon character (Pikachu) or a specific painting style (Van Gogh).

The tricky part is: you don't want to make the artist forget everything else. You still want them to be able to draw Mickey Mouse, Mario, or landscapes just fine.

This paper introduces a new, clever way to "unlearn" these specific concepts without having to retrain the whole artist from scratch. Here is how it works, broken down with simple analogies:

The Problem: The "Name Tag" Mistake

Previous methods tried to make the artist forget by looking at name tags.

How it worked: If you wanted the artist to forget "Pikachu," the computer would look at prompts like "a photo of Pikachu" or "a picture of Pikachu." It would find the "Pikachu" part in the text and try to delete that specific text pattern from the artist's brain.
The Flaw: This is like trying to stop someone from recognizing a friend by only banning their name. If you tell the artist, "Draw a yellow, electric mouse with red cheeks," the artist says, "Oh, I don't know what 'Pikachu' is, but I know how to draw that description!" The artist still draws Pikachu, even though you banned the name. The old methods were too focused on the words and missed the idea.

The Solution: Watching the "Sketching Process"

The authors, Saemi Moon and her team, realized that the artist doesn't just look at the name tag; they look at the actual sketching process.

The New Idea: Instead of looking at the text prompt, they watched the artist's internal "brushstrokes" (called cross-attention activations) while the artist was actually drawing the image.
The Analogy: Imagine the artist is drawing a picture.
- Old Method: You look at the order form and say, "Delete the word 'Pikachu'."
- New Method (PURE): You watch the artist's hand as they draw the yellow ears and red cheeks. You see exactly how the artist's hand moves to create those specific features. You then gently guide the hand away from that specific motion whenever it tries to draw those features again.

Because the artist's hand movements (the activations) are what actually create the image, blocking those specific movements works much better than just blocking the words. It stops the artist from drawing Pikachu, even if you describe it in a thousand different ways without using the name.

How They Did It (The "One-Time Fix")

The paper proposes a method called PURE. It's "closed-form," which is a fancy way of saying it's a one-time, quick math fix rather than a long, slow training process.

The Observation: They let the artist draw a few "Pikachu" prompts and recorded the specific hand movements (activations) used in every layer of the drawing process.
The Map: They created a "map" of these movements. This map shows exactly what the artist does when drawing the thing you want to forget.
The Edit: They applied a single mathematical "filter" to the artist's brain. This filter says: "If your hand tries to move in the 'Pikachu' direction, stop. But if you try to move in the 'Mario' or 'Landscape' direction, keep going."
The Result: The artist instantly forgets how to draw the target concept but keeps all their other skills perfectly intact.

Why It's Better

The paper tested this against other methods using a "Holistic Unlearning Benchmark" (a test with 10 different concepts like styles, celebrities, and cartoon characters).

Better at Forgetting: PURE stopped the artist from drawing the forbidden concepts even when people used tricky, rephrased descriptions (like "a yellow electric rodent").
Better at Remembering: Unlike other methods that accidentally made the artist forget other things (like making them bad at drawing Mickey Mouse when trying to ban Pikachu), PURE kept the artist's other skills sharp.
No Extra Cost: Because it's a quick math edit, it doesn't slow down the artist or require expensive retraining.

In a Nutshell

Think of the old methods as trying to erase a memory by crossing out a word in a dictionary. The new method (PURE) is like teaching the artist to stop making a specific hand gesture. Since the hand gesture is what actually creates the picture, this approach is much harder to trick and leaves the artist's other talents completely untouched.

Technical Summary: Concept Unlearning via Cross-Attention Activation Projection for Diffusion Models

Problem Statement

Concept unlearning aims to remove specific target concepts (e.g., copyrighted styles, celebrity likenesses, or unsafe imagery) from pretrained text-to-image diffusion models without the computational cost of retraining from scratch. Effective unlearning must satisfy two simultaneous requirements: forgetting the target concept even under paraphrased or adversarial prompts, and retaining the model's ability to generate all other unrelated concepts.

Existing closed-form methods, which apply a single deterministic edit to model weights without gradient-based fine-tuning, face a critical limitation. These methods (e.g., UCE, CURE) construct the "forget basis" using text-encoder embeddings derived from short anchor prompts (e.g., "Van Gogh"). The authors argue that text embeddings describe the user's prompt but fail to capture how the model actually utilizes that prompt during the denoising process. Consequently, paraphrased prompts that evoke the target concept without matching the anchor templates often bypass the unlearning edit, leading to "target leakage."

Methodology: PURE

The paper proposes PURE (Projection in U-Net Rendering for Erasure), a closed-form unlearning method that shifts the basis construction from the text-encoder embedding space to the cross-attention activation space within the U-Net.

Core Hypothesis:
Cross-attention activations ( $h^\ell$ ) at layer $\ell$ reflect how the denoiser uses the text condition under the current latent state and timestep. Unlike static text embeddings, these activations capture the specific concept evidence actually used during image generation, offering better generalization to diverse prompt phrasings.

Algorithmic Procedure:

Activation Capture: For a set of forget anchors ( $\mathcal{A}_f$ ) and retain anchors ( $\mathcal{A}_r$ ), the model runs a short denoising trajectory. At each layer $\ell$ and timestep, the post-attention activations are extracted, spatially mean-pooled, and aggregated into matrices $H^\ell_F$ and $H^\ell_R$ .
Subspace Construction: A Singular Value Decomposition (SVD) is performed on $H^\ell_F$ and $H^\ell_R$ for each layer. The top singular vectors, up to a cumulative variance threshold ( $\tau_F, \tau_R$ ), form the forget basis ( $V^\ell_F$ ) and retain basis ( $V^\ell_R$ ).
Linear Edit: The method applies a single linear update to the cross-attention key ( $W^\ell_K$ $W_{K}^{ℓ}$ ) and value ( $W^\ell_V$ $W_{V}^{ℓ}$ ) projection matrices. The update operator $E^\ell$ $E^{ℓ}$ is defined as:
$E^\ell = I - P^\ell_F (I - P^\ell_R)$
where $P^\ell_F = V^\ell_F (V^\ell_F)^\top$ $P_{F}^{ℓ} = V_{F}^{ℓ} (V_{F}^{ℓ})^{⊤}$ and $P^\ell_R = V^\ell_R (V^\ell_R)^\top$ $P_{R}^{ℓ} = V_{R}^{ℓ} (V_{R}^{ℓ})^{⊤}$ .
The weights are updated via left-multiplication: $W^\ell \leftarrow E^\ell W^\ell$ $W^{ℓ} \leftarrow E^{ℓ} W^{ℓ}$ .
- Mechanism: The term $(I - P^\ell_R)$ removes components aligned with retain concepts, ensuring $P^\ell_F$ only acts on the orthogonal subspace. Subtracting this from the identity ensures that inputs aligned with retain concepts pass through unchanged.
- Structural Shift: Unlike prior methods that multiply from the right in text space, PURE multiplies from the left in activation space ( $d_\ell \times d_\ell$ ), necessitating a per-layer edit.

Key Contributions

Activation-Based Basis: The paper identifies that constructing the forget basis from cross-attention activations rather than text embeddings significantly improves the detection of target concepts in diverse natural and adversarial prompts.
Probing Experiment: A binary classification experiment demonstrates that an activation-based basis achieves roughly five times higher recall on natural paraphrased prompts compared to a text-based basis, confirming that activation features better capture generation-relevant concept evidence.
Closed-Form Efficiency: PURE maintains the advantages of closed-form editing: it requires no gradient training, no auxiliary loss functions, no hyperparameter tuning of learning rates, and adds no inference-time cost after the edit is applied.
Layer-Specific Editing: The method performs a per-layer SVD and edit, acknowledging that concept representation varies across the U-Net depth.

Experimental Results

The authors evaluated PURE on the Holistic Unlearning Benchmark (HUB), covering ten concepts across four categories: Artistic Style, Intellectual Property (IP), Celebrity, and NSFW.

Forget-Retain Trade-off: PURE achieves the best harmonic-mean score across all categories compared to five baselines (ESD, MACE, RECELER, UCE, CURE).
Target Suppression: PURE significantly reduces target leakage under paraphrased and adversarial prompts compared to CURE, which often fails to suppress paraphrases.
Retention: PURE preserves "retain" concepts (neighboring concepts in the same category) better than aggressive training-based methods (which damage retention) and better than text-space closed-form methods (which often sacrifice retention for suppression in Style categories).
Ablation Studies:
- Anchor Size: Increasing the number of natural forget anchors improves the text basis's suppression but degrades retention. The activation basis maintains strong suppression with minimal retention loss even as anchor sets grow.
- Denoising Steps: Using a single denoising step ( $T=1$ ) severely degrades performance, indicating that capturing the concept signal across the trajectory is essential.
- Latent Averaging: Averaging over multiple random latents ( $n_{lat}$ ) stabilizes the basis toward concept-level directions rather than noise-specific artifacts.

Significance and Claims

The paper claims that PURE represents a fundamental shift in how closed-form unlearning is approached. By moving the basis construction from the static text-encoder space to the dynamic cross-attention activation space, the method overcomes the primary failure mode of existing closed-form techniques: the inability to handle prompt variations that fall outside the anchor templates.

The authors assert that PURE provides a robust, efficient, and state-of-the-art solution for concept unlearning, offering a superior balance between erasing unwanted concepts and preserving the model's general utility without the computational overhead of retraining. The work suggests that the "forget direction" is more accurately defined by the model's internal rendering process than by the input text alone.