Universal Prompt Optimizer for Safe Text-to-Image Generation (2024)

Abstract: Text-to-Image (T2I) models have shown great performance in generating images based on textual prompts. However, these models are vulnerable to unsafe input to generate unsafe content like sexual, harassment and illegal-activity images. Existing studies based on image checker, model fine-tuning and embedding blocking are impractical in real-world applications. Hence, we propose the first universal **p**rompt **o**ptimizer for **s**afe T2**I** (**POSI**) generation in black-box scenario. We first construct a dataset consisting of toxic-clean prompt pairs by GPT-3.5 Turbo. To guide the optimizer to have the ability of converting toxic prompt to clean prompt while preserving semantic information, we design a novel reward function measuring toxicity and text alignment of generated images and train the optimizer through Proximal Policy Optimization. Experiments show that our approach can effectively reduce the likelihood of various T2I models in generating inappropriate images, with no significant impact on text alignment. It is also flexible to be combined with methods to achieve better performance. Our code is available at [https://github.com/wzongyu/POSI](https://github.com/wzongyu/POSI).

## introduction text-to-image (t2i) generation has gained substantial attention, leading to the emergence of many powerful models like glide @xcite , imagen @xcite , dall-e 2 @xcite , stable diffusion (sd) @xcite and vq-diffusion @xcite . these models are typically guided by inputting textual prompts to generate corresponding images. their ability to generate high-quality images from textual descriptions can facilitate various real-world applications, such as book illustrations, brand identity design, and design of game environments. despite their wide adoption, t2i models are also used by malicious users to generate unsafe content like sexual, harassment and illegal-activity images @xcite . although t2i models have been developed to generate safe content through filtering training data foot_0 or robust learning @xcite in the developing stage, recent works @xcite @xcite have shown that t2i models are still vulnerable to prompt perturbations, which make these models disrupted to generate inappropriate content. therefore, further defensive methods have been proposed, such as rejecting the generation of toxic images by detection @xcite , guiding the model to generate safe content through embedding blocking @xcite or fine-tuning @xcite . though methods to some extent alleviate the generation of harmful content, directly rejecting images can affect user usability, and optimizing the model would require obtaining the internal structure of t2i models and could be time-consuming, which lacks universality in real-world scenarios. recent works @xcite @xcite show that prompt engineering has been used to improve the performance of prompt-based models. some works @xcite enhance the capability of large language models through prompt modification in text generation, math problem solving, code generation and etc. other works @xcite use prompt optimization to improve t2i generation. prompt optimization is regarded as a general method to improve prompt-based models without changing the parameters of the corresponding models. thus, it is promising to improve the safety of t2i models through prompt engineering, i.e., revising the toxic prompt so that t2i models can generate safe images that preserve the appropriate portion of the user's prompt (preserve the appropriate semantics of the toxic prompt). figure 1 shows the case of a toxic prompt with the corresponding modified one by our fine-tuned model. as shown in the figure, with the revised prompt, the generated image does not contain unsafe content meanwhile is semantically close to the image generated by the toxic prompt. however, there is no work in this promising direction. therefore, in this paper, we study a novel problem of safe t2i generation with prompt engineering. we propose a prompt optimizer called posi that can guide the t2i models to generate safe and semantic-preserving contents without obtaining the structure of t2i model. there are several challenges in developing the safe optimizer: (1) the optimizer should be universal and not require access to the parameters of t2i models; (2) a corresponding unsafe-safe dataset is needed for training; t2i generation. t2i generation aims to generate high-quality images based on text descriptions. various models such as vae (kingma and welling, 2014), arm (van den oord et al., 2016), flowbased models (kingma and dhariwal, 2018) and gan (goodfellow et al., 2020) are proposed in image generation and made great process to this field. however, they suffer from limitations like poor image quality and missing or weak prompt-following ability. recently, diffusion models (dms) such as dall-e (ramesh et al., 2021), imagen (saharia et al., 2022), and sd @xcite have made exciting strides in t2i generation. these models significantly improved the performance of generating high-quality images from arbitrary text descriptions @xcite . ## proposed framework our framework is inspired by promptist @xcite . the difference is that our framework aims to produce safe prompts for t2i generation. after a user inputs a toxic prompt for the t2i generation, posi automatically outputs the modified prompt to avoid generating inappropriate images while pre-serving the appropriate portion of the user's prompt (i.e., maintaining text alignment). an illustration of the proposed framework is shown in figure 2 . due to the absence of a publicly available toxicclean prompts pair dataset, we first produce a set of toxic-clean prompt pairs in section 3.1. then we use them to conduct supervised fine-tuning (sft) in section 3.2 to give the model the basic ability to turn toxic prompts into clean prompts. sft can be considered a warm-up phase, hence the effectiveness of the supervised fine-tuned model is generally moderate. to enhance the model's performance, we further perform proximal policy optimization in section 3.4 to maximize the reward score we design in section 3.3. it can reduce the inappropriateness of the generated images while not significantly affecting the text alignment. next, we give the details. ## conclusion in this work, we study a novel problem of safe image generation via automatic prompt optimization. we propose a novel framework which can revise a toxic prompt to generate safe and semanticpreserving images for black-box t2i models. experimental results demonstrated the effectiveness of the proposed framework. in addition, our approach has good transferability and is flexible to be plugged into various t2i models.