Bias mitigation in graph diffusion models

Aquí presento un resumen técnico detallado del artículo "Bias Mitigation in Graph Diffusion Models" (Mitigación de Sesgos en Modelos de Difusión de Grafos), publicado como ponencia en ICLR 2025.

1. El Problema: Sesgos en Modelos de Difusión de Grafos

Los autores identifican dos problemas críticos que degradan la calidad de generación en los modelos de difusión aplicados a grafos (como GDSS, MOOD, GSDM, HGDM):

• Sesgo de Inicio Inverso (Reverse-Starting Bias):

  • Causa: En la práctica, debido a limitaciones en la escala de los datos y la capacidad de aprendizaje de las redes, los modelos de difusión de grafos utilizan estrategias conservadoras en el proceso de difusión hacia adelante (forward diffusion). Esto hace que la distribución de perturbación máxima ($x_T$) no alcance la distribución Gaussiana estándar.
  • El conflicto: Sin embargo, durante el muestreo inverso (reverse sampling), los modelos siempre comienzan asumiendo una distribución Gaussiana estándar ($N(0, I)$).
  • Consecuencia: Existe una discrepancia fundamental entre el punto final de la difusión hacia adelante y el punto de inicio de la difusión inversa, lo que introduce un sesgo significativo que afecta la calidad de la generación.

• Sesgo de Exposición (Exposure Bias):

  • Causa: Durante el entrenamiento, el modelo ve muestras ruidosas generadas a partir de datos limpios reales. Durante la inferencia (muestreo), el modelo genera muestras ruidosas basándose en sus propias predicciones anteriores, que contienen errores acumulados.
  • Consecuencia: Esta discrepancia entre la distribución de entrenamiento y la de inferencia provoca que los errores se acumulen y propaguen a lo largo de los pasos de muestreo, degradando progresivamente la calidad de la muestra final.

2. Metodología Propuesta: Enfoque S++

El artículo propone un enfoque integral llamado S++ que mitiga ambos sesgos sin modificar la arquitectura de la red ni introducir nuevos componentes de aprendizaje (como generadores o discriminadores adicionales). La metodología consta de dos pilares principales:

A. Alineación del Inicio Inverso mediante Muestreo de Langevin

Para abordar el Sesgo de Inicio Inverso, los autores no fuerzan al modelo a aprender una distribución Gaussiana completa (lo cual es difícil en grafos), sino que ajustan el punto de partida del muestreo inverso para que coincida con la distribución de perturbación máxima real del proceso hacia adelante.

• Mecanismo: Utilizan un algoritmo de Muestreo de Langevin guiado por la red de puntuación (score network) preentrenada $s_{\theta, T}(\cdot)$ en el tiempo final $T$.
• Proceso: En lugar de iniciar directamente con ruido Gaussiano, se ejecutan $M$ pasos de Langevin para transicionar desde la distribución Gaussiana estándar hacia la distribución de alta densidad de la perturbación máxima real ($q(x_T)$).
• Beneficio: Esto alinea el punto de inicio con la realidad del modelo, mitigando el sesgo de inicio. Además, al situar el inicio en un estado de "ruido bajo" (donde el modelo es más robusto a errores de predicción), se reduce la propagación del sesgo de exposición.

B. Mecanismo de Corrección de Puntuación (Score Correction)

Para abordar el Sesgo de Exposición y refinar la distribución de estado estacionario, proponen una corrección basada en la diferencia de puntuaciones.

