Resumen Técnico: Modelos de Difusión Ajustados con Metropolis (MADM)

Enunciado del Problema

El muestreo a partir de modelos de difusión basados en puntuación (score-based) incurre inherentemente en un sesgo que surge de dos fuentes: el error de aproximación de la función de puntuación aprendida y el error de discretización temporal introducido por los integradores numéricos. En el marco estándar Predictor-Corrector, un paso predictor (a menudo una EDO de Flujo de Probabilidad) mueve una muestra entre niveles de ruido, seguido de un paso corrector (típicamente el Algoritmo de Langevin No Ajustado, ULA) para refinar la muestra en el nivel de ruido actual.

Aunque ULA utiliza la función de puntuación para proponer movimientos, es en sí mismo un muestreador sesgado porque discretiza un proceso de difusión de Langevin continuo. En consecuencia, el "corrector" en el esquema Predictor-Corrector introduce su propio sesgo de discretización, fallando en dejar la distribución marginal objetivo invariante. Los mecanismos estándar de aceptación-rechazo de Metropolis-Hastings (MH) o Barker, que corrigen dicho sesgo en las Cadenas de Markov de Monte Carlo (MCMC), no pueden aplicarse directamente a los modelos de difusión. Esto se debe a que MH requiere evaluar la razón de las densidades objetivo, $p_t(\tilde{x})/p_t(x)$, lo cual es intratable en modelos de difusión donde solo está disponible la función de puntuación $s_\phi(x, t) \approx \nabla_x \log p_t(x)$.

Metodología

Los autores proponen Modelos de Difusión Ajustados con Metropolis (MADM), un marco que reemplaza los correctores ULA sesgados por pasos de Langevin corregidos mediante aceptación-rechazo. La innovación central es un método para calcular la probabilidad de aceptación necesaria utilizando únicamente la función de puntuación, evitando la necesidad de la razón de densidades intratable.

1. Identidad de la Razón de Densidades

Los autores establecen que el logaritmo de la razón de densidades objetivo puede expresarse como una integral de línea de la función de puntuación a lo largo del camino entre el estado actual $x$ y la propuesta $\tilde{x}$:
$$\log \left( \frac{p_t(\tilde{x})}{p_t(x)} \right) = \int_0^1 \langle s(x + u(\tilde{x} - x), t), \tilde{x} - x \rangle \, du$$
Esta identidad permite estimar la razón de densidades utilizando únicamente evaluaciones de la puntuación.

2. Corrección Exacta mediante Fábrica de Bernoulli

Para lograr una corrección exacta sin evaluar la densidad, los autores proponen un algoritmo basado en una fábrica de Bernoulli de dos monedas (Gonçalves et al., 2017).

• Supuesto: Existe una cota $C(x, \tilde{x})$ tal que el valor absoluto del integrando en la integral de línea está acotado por $C$.
• Mecanismo: El algoritmo construye un estimador no sesgado $W$ del exponencial de la integral (escalado por la cota) utilizando una serie de potencias truncada por Poisson.
• Procedimiento: El algoritmo utiliza un procedimiento de "dos monedas" para simular una decisión de aceptación-rechazo de Barker. Primero rechaza con una probabilidad fija $\alpha'$, y si no es rechazado, lanza una moneda con probabilidad $W$. Si la moneda tiene éxito, la propuesta se acepta; de lo contrario, la ronda se reinicia.
• Resultado: Este procedimiento produce un corrector de Langevin ajustado exactamente a Barker que deja la distribución objetivo $p_t$ invariante, siempre que la puntuación sea exacta y se cumpla el supuesto de la cota.

3. Aproximación Eficiente mediante la Regla de Simpson

Reconociendo que el método exacto de fábrica de Bernoulli requiere un número aleatorio y potencialmente grande de evaluaciones de puntuación, los autores proponen una aproximación escalable.

