Resumen Técnico: Un Marco Unificado para el Olvido en Modelos de Difusión con Divergencia f

1. Enunciado del Problema

Los modelos de difusión (DM) de texto a imagen (T2I), entrenados con conjuntos de datos a gran escala extraídos de internet (por ejemplo, LAION-5B), han demostrado capacidades generativas notables, pero también heredan propiedades indeseables, incluida la capacidad de generar contenido No Seguro para el Trabajo (NSFW), estilos artísticos protegidos por derechos de autor e información personal. Aunque reentrenar los modelos en conjuntos de datos curados es la solución teórica para eliminar conceptos específicos, resulta prohibitivo desde el punto de vista computacional y a menudo impráctico debido a la eliminación de datos o a las limitaciones de recursos.

Los algoritmos de olvido existentes suelen basarse en un paradigma donde la generación del modelo para un concepto "objetivo" (que debe ser eliminado) se alinea con un concepto "ancla" (por ejemplo, un concepto nulo, una superclase o un concepto semánticamente relacionado). Esta alineación se logra casi exclusivamente minimizando el Error Cuadrático Medio (MSE) entre las salidas del desruidizador del objetivo y del ancla. Teóricamente, esta minimización del MSE corresponde a minimizar la divergencia de Kullback-Leibler (KL) entre dos distribuciones gaussianas condicionadas a los conceptos respectivos. Sin embargo, este enfoque limita el paisaje de optimización a una única métrica de divergencia, lo que puede ser subóptimo para equilibrar la eficacia del olvido con la preservación de la fidelidad generativa para conceptos no objetivo.

2. Metodología

Los autores proponen f-DMU, un marco unificado que generaliza el objetivo de olvido desde la divergencia KL estándar hacia cualquier divergencia f. El marco se basa en la observación de que el problema de olvido puede formularse como la minimización de la divergencia entre las distribuciones de probabilidad del modelo original $\Phi$ (condicionado al concepto ancla $c$) y el modelo olvidado $\hat{\Phi}$ (condicionado al concepto objetivo $c^*$) en pasos de tiempo específicos del proceso de difusión.

El objetivo central se define como:
$$\min_{\hat{\Phi}} \mathbb{E}_{c^*, c, t} \left[ D_f \left( p_{\Phi}(x_{t-1}|x_t, c) \parallel p_{\hat{\Phi}}(x_{t-1}|x_t, c^*) \right) \right]$$

El marco ofrece dos enfoques distintos para resolver este objetivo, dependiendo de las propiedades de la divergencia f elegida:

A. Pérdidas de Forma Cerrada

Para un subconjunto de divergencias f donde existe una expresión de forma cerrada para distribuciones gaussianas (específicamente cuando las varianzas están fijas, una práctica estándar en DM), los autores derivan funciones de pérdida específicas. Los casos clave incluyen:

• Distancia de Hellinger al Cuadrado ($H^2$): Produce una función de pérdida que implica un decaimiento exponencial del MSE: $J_{H^2} \propto -\exp(-\text{MSE})$. Esto resulta en gradientes acotados que disminuyen a medida que aumenta el MSE, mitigando el impacto de los valores atípicos y evitando cambios abruptos en el modelo.
• Divergencia $\chi^2$ de Pearson: Produce una función de pérdida que implica un crecimiento exponencial del MSE: $J_{\chi^2} \propto \exp(\text{MSE})$. Esto resulta en gradientes no acotados y de crecimiento rápido para errores grandes, lo que conduce a actualizaciones más agresivas.
• Divergencia de Jeffreys: Se reduce a una formulación estándar de MSE bajo suposiciones específicas.

B. Pérdidas Variacionales

Para divergencias f que carecen de una expresión de forma cerrada (por ejemplo, Variación Total, o cuando se desea una formulación más general), los autores utilizan la representación variacional de las divergencias f. Esto transforma el problema en un juego de optimización min-max:
$$\min_{\hat{\Phi}} \max_{T} \mathbb{E} \left[ \mathbb{E}_{p_{\Phi}}[T(\Phi)] - \mathbb{E}_{p_{\hat{\Phi}}}[f^*(T(\hat{\Phi}))] \right]$$
Aquí, $T$ actúa como un discriminador (parametrizado por una red neuronal) que estima la divergencia, mientras que $\hat{\Phi}$ actúa como el generador que intenta minimizarla. Este enfoque permite que el marco aplique cualquier divergencia f a la tarea de olvido.

3. Análisis Teórico

El artículo proporciona un análisis teórico riguroso del marco propuesto:

