Unified Safe In-context Image Generation in Multimodal Diffusion Transformers via Restricting Unsafe Information Flows

2026-06-04

AI Summary

Existing multimodal diffusion Transformers lack a unified and effective safety mechanism for image generation and editing, often failing to prevent the synthesis of harmful content. This work proposes a training-free Unified Visual Regulator (UVR), which, for the first time, identifies a task-agnostic harmful semantic priming phase within multimodal attention. Leveraging this insight, UVR dynamically analyzes information flow and proactively modulates relevant attention heads at an early stage to precisely suppress the propagation of unsafe signals. Requiring no additional training, the method seamlessly integrates into both image synthesis and editing pipelines. It achieves 91% and 77% harmful content mitigation rates respectively, substantially outperforming existing approaches while preserving high visual fidelity.

Abstract

Diffusion transformers (DiTs) equipped with multimodal attention (MM-Attn) have become a dominant paradigm for image generation. However, preventing the generation of harmful content remains a critical challenge, particularly in image-to-image (I2I) editing tasks. Existing safety mechanisms are primarily designed for text-to-image (T2I) synthesis or U-Net-based architectures, which limits their effectiveness for unified safety mitigation in DiT-based frameworks. To bridge this gap, we propose Unified Visual Safety Regulator (UVR), a training-free safe generation framework that regulates unsafe semantics in generated images. UVR is grounded in an analysis of attention dynamics from the perspective of information flow in MM-Attn. We identify a task-independent start-up stage, during which unsafe semantics in output patches rapidly emerge and can be accurately localized, followed by task-specific semantic amplification and interference stages, where harmful signals are further propagated and entangled with benign content. Based on these observations, UVR mitigates unsafe generation through unified, targeted attention modulation and explicit restriction of harmful information flow over the identified unsafe output patches. Experiments across various concepts show that UVR achieves state-of-the-art safety performance by achieving 91% and 77% erase rate in image synthesis and editing tasks, while preserving visual quality and fidelity with minimal degradation. Code is available at https://github.com/deng12yx/UVR.

Problem

• safe image generation
• multimodal diffusion transformers
• unsafe content prevention
• image-to-image editing
• visual safety

Innovation

• Diffusion Transformers
• Multimodal Attention
• Information Flow
• Safe Image Generation
• Attention Modulation