SafeDiffusion-R1：面向扩散模型安全后训练的在线奖励引导 SafeDiffusion-R1: Online Reward Steering for Safe Diffusion Post-Training

训练分四步。(1)预计算安全方向 $v_{safe}$:CLIP对安全/不安全描述求均值差并归一化,仅训练开始时算一次。(2)轨迹采样:用 $\pi_\theta$ 经DDIM(50步、guidance 7.5)对每个prompt生成K=4张图,按DDIM高斯转移记录每步对数概率,求和得轨迹 $\log p_\theta(z_0|c)$。(3)几何重定向奖励:CLIP编码图像与prompt得 $z_I$、$z_T$;当安全分数 $s_{text}(c)=z_T^\top v_{safe}<0$ 时,把 $z_T$ 沿 $v_{safe}$ 推 $\alpha=0.5$ 步得 $z'_T$,奖励 $r=z_I \cdot z'_T$;否则 $r=z_I \cdot z_T$。(4)GRPO策略更新:组内z-score归一化得优势 $A_{i,k}$,用PPO clip目标更新UNet,$\epsilon=0.0001$、$\beta_{KL}=0.5$、AdamW常数lr $1\times 10^{-5}$、batch 4/GPU、300 epoch、8张AMD MI210。

技术新颖性体现在三个层面。第一,在RL算法层,首次将GRPO系统性地用于扩散模型安全对齐;DDPO沿用PPO并需额外value network,GRPO用组内归一化替代,把'安全prompt奖励方差大'转化为组内相对比较,从机制上避免'不安全prompt绝对奖励过大导致全局过度抑制'。第二,在奖励信号层,把'训练安全分类器'彻底替换为'用CLIP嵌入几何做steering',完全不需要参数化分类器;与SafeCLIP(对比学习)、SafeCLIP+LLaVA(VLM penalty)等仍依赖判别式监督的方案有本质区别。第三,在训练数据层,完全不需要'安全-不安全配对',仅需一组对比性safe/unsafe文本描述定义方向 $v_{safe}$ 即可;Table 3的OOD实验显示,仅用nudity prompt训练后,模型在七类不安全问题(Sexual 11.60% vs 60.2%、Violence 17.33% vs 46.3%)均拿到SOTA,这是几何重定向'一次定义、多类迁移'特性的直接证据。

独立审视有四个值得改进的方向。第一,steering强度 $\alpha=0.5$ 是单一标量超参,但不同prompt类别对几何重定向容忍度不同,可用'自适应 $\alpha$'或基于prompt embedding模长/语义相似度的schedule替代,Fig. 5显示 $\alpha=0.8$–$1.0$ 已开始过修正。第二,当前安全方向 $v_{safe}$ 在训练开始一次性计算并固定,但CLIP embedding空间可能随训练漂移,可在训练中定期重估或用EMA更新,让方向跟踪模型当前状态。第三,GRPO只用K=4组内样本做优势估计,标准差估计噪声大;加大K到8或16、或引入critic-style baseline可进一步稳定;同时缺乏与DDPO、DPOK等PPO/正则化方法的直接对比,无法严格隔离'GRPO相对PPO'与'几何重定向相对监督奖励'各自贡献。第四,Unsafe Anchor虽nudity最强但跨类泛化退化,这一矛盾论文未深入分析;'Unsafe Anchor + 安全prompt scaling'组合值得后续研究。

可延伸五条研究方向。第一,把'单一nudity方向'扩展为'多个安全方向联合优化',每个方向用一个对比描述集合定义(暴力、隐私、骚扰等),通过向量加法/球面插值合成复合方向;Table 3显示默认变体七类接近最优但Unsafe Anchor跨类退化,提示'多方向联合'可能更鲁棒。第二,接入VLM-as-judge:在steering reward之外加LLaVA/Qwen-VL对生成图像做语义安全判断,把'几何奖励'与'语义奖励'融合,应对CLIP对长尾概念识别不准的问题。第三,把训练阶段的steering推广到推理阶段,在DDIM去噪每步做一次prompt方向修正,作为'inference-time steering',让训练与推理完全分离,变成即插即用模块,适配任意已发布SD模型。第四,扩展到SDXL/SD3/Flux及视频/3D扩散模型,验证几何重定向在高维潜空间是否仍然有效;SDXL更高维CLIP文本嵌入可能改变安全方向线性可分性,是开放问题。第五,理论分析CLIP嵌入空间'安全方向'的线性可分性,给出steering success的正式条件与适用边界。

复现评估分四方面。开源情况:论文给出GitHub(MAXNORM8650/SafeDiffusion-R1)与项目主页,代码与checkpoint发布需以仓库实际状态为准;算法核心(GRPO loss、steering reward、DDIM轨迹记录)在Algorithm 1与公式(7)–(11)中有详细描述,具备独立实现基础。数据可获得性:1900条nudity负prompt(Grok生成)、7100条GenEval prompt、30000+ SafetyDPO prompt均可获取;HPSv2、SD v1.4权重均开源。算力需求:每轮约72 GPU小时(8×AMD MI210 64GB),单卡A100/H100训练需约一周;显存峰值约40–50GB/卡。复现难度:涉及GRPO在扩散模型上的工程细节——DDIM轨迹每步对数概率计算、组内归一化优势向量化、PPO clip+KL双loss实现,以及CLIP嵌入steering(单位向量、负投影判断、重定向嵌入归一化),工程复杂度高于普通SFT;具备扩散模型训练经验的团队可在3–4周内复现主要数字。