PaperlayerView exact revision on arXiv

Source-linked AI summary

SafeDiffusion-R1: Online Reward Steering for Safe Diffusion Post-Training

Komal Kumar, Ankan Deria, Abhishek Basu, Fahad Shamshad, Hisham Cholakkal, Karthik Nandakumar

arXiv:2605.18719v1cs.CV

TL;DR

SafeDiffusion-R1 tackles unsafe associations and the data and model-degradation limitations of existing diffusion safety alignment methods. It uses online GRPO with a CLIP-based geometric steering reward to learn from benign and unsafe prompts without dedicated reward-model fine-tuning. The framework is reported to improve safety alignment and compositional utility while achieving broad out-of-domain generalization without supervised paired data.

  • Problem

    Text-to-image diffusion models can learn unsafe associations, while existing alignment methods depend on static supervised or offline data and are limited during concept unlearning.

  • Method

    The method combines online GRPO on evolving model outputs with a geometry-aware CLIP steering reward that represents safety as an embedding-space direction.

  • Results

    The framework consistently outperforms supervised fine-tuning and offline alignment on safety benchmarks while preserving generation quality on benign concepts.

  • Takeaways & Limitations

    SafeDiffusion-R1 broadly reduces inappropriate content and improves compositional utility with online geometric steering, without supervised data or a learned reward model.

Abstract

from arXiv · show

Diffusion models have been widely studied for removing unsafe content learned during pre-training. Existing methods require expensive supervised data, either unsafe-text paired with safe-image groundtruth or negative/positive image pairs, making them impractical to scale. Furthermore, offline reinforcement learning and supervised fine-tuning approaches that generate synthetic data offline suffer from catastrophic forgetting, degrading generation quality. We propose a novel online reinforcement learning framework that addresses both data scarcity and model degradation through post-training with Group Relative Policy Optimization (GRPO) on both negative and positive text prompts. To eliminate the need for fine-tuning specialized safe/unsafe reward models, we introduce a \textit{steering reward mechanism} that exploits an inherent property of CLIP embeddings: steering text representations toward positive safety directions and away from negative ones in the embedding space. Our online-policy approach enables the model to learn from diverse prompts, including explicit unsafe content, without catastrophic forgetting. Extensive experiments demonstrate that our method reduces inappropriate content to 18.07\% (vs. 48.9\% for SD v1.4) and nudity detections to 15 (vs. 646 baseline) while improving compositional generation quality from 42.08\% to 47.83\% on GenEval. Remarkably, these safety gains generalize to out-of-domain unsafe prompts across seven harm categories, achieving state-of-the-art performance without supervised paired data or reward tuning. Github: https://github.com/MAXNORM8650/SafeDiffusion-R1.

1 Introduction

SafeDiffusion-R1 addresses unsafe associations in text-to-image diffusion models with online policy optimization and geometry-aware steering rewards. The framework targets safety alignment without static supervised data or dedicated reward-model fine-tuning while preserving generation quality.

  • Text-to-image models can internalize unsafe associations from web data, creating safety risks when explicit or inappropriate prompts trigger them.
  • Existing post-training methods rely on curated safe/unsafe examples or fixed pre-generated data, limiting adaptation during concept unlearning.
  • SafeDiffusion-R1 formulates safety alignment as online policy optimization with GRPO, coupling safety learning to the model’s evolving generative distribution.
  • Its geometry-aware steering reward represents safety as a direction in CLIP embedding space, avoiding dedicated safe/unsafe reward models.
  • The framework is reported to outperform supervised fine-tuning and offline alignment on safety benchmarks while preserving generation quality on benign concepts.

2 Related Work

Diffusion safety research has progressed from output filtering toward parameter-level harmful-concept removal and reinforcement-learning-based alignment. Diffusion RL remains challenging because multi-step denoising creates long-horizon credit-assignment problems, while online optimization introduces reward-scale instability.

  • Output-only NSFW filters screen generated images but leave the diffusion model unchanged and can be bypassed with direct access.
  • Parameter-editing methods seek to remove harmful concepts from diffusion models without full retraining.
  • Reinforcement learning aligns generative models with objectives that supervised losses alone may not capture effectively.
  • Diffusion-model RL is harder than autoregressive language-model RL because multi-step denoising requires long-horizon credit assignment across timesteps.
  • Online policy optimization trains on current model outputs, but unsafe-prompt rewards can have higher magnitude and variance, causing standard PPO-style updates to overcorrect.
  • GRPO addresses this instability by normalizing advantages within groups of generations from the same prompt.

3 Methodology

