SafeDiffusion-R1: Direcionamento de Recompensa Online para Pós-Treinamento de Difusão Segura

Resumo

Modelos de difusão têm sido amplamente estudados para remover conteúdo indesejado aprendido durante o pré-treinamento. Métodos existentes exigem dados supervisionados caros, seja texto inseguro pareado com groundtruth de imagem segura ou pares de imagem negativa/positiva, tornando-os impraticáveis para escalar. Além disso, abordagens de aprendizado por reforço offline e ajuste fino supervisionado que geram dados sintéticos offline sofrem de esquecimento catastrófico, degradando a qualidade da geração. Propomos um novo framework de aprendizado por reforço online que aborda tanto a escassez de dados quanto a degradação do modelo por meio do pós-treinamento com Otimização de Política Relativa em Grupo (GRPO) em prompts de texto negativos e positivos. Para eliminar a necessidade de ajustar modelos de recompensa especializados seguros/inseguros, introduzimos um mecanismo de recompensa de direcionamento que explora uma propriedade inerente dos embeddings CLIP: direcionar representações de texto para direções de segurança positivas e afastá-las de direções negativas no espaço de embeddings. Nossa abordagem de política online permite que o modelo aprenda a partir de prompts diversos, incluindo conteúdo explícito indesejado, sem esquecimento catastrófico. Experimentos extensos demonstram que nosso método reduz o conteúdo inapropriado para 18,07% (vs. 48,9% para SD v1.4) e detecções de nudez para 15 (vs. 646 linha de base), enquanto melhora a qualidade de geração composicional de 42,08% para 47,83% no GenEval. Notavelmente, esses ganhos de segurança generalizam-se para prompts inseguros fora do domínio, abrangendo sete categorias de danos, alcançando desempenho de estado da arte sem dados pareados supervisionados ou ajuste de recompensa. GitHub: https://github.com/MAXNORM8650/SafeDiffusion-R1.

English

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 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.