Training-Free Safe Text Embedding Guidance for Text-to-Image Diffusion Models

AI summary

Overview

Research area: Safe generation for text-to-image diffusion models — specifically, preventing unsafe outputs (nudity, violence, copyrighted artist style) without retraining the model.

Technical level: Advanced. The paper combines score-based diffusion theory, classifier guidance formulations, and text-embedding optimization, with a theoretical theorem and a 2D toy analysis.

Scope: The paper introduces Safe Text embedding Guidance (STG), a training-free method that perturbs text embeddings during sampling using a safety function evaluated on the expected final denoised image, and compares it against training-based and training-free baselines on nudity, violence, and artist-style removal.

What This Paper Is About

Text-to-image diffusion models are trained on large-scale web-crawled data that contains inappropriate or biased content, so malicious prompts can produce harmful images. Existing fixes either require expensive fine-tuning on curated safe data or manipulate prompts/embeddings without ever checking the images the model actually produces. STG closes that gap by nudging the text embedding at sampling time based on a safety score computed on the model's own expected final image, so unsafe concepts are suppressed while the prompt's core meaning is preserved.

Key Contributions

1. A training-free safe guidance method (STG) that updates the text embedding as c ← c + ρ∇_c g_t(x_t, c) during diffusion sampling, where g_t is a time-dependent safety function evaluated on the expected denoised image. No model parameters are modified.

2. A theoretical analysis (Theorem 1) showing that STG's adjusted score decomposes into the original text-conditional score plus a safe guidance term of order ρ, with O(ρ²) residual. STG's implied safe probability is proportional to exp(ρ ∇_c g_t(xt, c)ᵀ ∇c log q_t(x_t | c)), which ties safety to the model's own likelihood.

3. A formal comparison of guidance placements: Safe Guidance (SG, analogous to Classifier Guidance), Safe Data Guidance (SDG, analogous to Universal Guidance, applied in perturbed data space), and STG (analogous to Diffusion Adaptive Text Embedding, applied in text embedding space). The paper argues that guiding text embeddings is less distorted by proxy safety functions than guiding data.

4. Empirical validation across three safety scenarios — nudity, violence, and artist-style removal — with both black-box and white-box red-teaming, plus generalization tests on FLUX, SDXL, SD3, PixArt-α, and LCM.

Main Findings

• STG achieves the lowest FID on COCO for nudity removal: FID 22.00 with CLIP 31.14, versus Base SD v1.4 at FID 23.22 / CLIP 31.96. STG is the only method in that table with a lower FID than the base model, though its CLIP score is slightly lower than the base model's. For comparison, SAFREE scored FID 28.39 / CLIP 30.27 and SDG scored FID 26.90 / CLIP 29.97.

• STG dominates the defense success rate / prior preservation trade-off curve for nudity (Ring-A-Bell and SneakyPrompt) and violence (Ring-A-Bell), occupying the upper-right region in Figure 3. Each experiment was repeated three times with different random seeds, and mean values are plotted with standard deviation error bars.

• Training-based DUO performs best among baselines on nudity but weakens on violence. The paper attributes this to the diverse categories of violence, which are harder to cover through training, whereas SDG and STG leverage test-time CLIP score guidance based on violence-related text.

• STG generalizes across backbones and schedulers. On Ring-A-Bell violence, base FLUX scored DSR 0.11 / FID-1K 56.58 / CLIP 32.67; STG at τ = 0.16 reached DSR 0.70 / PP 0.60 / FID-1K 57.77 / CLIP 32.00. Base SD3 went from DSR 0.12 to 0.68 at τ = 0.16. Base SDXL went from DSR 0.04 to 0.77 at τ = 0.16. Base LCM went from DSR 0.02 to 0.80 at τ = 0.16.

• STG works as an add-on to training-based defenses. When applied on top of DUO with fixed parameters, STG outperformed other training-free baselines under both black-box (Ring-A-Bell) and white-box (Concept Inversion) violence red-teaming.

• The 2D toy experiment shows SDG is sensitive to misspecified safety functions. With the ideal safety function g*(x₀) = exp(−½‖x₀ − μ‖²), SDG guides correctly; with the shape-distorted g̃(x₀) = exp(−½‖x₀ − μ‖⁴), SDG produces more biased samples with larger KL divergence, while STG remains robust under both safety functions.

• Artist-style removal results are mixed: On "Van Gogh," SDG reached the highest LPIPS_e (0.43) with the lowest LPIPS_u (0.09) and ACC_e 0.30, while SAFREE had LPIPS_e 0.39 and LPIPS_u 0.25. On "Kelly McKernan," SAFREE had the highest LPIPS_e (0.47) but also the highest LPIPS_u (0.46). The STG rows of Table 4 are cut off in the provided content, so those specific numbers are not reported here.

