Latent Guard: a Safety Framework for Text-to-image Generation

TL;DR

Latent Guard addresses the safety challenges of text-to-image generation by shifting from direct input classification to identifying blacklisted concepts within a learned latent space atop the text encoder. It introduces a trainable Embedding Mapping Layer and a contrastive learning regime to align prompt representations with concept embeddings, enabling fast, test-time webhook-like blocking of unsafe prompts without retraining the diffusion model. The authors construct CoPro, a large dataset of unsafe and safe prompts around 723 concepts, and demonstrate strong performance and generalization on in-distribution and out-of-distribution concepts, outperforming several baselines across multiple evaluation metrics. The approach offers practical benefits in deployment due to low computational overhead, adaptability to changing safety requirements, and compatibility with existing T2I pipelines, while acknowledging limitations related to concept list coverage and potential distribution shifts. Overall, Latent Guard provides a robust, scalable framework for proactive content safety in text-to-image systems and lays groundwork for open-set, adaptable safety mechanisms in generative AI.

Abstract

With the ability to generate high-quality images, text-to-image (T2I) models can be exploited for creating inappropriate content. To prevent misuse, existing safety measures are either based on text blacklists, which can be easily circumvented, or harmful content classification, requiring large datasets for training and offering low flexibility. Hence, we propose Latent Guard, a framework designed to improve safety measures in text-to-image generation. Inspired by blacklist-based approaches, Latent Guard learns a latent space on top of the T2I model's text encoder, where it is possible to check the presence of harmful concepts in the input text embeddings. Our proposed framework is composed of a data generation pipeline specific to the task using large language models, ad-hoc architectural components, and a contrastive learning strategy to benefit from the generated data. The effectiveness of our method is verified on three datasets and against four baselines. Code and data will be shared at https://latentguard.github.io/.

Figures (8)

• Figure 1: Recent text-to-image generators are composed of a text encoder and a diffusion model. Their deployment without appropriate safety measures creates risks of misuse (left). We propose Latent Guard (right), a safety method designed to block malicious input prompts. Our idea is to detect the presence of blacklisted concepts on a learned latent space on top of the text encoder. This allows to detect blacklisted concepts beyond their exact wording, extending to some adversarial attacks too ("<ADV>"). The blacklist is adaptable at test time, for adding or removing concepts without retraining. Blocked prompts are not processed by the diffusion model, saving computational costs.
• Figure 2: Overview of Latent Guard. We first generate a dataset of safe and unsafe prompts centered around blacklisted concepts (left). Then, we leverage pretrained textual encoders to extract features, and map them to a learned latent space with our Embedding Mapping Layer (center). Only the Embedding Mapping Layer is trained, while all other parameters are kept frozen. We train by imposing a contrastive loss on the extracted embedding, bringing closer the embeddings of unsafe prompts and concepts, while separating them from safe ones (right).
• Figure 3: Embedding Mapping Layer. We combine MLPs and multi-head cross-attention to extract embeddings used for contrastive training.
• Figure 4: Inference. We block the T2I prompt if at least one cosine similarity between concepts and prompts embedding is larger than a pre-defined threshold.
• Figure 5: CoPro generation. For $\mathcal{C}$ concepts, we sample unsafe $\mathcal{U}$ prompts with an LLM as described in Section \ref{['sec:data-generation']}. Then, we create Synonym prompts by replacing $c$ with a synonym, also using an LLM, and obtaining $\mathcal{U}^\text{syn}$. Furthermore, we use yang2023mmadiffusion to replace $c$ with an <ADV> Adversarial text ($\mathcal{U}^\text{adv}$). Safe prompts $\mathcal{S}$ are obtained from $\mathcal{U}$. This is done for each ID and OOD data.