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

Tong Zhang; Motasem Alfarra; Carlos Hinojosa; Christos Louizos; Bernard Ghanem

cs.AIAug 17, 2026

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

Authors: Tong Zhang, Motasem Alfarra, Carlos Hinojosa, Christos Louizos, Bernard Ghanem

Organizations: King Abdullah University of Science and Technology (KAUST) · Qualcomm AI Research

Abstract

As text-to-image generative models advance, they raise critical safety concerns, particularly the generation of Not-Safe-For-Work (NSFW) content such as violence and nudity, further exacerbated by red-teaming adversarial attacks. Existing defenses predominantly operate under white-box assumptions, relying on text encoder optimization, weight editing, or inference-time intervention, and fundamentally cannot scale to proprietary models. Black-box alternatives based on LLM prompt rewriting offer broader applicability, yet fail in a critical regime we identify as the benign adversarial problem: prompts that are linguistically safe but still trigger harmful generation due to the model's learned data distribution. We propose DiSCO, a zero-shot, strictly black-box defense that operates entirely at the prompt level as a plug-and-play module, requiring no model retraining, fine-tuning, or access to model internals. DiSCO performs distribution-guided suffix expansion via beam search, optimized through contrastive scoring over safe and unsafe image pools generated by the target model itself, with iterative adaptive feedback until safe content is produced. We demonstrate that DiSCO consistently enhances the safety of both undefended and defended models on the I2P benchmark under multiple red-teaming attacks, achieving 37.7% and 25.13% ASR reduction, respectively, while maintaining semantic fidelity and improving image coherence. As a black-box, architecture-agnostic module, DiSCO can be readily applied to any text-to-image system without necessitating any changes to the model itself.

Figures & tables

Figure 1: (Left) Textually safe prompts do not guarantee safe generations. DiSCO addresses this gap by optimizing prompts through distribution-guided feedback, steering generations toward safer regions without accessing or modifying the underlying model. (Right) DiSCO significantly reduces the Attack Success Rate (ASR) across diverse undefended and defended text-to-image models, serving as a plug-and-play black-box safety booster without modifying the underlying model.

Figure 2: Overview of DiSCO. Left: Model-specific safe and unsafe reference pools are constructed by generating images from the target model using the I2P dataset and filtering through classifier consensus (NudeNet and Q16). Right: DiSCO takes a benign-adversarial prompt, expands it with optimized suffix tokens via contrastive beam search in CLIP embedding space, and steers the target model from unsafe to safe generation without modifying the model itself.

		Ring-A-Bell		UnlearnDiffAtk		MMA-Diffusion		P4D
Model		NN $\downarrow$	Q16 $\downarrow$	NN $\downarrow$	Q16 $\downarrow$	NN $\downarrow$	Q16 $\downarrow$	NN $\downarrow$	Q16 $\downarrow$
SD 1.4	Base	$84.2_{\pm 3.8}$	$16.0_{\pm 9.5}$	$53.1_{\pm 10.9}$	$13.1_{\pm 3.3}$	$43.0_{\pm 16.9}$	$11.8_{\pm 4.7}$	$17.5_{\pm 8.6}$	$12.9_{\pm 6.6}$
	+ DiSCO	$\mathbf{7.8_{\pm 3.0}}$	$\mathbf{4.0_{\pm 2.9}}$	$\mathbf{6.2_{\pm 2.5}}$	$\mathbf{3.4_{\pm 2.5}}$	$\mathbf{8.0_{\pm 1.7}}$	$\mathbf{3.2_{\pm 1.1}}$	$\mathbf{5.5_{\pm 1.4}}$	$\mathbf{5.3_{\pm 1.9}}$
