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

Diffusion models excel at generating high-quality images but can memorize and reproduce harmful concepts when prompted, creating vulnerabilities in unlearned models that necessitate robust defenses. The core finding is that supposedly erased concepts persist as a coherent, interpretable linear subspace within the token embedding space, which allows for both novel jailbreaking attacks and corresponding plug-and-play defenses.

The Core Mechanism: Interpretable Residual Subspaces

The research identifies that harmful concepts are not completely erased but persist within a coherent, interpretable linear subspace of the token embedding space. This means these concepts can be extracted through linear combinations of existing vocabulary tokens, providing a human-understandable explanation for why unlearned models remain vulnerable. The paper demonstrates this structure through the SubAttack method, which learns an orthogonal set of attack token embeddings that span this residual subspace. These learned embeddings are not opaque vectors but are instead interpretable as recognizable words or semantic units, such as slave and hips for a nudity concept, revealing the textual associations retained by the model.

SubAttack: The Interpretable Jailbreaking Attack

SubAttack is a novel jailbreaking attack designed to read out this residual subspace. It proceeds in two main stages:

1. It learns an initial single-token embedding as a non-negative linear combination of existing token embeddings (Equation 2), parameterized by an MLP network.

2. It iteratively learns a set of orthogonal attack token embeddings, denoted as the set of vectors where orthogonality is enforced to promote diversity and improve attack effectiveness. This is achieved through a deflation process: after identifying one embedding, subsequent embeddings are learned such that they are orthogonal to vatt,k by subtracting the projection onto the previously found embedding (Equation 3).

SubDefense: The Plug-and-Play Defense Mechanism

The identified subspace directly inspires SubDefense, a lightweight defense strategy designed to suppress this residual concept. The defense is implemented by projecting out the learned attack subspace. Specifically, each token embedding in the vocabulary is updated according to Equation (4): vdef,i = v i − ProjVatt(v i), where Vatt represents the matrix spanning the learned attack vectors. This operation effectively removes the harmful concept from unlearned models by projecting onto the nullspace of learned subspace attacks, offering a defense that is both robust and versatile across various unlearned models and attack vectors.

Empirical Effectiveness and Transferability

Extensive experiments confirm that SubAttack is more effective than prior methods, demonstrating stronger empirical performance of efficiency and effectiveness and superior transferability across text prompts, initial noises, and unlearned models. The findings show that the attack token embeddings learned by SubAttack are robustly transferable between different unlearned LDMs (e.g., ESD to FMN) and even back to the original Stable Diffusion model (SD), suggesting that residual associations are inherited from the original SD model rather than independently formed. Furthermore, SubDefense is shown to be stronger robustness than existing defenses while better preserving safe generation quality, maintaining higher FID and CLIP scores on datasets like MSCOCO-10k.

Insights into Concept Persistence

The analysis of the learned embeddings provides crucial insights into how unlearning methods fail. The paper reveals that unlearning reduces surface-level cues but does not eliminate deeper associations. For concepts like church, the model retains explicit components, such as names (mary) and places (abbey), while for styles like Van Gogh, it retains both explicit words (vincent, gogh) and implicit ones (art, artist). This suggests that current unlearning methods retain more explicit associations with the target concept when applied to styles compared to NSFW or object concepts. The iterative learning process also shows a gradual decrease in the sparsity of the learned coefficients, indicating that later attacks require more complex associations to the target concept.

Ablations and Future Directions

The study includes detailed ablations on attack parameters, such as the number of attack tokens (using K=5 for efficiency), vocabulary size (showing optimal balance around 5000 tokens), and the orthogonality constraint, which is shown to consistently improve ASR across multiple unlearned models. While SubDefense demonstrates a clear trade-off between robustness and utility—where increasing blocked tokens beyond 100 leads to degradation in CLIP score and FID—the work suggests that future research should explore feature representations beyond the linear structure, design adaptive methods for selecting specific tokens to block, and investigate joint visual–textual embeddings to better defend against multimodal jailbreaks. The core principle of identifying and nullifying harmful semantic directions is proposed as a promising avenue for extending SubDefense beyond CLIP-based architectures.

The gist: Erased concepts persist as an interpretable residual subspace in the token embedding space, which can be read out by SubAttack to generate harmful content and removed by SubDefense via orthogonal projection.

