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.
翻译:文本到图像(T2I)模型的兴起日益引发对生成风险内容(如色情、暴力和受版权保护的图像)的担忧,凸显了在模型内部建立有效防护机制的必要性。尽管现有方法已提出从T2I模型中消除风险概念,但它们主要针对早期的U-Net架构设计,使得基于最先进扩散变换器的T2I模型缺乏充分保护。这一缺陷源于根本性的架构转变:扩散变换器(DiTs)通过联合注意力机制将语义注入与视觉合成纠缠在一起,使得在生成过程中隔离和擦除风险内容变得困难。为弥补这一不足,我们研究了语义概念在DiTs中的表征方式,发现注意力头展现出对特定概念的高度敏感性。这一特性使检测和抑制风险内容成为可能。基于此发现,我们提出AHV-D\&S——一种无需训练、推理阶段生效的DiTs图像生成防护方法。具体而言,AHV-D\&S将每个文本标记在全部注意力头上的敏感性量化为注意力头向量(AHV),作为检测风险生成倾向的判别特征。在推理阶段,我们提出基于动量的策略动态追踪跨去噪步骤的逐标记AHV,并设计敏感性引导的自适应抑制策略,基于头特定风险分数抑制已识别风险标记的注意力权重。大量实验表明,AHV-D\&S能有效抑制色情、受版权保护的风格及各类有害内容,同时保持视觉质量,并进一步展现出对对抗性提示的强大鲁棒性及在不同DiT基T2I模型间的可迁移性。