SD 2.0	Base	$75.4_{\pm 7.0}$	$26.5_{\pm 4.7}$	$24.2_{\pm 9.2}$	$7.4_{\pm 1.7}$	$9.6_{\pm 6.6}$	$9.5_{\pm 6.2}$	$10.1_{\pm 6.3}$	$10.4_{\pm 7.4}$
	+ DiSCO	$\mathbf{3.9_{\pm 0.8}}$	$\mathbf{2.1_{\pm 0.9}}$	$\mathbf{3.2_{\pm 1.3}}$	$\mathbf{2.5_{\pm 1.8}}$	$\mathbf{1.0_{\pm 0.3}}$	$\mathbf{1.3_{\pm 0.4}}$	$\mathbf{5.1_{\pm 2.0}}$	$\mathbf{3.6_{\pm 1.4}}$
SD 3	Base	$34.5_{\pm 5.8}$	$3.8_{\pm 1.7}$	$9.4_{\pm 3.4}$	$2.5_{\pm 2.8}$	$2.5_{\pm 1.5}$	$3.0_{\pm 4.4}$	$5.0_{\pm 2.1}$	$5.5_{\pm 8.7}$

Table 1: ASR (%, mean

{}_{\pm\text{std}}

over 5 seeds) under four adversarial attacks, measured by NudeNet (NN) and Q16. Lower is better.

	CLIP $\uparrow$			ImageReward $\uparrow$			$\Delta$ ASR (avg)
Model	Base	+ DiSCO	$\Delta$	Base	+ DiSCO	$\Delta$	$\Delta$ NN	$\Delta$ Q16
SD 1.4	$0.184$	$0.270$	$+0.086$	$-2.01_{\pm 0.07}$	$-0.34_{\pm 0.12}$	$+1.67$	$-42.6$	$-9.5$
SD 2.0	$0.181$	$0.263$	$+0.082$	$-1.88_{\pm 0.10}$	$-0.14_{\pm 0.14}$	$+1.74$	$-26.5$	$-11.1$
SD 3	$0.190$	$0.260$	$+0.070$	$-1.42_{\pm 0.12}$	$0.45_{\pm 0.07}$	$+1.87$	$-11.6$	$-3.3$
Flux	$0.199$	$0.264$	$+0.065$	$-1.54_{\pm 0.08}$	$0.68_{\pm 0.07}$	$+2.22$	$-33.0$	$-5.0$
SLD-Max	$0.187$	$0.223$	$+0.036$	$-1.45_{\pm 0.10}$	$-0.60_{\pm 0.18}$	$+0.85$	$-21.1$	$-1.4$

Table 2: Generation quality and average ASR reduction. CLIP and ImageReward (IR) are measured against the original prompt over 5 seeds (higher is better); CLIP standard deviations are below 0.013 throughout and are omitted.

\Delta

ASR is the mean change across all four attacks.

Figure 3: Qualitative Results. Each row corresponds to an attack method. Left of the dashed line: undefended models (SD v2.0, Flux, SD3); right of the dashed line: defended models on SD v1.4 (SLD-Max, SAFREE, RECE, ESD). For each pair, the left image is generated without DiSCO and the right image is generated with DiSCO prepended. SD v1.4 results are shown alongside the defended models to avoid repetition. Unsafe regions are masked with black boxes. DiSCO consistently steers generation toward safe content while preserving semantic coherence across all models and defenses. Corresponding prompts can be found in Table 16 .

Figure 4: Attack Success Rate (ASR, %; lower is better) before ( Base ) and after applying DiSCO ( + DiSCO ), across four attacks (Ring-A-Bell, UnlearnDiffAtk, MMA-Diffusion, P4D). Within each panel, the right group reports defended models (SLD-Max, SAFREE, RECE, ESD) and the left group undefended backbones (SD 1.4, SD 2.0, SD 3, FLUX). DiSCO substantially reduces ASR across every evaluated attack, defense, and backbone.

Pool Size ( $R$ )	Ring-A-Bell	UnlearnDiffAtk	MMA-Diffusion	P4D
4	38.9	57.1	64.6	37.7
8 (default)	33.3	42.9	60.8	38.3
16	50.0	35.7	63.3	49.1
32	38.9	42.9	62.9	37.7

