Safeguarding Text-to-Image Generation via Inference-Time Prompt-Noise Optimization

TL;DR

The paper tackles unsafe image generation in text-to-image diffusion models by introducing Prompt-Noise Optimization (PNO), a training-free inference-time method that jointly optimizes the continuous prompt embedding $c$ and the injected noise trajectory during DDIM sampling to minimize a toxicity objective while preserving prompt-image alignment. It formulates the problem as minimizing $\mathcal{L}_{tox}(x_0) + \lambda \mathcal{L}_{reg}(x_T, z_1, \dots, z_T)$ subject to $x_0 = \text{DDIM}(c, x_T, z_1, \dots, z_T)$ and solves it with gradient-based updates, enabling robust defense against adversarial prompts without model fine-tuning. Extensive experiments on I2P and Ring-a-bell show state-of-the-art safety performance and resilience, with PNO maintaining competitive image quality (CLIP, HPSv2, PickScore) and a favorable safety-alignment Pareto frontier. The approach is flexible, extensible to other safety evaluators, and requires modest inference-time overhead, making it a practical tool for trustworthy diffusion-based generation.

Abstract

Text-to-Image (T2I) diffusion models are widely recognized for their ability to generate high-quality and diverse images based on text prompts. However, despite recent advances, these models are still prone to generating unsafe images containing sensitive or inappropriate content, which can be harmful to users. Current efforts to prevent inappropriate image generation for diffusion models are easy to bypass and vulnerable to adversarial attacks. How to ensure that T2I models align with specific safety goals remains a significant challenge. In this work, we propose a novel, training-free approach, called Prompt-Noise Optimization (PNO), to mitigate unsafe image generation. Our method introduces a novel optimization framework that leverages both the continuous prompt embedding and the injected noise trajectory in the sampling process to generate safe images. Extensive numerical results demonstrate that our framework achieves state-of-the-art performance in suppressing toxic image generations and demonstrates robustness to adversarial attacks, without needing to tune the model parameters. Furthermore, compared with existing methods, PNO uses comparable generation time while offering the best tradeoff between the conflicting goals of safe generation and prompt-image alignment.

Figures (13)

• Figure 1: The workflow of Prompt-Noise Optimization (PNO). (Left) demonstrates the use case of PNO, where the user provides a potentially toxic prompt to the model, and the model generates an image that is evaluated by a toxicity score, which is used to update the noise trajectory and prompt embedding. (Right) shows the detailed process of PNO, where the optimization process jointly optimizes the prompt embedding $c$ and the noise trajectory $\{\bm{x}_T, \bm{z}_T, \dots, \bm{z}_1\}$ to minimize the toxicity score of the generated image.
• Figure 2: Tradeoff between CLIP Score $\uparrow$ and toxicity $\downarrow$ PNO (with different learning rates specified in the parenthesis) offers superior tradeoffs between image safety and prompt alignment, when compared with state-of-the-art T2I safety mechanisms.
• Figure 3: Illustration of the optimization landscape of the Prompt-Noise Optimization process, plotted over the prompt embedding space. Higher scores (lighter background) indicate safer outputs. Jointly optimized embedding stays closest to the original prompt, while direct prompt modification causes greatest deviation.
• Figure 4: Percentage of non-toxic outputs $\uparrow$ on I2P Dataset: Q16 Evaluations. The center of the circle represents all generated images are toxic, while the outer most frontier means all generations are safe. PNO achieves almost 100% safe percentage, outperforming state-of-the-art baselines.
• Figure 5: Demonstration of PNO iterations. (Upper) and (Lower) are images generated from different prompts. Specific prompts included in Sec. \ref{['subsec:supp_demoiter']}. PNO is able to substantially reduce image toxicity at each step.

Theorems & Definitions (1)

• Lemma 1: wainwright2019high