SafeRedir: Prompt Embedding Redirection for Robust Unlearning in Image Generation Models

AI summary

Overview

Research area: Safety and machine unlearning for text-to-image diffusion models (computer vision / generative AI safety).

Technical level: Advanced — the paper assumes familiarity with latent diffusion, prompt embeddings, cross-attention, LoRA-style adapters, and embedding-space perturbation methods.

Scope: This paper introduces SafeRedir, an inference-time, plug-and-play framework that intercepts prompt embeddings and performs token-level semantic redirection to suppress unsafe concepts in image generation models without modifying the underlying diffusion weights.

Note on completeness: the provided paper content is truncated partway through Section 4.4 (Training Objective). The main experiment section (Section 5), including the headline quantitative comparison of SafeRedir against baselines, is not present in the supplied text. Claims about SafeRedir's comparative performance below are drawn only from the abstract and introduction; specific SafeRedir benchmark numbers beyond Tables 1, 3, and 4 are not reported in the available content.

What This Paper Is About

Image generation models such as Stable Diffusion can memorize and reproduce undesirable content — explicit nudity, violent imagery, and copyrighted artistic styles — from their training data. Existing fixes either filter text or images after the fact (fragile against paraphrasing) or modify model weights (expensive and prone to degrading benign image quality). SafeRedir's goal is to erase harmful concepts at inference time by nudging only the unsafe parts of a prompt's embedding toward a "safe" region, leaving the diffusion backbone completely untouched.

Key Contributions

1. A plug-and-play prompt-level redirection framework. SafeRedir performs semantic unlearning without modifying the underlying diffusion model — no changes to the U-Net, text encoder, or VAE — operating entirely through prompt embedding hooks at inference time.

2. A latent-aware multi-modal safety classifier. Unsafe generation trajectories are detected by jointly analyzing text embeddings, image latents at each diffusion step, and the diffusion timestep, rather than relying on any single modality.

3. A token-level redirection module. Once unsafe content is detected, a redirector implements latent-conditioned, token-wise semantic guidance by combining a learned shift direction, a token soft mask, and an adaptive per-token scaling factor for precise, minimal intervention.

4. An empirical study of existing unlearning failures. The paper benchmarks representative methods spanning the text encoder, cross-attention, self-attention, and feed-forward layers, and documents four persistent limitations (incomplete forgetting, quality degradation, prompt-manipulation vulnerability, and poor transferability) that motivate the design.

Main Findings

• Incomplete forgetting is pervasive. Across three unlearning tasks — NSFW (Nudity), Van Gogh style, and Church — existing methods left residual traces. UCE and ConceptPrune still produced NSFW content in some cases; MACE, UCE, and ConceptPrune retained Van Gogh style cues; and the majority of evaluated methods preserved recognizable church structures (Observation 1).

• Unlearning degrades benign generation quality. On Van Gogh, all evaluated methods had high CSDR (minimum 16.83, where lower is better). On Nudity, CSDR was lower (minimum 7.42) but still non-negligible. FID ranged from 18.73 to 58.92 across both tasks, and PDR ranged from 60.96% (Receler) to 90.24% (MACE) on Nudity (Table 1).

• Prompt manipulation easily bypasses unlearning. Minor rephrasings, synonym substitutions, and implicit expressions reactivate forbidden concepts even when specific keywords such as "nude" appear forgotten — a consequence of the keyword-based paradigm used by most existing methods (Observation 3).

• Most unlearning methods transfer poorly. Of the seven methods compared in Table 2 (ESD, AdvUnlearn, MACE, UCE, ConceptPrune, Receler, SafeGen), only AdvUnlearn was marked as transferable; the rest edit U-Net cross-attention, self-attention, or FFN components and require per-model adaptation (Observation 4).

• Multi-modal detection beats single-modality detection. On the IGMU dataset, text-only detection achieved 89.17% overall / 97.20% any-step success, latent-only 66.98% / 93.76%, text-and-latent 99.46% / 99.60%, and text-and-latent-and-timestep 99.73% / 99.98%. On the adversarial MMA dataset, text-only reached only 51.68% overall (93.50% any-step), latent-only 42.37% (90.12%), text-and-latent 64.29% (99.22%), and the full tri-modal setting 74.72% (99.60%) (Table 3).

