DiSCO: Defending text-to-image generation through distribution-guided contrastive prompt optimization

The Problem: Black-Box Safety for Image Generators

Text-to-image models like DALL-E and Stable Diffusion can produce harmful content—violence, nudity, and other NSFW material—especially when attacked with adversarial prompts. Existing defenses require white-box access (modifying the model's weights or internals) or rely on LLM-based prompt rewriting, which fails on a subtle but dangerous class of attacks: prompts that sound innocent but still trigger harmful outputs because of biases learned from training data. This is the "benign adversarial problem," and current LLM rewriting misses it entirely.

The authors propose DiSCO, a black-box defense that runs as a pure prompt-optimization layer with no access to model internals, weights, or training data.

How DiSCO Works

DiSCO expands the user's input prompt by appending safe suffixes via beam search. Rather than rewriting the prompt with an LLM, it:

1. Generates candidate expansions using beam search to add multiple possible suffix continuations to the original prompt.

2. Tests each expansion by sending it to the target image model and collecting both safe and unsafe images as reference pools. For each candidate suffix, the system generates sample images and scores them using contrastive learning—comparing how similar they are to known safe versus unsafe content.

3. Iterates with adaptive feedback. The contrastive scoring identifies which suffix best steers the model toward safe outputs. The system repeats until the final image passes safety checks.

The key insight is using the target model itself as an oracle: rather than assume an external LLM knows what the model will do, DiSCO learns from actual outputs what prompt modifications work.

Why This Matters

DiSCO achieves 37.7% attack success rate (ASR) reduction on undefended models and 25.13% ASR reduction on models that already have defenses, evaluated on the I2P benchmark (a standard red-teaming dataset). Critically, it does this while preserving semantic fidelity—the intended meaning of the prompt survives—and often improves image coherence.

The method is genuinely plug-and-play: it works on proprietary models (OpenAI's, closed-source competitors) without querying internals, requires no retraining or fine-tuning, and adds minimal inference overhead. This is a substantial advantage over white-box defenses, which cannot ship as third-party modules and require vendor cooperation.

The authors demonstrate that their approach handles the benign adversarial regime that simple LLM rewriting cannot: prompts like "a woman in minimal clothing" that are grammatically harmless but exploit learned model biases. By iteratively testing real outputs from the target model, DiSCO discovers and adjusts for these hidden failure modes.

The trade-off is computational cost—generating multiple candidate images per query adds latency—but for safety-critical applications (content moderation at deployment time), this is often acceptable.