💻 computer science

Modular Energy Steering for Safe Text-to-Image Generation with Foundation Models

This paper proposes a modular, training-free inference-time framework that leverages gradient feedback from frozen vision-language foundation models as energy-based supervisory signals to steer text-to-image generation toward safety without compromising image quality or model scalability.

Original authors: Yaoteng Tan, Zikui Cai, M. Salman Asif

Published 2026-04-03

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Yaoteng Tan, Zikui Cai, M. Salman Asif

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 you have a magical, super-talented artist named Diffusion. This artist can paint any picture you describe in words. If you say, "A cat on a skateboard," they paint it instantly. If you say, "A castle made of clouds," they paint that too.

But there's a problem: Diffusion is too talented. If you ask for something dangerous, illegal, or just plain inappropriate (like a picture of a famous celebrity without their permission, or something explicit), Diffusion will happily paint it, too. They don't have a "moral compass" built-in; they just follow orders.

The Old Ways: The "Brute Force" Approach

Previously, people tried to fix this in two messy ways:

The "Re-Training" Method: They tried to teach Diffusion new rules by showing it thousands of "bad" pictures and saying, "Don't do this!" It's like trying to retrain a dog that has already learned to fetch. It often makes the dog forget how to fetch good things, or it just doesn't learn the new rules well enough.
The "Bouncer" Method: They put a security guard at the door (before the artist starts) and another at the exit (after the painting is done). But clever tricksters (hackers) can whisper secret codes to the artist that the bouncer doesn't understand, bypassing the security.

The New Idea: The "Smart Spotter" (Modular Energy Steering)

This paper introduces a clever, new way to control the artist without retraining them and without needing a massive library of bad pictures.

Think of it like this:

1. The Artist (The Generator)
Diffusion is the artist. They are frozen in time; we don't change their brain or their training. They just keep painting based on their original skills.

2. The Spotter (The Foundation Model)
We bring in a second character: The Spotter. This is a different AI (like CLIP or a VLM) that is already an expert at understanding pictures. It knows what a "naked person" looks like, or what "Elon Musk" looks like. It's like a seasoned art critic who has seen millions of images.

3. The Dance (The Process)
The artist doesn't paint the whole picture in one go. They paint it in steps, starting with a blurry mess and slowly refining it into a clear image.

Step 1: The artist paints a blurry version of the image.
Step 2: Before the artist moves to the next step, The Spotter looks at that blurry version and asks, "Hey, does this look like something we shouldn't be making?"
Step 3: If the Spotter says, "Yes, that looks like a celebrity," it gently pushes the artist's hand in a different direction. It's like a gentle nudge saying, "No, go this way instead."
Step 4: The artist continues painting, now slightly steered away from the bad idea, but still following your original request.

The "Energy" Analogy

The paper calls this "Energy Steering." Imagine the artist is a ball rolling down a hill.

The Goal: You want the ball to roll to a specific spot (a safe, beautiful image).
The Danger: There are "potholes" (unsafe concepts like nudity or specific people) on the hill.
The Old Way: You try to pave over the potholes (retraining), which is hard and might ruin the road.
The New Way: You use a magnet (The Spotter). As the ball rolls, the magnet senses the pothole and gently pulls the ball away from it, guiding it safely to the destination. The road itself never changes; you just added a smart guide.

Why is this cool?

It's Plug-and-Play: You don't need to rebuild the artist. You just plug in the Spotter.
It's Flexible: If you want to block "Elon Musk" today, you tell the Spotter. If you want to block "Guns" tomorrow, you just update the Spotter's list. No need to retrain the artist.
It's Strong: Because the Spotter is looking at the image while it's being made, it catches bad ideas early, even if you try to trick the artist with weird words.
It Keeps Quality: Since we aren't messing with the artist's brain, the pictures still look amazing. They don't get blurry or weird just because we are being safe.

In a Nutshell

This paper is about giving a super-powerful AI artist a smart, real-time supervisor. The supervisor doesn't change how the artist thinks; it just whispers, "Hey, that looks a bit risky, let's try a different angle," at every single step of the painting process. This keeps the art safe, high-quality, and flexible, without needing to rebuild the whole system.

1. Problem Statement

Text-to-Image (T2I) generative models, while powerful, pose significant safety risks, including the generation of Not-Safe-For-Work (NSFW) content, copyrighted material, and non-consensual depictions of real individuals. Existing safety mechanisms suffer from critical limitations:

Input/Output Filters: Prone to high false positives and easily bypassed by adversarial prompts.
Model Fine-tuning/Unlearning: Requires curated "forget" and "retain" datasets, often degrading the model's general generation capabilities and scalability.
Existing Inference-Time Steering: Often limited to prompt embedding refinement (suboptimal) or relies on curated image datasets for guidance.

The core challenge is to develop a training-free, dataset-free, and scalable safety mechanism that effectively suppresses specific unwanted concepts (e.g., nudity, specific celebrities) without modifying the underlying generator's weights or sacrificing generation quality.

2. Methodology