• Concepto: Se entrena una "red de puntuación pseudo" ($s_{\psi, t}$) utilizando un lote de muestras generadas por la red original preentrenada ($s_{\theta, t}$). Dado que las muestras generadas tienen sesgos inherentes, la red pseudo aprende estos sesgos.
• Diferencia de Puntuación: Se calcula la diferencia entre la puntuación original y la pseudo: $\Delta s = s_{\theta, t} - s_{\psi, t}$. Esta diferencia contiene información sobre la estructura de los datos originales ($x_0$) que ha sido perdida o distorsionada.
• Corrección: Se ajusta la puntuación predicha utilizando esta diferencia como una señal de guía:
  $$\hat{s}_{\theta, t} = \frac{s_{\theta, t} + \lambda(s_{\theta, t} - s_{\psi, t})}{\omega}$$
  Donde $\lambda$ controla la dirección (ángulo) y $\omega$ ajusta la magnitud.
• Ventaja: Este método empuja la puntuación corregida hacia la puntuación óptima real, reduciendo el error de predicción y mitigando el sesgo de exposición sin reentrenar la red principal.

3. Contribuciones Clave

1. Análisis Sistemático: Son los primeros en identificar y abordar sistemáticamente tanto el sesgo de inicio inverso como el sesgo de exposición específicamente en el contexto de modelos de difusión de grafos.
2. Solución sin Modificación de Red: El método S++ es "plug-and-play". No requiere cambiar la arquitectura de la red, ni sus parámetros, ni añadir nuevos módulos de aprendizaje (como GANs o discriminadores). Funciona sobre modelos preentrenados existentes.
3. Eficiencia y Robustez: Demuestran que alinear el inicio y corregir la puntuación permite reducir significativamente el número de pasos de muestreo necesarios sin perder calidad, mejorando la eficiencia computacional.
4. Validación Extensa: El enfoque ha sido validado en múltiples modelos (GDSS, GSDM, HGDM, MOOD), diversos datasets (grafos genéricos, moléculas QM9, ZINC250k) y tareas (generación de grafos, optimización de propiedades moleculares).

4. Resultados Experimentales

Los experimentos muestran que S++ supera consistentemente a los modelos base (baselines) y a otros métodos avanzados:

• Generación de Grafos Genéricos: En datasets como Community-small, Enzymes y Grid, S++ logra métricas de estado del arte (SOTA) en medidas de distribución (MMD para grados, coeficientes de agrupamiento y órbitas), superando a los modelos base sin correctores y a menudo superando a los modelos base con correctores tradicionales.
• Generación Molecular: En QM9 y ZINC250k, S++ mejora drásticamente la validez de las moléculas generadas y reduce la distancia Fréchet ChemNet (FCD), indicando una mayor similitud con la distribución real.
• Descubrimiento de Fármacos: En tareas de optimización de propiedades (docking score, QED, SA), la variante MOOD-S++ logra mejores ratios de "hits" novedosos y mejores puntuaciones de acoplamiento que el modelo base MOOD.
• Eficiencia: S++ logra mantener un alto rendimiento incluso con un número reducido de pasos de difusión (ej. $N=100$), mientras que los modelos base sufren una degradación severa de calidad.

5. Significado e Impacto

Este trabajo es fundamental porque:

• Desbloquea el potencial de los modelos de difusión de grafos: Al eliminar los sesgos inherentes que limitan la calidad de generación, permite que estos modelos operen cerca de su límite teórico.
• Ofrece una solución generalizable: Al no depender de arquitecturas específicas (espaciales, espectrales, hiperbólicas), el método S++ puede integrarse en cualquier modelo de difusión de grafos existente.
• Cambia el paradigma de corrección: A diferencia de métodos anteriores que requieren reentrenamiento costoso o componentes adicionales, S++ demuestra que la corrección de sesgos puede lograrse puramente mediante estrategias de muestreo y post-procesamiento de puntuaciones inteligentes.
• Diferenciación con Imágenes: El artículo aclara que los sesgos en grafos son más agudos que en imágenes debido a la estructura de datos y la escala, invalidando soluciones directas de visión por computadora y proponiendo un enfoque adaptado a la naturaleza de los grafos.

En resumen, S++ representa un avance significativo hacia la madurez de los modelos generativos de grafos, ofreciendo una vía práctica y eficiente para lograr generación de alta calidad y robusta.