通过旋转位置编码精准抑制扩散模型中的不安全内容。
SafeRoPE: Risk-specific Head-wise Embedding Rotation for Safe Generation in Rectified Flow Transformers
- 按注意力头划分不安全语义子空间,定位关键风险头。
- 引入潜在风险评分与头级旋转扰动,实现细粒度控制。
- 轻量高效,适合文本生成图像的Transformer架构。
基于修正流变压器(如SD3、FLUX)的文生图模型虽具备高生成保真度,但对多标记交互引发的不安全语义仍易受攻击。现有方法依赖微调或注意力调制进行概念去学习,但计算开销大且仅适配于U-Net结构的去噪器,难以直接用于基于Transformer的扩散模型(如MMDiT)。本文深入分析MMDiT的注意力机制,发现不安全语义集中于可解释的低维头级子空间,少数安全关键头负责提取不安全特征。进一步观察到:对查询与键向量施加旋转位置编码(RoPE)扰动,可有效改变生成图像中的特定概念。基于此,提出SafeRoPE——一种面向MMDiT的轻量级、细粒度安全生成框架。该方法首先通过分解安全关键头内的不安全嵌入,构建头级不安全子空间,并通过投影计算每个输入向量的潜在风险评分(LRS)。随后引入头级RoPE扰动,在不损害良性内容与图像质量的前提下抑制不安全语义。SafeRoPE结合头级LRS与RoPE扰动,实现风险相关的头级查询与键向量嵌入旋转,精确压制不安全输出同时保持生成保真度。大量实验表明,SafeRoPE在平衡有害内容消除与生成效用保留方面达到当前最优性能。代码已开源:https://github.com/deng12yx/SafeRoPE。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。