💻 computer science

The Linear Geometry of Interpretable Tokens: Jailbreaking Attacks and Defenses for Unlearned Diffusion Models

This paper reveals that erased concepts in unlearned diffusion models persist as coherent linear subspaces within token embeddings, leading to the development of SubAttack, a powerful jailbreaking method that exploits this structure, and SubDefense, a lightweight mechanism that projects out the subspace to robustly suppress residual harmful concepts while preserving generation quality.

Original authors: Siyi Chen, Yimeng Zhang, Sijia Liu, Qing Qu

Published 2026-08-17

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Siyi Chen, Yimeng Zhang, Sijia Liu, Qing Qu

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 a world where computers can paint pictures just by listening to your words. You say, "a cat wearing a hat," and poof, a perfect image appears. These are called Diffusion Models, and they are like digital artists that learned to create by studying millions of photos. But sometimes, these artists get a little too good at remembering things they shouldn't. If you ask them to draw something dangerous or copyrighted, they might do it anyway. To fix this, scientists invented a process called "Unlearning." Think of it like a teacher trying to make a student forget a specific fact, like the capital of a country, without making them forget how to do math or read. The goal is to "erase" the bad idea from the computer's brain while keeping its ability to draw everything else. But here's the tricky part: even after the teacher says, "Forget that fact," the student might still have a secret way of remembering it, hidden in a way we can't easily see. This paper dives into that mystery, asking: If we think we've erased a bad idea, where does it actually hide, and how can we finally get rid of it?

