Recent Text-to-Image (T2I) models based on rectified-flow transformers (e.g., SD3, FLUX) achieve high generative fidelity but remain vulnerable to unsafe semantics, especially when triggered by multi-token interactions. Existing mitigation methods largely rely on fine-tuning or attention modulation for concept unlearning; however, their expensive computational overhead and design tailored to U-Net-based denoisers hinder direct adaptation to transformer-based diffusion models (e.g., MMDiT). In this paper, we conduct an in-depth analysis of the attention mechanism in MMDiT and find that unsafe semantics concentrate within interpretable, low-dimensional subspaces at head level, where a finite set of safety-critical heads is responsible for unsafe feature extraction. We further observe that perturbing the Rotary Positional Embedding (RoPE) applied to the query and key vectors can effectively modify some specific concepts in the generated images. Motivated by these insights, we propose SafeRoPE, a lightweight and fine-grained safe generation framework for MMDiT. Specifically, SafeRoPE first constructs head-wise unsafe subspaces by decomposing unsafe embeddings within safety-critical heads, and computes a Latent Risk Score (LRS) for each input vector via projection onto these subspaces. We then introduce head-wise RoPE perturbations that can suppress unsafe semantics without degrading benign content or image quality. SafeRoPE combines both head-wise LRS and RoPE perturbations to perform risk-specific head-wise rotation on query and key vector embeddings, enabling precise suppression of unsafe outputs while maintaining generation fidelity. Extensive experiments demonstrate that SafeRoPE achieves SOTA performance in balancing effective harmful content mitigation and utility preservation for safe generation of MMDiT. Codes are available at https://github.com/deng12yx/SafeRoPE.
Figures & tables
Figure 2: Illustration of RoPE in FLUX.1. Each token contains three predefined positional vectors (m,n,q) , and RoPE applies their associated rotations to the corresponding segments of the query and key vectors.
Figure 3: Differential impact of random perturbations to RoPE text positional IDs in FLUX.1-dev across explicit, violence, style, and benign prompts.
Figure 4: Overview of how SafeRoPE identifies safety-critical heads and applies risk-aware rotations. (a) Head Discrimination Score (HDS) used to identify safety-critical heads; higher scores indicate heads that more strongly differentiate unsafe from safe token projections onto the estimated unsafe subspace. (b) SafeRoPE pipeline: (1) SVD-based construction of head-wise unsafe subspaces Ur,b,h (2) Latent Risk Score (LRS) computed by projecting token features onto these subspaces; (3) LRS-guided orthogonal rotations applied only to safety-critical heads to suppress unsafe activations while preserving benign semantics.
Figure 5: The comparison of LRS-guided random rotation perturbations applied to different head groups, where ” H∗ ” denotes the safety-critical heads
Nude
Bloody
VanGogh
Pikachu
CLIP ↑
VQA ↑
FID ↓
UR ↓
CLIP ↑
VQA ↑
FID ↓
UR ↓
CLIP ↑
VQA ↑
FID ↓
UR ↓
CLIP ↑
VQA ↑
FID ↓
UR ↓
Unsafe-1k
I2P
on FLUX.1-dev
ESD
31.09
86.2
76.62
18.57
9.2
31.17
86.8
76.67
32.87
31.37
86.7
76.03
30.13
31.40
87.5
75.31
18.22
SLD
31.82
88.9
76.58
21.17
8.2
31.85
88.3
76.68
43.83
31.89
88.5
76.28
57.53
31.87
88.7
75.33
14.11
UCE
31.31
87.5
76.83
23.00
7.8
31.32
85.4
76.75
25.17
31.39
87.7
76.71
26.02
31.28
87.2
76.37
14.29
Table 1: Cross-concept evaluation of SafeRoPE against baseline methods. SafeRoPE consistently outperforms baselines by achieving safety performance while preserving original generation quality. Furthermore, the learned rotation matrices exhibit cross-concept generalization, maintaining robust efficacy even when transferred to unseen or mismatched domains.
Figure 6: Qualitative comparison of different methods for concept erasure. Corresponding to Table 1 , this figure visualizes results across various target concepts alongside benign prompts. SafeRoPE effectively removes the undesired concepts while maintaining high visual fidelity and semantic consistency across diverse scenarios.
CLIP ↑
VQA ↑
Unsafe-1k ↓
I2P ↓
Shr-NS
31.1
85.52
24.24
9.32
Shr-S
31.2
87.54
29.00
7.1
Indep
31.1
86.31
26.32
8.23
Rank-Low
31.3
89.2
34.00
10.43
Rank-High
31.2
87.55
21.58
11.14
SafeRoPE
31.3
88.73
15.4
6.95
Table 2: Ablation study on key design choices in SafeRoPE. We analyze the effects of rotation sharing strategies and rotation rank on the safety–fidelity trade-off. All symbol definitions and abbreviations are detailed in Section 4.2.3 .
Figure 7: Qualitative ablation comparison under different ablation settings demonstrate that SafeRoPE achieves a more balanced trade-off between safety and generation utility.
Figure 8: Comparison between CA in LDM and MMSA in MMDiT.
Figure 9: Visualization of cross-modal unsafe subspace activation in Single-DiT. Gray noise denotes the spatial locations of image tokens classified as unsafe, identified by LRS >0.7 computed using the unsafe subspace constructed from text embeddings.
Figure 10: Comparison of unsafe-token detection across full heads and safety-critical heads. Unsafe tokens are identified when projection scores onto head-wise unsafe subspaces exceed 0.7. Each subplot shows head counts (y-axis) versus detected unsafe tokens (x-axis) for explicit and safe prompts. Safety-critical heads yield sharper separation between prompt types, confirming detection accuracy.
Hyperparameter
Value
Learning Rate
1×10−3
LR Warmup Cycles
1
LR Scheduler Power
1
Optimizer
AdamW
Adam β1
0.9
Adam β2
0.999
Table 3: Training Hyperparameters Used in SafeRoPE.
Figure 11: Effect of positional ID perturbations on Flux generation
Figure 12: Effectiveness of RoPE text position ID perturbations for safety alignment.
CLIP ↑
VQA ↑
Unsafe-1k ↓
I2P ↓
Method
food
scenery
person
other
food
scenery
person
other
non-hard [%]
hard [%]
non-hard [%]
hard [%]
EraseAnything
30.9
31.1
31.2
30.5
85.8
85.8
89.7
86.7
55.53
21.46
19.67
7.49
SafeRoPE
31.1
31.4
31.1
30.6
89.8
88.1
87.7
87.0
36.67
15.4
12.48
7.0
FLUX.1-dev
31.4
31.5
31.3
30.8
87.2
86.0
88.4
85.9
73.72
38.79
25.18
10.30
Table 4: Quantitative comparison on category-specific alignment (CLIP, VQA; higher is better) and safety (Unsafe-1k/I2P; lower is better). Percentages are shown without the % sign for alignment; units are indicated in the headers.
Figure 13: Qualitative comparison on unsafe prompts from Unsafe-1K and I2P..
State-of-the-art flow based text-to-image (T2I) models exhibit remarkable generative abilities but remain vulnerable to producing unsafe content. Prior safety efforts range from concept erasure and prompt filtering to classifier-based gating. However, simple techniques like parameter efficient adaptations of the models easily bypass such guardrails. We introduce a unique principled approach that achieves safety by regulating the model's attention dynamics through inference-time introspection, exhibiting intrinsic robustness. Our method analyzes and rebalances attention activations throughout image synthesis, steering generations away from unsafe concepts while preserving semantic alignment. This introspective control ensures safety of deployed models. Across standard and adversarial safety benchmarks, our approach achieves remarkable safety scores while maintaining or even improving alignment and perceptual quality. Our results reveal that attention-space regulation offers a considerably more promising path to safer diffusion transformer based image generation than the existing concept erasing mechanism.Our code can be accessed at https://basim-azam.github.io/iam/
Basim Azam, Hossein Rahmani, Naveed Akhtar
The University of Melbourne, Australia · Lancaster University, United Kingdom
Diffusion Transformers have become a powerful backbone for text-to-image generation, but their layered and cross-modal generation process makes safety control fundamentally different from prompt-level filtering or output-level detection. Harmful semantics may be weakly expressed in text representations, progressively bound to visual latents, and finally entangled with rendering dynamics. As a result, safety steering at a fixed layer can be unstable, and a steering mechanism learned from known risks may not transfer reliably to a shifted target risk domain. We propose SafeDIG, a safety steering framework that formulates DiT safety adaptation as position-aware sparse feature transfer. SafeDIG first constructs Sparse Autoencoders over functionally distinct DiT intervention positions and uses robustness-aware pre-training routing to prioritize intervention sites that are expected to remain stable under source-target risk shift. It then separates transferable safety features from domain-specific activation geometry by freezing the SAE encoder as a reusable sparse safety dictionary and adapting only the decoder to the target-domain activation manifold. During inference, SafeDIG combines Blend and Repel operations to steer unsafe activations toward transferred safety manifolds or away from harmful sparse directions. Experiments on FLUX.1 Dev and Stable Diffusion 3.5 Large show that SafeDIG consistently reduces target-domain and overall unsafe generation rates while preserving source-domain safety and image quality.
Zihao Xue, Yan Wang, Zhen Bi +7
1Huzhou Normal University · 2Alibaba Group · University of Science and Technology of China +2
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.
Xiang Yang, Feifei Li, Mi Zhang +4
Fudan University, Shanghai, China · East China University of Science and Technology, Shanghai, China · Shanghai University of Electric Power, Shanghai, China