The authors propose a Modular Energy Steering Framework that operates at inference time. The core insight is that frozen, off-the-shelf Vision-Language Models (VLMs) like CLIP or large VLMs already encode rich semantic knowledge that can be repurposed as semantic energy estimators to guide the generation process.

A. Energy-Based Formulation

The method frames safety control as an energy-based sampling problem. Instead of modifying the generator, it introduces a safety guidance term derived from the gradient of a safety energy function.

Objective: Modify the generative score field to penalize the presence of an unsafe target concept $y^*$ .
Formulation: The conditional score is adjusted as:
$\nabla_{x_t} \log p(x_t | c \neq y^*) \approx \nabla_{x_t} \log p(x_t | c) - \lambda \nabla_{x_t} \log p(y^* | x_t)$
Here, $E_{safety}(x) = \log p(y^* | x)$ acts as the energy penalty, and $\lambda$ controls the steering strength.

B. Foundation Models as Energy Estimators

The framework utilizes frozen foundation models to estimate the probability $p(y^* | x)$ at each sampling step without training:

Clean Latent Estimation: At each diffusion step $t$ , the method estimates the clean latent data $\hat{x}_{0|t}$ from the noisy latent state.
Semantic Energy Calculation:
- CLIP-based: Uses a Log-Mean-Exponential (LME) function to aggregate similarity scores between the image and a set of positive (unsafe) and negative (safe) text prompts. The energy is the difference between the positive and negative logit scores.
- VLM-based: Formulates safety checking as a next-token prediction task. The VLM is queried with a binary question (e.g., "Does this image contain nudity?"). The logits for "Yes" vs. "No" tokens are used to compute the energy gradient.
Gradient Injection: The gradient of this energy estimate is backpropagated through the frozen foundation model to the image latent space, providing a steering signal that pushes the generation away from the unsafe region.

C. Algorithm Flow (Algorithm 1)

The process integrates into the standard ODE/SDE sampling loop:

Initialize latent noise.
For each timestep, estimate the clean latent $\hat{x}_{0|t}$ .
Pass $\hat{x}_{0|t}$ to the plug-in estimator (CLIP/VLM) to compute the safety logit $p_t$ .
If $p_t$ indicates a high probability of the unsafe concept, apply a gradient update to the latent state using the safety guidance term.
Proceed to the next timestep.

3. Key Contributions

Training-Free & Dataset-Free: The method requires no model weight updates and no domain-specific image datasets. It only requires a user-defined textual "blacklist" of concepts to suppress.
Modular Energy Steering: It formulates safety as an energy-based sampling problem, allowing for principled integration with both Diffusion Models (DMs) and Flow-Matching Models (FMs).
Scalability: The framework naturally scales to multi-concept steering. Adding new targets (e.g., new celebrities) only requires expanding the text prompt set in the estimator, avoiding the catastrophic forgetting issues associated with continual learning.
Robustness: It demonstrates state-of-the-art robustness against adversarial red-teaming prompts designed to bypass safety filters.

4. Experimental Results

The authors evaluated the framework on Stable Diffusion v1.4 and Stable Diffusion v3 (Flow Matching).

NSFW Suppression:
- Achieved significantly lower Attack Success Rates (ASR) across four red-teaming benchmarks (P4D, Ring-A-Bell, MMA-Diffusion, UnlearnDiffAtk) compared to prior methods like SAFREE, SafeDenoiser, and CURE.
- For example, on the MMA-Diffusion benchmark, the proposed method reduced ASR to 0.155 (CLIP) and 0.115 (VLM) compared to the original model's 0.952.
Identity Steering:
- Successfully suppressed specific celebrity identities (e.g., Elon Musk, Taylor Swift) while preserving the generation of non-targeted identities.
- Demonstrated effective joint steering of up to 9 identities simultaneously without degrading the suppression of non-targets.
Quality Preservation:
- Evaluated on the COCO-30K dataset. The method maintained high CLIP scores (text-image alignment) and competitive FID scores, indicating that the safety steering did not degrade the quality of benign, non-targeted generations.
Scalability to Advanced Models:
- The framework was successfully applied to SD-v3 (a Flow Matching model with DiT architecture), proving its architecture-agnostic nature.

5. Significance and Implications

Practical Deployment: By eliminating the need for retraining or large curated datasets, this approach offers a highly practical solution for deploying safe T2I models in dynamic environments where safety requirements (e.g., new celebrity bans) change frequently.
Leveraging Foundation Models: It pioneers the use of frozen foundation models not just as classifiers, but as semantic energy estimators that provide real-time, gradient-based feedback to generative processes.
Trade-off Management: While the method introduces a modest computational overhead (approx. 2.7s/image for CLIP, 10s for VLM on SD-v1.4), the authors argue this is a favorable trade-off for achieving robust, training-free safety control.

In conclusion, this paper presents a principled, modular, and highly effective framework for safe text-to-image generation that leverages the semantic understanding of existing foundation models to steer generation away from unsafe concepts without compromising the underlying model's utility.