Methodology in Plain English

The researchers start from a simple observation: unsafe images usually trace back to text prompts containing explicit or implicit unsafe concepts. So instead of retraining the model or editing prompts, they modify the text embedding the model conditions on.

The challenge is that safety detectors (like NudeNet) only judge finished images, not the noisy intermediate latents the diffusion process works with. The team solves this by using Tweedie's formula to estimate the final clean image from any noisy intermediate step, then evaluating the safety function on that estimate.

From there, the method computes the gradient of the safety score with respect to the text embedding and takes a small step in that direction — gradient ascent on safety. This updated embedding is fed into the score network for that timestep. Two extra hyperparameters, an update threshold τ and an update step ratio γ, control when and how often the safety update runs, trading computation for safety strength. A scale hyperparameter ρ controls overall guidance strength.

To understand why this works, they prove a theorem showing that the modified score function splits into the original text-conditional score plus a safety term, meaning the model stays close to its original distribution while drifting toward safer outputs. They also build a tractable 2D Gaussian toy example where the true safe distribution is known, letting them compare SDG and STG when the safety function is ideal versus shape-distorted.

Why This Matters

Impact on research: The paper provides a theoretical account of how text-embedding-space guidance relates to the underlying model distribution, addressing a gap the authors identify in prior training-free methods, which they say do not directly use intermediate or final diffusion states and lack a clear theoretical foundation. It also shows that training-free and training-based defenses are complementary rather than competing.

Real-world applications:

• Content moderation for image-generation platforms that need to block nudity or violence on demand without retraining deployed models.
• Copyright and intellectual-property compliance, via the artist-style removal task, which matters for commercial deployments sensitive to style mimicry.
• Deployment across heterogeneous model fleets, since STG worked on FLUX, SDXL, SD3, PixArt-α, and the fast LCM model without modification.
• Adaptable safety policies, because the safety function g is pluggable — NudeNet for nudity, CLIP-based scoring for violence or artist style — letting operators swap definitions of "safe" without retraining.

Industry relevance: The method requires no curated safety-annotated datasets and no fine-tuning compute, and it only needs access to mean predicted images at intermediate timesteps, which most diffusion frameworks already expose. That makes it cheap to bolt onto existing pipelines, including accelerated few-step samplers like LCM.

Future Directions

• Resolving the truncated artist-style results. Table 4's STG rows are incomplete in the provided content, so the paper's full claim about outperforming baselines on artist-style removal cannot be verified from what is shown here.

• Tuning the τ and γ hyperparameters. The reported results show a clear DSR/PP/FID trade-off as τ moves from 0.20 to 0.16, but the paper does not report an automatic way to pick these per deployment.

• Extending beyond the three tested safety categories. The paper covers nudity, violence, and artist-style removal; whether the same pluggable safety function works for bias, hate symbols, or other harms is not reported.

• Closing the residual gap in adversarial robustness. The paper notes that filtering-based methods are vulnerable to adversarial prompts, and its own tests include Ring-A-Bell, SneakyPrompt, and Concept Inversion — but the highest DSR values reported (e.g., 0.80 on LCM) still leave a meaningful fraction of unsafe outputs.

Target Audience

Researchers and engineers working on diffusion-model safety, content moderation, or controllable generation who already understand score-based diffusion and classifier guidance. Practitioners deploying text-to-image systems at scale will find the training-free and backbone-agnostic properties most useful. Readers looking for a beginner-level introduction to diffusion safety will need background in score functions, classifier-free guidance, and text embeddings to follow the theoretical sections.

Authors’ abstract

Text-to-image models have recently made significant advances in generating realistic and semantically coherent images, driven by advanced diffusion models and large-scale web-crawled datasets. However, these datasets often contain inappropriate or biased content, raising concerns about the generation of harmful outputs when provided with malicious text prompts. We propose Safe Text embedding Guidance (STG), a training-free approach to improve the safety of diffusion models by guiding the text embeddings during sampling. STG adjusts the text embeddings based on a safety function evaluated on the expected final denoised image, allowing the model to generate safer outputs without additional training. Theoretically, we show that STG aligns the underlying model distribution with safety constraints, thereby achieving safer outputs while minimally affecting generation quality. Experiments on various safety scenarios, including nudity, violence, and artist-style removal, show that STG consistently outperforms both training-based and training-free baselines in removing unsafe content while preserving the core semantic intent of input prompts. Our code is available at https://github.com/aailab-kaist/STG.