Technical Summary: SafeDiffusion-R1

Problem Statement

Text-to-image (T2I) diffusion models, trained on large-scale web data, have democratized visual content generation but inadvertently internalized unsafe and explicit associations. Existing safety interventions face significant limitations:

1. Data Scarcity and Cost: Methods relying on supervised fine-tuning (SFT) or offline reinforcement learning (RL) require expensive, curated datasets of safe/unsafe image-text pairs or negative/positive image pairs, which are difficult to scale.
2. Catastrophic Forgetting: Offline approaches that generate synthetic data or rely on fixed datasets often suffer from catastrophic forgetting, degrading the model's general generative capabilities and compositional quality.
3. Static Supervision: Current paradigms (SFT and offline RL) optimize against fixed examples or pre-generated data. They fail to adapt their training signals to the model's evolving generative behavior, making them insufficient for tracking and suppressing unsafe content that emerges dynamically during training.
4. Reward Model Overhead: Many safety alignment methods require training or fine-tuning specialized reward models to classify images as safe or unsafe, introducing additional computational overhead.

Methodology

The authors propose SafeDiffusion-R1, an online reinforcement learning framework designed to align diffusion models with safety constraints without supervised paired data or specialized reward model fine-tuning. The framework integrates two core components:

1. Online Policy Optimization via GRPO

The method employs Group Relative Policy Optimization (GRPO) as the training algorithm. Unlike standard PPO or offline RL, GRPO operates on-policy:

• Dynamic Sampling: The model continuously generates image samples from both benign and unsafe prompts during training.
• Group-Based Advantage: For each prompt, $K$ independent samples are generated. Advantages are normalized within the group ($A_{i,k} = \frac{r_{i,k} - \bar{r}_i}{\sigma_i + \delta}$) rather than globally. This mitigates reward scale mismatches between benign and unsafe prompts, preventing over-correction and ensuring stable updates relative to the current policy's distribution.
• On-Policy Learning: By coupling safety optimization with the model's evolving sampling distribution, the method avoids distribution mismatch and preserves general generative capabilities while progressively unlearning unsafe concepts.

2. Geometry-Aware Steering Reward

To eliminate the need for a separate safe/unsafe classifier or reward model, the authors introduce a steering reward mechanism that exploits the geometric properties of CLIP embeddings:

• Safety Direction Vector ($v_{safe}$): A unit vector is computed in the text embedding space as the normalized difference between the mean embeddings of a small set of safe text anchors ($T_{safe}$) and unsafe text anchors ($T_{unsafe}$):
  $$v_{safe} = \frac{\bar{z}_{safe} - \bar{z}_{unsafe}}{\|\bar{z}_{safe} - \bar{z}_{unsafe}\|_2}$$
• Conditional Steering: For an input prompt $c$, the text embedding $z_T$ is projected onto $v_{safe}$. If the prompt is identified as unsafe (negative projection), the embedding is geometrically steered toward the safe direction:
  $$z'_T = \text{normalize}(z_T + \alpha v_{safe})$$
  where $\alpha$ is a steering strength hyperparameter.
• Reward Computation: The reward is calculated as the cosine similarity between the image embedding ($z_I$) and the steered text embedding ($z'_T$) for unsafe prompts, or the original text embedding for safe prompts. This reshapes the optimization signal to encourage alignment with safe semantic attributes without explicitly penalizing the model for "understanding" the unsafe prompt.

Key Contributions

1. Online Policy Formulation: The paper formulates safety alignment for T2I diffusion models as an online policy optimization problem, introducing a GRPO-based framework that couples safety learning with the model's evolving generative distribution to prevent catastrophic forgetting.
2. Steering Reward Mechanism: A novel geometry-aware reward is proposed that represents safety as a direction in CLIP embedding space. This enables concept suppression without training dedicated safe/unsafe reward models or relying on supervised paired data.
3. Empirical Validation: Extensive experiments demonstrate that the framework outperforms supervised fine-tuning and offline alignment methods on safety benchmarks while preserving or improving compositional generation quality.

Experimental Results

The method was evaluated on Stable Diffusion v1.4 using the I2P benchmark (nudity detection) and GenEval (compositional utility).

• Safety Performance:
  • Nudity Detection: SafeDiffusion-R1 reduced nudity detections to 15 (compared to 646 for the baseline SD v1.4 and 66 for the RECE baseline).
  • Inappropriate Content: The method reduced the overall inappropriate content rate to 18.07% (vs. 48.9% for SD v1.4). Notably, it achieved strong performance across seven harm categories (Hate, Harassment, Violence, Self-harm, Sexual, Shocking, Illegal activity) despite being trained primarily on nudity-focused prompts, demonstrating strong Out-of-Distribution (OOD) generalization.
• Utility Preservation:
  • Compositional Quality: Unlike baselines that degrade utility, SafeDiffusion-R1 improved the GenEval score from 42.08% (SD v1.4) to 47.83%.
  • Visual Fidelity: Qualitative analysis shows the model preserves structural integrity, facial details, and lighting consistency, avoiding the distortion, oversmoothing, or semantic incoherence observed in methods like EraseDiff or ESD.
  • CLIP Score: The model maintained CLIP-T scores close to the baseline (0.311 vs. 0.313), indicating preserved instruction following.

Significance and Claims

The paper claims that SafeDiffusion-R1 addresses the critical trade-off between safety and utility in diffusion models. By utilizing online GRPO and geometric steering, the approach:

• Eliminates Data Dependencies: It removes the need for expensive, curated supervised safety datasets.
• Avoids Reward Model Overhead: It bypasses the computational cost of training specialized reward classifiers.
• Ensures Stability: The on-policy formulation and group-relative normalization prevent catastrophic forgetting and distribution collapse.
• Generalizes Robustly: The method achieves state-of-the-art safety performance that generalizes to OOD unsafe prompts across diverse harm categories, even when trained on a narrow set of unsafe concepts.

The authors position this work as a practical strategy for post-training modification that directly adjusts pre-trained models to suppress unsafe concepts without retraining from scratch, remaining compatible with publicly released systems like Stable Diffusion.