Resumen Técnico: SafeDiffusion-R1

Enunciado del Problema

Los modelos de difusión de texto a imagen (T2I), entrenados con datos web a gran escala, han democratizado la generación de contenido visual, pero han internalizado inadvertidamente asociaciones inseguras y explícitas. Las intervenciones de seguridad existentes enfrentan limitaciones significativas:

1. Escasez y Costo de Datos: Los métodos que dependen de ajuste fino supervisado (SFT) o aprendizaje por refuerzo (RL) fuera de línea requieren conjuntos de datos costosos y curados de pares imagen-texto seguros/inseguros o pares de imágenes negativos/positivos, los cuales son difíciles de escalar.
2. Olvido Catastrófico: Los enfoques fuera de línea que generan datos sintéticos o dependen de conjuntos de datos fijos a menudo sufren de olvido catastrófico, degradando las capacidades generativas generales del modelo y la calidad composicional.
3. Supervisión Estática: Los paradigmas actuales (SFT y RL fuera de línea) optimizan contra ejemplos fijos o datos pregenerados. No logran adaptar sus señales de entrenamiento al comportamiento generativo evolutivo del modelo, lo que los hace insuficientes para rastrear y suprimir contenido inseguro que emerge dinámicamente durante el entrenamiento.
4. Sobrecarga del Modelo de Recompensa: Muchos métodos de alineación de seguridad requieren entrenar o ajustar fino modelos de recompensa especializados para clasificar imágenes como seguras o inseguras, introduciendo una sobrecarga computacional adicional.

Metodología

Los autores proponen SafeDiffusion-R1, un marco de aprendizaje por refuerzo en línea diseñado para alinear los modelos de difusión con restricciones de seguridad sin datos pareados supervisados ni ajuste fino de modelos de recompensa especializados. El marco integra dos componentes principales:

1. Optimización de Política en Línea mediante GRPO

El método emplea Optimización de Política Relativa por Grupos (GRPO) como algoritmo de entrenamiento. A diferencia del PPO estándar o el RL fuera de línea, GRPO opera en política:

• Muestreo Dinámico: El modelo genera continuamente muestras de imágenes a partir de prompts benignos e inseguros durante el entrenamiento.
• Ventaja Basada en Grupos: Para cada prompt, se generan $K$ muestras independientes. Las ventajas se normalizan dentro del grupo ($A_{i,k} = \frac{r_{i,k} - \bar{r}_i}{\sigma_i + \delta}$) en lugar de globalmente. Esto mitiga las discrepancias en la escala de recompensas entre prompts benignos e inseguros, previniendo la sobre-corrección y asegurando actualizaciones estables en relación con la distribución de la política actual.
• Aprendizaje en Política: Al acoplar la optimización de seguridad con la distribución de muestreo evolutiva del modelo, el método evita la discrepancia de distribución y preserva las capacidades generativas generales mientras desaprende progresivamente conceptos inseguros.

2. Recompensa de Dirección Consciente de la Geometría

Para eliminar la necesidad de un clasificador separado seguro/inseguro o un modelo de recompensa, los autores introducen un mecanismo de recompensa de dirección que explota las propiedades geométricas de las incrustaciones de CLIP:

• Vector de Dirección de Seguridad ($v_{safe}$): Se calcula un vector unitario en el espacio de incrustaciones de texto como la diferencia normalizada entre las incrustaciones medias de un pequeño conjunto de anclajes de texto seguros ($T_{safe}$) y anclajes de texto inseguros ($T_{unsafe}$):
  $$v_{safe} = \frac{\bar{z}_{safe} - \bar{z}_{unsafe}}{\|\bar{z}_{safe} - \bar{z}_{unsafe}\|_2}$$
• Dirección Condicional: Para un prompt de entrada $c$, la incrustación de texto $z_T$ se proyecta sobre $v_{safe}$. Si el prompt se identifica como inseguro (proyección negativa), la incrustación se dirige geométricamente hacia la dirección segura:
  $$z'_T = \text{normalize}(z_T + \alpha v_{safe})$$
  donde $\alpha$ es un hiperparámetro de fuerza de dirección.