• Método: La integral de línea se aproxima de manera determinista utilizando la regla 1/3 de Simpson (una regla de cuadratura de Newton-Cotes). Esto requiere evaluar la puntuación en el estado actual, en la propuesta y en el punto medio.
• Aceptación: La estimación resultante se introduce en la probabilidad de aceptación estándar de Metropolis-Hastings.
• Propiedades: Aunque esto introduce un pequeño sesgo debido a la truncación de la cuadratura, la aproximación logra una precisión de orden $O(h^{5/2})$ en el tamaño del paso $h$. Es computacionalmente económica, requiriendo solo una evaluación adicional de puntuación en el punto medio por propuesta.

Contribuciones Clave

1. Marco MADM: Introducción de un marco de corrector ajustado con Metropolis para modelos de difusión que reemplaza sistemáticamente los pasos ULA sesgados.
2. Método Exacto: El primer método exacto para ajustar correcciones de Langevin en modelos de difusión utilizando una fábrica de Bernoulli de dos monedas, requiriendo únicamente la función de puntuación y una cota sobre la integral de línea.
3. Aproximación Escalable: Una aproximación práctica mediante la regla de Simpson que logra una precisión de alto orden ($O(h^{5/2})$) con un costo marginal cercano a cero (una evaluación extra de puntuación).
4. Validación Empírica: Demostración de una mejora en la calidad de las muestras tanto en conjuntos de datos sintéticos como en imágenes.

Resultados

Los autores evalúan MADM en conjuntos de datos sintéticos 2D y en benchmarks estándar de imágenes (CIFAR-10, FFHQ, AFHQv2, ImageNet-64) utilizando modelos preentrenados de Karras et al. (2022).

• Conjuntos de Datos Sintéticos: En conjuntos de datos 2D (espiral, embudo, triángulo de Sierpinski, molinillo), el muestreador ancestral estándar (equivalente a un predictor con un corrector ULA) produjo muestras con valores atípicos significativos fuera del soporte de la distribución de datos verdadera. Los pasos correctores de MADM movieron con éxito estas muestras de vuelta al soporte de los datos, eliminando los valores atípicos.
• Conjuntos de Datos de Imágenes:
  • Para modelos incondicionales (CIFAR-10, FFHQ, AFHQv2), aplicar un paso corrector MADM por paso predictor produjo mejoras consistentes, aunque modestas, en la Distancia de Inception Fréchet (FID) en comparación con las líneas base deterministas (resolutores de EDO Heun y Euler).
  • Para el modelo condicional por clase ImageNet-64, las ganancias fueron más pronunciadas. Usar MADM con un predictor Euler produjo mejoras mayores en FID que usar un predictor Heun de orden superior sin corrección. Esto sugiere que los ajustes MH son más efectivos para mitigar errores de discretización que simplemente aumentar el orden del resolutor de EDO cuando el número de pasos es grande.
• Costo Computacional: Se encontró que el método exacto (Algoritmo 2) era computacionalmente costoso debido al número aleatorio de evaluaciones de puntuación requeridas. En consecuencia, los experimentos utilizaron principalmente la aproximación de la regla de Simpson (Algoritmo 3), que proporcionó ganancias consistentes con una sobrecarga manejable.

Significado y Afirmaciones

El artículo afirma proporcionar el primer método exacto para ajustar correcciones de Langevin en modelos de difusión, abordando el sesgo inherente de los correctores ULA. Al aprovechar la función de puntuación para estimar razones de densidades mediante integrales de línea, MADM permite el uso de criterios de aceptación rigurosos de MCMC en modelado generativo sin requerir acceso a la densidad de datos intratable.

Los autores posicionan a MADM como una extensión natural del marco Predictor-Corrector que mejora la calidad de las muestras. Señalan que, aunque el método exacto es teóricamente significativo, la aproximación de la regla de Simpson ofrece una solución práctica y escalable que mejora consistentemente el FID en benchmarks de imágenes. El trabajo asume una estimación perfecta de la puntuación, reconociendo esto como una limitación inherente de los modelos de difusión y no como un defecto específico de su método. Se sugiere trabajo futuro centrado en desarrollar variantes eficientes del método exacto o de sus subrutinas.