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

Most defenses against unsafe image generation assume you can reach inside the model — editing weights, steering attention, or intercepting activations. That assumption breaks the moment you're dealing with a proprietary API. DiSCO sidesteps it entirely by working only at the prompt level, appending an optimized suffix that steers the generator toward safe outputs without touching anything internal.

The core insight is what the paper calls the benign adversarial problem: a prompt can pass every text-level safety check and still reliably produce harmful images, because safety depends on how the prompt maps to the model's learned visual distribution, not on whether the words look dangerous. LLM-based prompt rewriting — the standard black-box fallback — misses this entirely. DiSCO addresses it by treating defense as a distributional alignment problem.

The mechanism is concrete. DiSCO builds two reference pools by querying the target model on ordinary prompts and filtering outputs through classifier consensus: images both NudeNet and Q16 agree are safe go in one pool, images both agree are unsafe go in the other. Ambiguous cases are discarded. Given a problematic prompt, DiSCO then runs beam search using LLaMA-3-8B to expand a suffix token by token, scoring each candidate by its CLIP-space alignment to the safe pool minus its alignment to the unsafe pool. The contrastive objective matters: safe-only scoring averaged 21.6% attack success rate in ablations, unsafe-only 20.6%, but the combined objective reached 15.6%, with the gap widest on Ring-A-Bell and MMA-Diffusion.

Across 32 system-attack combinations and five seeds, average attack success rate drops from 23.6% to 2.4% under NudeNet and from 8.3% to 1.7% under Q16. Crucially, semantic fidelity doesn't degrade — CLIP alignment to the original prompt actually increases across all four tested architectures, and ImageReward improves too. The suffix expansion preserves prompt meaning: cosine similarity between original and suffixed prompt embeddings stays between 0.84 and 0.92 across all settings.

The computational cost is real. Under default settings, each prompt requires 64 candidate image generations for beam scoring. On an A100, that runs roughly 3.5 minutes per prompt for SD-based models. But the paper demonstrates a practical shortcut: scoring at 4 diffusion steps instead of 50 cuts that to about 34 seconds with negligible effect on defense quality, since CLIP operates on high-level semantics that emerge early in denoising. The optimized suffix can also be cached for repeated use.

One honest limitation surfaces in the ablations: suffix expansion occasionally injects tokens tied to visually dominant training concepts — popular characters, common animals — that can override the original prompt's intent. The contrastive objective reduces this but doesn't eliminate it. The paper flags explicit semantic preservation constraints as future work rather than a solved problem.