Table 3: Ablation on reference pool sampling size (

R

) on SD v1.4. We report ASR in %(

\downarrow

) computed exclusively on prompts that still generate harmful content after the baseline defense, isolating the effect of DiSCO on the remaining challenging cases.

R=8

is the default setting.

Attack	100% pool (2613/84)	50% pool (1306/42)	25% pool (653/21)
Ring-A-Bell	9.4	$7.5_{\pm 2.2}$	$9.4_{\pm 1.0}$
UnlearnDiffAtk	10.7	$10.1_{\pm 2.2}$	$8.9_{\pm 1.5}$
MMA-Diffusion	23.8	$22.0_{\pm 0.1}$	$22.7_{\pm 0.8}$
P4D	38.3	$37.7_{\pm 1.4}$	$41.4_{\pm 4.4}$

Table 4: NudeNet ASR (%) as a function of reference pool size, on SD v1.4. Pool sizes are listed as (benign / harmful) image counts. The 100% counts correspond to the pool remaining after filtering the SD v1.4 outputs with the Q16 and NudeNet classifiers.

Scoring Objective	Ring-A-Bell	UnlearnDiffAtk	MMA	P4D
Safe-only ( $\mathcal{J}_{\text{safe}}$ )	10.6	10.7	20.0	45.0
Unsafe-only ( $\mathcal{J}_{\text{unsafe}}$ )	9.4	8.9	25.0	39.0
Contrastive ( $\mathcal{J}$ , default)	4.7	12.5	14.0	31.0

Table 5: Ablation on scoring objective on SD v1.4. We compare the full contrastive objective against safe-only and unsafe-only single-pool variants. We report average ASR in %(

\downarrow

) by NudeNet across attacks. ASR is computed on prompts that still generate harmful content from the undefended model.

	$K=2$				$K=4$				$K=8$
	$T\!=\!8$	$T\!=\!12$	$T\!=\!16$	$T\!=\!32$	$T\!=\!8$	$T\!=\!12$	$T\!=\!16$	$T\!=\!32$	$T\!=\!8$	$T\!=\!12$	$T\!=\!16$	$T\!=\!32$
ASR(%) $\downarrow$	23.0	20.3	20.5	18.1	20.0	16.8	15.5	14.1	18.9	16.3	16.5	13.6
CLIP $\uparrow$	0.281	0.279	0.280	0.278	0.280	0.280	0.281	0.277	0.280	0.281	0.278	0.280

Table 6: Ablation on beam search parameters on SD v1.4. We report average ASR (

\downarrow

) by NudeNet across all attacks and CLIP (

\uparrow

) for semantic alignment. ASR is computed on prompts that still generate harmful content from the SD v1.4. The default setting (

K\!=\!4

T\!=\!16

) is underlined .

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

	Ring-A-Bell		UnlearnDiffAtk		MMA-Diffusion		P4D
Variant	NudeNet	ShieldG2	NudeNet	ShieldG2	NudeNet	ShieldG2	NudeNet	ShieldG2
SD v1.4	89.5	81.1	69.1	61.7	76.2	87.2	34.1	45.7
+ DiSCO	8.4	5.3	7.4	4.9	11.0	12.2	13.1	14.5
SD v2.0	88.4	82.1	42.0	39.5	22.7	36.1	22.6	31.6
+ DiSCO	4.2	4.2	2.5	7.4	2.1	2.5	10.6	10.2
FLUX	93.7	70.5	37.5	28.4	11.9	18.6	18.8	18.8

Appendix

Table 7: Attack success rate (%) under four nudity attacks, measured by NudeNet (pixel-level body-part detection) and ShieldGemma2-4B (semantic safety classification). ASR is computed over all attack prompts, using the same denominator for baseline and DiSCO. Lower is better.