SafeDiffusion-R1 trains text-to-image diffusion models online using geometric safety steering and GRPO. The framework redirects unsafe prompts toward safe embedding directions, then optimizes sampled denoising trajectories with normalized relative advantages and clipped policy updates.

  • 3 Methodology: SafeDiffusion-R1 combines steering rewards, GRPO, and denoising trajectory optimization for safe reinforcement learning of text-to-image diffusion models.The framework is designed to train on diverse prompt distributions, including unsafe content.
  • 3 Methodology: A safety direction in CLIP embedding space is constructed from safe and unsafe text descriptions to represent movement toward safer concepts.The direction is computed once during initialization and remains fixed throughout training.
  • 3 Methodology: Unsafe prompts are detected by projecting their text embeddings onto the safety direction, with negative cosine similarity indicating alignment with unsafe concepts.The projection score lies in [−1, 1] because the text and safety embeddings are normalized.
  • 3 Methodology: For unsafe prompts, the reward uses a text embedding geometrically steered toward safety instead of the original text embedding.The steering strength α controls how strongly negative prompts are redirected.
  • 3 Methodology: The policy generates multiple image samples through DDIM trajectories, computes steering rewards, and normalizes advantages within each prompt’s generation group.Group normalization makes updates depend on relative rewards and improves stability across prompts with different reward scales.
  • 3 Methodology: Clipped PPO-style GRPO updates constrain policy changes, while a KL term helps prevent catastrophic forgetting.The clipping operation limits large policy updates during optimization.

4 Experiments

Experiments evaluate SafeDiffusion-R1 across safety suppression, out-of-domain generalization, compositional utility, image-quality locality, and qualitative fidelity. The method achieves strong safety improvements while preserving or improving several generation-quality measures, though FID remains higher than most baselines.

  • Experimental setup: SafeDiffusion-R1 is trained with GRPO using HPSv2 embeddings to construct a steering direction, with α = 0.5 used throughout experiments.Training uses prompts rather than paired images, including nudity-focused negative prompts and broader safety and compositional prompt sets.
  • Safety evaluation: 15 total nude-body-part detections are achieved on I2P, compared with 646 for base Stable Diffusion v1.4.NudeNet uses a 0.6 threshold, and lower detection counts indicate better safety.
  • Safety evaluation: 18.07% overall inappropriate content is achieved on I2P, compared with 48.9% for base Stable Diffusion v1.4.SafeDiffusion-R1 also reaches 11.60% in Sexual, 17.33% in Violence, and 15.86% in Self-harm, despite primarily using nudity-focused training prompts.
  • Compositional utility: 47.83% GenEval accuracy is achieved with GenEval+Nudity prompts, while nudity-only training reaches 44.12%.The largest gains occur in two-object composition and position and color-attribute relational reasoning; the benchmark contains 553 prompts and 2,212 generated images.
  • Image quality and locality: GRPO post-trained models maintain CLIP-T close to baseline but exhibit higher FID than most baselines on the locality evaluation.The authors attribute the FID increase primarily to synthetic post-training samples differing from the COCO reference distribution.
  • Qualitative results: Qualitative comparisons show safe, semantically coherent, and visually faithful outputs across challenging explicit or sensitive prompts.Several baselines suppress unsafe content while introducing distortion, blur, unnatural textures, oversmoothing, or weakened structural consistency.

5 Conclusion

SafeDiffusion-R1 combines online GRPO with geometric steering rewards to train safe text-to-image diffusion models without supervised data or learned reward models. The framework supports diverse prompts, preserves model capabilities, improves compositional utility, and reduces inappropriate content with strong out-of-domain generalization.

  • SafeDiffusion-R1 uses online GRPO with geometric steering rewards for safe reinforcement learning of text-to-image diffusion models.
  • The steering reward enables training on diverse prompt distributions, including unsafe content, by redirecting unsafe prompts toward safe alternatives before reward computation.
  • The framework requires no supervised data or learned reward model while improving compositional utility and broadly reducing inappropriate content.
  • The supplementary material adds training algorithms, implementation details, ablations, and qualitative results.

A GRPO Training Algorithms

The training procedure separates online trajectory generation and group-relative advantage computation from clipped policy-gradient optimization over stored trajectories.

  • Algorithm 2 generates trajectories and computes group-relative advantages during sampling and reward evaluation.
  • Together, the algorithms implement the two-phase GRPO training procedure.
  • Algorithm 3 applies clipped policy-gradient updates to the stored trajectories.

B.1 Model Architecture and LoRA Configuration

