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

TL;DR

SafeRedir tackles safety in image generation by enabling robust unlearning without modifying diffusion backbones. It performs detection and token-level embedding redirection at inference time, redirecting unsafe prompts toward safe semantic regions while preserving benign content. The framework demonstrates state-of-the-art forgetting, strong semantic and perceptual preservation, and notable robustness to adversarial prompts across diverse backbones and unlearning tasks. Its plug-and-play, model-agnostic design promises scalable deployment in real-world generative AI systems, with code and data available for reproducibility.

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.

Figures (19)

• Figure 1: A demo case of SafeRedir. Given the prompt $p$="A naked woman sits on a rock by a waterfall", a standard diffusion pipeline (left) directly encodes the prompt and generates images $I$ containing explicit content. In contrast, SafeRedir (right) intercepts the prompt embedding, performs token-level semantic redirection to filter unsafe concepts, and injects the updated embedding into the denoising process. The resulting images $I'$ preserve the benign semantic while ensuring safe, well-clothed outputs. Sensitive parts are covered by .
• Figure 2: Generated images of unlearning methods on three tasks. Each column corresponds to a mainstream unlearning method, and each row presents a sample for a specific task: top row:NSFW (Nudity) unlearning; middle row:Van Gogh style removal; bottom row:Church object removal.
• Figure 3: (a) Images generated by leveraging unlearned models in response to unsafe prompts containing sensitive or explicit words (e.g., "naked", "nude", "Van Gogh style") . (b) Images generated from paraphrased or adversarial prompts that evade unlearning and lead to the regeneration of sensitive content.
• Figure 4: SafeRedir inference pipeline for safety-aware text-to-image generation. The framework intercepts user prompts and injects token-wise semantic guidance during the denoising process. Unsafe semantic elements (e.g., "naked person") are automatically redirected in the prompt embedding space at each denoising step $t$, resulting in sanitized and semantically coherent outputs. For safe prompts, the original generation trajectory is preserved.
• Figure 5: Selective semantic redirection. Prompt embeddings for unsafe and safe content form distinct clusters separated by a safe boundary. SafeRedir minimally shifts only unsafe embeddings into the safe region using $\alpha \cdot \tilde{\Delta}$, leaving benign prompts unchanged. Solid arrows indicate effective redirection; dashed arrows indicate ineffective directions or scales.