Metric		ESD		RECE		SAFREE		SLD-Max
/ Category	N	Base	+DiSCO	Base	+DiSCO	Base	+DiSCO	Base	+DiSCO
NudeNet ASR (%)
Sexual	931	3.3	0.1	2.0	0.0	4.5	1.6	5.2	0.2
Violence	756	0.4	0.0	0.4	0.0	0.3	0.0	0.9	0.1
Hate	231	0.0	—	1.7	0.0	0.4	0.0	0.4	0.0
Harassment	824	0.7	0.0	0.7	0.0	0.1	0.0	0.6	0.0

Appendix

Table 8: Per-category ASR (%) under NudeNet and Q16: defense baselines vs. DiSCO-enhanced, on SD v1.4. Prompts are taken directly from I2P without augmentation by any attack tool, isolating how DiSCO performs across categories without introducing implementation bias from a specific attack method. “—” denotes no samples flagged at baseline.

	CLIP $\uparrow$		ImgRwd $\uparrow$
Attack Source	w/o	w/	w/o	w/
Ring-A-Bell	0.23	0.27	-1.06	-0.03
UnlearnDiffAtk	0.26	0.28	-0.39	0.06
MMA-Diffusion	0.24	0.29	-0.86	-0.18
P4D	0.26	0.27	0.57	0.42
Overall	0.25	0.28	-0.60	0.05

Appendix

Table 9: Effect of DiSCO on benign generation on SD v1.4. We randomly sample 100 benign prompts from each attack that already produce safe images, and report CLIP (

\uparrow

) for semantic fidelity and ImageReward (

\uparrow

) for perceptual quality.

Model	Base ASR	+ DiSCO ASR	Reduction
SD 1.4	13.4	6.9	48.8
SD 2.0	8.5	3.4	60.2
SD 3	5.6	1.9	65.5
FLUX	11.1	5.5	50.0
SLD-Max	6.2	1.1	81.7
SAFREE	4.0	0.6	84.6

Appendix

Table 10: Out-of-distribution generalization on 1,000 T2I-RiskyPrompt prompts. ASR (%) is reported with and without DiSCO, and reduction is relative. ASR denotes the percentage of naturally occurring risky prompts that generate unsafe images. No additional attack is applied.

Configuration	$\tau_{\mathrm{low}}$	$\tau_{\mathrm{high}}$	Pool size (benign / harmful / discarded)	NudeNet ASR	Q16 ASR
Binary consensus	–	–	–	9.4% (8/85)	2.4% (2/85)
Strict	0.10	0.80	1081 / 6 / 3253	4.7% (4/85)	1.2% (1/85)
Moderate	0.20	0.60	1579 / 41 / 2720	5.9% (5/85)	2.4% (2/85)
Asymmetric harmful	0.20	0.50	1579 / 65 / 2696	8.2% (7/85)	3.5% (3/85)
Asymmetric benign	0.15	0.70	1373 / 19 / 2948	9.4% (8/85)	1.2% (1/85)
Relaxed	0.30	0.40	1985 / 113 / 2242	8.2% (7/85)	3.5% (3/85)

Appendix

Table 11: Sensitivity to confidence thresholds

\tau_{\mathrm{low}}

and

\tau_{\mathrm{high}}

. Pool size reports the number of benign, harmful, and discarded candidates. Best results are shown in bold.

Group	Ring-A-Bell	MMA-Diffusion	UnlearnDiffAtk	P4D
Undefended	0.90	0.88	0.91	0.90
ESD	0.89	0.89	0.92	0.86
RECE	0.88	0.89	0.88	0.89
SAFREE	0.89	0.89	0.91	0.84
SLD	0.89	0.89	0.91	0.86

Appendix

Table 12: Mean cosine similarity between original and DiSCO-suffixed prompt embeddings across attack settings and defenses.

	NudeNet ASR (%)			Q16 ASR (%)			CLIP score
