What Concepts Lie Within? Detecting and Suppressing Risky Content in Diffusion Transformers

2026-05-11

AI Summary

This work addresses the challenge of content safety in Diffusion Transformers (DiTs), where tight coupling between text and image generation via cross-attention renders conventional safety mechanisms ineffective against harmful outputs such as sexually suggestive, violent, or copyright-infringing content. The study reveals for the first time that individual attention heads in DiTs exhibit concept-specific sensitivity, and leverages this insight to propose AHV-D&S—a training-free, inference-time defense framework that dynamically detects and adaptively suppresses risky content. By quantifying Attention Head Vectors (AHVs), tracking them across denoising steps with momentum-based smoothing, and applying head-specific risk scores to modulate attention weights, AHV-D&S effectively blocks diverse harmful content while preserving image fidelity, demonstrating robustness to adversarial prompts and strong transferability across DiT architectures.

Abstract

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.

Problem

• Diffusion Transformers
• risky content
• text-to-image generation
• content safety
• semantic concepts

Innovation

• Diffusion Transformer
• Attention Head Vector
• Risky Content Suppression
• Training-free Safeguard
• Concept-specific Sensitivity