• Latent-only detection is unreliable early in denoising. Latent-based detectors perform near random guessing (approximately 50%) at early diffusion steps when latents resemble Gaussian noise, and even on the all-adversarial MMA dataset early-step accuracy rarely exceeded 70%, motivating the inclusion of text and timestep signals.

• Adaptive, masked, token-level redirection outperforms global strategies. Comparing redirection configurations (Table 4): a simple pairwise safe–unsafe embedding difference (emb1) with no scaling or mask gave 14.27 FSR, 0.51 LPIPS, 80.00 PDR, 285.65 FID. A latent-aware generated embedding (emb2) alone improved this to 60.40 FSR, 0.26 LPIPS, 97.20 PDR, 109.33 FID. Adding a fixed scaling factor raised forgetting (88.00 FSR at α = 1.5; 96.80 at α = 2.0; 99.87 at α = 3.0) but degraded quality (FID rose to 138.28 and then 200.72, LPIPS to 0.39). Adding a mask alone gave 55.60 FSR, 0.23 LPIPS, 96.40 PDR, 101.01 FID; combining mask with scaling gave 83.07 FSR at α = 1.5, 92.40 at α = 2.0, with FID of 103.97 and 123.95 respectively.

• Direction is easier to learn than magnitude. Inspired by adversarial attack research, SafeRedir predicts only the normalized direction of the shift vector rather than its absolute value, on the intuition that directionality generalizes better in high-dimensional embedding spaces.

• Latent context disambiguates redirection. The paper argues that text-only redirection can be ambiguous for prompts that appear benign but lead to unsafe generation, and that incorporating the image latent enables more accurate redirection when textual cues are insufficient.

• Abstract-level performance claims. The abstract states that SafeRedir achieves effective unlearning, high semantic and perceptual preservation, robust image quality, and enhanced resistance to adversarial attacks, and that it transfers across compatible Stable Diffusion-family variants and existing unlearned models. The specific benchmark values supporting these claims are not present in the available content.

Methodology in Plain English

The approach rests on a geometric intuition: safe prompt embeddings and unsafe prompt embeddings tend to occupy separable regions of the embedding space, divided by what the paper calls a "safe boundary." The challenge is to move only unsafe embeddings across that boundary, by a minimal amount, while leaving safe embeddings untouched.

SafeRedir works in two stages, identify then redirect.

Stage one: detection. A lightweight classifier examines three things at every denoising step — the current image latent (what the model is actually generating), the prompt embedding (what the user typed), and the diffusion timestep (how far along generation is). Each is encoded separately and then fused through concatenation and multi-scale cross-attention into a unified representation. This tri-modal setup matters because text alone misses images that turn unsafe despite benign wording, latents alone carry little signal early in denoising, and the timestep tells the system when risk is developing.

Stage two: redirection. When an unsafe state is detected, the framework computes a correction applied to the prompt embedding. Three token-level quantities govern it:

1. A shift vector describing the direction to move each token's embedding. A small prediction head (a cross-attention-augmented MLP with a LoRA-style low-rank branch) outputs this vector, which is then normalized so only the direction is used.
2. A soft mask that limits the correction to the tokens actually responsible for unsafe semantics. Supervision for this mask comes from a pseudo-ground-truth signal: tokens whose safe and unsafe prompt embeddings differ most in cosine similarity are the ones targeted for correction, while tokens with small differences are treated as shared benign context.
3. An adaptive scaling factor in the range [0, 1], predicted per token by a two-layer MLP with sigmoid activation, controlling how far to move along the chosen direction.

The corrected embedding is a masked, normalized, and scaled addition to the original, preserving per-token norms so the intervention stays directionally consistent and scale-aware.

Training. The system is trained on aligned safe–unsafe prompt pairs that share benign context and differ only in the unsafe concept. Supervision combines cross-entropy with label smoothing for the safety classifier, MSE and cosine losses aligning the predicted shift with the reference difference between safe and unsafe embeddings, binary cross-entropy for the mask, and a final embedding alignment loss pulling the redirected embedding toward the aligned safe embedding. The scaling factor is not supervised directly — it is learned implicitly through its contribution to the final redirected embedding.

Deployment. Everything happens in the embedding space through inference-time hooks, so no weights are touched and no retraining is needed when swapping between compatible Stable Diffusion variants.

Why This Matters