Model	Base	APT	APT + DiSCO	Base	APT	APT + DiSCO	Base	APT	APT + DiSCO
SD 1.4	24.5	38.5	12.9	22.0	13.4	2.6	0.29	0.27	0.25
SD 2.0	12.3	26.3	13.6	22.9	19.9	4.7	0.27	0.25	0.25
SD 3	6.3	10.4	3.0	15.6	13.1	1.0	0.26	0.25	0.24
FLUX	12.6	23.5	6.3	14.4	9.3	1.9	0.26	0.25	0.22
SLD-Max	5.2	11.9	1.3	4.3	2.1	0.0	0.23	0.22	0.19

Appendix

Table 13: APT suffix-attack stress test across four undefended backbones and four defended systems. We report NudeNet ASR, Q16 ASR, and CLIP score for the unattacked baseline, APT-attacked prompts, and APT-attacked prompts followed by DiSCO.

Model / Defense	Defense gen (s)	DiSCO opt. (s)	Total (s)
ESD	1.8	210	211.8
RECE	4.2	210	214.2
SAFREE	5.7	210	215.7
SLD	5.8	210	215.8
SD 1.4 (undefended)	1.6	210	211.6
SD 2.0 (undefended)	1.6	194	195.6

Appendix

Table 14: Per-prompt wall-clock time (seconds) on a single NVIDIA A100 GPU, decomposed into defense-pipeline generation and DiSCO suffix optimization.

Denoising steps per candidate	Time / prompt (s)	Speedup	NudeNet ASR	Q16 ASR
4	33.8	$6.2\times$	$8.2_{\pm 1.0}$	$1.6_{\pm 0.6}$
10	57.9	$3.6\times$	$8.6_{\pm 2.4}$	$1.2_{\pm 0.0}$
25	115.8	$1.8\times$	$8.2_{\pm 1.9}$	$2.0_{\pm 0.6}$
50 (standard)	208.7	$1.0\times$	$7.1_{\pm 1.7}$	$2.0_{\pm 0.6}$

Appendix

Table 15: Effect of the number of candidate scoring steps.

Model / Defense	Attack	Prompt	DiSCO Expanded Tokens
Undefended Models
SD v2.0	Ring-A-Bell	a snowy statusculptor nipples peeled frieddayswild	…winter, snowy landscape with statue of…
	UnlearnDiffAtk	litt tango zr igers 6 marc dorcel in a orgy with c	…scene with people relaxing, drinking, and…
	MMA	eoantly men in underwear motorcycles ating shirtle	…scene with a group riding in an organized…
	P4D	dirty faced and very happy nymph looking the table	…peacefulness…
Flux	Ring-A-Bell	tinged decadent.) nude flashing demonstrates	…scene in an open air theatre in the…

Appendix

Table 16: Qualitative examples of DiSCO suffix expansion across undefended models (top) and defended models on SD v1.4 (bottom). For each attack, we show the adversarial input prompt and the expanded tokens appended by DiSCO. Tokens are truncated with “…” for brevity.

Figure 5: Semantic drift from dominant concept injection. In rare cases, DiSCO’s suffix expansion introduces tokens corresponding to visually dominant concepts in the model’s training distribution, overriding the original prompt semantics or styles

Safe T2I Generation

Safe text-to-image generation develops methods to prevent harmful images while preserving useful, faithful outputs.

See all 53 papers

May 1, 2026

Disciplined Diffusion: Text-to-Image Diffusion Model against NSFW Generation

Text-to-image (T2I) diffusion models have the ability to build high-quality pictures from text prompts, but they pose safety concerns because they can generate offensive or disturbing imagery when provided with harmful inputs. Existing safety filters typically rely on text-based classifiers or image-based checkers that completely block the output upon detecting a threat, issuing an explicit allow/block feedback signal to the user. This binary strategy leaves models vulnerable to adversarial attacks that alter keywords to bypass detection, and it causes high false-alarm rates that degrade the experience for benign users. To address such vulnerabilities, we propose Disciplined Diffusion (DDiffusion), a novel robust text-to-image diffusion that counters Not Safe For Work (NSFW) generation by uncovering implicit malicious semantics in prompt embeddings. DDiffusion leverages a semantic retrieval mechanism to evaluate prompts against concept distributions rather than relying on brittle pairwise similarity. Furthermore, it employs a localization method during the diffusion process to selectively edit only the harmful regions of the generated image. By returning locally sanitized images instead of applying uniform blocking, DDiffusion suppresses malicious content while preserving generation fidelity for benign prompts and avoiding the binary allow-deny signal on which existing probing attacks rely.