As a fastidious and diligent AI researcher, I have thoroughly reviewed this paper, The Linear Geometry of Interpretable Tokens: Jailbreaking Attacks and Defenses for Unlearned Diffusion Models. The core contribution is establishing that concepts supposedly erased by unlearning methods persist as an interpretable linear subspace in the token embedding space.

Here are the specific improvements to AI systems based on this research, categorized by application:

)Specific Improvements and Capabilities of Enhanced AI Systems:

1. """

A. Subspace-Aware Robust Unlearning Frameworks (SubSpace-MU):

The system will integrate a diagnostic step before or after concept erasure. Instead of simply fine-tuning the UNet weights, it will first perform an analysis on the resulting token embedding space to identify residual subspaces using a method analogous to SubAttack (learning orthogonal attack tokens).

Capability: This allows for concept-aware unlearning. If a residual subspace is found, the system can either refine its unlearning process by specifically targeting those directions or trigger a subsequent defense mechanism.

2. """

B. Adaptive, Interpretable Defense Layer (SubDefense Plug-in):

The core defense mechanism will be a lightweight, plug-and-play projection layer (SubDefense). This layer will take the learned residual subspace vectors (from SubAttack) and project them out of the model's token embedding space before generation.

Capability: This provides dynamic, targeted robustness. Unlike naive blocking of all related tokens, SubDefense blocks only the specific semantic directions identified as residual hidden words, preserving generation quality for safe concepts while effectively neutralizing the persistent harmful concept subspace.

3. """

C. Interpretable Jailbreaking Attack Generation (SubAttack Attacker):

The system will utilize a learned set of orthogonal attack token embeddings to construct jailbreak prompts that are semantically grounded in human-interpretable concepts (e.g., slave + hips).

Capability: This enables the creation of highly effective, transferable, and interpretable attacks. The AI can systematically probe an unlearned model's vulnerabilities by generating attacks that leverage the specific textual components (tokens) that remain associated with the erased concept, rather than relying on opaque adversarial perturbations.

4. """

D. Concept Persistence Diagnostic Tool:

A diagnostic module will be implemented to analyze the latent space of any diffusion model after an unlearning attempt. This tool will quantify the residual harm by measuring the CLIP similarity between learned attack embeddings and original concept embeddings (as shown in Table 1).

Capability: This provides a quantitative metric for assessing the success or failure of unlearning methods, moving beyond simple ASR to show exactly how much of a concept remains embedded in the model's structure.

5. """

E. Robustness-Aware Model Selection and Fine-Tuning:

The research suggests that the effectiveness of defenses (like SubDefense) can be tuned by blocking a specific number of tokens, balancing ASR reduction against generation quality degradation (Table 22).

Capability: The system will incorporate an adaptive tuning mechanism where the optimal number of blocked tokens is determined based on a trade-off function between maximizing ASR reduction and minimizing FID/CLIP score degradation for the target concept.

6. """

F. Cross-Model Transferability Verification Module:

A module will be built to verify the transferability of learned attack tokens across different unlearned models (e.g., ESD to FMN).

Capability: This allows researchers and deployers to rapidly assess the generalizability of a discovered vulnerability by testing it on a diverse set of production or research models without needing extensive, manual testing for every new model variant.

"""

Diffusion models excel at generating high-quality images but can memorize and reproduce harmful concepts when prompted. Although fine-tuning methods have been proposed to unlearn a target concept, they struggle to fully erase it while maintaining generation quality on other concepts, leaving models vulnerable to jailbreak attacks. Existing jailbreak methods demonstrate this vulnerability but offer limited insight into how unlearned models retain harmful concepts, limiting progress on effective defenses. In this work, we show that a coherent, interpretable, attack-accessible linear residual of the erased concept can be recovered in the token embedding space, and that both an attack and a defense follow from this structure. We introduce SubAttack, a novel jailbreaking attack that reads out this subspace by learning an orthogonal set of attack token embeddings, each being a linear combination of human-interpretable textual elements, revealing that unlearned models still retain the target concept through related textual components. Furthermore, our attack is also more powerful and transferable across text prompts, initial noises, and unlearned models than prior attacks. Conversely, projecting out the same subspace yields SubDefense, a lightweight plug-and-play defense mechanism that suppresses the residual concept in unlearned models. SubDefense provides stronger robustness than existing defenses while better preserving safe generation quality. Extensive experiments across multiple unlearning methods, concepts, and attack types demonstrate that our approach advances both understanding and mitigation of vulnerabilities in diffusion unlearning.