SafeRoPE: Risk-specific Head-wise Embedding Rotation for Safe Generation in Rectified Flow Transformers

• URL: http://arxiv.org/abs/2604.01826v1
• Date: Thu, 02 Apr 2026 09:37:32 GMT
• Title: SafeRoPE: Risk-specific Head-wise Embedding Rotation for Safe Generation in Rectified Flow Transformers
• Abstract summary: We present SafeRoPE, a safe generation framework for transformer-based diffusion models (e.g., MMDiT)<n>We show 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.<n>We then introduce head-wise RoPE perturbations that can suppress unsafe semantics without degrading benign content or image quality.
• Score: 20.61907377210064
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: 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.

Related papers

• Pulling The REINS: Training-Free Safety Alignment of Video Diffusion Models via Representation Steering [20.152820762914214]
  We present REINS, a training-free method that aligns video diffusion models at inference time by steering internal representations toward safe generation.<n>Our key finding is that safety-relevant structure is linearly encoded in the hidden-state activations of video diffusion transformers.<n>At inference, adding this direction to hidden states at an intermediate transformer layer redirects generation from harmful content to semantically related safe alternatives, with no weight updates, no concept enumeration, and negligible computational overhead.
  arXiv  Detail & Related papers  (2026-06-15T20:03:09Z)
• Unified Safe In-context Image Generation in Multimodal Diffusion Transformers via Restricting Unsafe Information Flows [20.386952794426833]
  Unified Visual Safety Regulator (UVR) is a training-free safe generation framework that regulates unsafe semantics in generated images.<n>UVR mitigates unsafe generation through unified, targeted attention modulation and explicit restriction of harmful information flow.<n>Experiments show that UVR achieves state-of-the-art safety performance by achieving 91% and 77% erase rate in image synthesis and editing tasks.
  arXiv  Detail & Related papers  (2026-06-05T03:43:04Z)
• Do Models Share Safety Representations? Cross-Model Steering for Safe Visual Generation [52.122731171289665]
  We introduce the first framework for cross-model safety steering.<n>Our pipeline never accesses unsafe data on the target side.<n>We evaluate our approach in text-to-image and text-to-video generation across diverse source-target model pairs.
  arXiv  Detail & Related papers  (2026-06-03T18:00:04Z)
• Robust and Generalizable Safety Steering for Text-to-Image Diffusion Transformers [28.791649151495008]
  We propose SafeDIG, a safety steering framework that formulates DiT safety adaptation as position-aware sparse feature transfer.<n> Experiments on FLUX.1 Dev and Stable Diffusion 3.5 Large show that SafeDIG consistently reduces target-domain and overall unsafe generation rates.
  arXiv  Detail & Related papers  (2026-05-28T15:00:18Z)
• Safety-Guided Flow (SGF): A Unified Framework for Negative Guidance in Safe Generation [2.5489046505746704]
  In this paper, we introduce a unified probabilistic framework using a Maximum Mean Discrepancy (MMD) potential for image generation tasks.<n>We leverage control-barrier functions analysis to justify the existence of a critical time window in which negative guidance must be strong.<n>We evaluate our unified framework on several realistic safe generation scenarios, confirming that negative guidance should be applied in the early stages of the denoising process.
  arXiv  Detail & Related papers  (2026-03-02T07:31:29Z)
• BarrierSteer: LLM Safety via Learning Barrier Steering [83.12893815611052]
  BarrierSteer is a novel framework that formalizes safety by embedding learned non-linear safety constraints directly into the model's latent representation space.<n>We show that BarrierSteer substantially reduces adversarial success rates, decreases unsafe generations, and outperforms existing methods.
  arXiv  Detail & Related papers  (2026-02-23T18:19:46Z)
• LSSF: Safety Alignment for Large Language Models through Low-Rank Safety Subspace Fusion [16.434293020863592]
  The safety mechanisms of large language models (LLMs) exhibit notable fragility, as even fine-tuning on datasets without harmful content may still undermine their safety capabilities.<n>We introduce LSSF, a novel safety re-alignment framework with underlineLow-Rank underlineSafety underlineSubspace underlineFusion.<n>Our proposed method exploits the low-rank characteristics of safety information in LLMs by constructing a low-rank projection matrix.
  arXiv  Detail & Related papers  (2026-01-19T03:59:12Z)
• SafeRedir: Prompt Embedding Redirection for Robust Unlearning in Image Generation Models [67.84174763413178]
  We introduce SafeRedir, a lightweight inference-time framework for robust unlearning via prompt embedding redirection.<n>We show that SafeRedir achieves effective unlearning capability, high semantic and perceptual preservation, robust image quality, and enhanced resistance to adversarial attacks.
  arXiv  Detail & Related papers  (2026-01-13T15:01:38Z)
• UpSafe$^\circ$C: Upcycling for Controllable Safety in Large Language Models [67.91151588917396]
  Large Language Models (LLMs) have achieved remarkable progress across a wide range of tasks, but remain vulnerable to safety risks such as harmful content generation and jailbreak attacks.<n>We propose UpSafe$circ$C, a unified framework for enhancing LLM safety through safety-aware upcycling.<n>Our results highlight a new direction for LLM safety: moving from static alignment toward dynamic, modular, and inference-aware control.
  arXiv  Detail & Related papers  (2025-10-02T16:43:33Z)
• Turning the Spell Around: Lightweight Alignment Amplification via Rank-One Safety Injection [47.347413305965006]
  Safety alignment in Large Language Models (LLMs) often involves mediating internal representations to refuse harmful requests.<n>Recent research has demonstrated that these safety mechanisms can be bypassed by ablating or removing specific representational directions.<n>We propose Rank-One Safety Injection (ROSI), a white-box method that amplifies a model's safety alignment by permanently steering its activations toward the refusal-mediating subspace.
  arXiv  Detail & Related papers  (2025-08-28T13:22:33Z)
• Shape it Up! Restoring LLM Safety during Finetuning [65.75757313781104]
  Finetuning large language models (LLMs) enables user-specific customization but introduces critical safety risks.<n>We propose dynamic safety shaping (DSS), a framework that uses fine-grained safety signals to reinforce learning from safe segments of a response while suppressing unsafe content.<n>We present STAR-DSS, guided by STAR scores, that robustly mitigates finetuning risks and delivers substantial safety improvements across diverse threats, datasets, and model families.
  arXiv  Detail & Related papers  (2025-05-22T18:05:16Z)
• Superficial Safety Alignment Hypothesis [15.215130286922564]
  We propose the Superficial Safety Alignment Hypothesis (SSAH), which posits that safety alignment teaches an otherwise unsafe model to choose the correct reasoning direction.<n>We identify four types of attribute-critical components: Safety Critical Unit (SCU), Utility Critical Unit (UCU), Complex Unit (CU) and Redundant Unit (RU)<n>Our findings show that freezing certain safety-critical components during fine-tuning allows the model to retain its safety attributes while adapting to new tasks.
  arXiv  Detail & Related papers  (2024-10-07T19:53:35Z)
• Ring-A-Bell! How Reliable are Concept Removal Methods for Diffusion Models? [52.238883592674696]
  Ring-A-Bell is a model-agnostic red-teaming tool for T2I diffusion models. It identifies problematic prompts for diffusion models with the corresponding generation of inappropriate content. Our results show that Ring-A-Bell, by manipulating safe prompting benchmarks, can transform prompts that were originally regarded as safe to evade existing safety mechanisms.
  arXiv  Detail & Related papers  (2023-10-16T02:11:20Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.