SafeRoPE:面向整流流 Transformer 安全生成的风险特定头级嵌入旋转
SafeRoPE: Risk-specific Head-wise Embedding Rotation for Safe Generation in Rectified Flow Transformers
针对 MMDiT 架构,提出在注意力头级别定位不安全语义子空间,并通过扰动 RoPE 的 query/key 嵌入实现无需微调的安全概念抑制,为扩散 Transformer 的安全生成提供了轻量级、可复用的干预手段。
原文依据 perturbing the Rotary Positional Embedding (RoPE) applied to the query and key vectors can effectively modify some specific concepts
摘要译文
由模型翻译 · 术语统一 · 请对照原文核对关键结论近期基于整流流 Transformer(如 SD3、FLUX)的文本到图像(T2I)模型实现了高生成保真度,但仍易受不安全语义影响,尤其是由多词元交互触发时。现有缓解方法主要依赖微调或注意力调制进行概念遗忘;然而,其高昂的计算开销和针对 U-Net 去噪器的设计阻碍了直接适配基于 Transformer 的扩散模型(如 MMDiT)。本文深入分析 MMDiT 中的注意力机制,发现不安全语义集中在头级可解释的低维子空间中,其中有限的安全关键头负责不安全特征提取。我们进一步观察到,扰动应用于查询和键向量的旋转位置嵌入(RoPE)能有效修改生成图像中的某些特定概念。受此启发,我们提出 SafeRoPE,一种面向 MMDiT 的轻量级细粒度安全生成框架。具体而言,SafeRoPE 首先通过在安全关键头内分解不安全嵌入来构建头级不安全子空间,并通过投影到这些子空间为每个输入向量计算潜在风险分数(LRS)。然后,我们引入头级 RoPE 扰动,可在不降低良性内容或图像质量的情况下抑制不安全语义。SafeRoPE 结合头级 LRS 和 RoPE 扰动,对查询和键向量嵌入执行风险特定的头级旋转,从而在保持生成保真度的同时精确抑制不安全输出。大量实验表明,SafeRoPE 在平衡有效缓解有害内容和保持实用性方面达到 SOTA 性能,适用于 MMDiT 的安全生成。代码见 https://github.com/deng12yx/SafeRoPE。
英文原文摘要 1900 字符
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.