Impact on research. The paper reframes image-generation unlearning as an inference-time intervention problem rather than a weight-editing problem, and it supplies a systematic empirical audit showing why keyword-level erasure, component editing, and parameter selection fall short. By treating unsafe semantics as distributed across representations rather than isolated in specific tokens or layers, it argues for embedding-space, latent-conditioned intervention as a more robust design point. The multi-modal detection ablation also provides useful evidence about the limits of latent-only and text-only safety signals.

Real-world applications:

• Content moderation for generative platforms. An external safety layer that can be attached to deployed image services without retraining or accessing proprietary weights.
• Copyright and artist-style protection. Suppressing protected artistic styles (the paper studies Van Gogh) at generation time rather than through post-hoc takedowns.
• Enterprise and parental controls. Customizable concept suppression for deployments where the operator cannot or will not modify the base model.
• Retrofitting already-deployed or personalized models. Because it works as a plug-in, it can be layered onto fine-tuned or personalized Stable Diffusion variants that would otherwise need separate adaptation.

Industry relevance. The paper claims transferability across compatible Stable Diffusion-family variants and existing unlearned models, which maps directly onto real deployment constraints: model providers rarely release weights, model families iterate quickly (the paper itself notes models are "frequently updated"), and per-model retraining is costly. A hook-based, low-overhead module is a more plausible product form than a weight-editing pipeline, though the paper reports only "limited deployment overhead" qualitatively rather than quantified latency figures in the available content.

Future Directions

• Extending beyond the Stable Diffusion family. The paper scopes transferability to variants sharing the same target SD architecture; whether the embedding-space geometry it exploits holds for fundamentally different backbones (other latent diffusion models, or non-diffusion generators mentioned such as VAEs, GANs, and normalizing flows) is untested in the available content.
• Broader robustness evaluation. The empirical study covers paraphrasing, synonym substitution, and adversarial prompts, but the paper itself cites multimodal jailbreaks as a weakness of prompt-side defenses generally. Whether SafeRedir resists multimodal prompt manipulation is not established in the supplied material.
• Quantifying deployment overhead. The paper claims limited deployment overhead but the provided content reports no latency, memory, or throughput numbers for the classifier and redirector operating at every denoising step.
• Scaling to more concepts and combined defenses. The study covers three concept types (NSFW, artistic style, object category). The threat model also notes that a governor "may optionally combine additional semantic filtering mechanisms," raising the open question of how SafeRedir composes with other defenses and whether multi-concept suppression holds at scale.

Target Audience

Researchers and engineers working on generative AI safety, machine unlearning, and diffusion model deployment — particularly those interested in inference-time interventions that avoid modifying model weights. It will also be valuable to practitioners building content-moderation or compliance layers for text-to-image systems, and to readers already familiar with adversarial attack and embedding-space perturbation literature, given the paper's explicit analogy between crossing a decision boundary and crossing its "safe boundary." Readers without background in latent diffusion and prompt-embedding mechanics will find Sections 4.3 and 4.4 dense, though the paper defers additional implementation details to its appendices.

Authors’ abstract

Image generation models (IGMs), while capable of producing impressive and creative content, often memorize a wide range of undesirable concepts from their training data, leading to the reproduction of unsafe content such as NSFW imagery and copyrighted artistic styles. Such behaviors pose persistent safety and compliance risks in real-world deployments and cannot be reliably mitigated by post-hoc filtering, owing to the limited robustness of such mechanisms and a lack of fine-grained semantic control. Recent unlearning methods seek to erase harmful concepts at the model level, which exhibit the limitations of requiring costly retraining, degrading the quality of benign generations, or failing to withstand prompt paraphrasing and adversarial attacks. To address these challenges, we introduce SafeRedir, a lightweight inference-time framework for robust unlearning via prompt embedding redirection. Without modifying the underlying IGMs, SafeRedir adaptively routes unsafe prompts toward safe semantic regions through token-level interventions in the embedding space. The framework comprises two core components: a latent-aware multi-modal safety classifier for identifying unsafe generation trajectories, and a token-level delta generator for precise semantic redirection, equipped with auxiliary predictors for token masking and adaptive scaling to localize and regulate the intervention. Empirical results across multiple representative unlearning tasks demonstrate that SafeRedir achieves effective unlearning capability, high semantic and perceptual preservation, robust image quality, and enhanced resistance to adversarial attacks. Furthermore, SafeRedir generalizes effectively across a variety of diffusion backbones and existing unlearned models, validating its plug-and-play compatibility and broad applicability. Code and data are available at https://github.com/ryliu68/SafeRedir.