发现扩散Transformer中注意力头对风险概念敏感,可无训练检测并抑制有害图像生成。
What Concepts Lie Within? Detecting and Suppressing Risky Content in Diffusion Transformers
- 通过注意力头向量量化文本词元敏感度,识别风险生成倾向。
- 在推理阶段动态追踪敏感度,按头级风险分数自适应抑制注意力权重。
- 无需训练即可有效防御性内容生成,兼容多种扩散Transformer模型。
文本到图像(T2I)模型的兴起引发对生成色情、暴力及受版权保护内容的担忧,亟需模型内嵌的有效防护机制。现有方法多针对早期U-Net架构,难以适配当前主流的扩散变压器(DiT)模型。其根本原因在于:DiT通过联合注意力将语义注入与视觉合成耦合,导致难以分离并清除风险内容。本文研究发现,注意力头对特定概念具有敏感性,这一特性可用于风险内容检测与抑制。基于此,提出AHV-D&S——一种无需训练的推理阶段防护方法。该方法通过注意力头向量(AHV)量化每个文本词元在所有注意力头中的敏感度,作为风险判别特征。推理时,采用动量策略动态追踪各去噪步骤中的词元级AHV,并设计敏感度引导的自适应抑制策略,根据头级风险评分降低风险词元的注意力权重。大量实验表明,AHV-D&S能有效抑制色情、版权风格及各类有害内容,同时保持高质量生成效果,并对对抗性提示具备强鲁棒性,且在不同DiT-based T2I模型间具有良好迁移能力。
原文摘要 · Abstract (English)
The rise of text-to-image (T2I) models has increasingly raised concerns regarding the generation of risky content, such as sexual, violent, and copyright-protected images, highlighting the need for effective safeguards within the models themselves. Although existing methods have been proposed to eliminate risky concepts from T2I models, they are primarily developed for earlier U-Net architectures, leaving the state-of-the-art Diffusion-Transformer-based T2I models inadequately protected. This gap stems from a fundamental architectural shift: Diffusion Transformers (DiTs) entangle semantic injection and visual synthesis via joint attention, which makes it difficult to isolate and erase risky content within the generation. To bridge this gap, we investigate how semantic concepts are represented in DiTs and discover that attention heads exhibit concept-specific sensitivity. This property enables both the detection and suppression of risky content. Building on this discovery, we propose AHV-D\&S, a training-free inference-time safeguard for image generation in DiTs. Specifically, AHV-D\&S quantifies each textual token's sensitivity across all attention heads as an Attention Head Vector (AHV), which serves as a discriminative signature for detecting risky generation tendencies. In the inference stage, we propose a momentum-based strategy to dynamically track token-wise AHVs across denoising steps, and a sensitivity-guided adaptive suppression strategy that suppresses the attention weights of identified risky tokens based on head-specific risk scores. Extensive experiments demonstrate that AHV-D\&S effectively suppresses sexual, copyrighted-style, and various harmful content while preserving visual quality, and further exhibits strong robustness against adversarial prompts and transferability across different DiT-based T2I models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。