The researchers in this paper discovered that when a diffusion model tries to "unlearn" a harmful concept (like nudity or a specific artist's style), it doesn't actually delete it. Instead, the concept hides in a secret, linear subspace within the computer's memory. To use a metaphor, imagine the computer's brain is a giant library of words. When we try to remove the book about "Nudity," the computer doesn't burn the book. Instead, it shreds the pages and hides the scraps inside a specific, invisible drawer in the library. The paper shows that these scraps are still organized in a neat, predictable line. The authors call this hidden drawer a "residual subspace."

To prove this, the team created a new tool called SubAttack. Instead of guessing random words to trick the computer (which is how previous hackers tried), SubAttack acts like a master librarian who knows exactly where that invisible drawer is. It learns a special set of "magic keys" (which are just combinations of normal words) that unlock the drawer. When they used these keys, the "unlearned" computer immediately started drawing the harmful images again, proving the concept was never truly gone. What's fascinating is that these keys are made of words we can understand. For example, to make the computer draw a "nude" person again, the keys might be a mix of words like "slave," "hips," and "skin." This shows that the computer is still connecting the bad idea to these related, hidden clues.

The paper also found that this secret drawer isn't just in one computer; it's in many different "unlearned" models. If you find the key for one model, it often works on others, too. This suggests that the way these computers "forget" is flawed in a very similar way across the board. They aren't deleting the memory; they are just hiding it in a way that follows a straight, predictable line.

But the story doesn't end with the hack. Because the researchers understood exactly how the concept was hiding, they built a shield called SubDefense. If SubAttack is the key that opens the drawer, SubDefense is a heavy metal plate that welds that drawer shut. It works by taking the computer's entire library of words and mathematically "projecting" them away from that secret drawer. It's like telling the librarian, "No matter what word you use, make sure it never points toward that specific hidden drawer."

The results were impressive. When they used SubDefense, the computer became much harder to trick. Even when hackers tried to use the old methods or the new SubAttack keys, the computer refused to draw the harmful images. At the same time, the computer didn't lose its ability to draw safe, beautiful pictures. In fact, the images it drew after the defense were just as high-quality as before. The paper suggests that while we can't yet erase the concept perfectly, we can at least block the specific path it uses to sneak back in. This gives us a new, clear way to understand why "unlearning" is so hard and offers a practical, plug-and-play fix to make these digital artists safer for everyone.

Technical Summary: The Linear Geometry of Interpretable Tokens in Unlearned Diffusion Models

Problem Statement
Diffusion models (DMs), particularly Text-to-Image (T2I) models, excel at generating high-quality content but pose safety risks by potentially memorizing and reproducing harmful concepts (e.g., nudity, copyrighted styles, or specific objects). Machine Unlearning (MU) methods have been developed to "erase" these concepts from pretrained models via fine-tuning without retraining from scratch. However, existing unlearning methods struggle to completely remove harmful concepts while maintaining generation quality on safe content. Consequently, unlearned models remain vulnerable to jailbreaking attacks, where adversaries manipulate prompts to regenerate the erased concepts.

Current jailbreaking approaches (e.g., adversarial prompt optimization or embedding-based attacks) often rely on non-interpretable, optimization-based perturbations. While effective, these methods offer limited insight into how harmful concepts persist within the model's internal state. This lack of interpretability hinders the development of robust, targeted defense strategies. Furthermore, lightweight, post-hoc defense mechanisms that operate on already-unlearned models remain underexplored compared to adversarial re-training approaches.

Methodology
The authors propose that erased concepts do not vanish entirely but persist as a coherent, interpretable linear subspace within the token embedding space of the CLIP text encoder used by Latent Diffusion Models (LDMs). Based on this hypothesis, they introduce two complementary methods: SubAttack and SubDefense.

SubAttack (Subspace Attack):
- Goal: To identify and exploit the residual subspace of a harmful concept.
- Mechanism: The method learns a set of orthogonal attack token embeddings $\{v_{att, k}\}$ that span the residual subspace. Each embedding is constrained to be a non-negative linear combination of existing vocabulary token embeddings ( $v_{att} = \sum \alpha_i v_i, \alpha_i \geq 0$ ).
- Optimization: An MLP network is trained to optimize the coefficients $\alpha_i$ to minimize the denoising loss when the generated image contains the target harmful concept.
- Deflation: To learn multiple orthogonal embeddings, the authors employ a deflation strategy. After learning $v_{att, k}$ , they project the entire vocabulary orthogonally to $v_{att, k}$ before learning the next embedding. This ensures diversity and captures the full subspace of the residual concept.
- Interpretability: Because the embeddings are linear combinations of existing tokens, the top-weighted tokens can be extracted and analyzed (e.g., via WordClouds) to reveal the semantic components (e.g., "slave," "hips," "nip") that the unlearned model still associates with the erased concept.
SubDefense (Subspace Defense):
- Goal: To neutralize the residual subspace identified by SubAttack.
- Mechanism: Once the attack subspace $V_{att}$ (spanned by the learned orthogonal embeddings) is identified, SubDefense projects the token embeddings of the unlearned model's vocabulary onto the nullspace of $V_{att}$ .
- Update Rule: For every token embedding $v_i$ in the vocabulary, the defense updates it as $v_{def, i} = v_i - \text{Proj}_{V_{att}}(v_i)$ .
- Plug-and-Play: This operation is lightweight and does not require retraining the diffusion model's weights. It can be applied to any unlearned model, including those that have undergone adversarial fine-tuning.

Key Contributions

Identification of Interpretable Residual Subspaces: The paper demonstrates that harmful concepts in unlearned DMs persist as a linear subspace in the token embedding space, composed of human-interpretable textual elements.
SubAttack: A novel, interpretable jailbreaking attack that learns orthogonal token embeddings to span this subspace. It outperforms existing baselines in Attack Success Rate (ASR) and demonstrates high transferability across different prompts, initial noises, and unlearned models.
SubDefense: A robust, lightweight defense mechanism that projects out the identified residual subspace. It improves robustness against diverse jailbreaking attacks (including adversarial prompts and embedding attacks) while better preserving safe generation quality compared to existing defenses like RECE.
Insight into Unlearning Limitations: The work reveals that unlearning often suppresses surface-level keywords but fails to eliminate deeper, implicit associations (e.g., "dumpster" for "garbage truck" or "slave" for "nudity"), which are captured by the linear subspace.

Experimental Results

Attack Performance: SubAttack achieved significantly higher ASR than baselines (UnlearnDiff and CCE) across multiple concepts (nudity, Van Gogh style, church) and unlearned models (ESD, FMN, UCE, SPM). For instance, on the "nudity" concept across ESD, FMN, UCE, and SPM, SubAttack achieved ASRs of 97.56%, 100%, 81.67%, and 74.89% respectively, outperforming CCE and UnlearnDiff.
Transferability: Attack embeddings learned on one unlearned model (e.g., ESD) transferred effectively to others (FMN, SPM, UCE) and even back to the original Stable Diffusion model, with ASRs consistently above 80%.
Defense Performance: SubDefense significantly reduced ASR for various attacks. For example, when defending UCE against SubAttack on the "Van Gogh" concept, ASR dropped from 98.33% to 29.44%. It also reduced ASR for UnlearnDiff and CCE attacks.
Quality Preservation: Unlike some defenses that degrade image quality, SubDefense maintained or improved FID and CLIP scores on safe prompts (MSCOCO-10k dataset) compared to the baseline RECE defense.
Complementarity: SubDefense was shown to be effective even when applied on top of adversarially fine-tuned models like STEREO, further reducing ASR from 21.67% to 6.67%, indicating it addresses a different (linear residual) vulnerability than adversarial training.

Significance and Claims
The paper claims that the persistence of harmful concepts in unlearned models is not random noise but a structured, linear phenomenon that can be both exploited and mitigated. By framing the problem through the lens of linear geometry in token embeddings, the authors provide a unified framework for understanding vulnerabilities (SubAttack) and constructing defenses (SubDefense).

The significance lies in shifting the paradigm from opaque, optimization-based attacks to interpretable, subspace-based analysis. This approach not only exposes deeper vulnerabilities in current unlearning techniques but also offers a practical, model-agnostic defense strategy that does not require retraining. The authors emphasize that while linear projection cannot capture nonlinear remnants (a limitation they acknowledge), addressing the linear residual subspace provides a critical first step toward more robust and interpretable safety mechanisms for generative AI.

Technical Summary: The Linear Geometry of Interpretable Tokens in Unlearned Diffusion Models

More like this