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

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 \textit{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.

한 줄 요약

DiSCO는 텍스트-이미지 생성 모델의 안전성을 향상시키는 블랙박스 방식의 프롬프트 최적화 모듈로, ASR을 23.6%에서 2.4%까지 감소시킨다.

핵심 기여도

**benign adversarial 문제**를 정식화: 언어적으로 안전한 프롬프트가 모델의 학습 분포로 인해 해로운 이미지를 생성할 수 있음.
**DiSCO 모듈** 제안: 트레이닝 없이 블랙박스 방식으로, 프롬프트 레벨에서 분포 가이드 기반 대비적 최적화를 수행.
**I2P 벤치마크에서 4가지 공격**에 대해 평균 ASR 23.6% → 2.4% (NudeNet 기준), 8.3% → 1.7% (Q16 기준) 감소.
**모델 아키텍처와 무관하게 적용 가능**하며, 기존 방어 방법과 결합 시 안전성 향상 효과 증가.

핵심 아이디어

기존 방어 방식은 모델 내부 파라미터 수정이나 텍스트 수준만 고려해, 모델의 학습된 시각적 분포와의 관계를 무시했다. DiSCO는 이 문제를 **분포 정렬 문제**로 접근해, 프롬프트를 최적화하여 모델의 안전한 생성 영역으로 유도한다. 이는 **모델 자체의 안전/비안전 이미지 풀**을 기반으로 대비적 점수를 계산하고, **빔 서치를 활용한 접미사 확장**을 통해 반복적 피드백을 주입하는 방식이다. 핵심 아이디어는 모델의 내재적 분포를 활용해, **언어적으로는 안전하나 시각적으로는 위험한 프롬프트**(benign adversarial)를 효과적으로 방어하는 것이다.

기술적 접근법

**DiSCO 모듈**: 트레이닝 없이 블랙박스 방식으로 작동.
**접미사 확장**: LLaMA-3-8B 기반의 가벼운 언어 모델을 사용, 빔 서치(beam width = 4, suffix length = 16)로 프롬프트를 확장.
**대비적 점수 계산**: CLIP 임베딩 공간에서 안전/비안전 이미지 풀(각각 8개 샘플)을 기반으로 계산.
**샘플링 기준**: NudeNet과 Q16 분류기의 일치 여부로 샘플 선택 (일치하지 않는 샘플은 제거).
**적응적 피드백**: 생성된 이미지의 해로운 정도에 따라 최적화 목적 함수를 조정.

주요 결과

**I2P 벤치마크**에서 DiSCO는 4가지 공격(Black-box, White-box 포함)에 대해 안전성을 향상.
**NudeNet 기준**: 평균 ASR 23.6% → 2.4% (37.7% 감소), **Q16 기준**: 8.3% → 1.7% (25.13% 감소).
**모델 아키텍처별 평가**: UNet 기반 (SD v1.4, SD v2.0) 및 DiT 기반 (Flux, SD 3) 모델 모두에서 효과.
**생성 품질 유지**: CLIP 점수와 ImageReward를 통해 의미적 일관성과 시각 품질이 유지되거나 향상됨.

의의 및 한계

DiSCO는 모델 재학습 없이도 블랙박스 방식으로 작동하며, **모델 아키텍처와 무관하게 적용 가능**하다는 점에서 실용적 가치가 높다. 특히, 기존 방어 방법과 결합 시 **안전성 향상 효과가 누적**되며, 다양한 공격 시나리오에서 안정적인 성능을 보인다. 그러나 **APT와 같은 강력한 적응형 공격**에 대해서는 완전한 방어가 어려울 수 있다. 예를 들어, SD 2.0에서 ASR은 26.3% → 13.6%로 감소했으나, **기본값 12.3%보다 높아** 일부 해로운 접미사가 완전히 중화되지 못한 것으로 보인다. 이는 DiSCO의 **적응적 공격에 대한 한계**를 드러내며, 보다 강력한 공격을 통해 안정성 한계를 정확히 평가할 필요가 있다.

실용적 활용

DiSCO는 **모든 텍스트-이미지 생성 시스템**에 즉시 적용 가능한 모듈로, 특히 **프로퍼티 모델**(사유 모델)이나 **사전 학습된 모델**에 적합하다. 기업이 모델 내부에 접근하지 못하는 경우에도 **안전성 향상을 위한 외부 모듈**로 활용할 수 있다. 또한, **기존 방어 방법과 결합**하여 보다 강력한 방어 체계를 구축하는 데 유용하다. 예를 들어, 인공지능 기반 콘텐츠 생성 플랫폼이나, 디지털 마케팅, 게임 산업 등에서 **NSFW 콘텐츠 생성 방지**에 활용 가능하다.