• Cálculo de Recompensa: La recompensa se calcula como la similitud del coseno entre la incrustación de la imagen ($z_I$) y la incrustación de texto dirigida ($z'_T$) para prompts inseguros, o la incrustación de texto original para prompts seguros. Esto reconfigura la señal de optimización para fomentar la alineación con atributos semánticos seguros sin penalizar explícitamente al modelo por "comprender" el prompt inseguro.

Contribuciones Clave

1. Formulación de Política en Línea: El artículo formula la alineación de seguridad para modelos de difusión T2I como un problema de optimización de política en línea, introduciendo un marco basado en GRPO que acopla el aprendizaje de seguridad con la distribución generativa evolutiva del modelo para prevenir el olvido catastrófico.
2. Mecanismo de Recompensa de Dirección: Se propone una recompensa novedosa consciente de la geometría que representa la seguridad como una dirección en el espacio de incrustaciones de CLIP. Esto permite la supresión de conceptos sin entrenar modelos de recompensa seguros/inseguros dedicados ni depender de datos pareados supervisados.
3. Validación Empírica: Experimentos extensivos demuestran que el marco supera al ajuste fino supervisado y a los métodos de alineación fuera de línea en benchmarks de seguridad, mientras preserva o mejora la calidad de la generación composicional.

Resultados Experimentales

El método se evaluó en Stable Diffusion v1.4 utilizando el benchmark I2P (detección de desnudez) y GenEval (utilidad composicional).

• Rendimiento de Seguridad:
  • Detección de Desnudez: SafeDiffusion-R1 redujo las detecciones de desnudez a 15 (en comparación con 646 para la línea base SD v1.4 y 66 para la línea base RECE).
  • Contenido Inapropiado: El método redujo la tasa general de contenido inapropiado al 18.07% (vs. 48.9% para SD v1.4). Cabe destacar que logró un rendimiento sólido en siete categorías de daño (Odio, Acoso, Violencia, Autolesión, Sexual, Impactante, Actividad ilegal) a pesar de haber sido entrenado principalmente con prompts centrados en la desnudez, demostrando una fuerte generalización fuera de distribución (OOD).
• Preservación de Utilidad:
  • Calidad Composicional: A diferencia de las líneas base que degradan la utilidad, SafeDiffusion-R1 mejoró la puntuación de GenEval del 42.08% (SD v1.4) al 47.83%.
  • Fidelidad Visual: El análisis cualitativo muestra que el modelo preserva la integridad estructural, los detalles faciales y la consistencia de la iluminación, evitando la distorsión, el suavizado excesivo o la incoherencia semántica observada en métodos como EraseDiff o ESD.
  • Puntuación CLIP: El modelo mantuvo puntuaciones CLIP-T cercanas a la línea base (0.311 vs. 0.313), indicando una preservación del seguimiento de instrucciones.

Significado y Afirmaciones

El artículo afirma que SafeDiffusion-R1 aborda el compromiso crítico entre seguridad y utilidad en los modelos de difusión. Al utilizar GRPO en línea y dirección geométrica, el enfoque:

• Elimina Dependencias de Datos: Elimina la necesidad de costosos conjuntos de datos supervisados de seguridad curados.
• Evita la Sobrecarga del Modelo de Recompensa: Evita el costo computacional de entrenar clasificadores de recompensa especializados.
• Garantiza Estabilidad: La formulación en política y la normalización relativa al grupo previenen el olvido catastrófico y el colapso de la distribución.
• Generaliza de Forma Robusta: El método logra un rendimiento de seguridad de vanguardia que se generaliza a prompts inseguros OOD en diversas categorías de daño, incluso cuando se entrena en un conjunto estrecho de conceptos inseguros.

Los autores posicionan este trabajo como una estrategia práctica para la modificación post-entrenamiento que ajusta directamente los modelos preentrenados para suprimir conceptos inseguros sin reentrenar desde cero, manteniéndose compatible con sistemas de lanzamiento público como Stable Diffusion.