DREAM: Scalable Red Teaming for Text-to-Image Generative Systems via Distribution Modeling

TL;DR

DREAM is a scalable framework that models the distribution of unsafe prompts in text-to-image systems, enabling efficient large-scale red teaming to improve safety by identifying diverse problematic prompts.

Contribution

It introduces a probabilistic distribution modeling approach for red teaming, improving scalability, diversity, and effectiveness over prior prompt-level optimization methods.

Findings

Achieves state-of-the-art success rate in identifying unsafe prompts.

Enhances diversity of problematic prompts compared to existing methods.

Demonstrates robustness across various T2I models and safety filters.

Abstract

Despite the integration of safety alignment and external filters, text-to-image (T2I) generative systems are still susceptible to producing harmful content, such as sexual or violent imagery. This raises serious concerns about unintended exposure and potential misuse. Red teaming, which aims to proactively identify diverse prompts that can elicit unsafe outputs from the T2I system, is increasingly recognized as an essential method for assessing and improving safety before real-world deployment. However, existing automated red teaming approaches often treat prompt discovery as an isolated, prompt-level optimization task, which limits their scalability, diversity, and overall effectiveness. To bridge this gap, in this paper, we propose DREAM, a scalable red teaming framework to automatically uncover diverse problematic prompts from a given T2I system. Unlike prior work that optimizes prompts individually, DREAM directly models the probabilistic distribution of the target system's problematic prompts, which enables explicit optimization over both effectiveness and diversity, and allows efficient large-scale sampling after training. To achieve this without direct access to representative training samples, we draw inspiration from energy-based models and reformulate the objective into a simple and tractable form. We further introduce GC-SPSA, an efficient optimization algorithm that provides stable gradient estimates through the long and potentially non-differentiable T2I pipeline. During inference, we also propose a diversity-aware sampling strategy to enhance prompt variety. The effectiveness of DREAM is validated through extensive experiments, demonstrating state-of-the-art performance across a wide range of T2I models and safety filters in terms of both prompt success rate and diversity. Our code is available at https://github.com/AntigoneRandy/DREAM