The method uses parameter-efficient LoRA adaptation of Stable Diffusion v1.4 and an online steering-reward pipeline that preserves the original unsafe prompt as model input while modifying reward computation.

  • Model Architecture and LoRA Configuration: Stable Diffusion v1.4 uses an approximately 860M-parameter UNet with rank-4 LoRA adapters, yielding approximately 2.4M trainable parameters.
  • Policy Optimization: Policy optimization uses AdamW with learning rate 1 × 10−5, gradient norm 1.0, and advantage clipping to [−5, 5].
  • Sampling Configuration: Training uses DDIM with 50 denoising steps, guidance scale 7.5, eta 1.0, and the last 40 denoising steps.
  • GRPO Sampling and Reward: Sampling generates K = 16 images per prompt in main experiments and uses the steered reward with prompts drawn from D.
  • Steering Reward: The steering reward uses α = 0.5, HPSv2, safe text anchors, and unsafe descriptions to define the safety direction.
  • Computational Setup: The steering direction is computed once before training, while training uses 8 AMD MI210 GPUs for approximately 300 epochs and 300K gradient updates.
  • Geometric Interpretation: Adding αv_safe and renormalizing rotates unsafe text embeddings toward safer representations, with larger α producing stronger steering.
  • Geometric Interpretation: The original unsafe prompt remains the model input, while the steered representation is used exclusively for reward computation.

C Qualitative Results

The method reduces inappropriate content across harm categories, preserves benign compositional generation quality, and avoids the collapse observed with negative-only rewards. Safety improvements also generalize beyond the nudity-focused training prompts.

  • Safety Results: Inappropriate rates decrease across Hate, Harassment, Violence, Self-harm, Sexual, Shocking, and Illegal Activity categories.
  • Safety Results: Categories not directly targeted by training prompts, including Hate, Violence, and Illegal Activity, also improve, indicating out-of-domain generalization.
  • Utility Preservation: GenEval improves from 42.08% to 47.83% while maintaining fidelity on single-object, two-object, color-attribute, and spatial-relation prompts.
  • Reward Ablation: Negative-only reward produces very low nudity rates but degrades CLIP-T to 23.31 and FID to 167.49, causing structureless generations dissimilar to text prompts.

D.2 Utility Degradation under Negative-Only Training

Negative-only CLIP training suppresses unsafe content but causes severe utility collapse. Steering and anchor-based rewards preserve image quality while improving safety suppression.

  • Negative-only training causes severe image degradation, including blurring, color collapse, and loss of semantic structure.
  • SafeCLIP with 7K positive prompts achieves good CLIP-T but has higher MeanUnsafe than steering reward, 0.246 vs. 0.002.
  • Positive-only SafeCLIP preserves utility but provides insufficient safety suppression, with MeanUnsafe = 0.816 for 100K positives.
  • Early SafeCLIP configurations show inconsistent suppression across broad I2P prompts, whereas anchor-based steering progressively improves safety suppression.

F LLaVA-Augmented Penalty Analysis

The LLaVA-augmented penalty improves safety over plain SafeCLIP but introduces image-quality degradation from its discrete penalty. SafeDiffusion-R1 achieves stronger suppression without this degradation.

  • On nudity-focused I2P prompts, SafeDiffusion-R1 consistently generates safe, high-fidelity images, while competing methods either retain explicit content or reduce realism.
  • SafeCLIP+LLaVA improves MeanUnsafe over plain SafeCLIP, 0.151 vs. 0.246, but introduces higher FID of 103.40.
  • SafeDiffusion-R1 achieves MeanUnsafe = 0.002 without the FID degradation associated with the LLaVA penalty.
  • Across checkpoints, steering progressively redirects outputs from explicit content to fully clothed, semantically appropriate images by step 600.

G.2 Utility Preservation Across Training

SafeDiffusion-R1 preserves compositional utility during safety training. Its outputs remain accurate on object counts, spatial relations, attributes, and representative GenEval tasks.

  • Object counts, spatial relations, and color attributes remain accurately rendered across training checkpoints.
  • SafeDiffusion-R1 produces semantically accurate and visually coherent generations across single-object, multi-object, spatial-relation, and color-binding tasks.
  • The method maintains or improves upon the SD v1.4 baseline on benign compositional prompts.

H Training Stability Analysis

Training remains stable under GRPO with reward normalization, KL monitoring, and gradient clipping. Negative-only penalties are an exception, producing early gradient spikes and utility degradation, while anchor variants expose safety–utility trade-offs.

  • Training Stability: Group-relative normalization reduces reward variance by normalizing advantages within each prompt’s generation group.This makes the optimization signal invariant to absolute reward scale across benign and unsafe prompts.
  • Training Stability: KL divergence remains below 0.05 throughout training with βKL = 0.5 and ε = 0.0001.The reported monitoring indicates stable on-policy optimization without model drift.
  • Training Stability: Gradient norms are clipped at 1.0 and remain stable across 300 epochs except for negative-only CLIP training.The negative-only variant exhibits early gradient spikes because it lacks a positive reward signal.
  • Reward Design: Negative-only training reaches the lowest unsafe score but causes utility collapse, with degraded FID and CLIP-T scores.The model generates degenerate images that match no prompt instead of safe, semantically appropriate content.
  • Reward Design: Steering reward outputs preserve visual quality and semantic alignment on benign prompts, unlike negative-only outputs.
  • Reward Design: Steering reward produces sharper, more compositionally accurate images and the lowest unsafe score among SafeCLIP configurations.
Loading 2605.18719v1…