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..
Fudan University, Shanghai, China · East China University of Science and Technology, Shanghai, China · Shanghai University of Electric Power, Shanghai, China