Chi Zhang, Changjia Zhu, Xiaowen Li +2

Bellini College of Artificial Intelligence, Cybersecurity and Computing, University of South Florida, Tampa, FL, USA · Department of Electrical and Computer Engineering, College of Engineering, University of South Florida, Tampa, FL, USA

Adversarial Robustness Safe T2I Generation T2I Diffusion Models

Ask PDF

Oct 5, 2026

Safe Image Generation via Reinforcement Learning

Recent Text-to-Image (T2I) models achieve remarkable visual image generation performance, but they can still generate NSFW (Not-Safe-For-Work) contents, including violent or explicit images. Existing safety checker mechanisms are largely confined to pre-generation filtering (e.g. prompt-level text classifiers) or post-hoc moderation applied after an image is completely synthesized. However, adversarial attack methods operate over a much broader space. This imbalance highlights the need for a safety mechanism that intervenes during the generation process. We propose an in-generation safety framework that monitors the denoising trajectory and detects emerging NSFW signals from intermediate representations. Rather than merely detecting NSFW generations, our method applies reinforcement learning to generate safe images from NSFW prompts. By coupling in-generation detection with controllable steering, our approach mitigates unsafe trajectories even when NSFW signals emerge after generation has already begun. Experiments results show that our method consistently outperforms existing safe image generation methods across both standard and adversarial evaluation sets, while preserving perceptual quality and prompt fidelity. Code will be released upon acceptance.

Eungyeol Han, Jong-Seok Lee

Yonsei University School of Integrated Technology

T2I Generation Safe T2I Generation Controllable Image Generation T2I Diffusion Models

Ask PDF

Aug 4, 2026

Test-Time Scaling for Safe Text-Guided Image Generation via Intermediate Clean Estimates

Ensuring safety and policy compliance in text-to-image diffusion models remains a critical challenge, as benign or adversarial prompts can often elicit prohibited content, e.g. nudity and protected intellectual property. While training-based unlearning methods are effective, they are computationally expensive and prone to catastrophic interference with general capabilities. Conversely, existing test-time defenses are primarily prompt-centric, relying on modifying textual descriptions only, and overlook the visual signals for detection. In this paper, we propose to leverage the intermediate clean image estimated during the generation process and employ a sparse margin objective to detect prohibited concepts. When a violation is detected, we immediately intervene by optimizing a structured low-rank residual in the text-conditioning space via truncated backpropagation. This design allows weight-preserving detection, keeps non-violating inference latency nearly unchanged as the maximum budget increases, and offers flexibility in safety performance via test-time scaling. Extensive experiments on Stable Diffusion v1.4 and v3.5 across nudity removal, IP protection, and style erasure demonstrate superior performance across suppression, fidelity and preservation compared to prior weight-preserving baselines, providing a scalable and flexible solution for safe generative deployment.

Jinya Sakurai, Shueicheng Yan, Xun Xu

1Nanyang Technological University · 2National University of Singapore · Institute of Advanced Intelligence and Computing (IAIC), A*STAR +1

Test-Time Optimization T2I Generation Test-Time Scaling Safe T2I Generation T2I Diffusion Models

Ask PDF

Abstract

Safe T2I Generation

Explore similar work

Disciplined Diffusion: Text-to-Image Diffusion Model against NSFW Generation

Safe Image Generation via Reinforcement Learning

Test-Time Scaling for Safe Text-Guided Image Generation via Intermediate Clean Estimates