• Dinámica de Gradientes: Los autores analizan las magnitudes de los gradientes de las pérdidas de forma cerrada. Demuestran que $H^2$ y $\chi^2$ exhiben comportamientos opuestos en comparación con el MSE/KL estándar. $H^2$ suprime los gradientes para errores grandes (valores atípicos), mientras que $\chi^2$ los amplifica. Esta ponderación por muestra es una propiedad inherente a la elección de la divergencia f, distinta de los programadores estándar de la tasa de aprendizaje.
• Propiedades de Convergencia: Utilizando la teoría de sistemas dinámicos no lineales, los autores prueban la estabilidad exponencial local de la formulación variacional alrededor de los puntos de equilibrio. Derivan límites sobre la tasa de convergencia basados en los valores propios de la matriz jacobiana. Crucialmente, muestran que la elección de la divergencia f impacta la velocidad de convergencia; específicamente, las divergencias con segundas derivadas más grandes de sus conjugados de Fenchel en el equilibrio (por ejemplo, $H^2$ y Jensen-Shannon) convergen teóricamente más rápido que otras como KL o $\chi^2$.
• Búsqueda de Modos vs. Cobertura de Modos: El artículo discute cómo diferentes divergencias afectan la distribución de las imágenes generadas. $H^2$ y Jensen-Shannon exhiben propiedades de "búsqueda de modos" de "nivel medio", mientras que KL y $\chi^2$ son de "cobertura de modos", asignando potencialmente una probabilidad más alta a las áreas entre los modos.

4. Resultados Experimentales

El marco se evaluó en modelos Stable Diffusion (v1.4, v1.5, v2.1, XL y 3) y FLUX, dirigiéndose a la eliminación de objetos (por ejemplo, R2D2), estilos artísticos (por ejemplo, Van Gogh, Monet) y desnudez.

• Estabilidad de Gradientes: La validación empírica confirma que $H$-DMU (utilizando $H^2$) mantiene gradientes acotados durante todo el entrenamiento, evitando la generación degenerada repentina de imágenes observada en los métodos basados en MSE durante las etapas intermedias de ajuste fino.
• Eficacia del Olvido vs. Preservación:
  • Forma Cerrada: $H$-DMU supera consistentemente al MSE estándar en la preservación de conceptos no objetivo (medido por la Distancia de Inception del Núcleo, KID) mientras logra una eliminación comparable o mejor de los conceptos objetivo. Por el contrario, $\chi^2$-DMU logra una eliminación más fuerte, pero a menudo a costa de la calidad generativa y la preservación del prior.
  • Variacional: Los métodos variacionales (por ejemplo, $H_v$-DMU) tienden a ser más "agresivos", logrando una eliminación rápida en menos iteraciones, pero a menudo resultan en puntuaciones KID más altas (menor calidad de imagen) debido a estimaciones de divergencia toscas en lotes pequeños.
• Robustez: En el contexto de la eliminación de desnudez, $H$-DMU demostró una robustez superior frente a prompts adversarios (por ejemplo, Ring-A-Bell, MMA-Diffusion) en comparación con métodos de última generación como CAbl, DoCo y SLD, generando significativamente menos imágenes NSFW bajo ataque.
• Eliminación de Múltiples Conceptos: El marco manejó con éxito la eliminación secuencial de 10 estilos artísticos, mostrando $H$-DMU compensaciones superiores entre eliminación y preservación en comparación con MSE.

5. Significado y Afirmaciones

El artículo afirma presentar el primer marco unificado para el olvido en modelos de difusión basado en divergencias f. Sus contribuciones principales son:

1. Generalización: Ir más allá de la suposición implícita de KL/MSE hacia una familia flexible de divergencias, permitiendo a los practicantes seleccionar el objetivo óptimo basándose en los objetivos específicos de la aplicación (por ejemplo, eliminación agresiva vs. preservación de alta fidelidad).
2. Fundamentación Teórica: Proporcionar un análisis teórico de la escalado de gradientes y las propiedades de convergencia, vinculando la elección de la divergencia con la estabilidad y velocidad algorítmica.
3. Utilidad Práctica: Demostrar que instancias específicas, particularmente la distancia de Hellinger al cuadrado de forma cerrada ($H$-DMU), ofrecen una opción predeterminada superior para la mayoría de las tareas de olvido al equilibrar la eliminación efectiva de conceptos con la preservación de las capacidades generativas del modelo.

Los autores enfatizan que su marco es agnóstico al modelo (aplicable también a modelos basados en flujo) y no requiere reentrenar desde cero, ofreciendo una alternativa computacionalmente eficiente y teóricamente sólida a las técnicas existentes de eliminación de conceptos.