Likelihood Matching for Diffusion Models

Resumen Técnico: Coincidencia de Verosimilitud para Modelos de Difusión

Planteamiento del Problema
Los modelos de difusión actuales predominantes, particularmente los Modelos Generativos Basados en Puntuación (SGM, por sus siglas en inglés), dependen de la Coincidencia de Puntuación (SM) como su objetivo de entrenamiento principal. Aunque es empíricamente efectivo, SM minimiza la discrepancia entre una función de puntuación parametrizada y la verdadera función de puntuación en varios niveles de ruido. Teóricamente, este enfoque solo minimiza un límite superior de la log-verosimilitud negativa de la distribución de datos original en lugar de la verosimilitud misma. Esta conexión indirecta puede conducir a una pérdida de eficiencia estadística en comparación con la Estimación de Máxima Verosimilitud (MLE), la cual es conocida por producir estimadores consistentes y eficientes. Además, el SM estándar se enfoca exclusivamente en la información del momento de primer orden (la puntuación/gradiente), potencialmente descuidando la información estructural de orden superior contenida en las densidades de transición.

Metodología
Los autores proponen la Coincidencia de Verosimilitud (LM), un marco que apunta directamente a la verosimilitud de la distribución de datos aprovechando la equivalencia entre la verosimilitud de los datos objetivo y la verosimilitud de la trayectoria del proceso de difusión inversa.

1. Fundamento Teórico:
   El método se basa en la propiedad de que la verosimilitud conjunta de la trayectoria de difusión inversa es equivalente a la verosimilitud de la distribución de datos original (salvo constantes relacionadas con el proceso de difusión hacia adelante). Los autores establecen que maximizar la verosimilitud exacta de la trayectoria del proceso inverso es equivalente a maximizar la log-verosimilitud de los datos.

2. Estimación de Cuasi-Máxima Verosimilitud (QMLE):
   Dado que las densidades de transición inversas exactas son intratables, los autores las aproximan utilizando un enfoque de Cuasi-Máxima Verosimilitud. Modelan cada densidad de transición inversa como una distribución Gaussiana donde la media y la covarianza condicional se ajustan al verdadero proceso inverso.

• Idea Clave: La media y la covarianza condicional del proceso inverso dependen no solo de la función de puntuación ($\nabla \log q_t$), sino también de la función Hessiana ($\nabla^2 \log q_t$).
• Parametrización: Los autores parametrizan tanto la puntuación como la Hessiana utilizando redes neuronales (específicamente U-Nets). Para gestionar la complejidad computacional, la Hessiana se aproxima frecuentemente utilizando una estructura de diagonal más bajo rango.

3. El Objetivo de Coincidencia de Verosimilitud:
   El objetivo propuesto, la Coincidencia de Verosimilitud (LM), se deriva minimizando la log-verosimilitud cuasi-negativa de las transiciones inversas.

• El objetivo unifica la coincidencia de puntuación y la coincidencia de covarianza.
• Incorpora un mecanismo de ponderación de verosimilitud implícito mediante la matriz de covarianza ajustada $(I + \sigma^2 H_t)^{-1/2}$, que ajusta automáticamente el peso del término de coincidencia de puntuación basado en la curvatura local.
• El entrenamiento utiliza un muestreador estocástico que selecciona aleatoriamente pasos temporales a lo largo de la trayectoria de difusión para calcular la pérdida, asegurando el uso eficiente de la información evolutiva de la SDE hacia adelante.

4. Muestreador Estocástico:
   Se introduce un nuevo procedimiento de muestreo que utiliza tanto la información de la puntuación como de la Hessiana aprendidas. El paso inverso se define como:
   $$\tilde{Y}_{t-1} = \hat{\mu}_{t-1|t}(\tilde{Y}_t) + \hat{\Sigma}_{t-1|t}^{1/2}(\tilde{Y}_t) Z_t$$
   donde la media $\hat{\mu}$ y la covarianza $\hat{\Sigma}$ son funciones de la puntuación y la Hessiana estimadas, respectivamente.

Contribuciones Clave

• Nuevo Objetivo de Entrenamiento: El artículo introduce un objetivo de entrenamiento de cuasi-verosimilitud (LM) que aproxima la log-verosimilitud de la trayectoria inversa, combinando efectivamente la coincidencia de puntuación y la coincidencia de covarianza con ponderación de verosimilitud implícita.
• Muestreo Mejorado por la Hessiana: Se deriva un muestreador estocástico que aprovecha tanto la información de la puntuación como de la Hessiana, ofreciendo una aproximación más precisa de la densidad de transición inversa que los muestreadores estándar de tipo DDPM.
• Garantías Teóricas:
  • Convergencia: Se proporcionan garantías de convergencia no asintótica en distancia de variación total. El límite de error escala como $O(d^3 \log^{4.5} T / T)$ para el error del paso inverso, con términos adicionales para los errores de estimación de la puntuación y la Hessiana.
  • Consistencia: El artículo demuestra la consistencia de la propuesta QMLE bajo un entorno paramétrico oráculo, mostrando que el estimador converge a los verdaderos parámetros de los datos a medida que aumentan el tamaño de la muestra y los pasos de difusión.
• Implementación Eficiente: El uso de aproximaciones de la Hessiana de bajo rango y la fórmula de Sherman-Morrison-Woodbury permite una computación eficiente de la verosimilitud y los pasos de muestreo.

Resultados Experimentales
Los autores validan el método tanto en conjuntos de datos sintéticos como en el mundo real:

• Datos Sintéticos: En mezclas de distribuciones Gaussianas y t, LM superó consistentemente a la Coincidencia de Puntuación (SM) estándar, particularmente al capturar modos bien separados. En una distribución "Swiss-roll" altamente enrollada (que representa variedades curvas), LM demostró un rendimiento superior, validando el beneficio de la covarianza informada por la Hessiana.
• Estimación de Parámetros: En un entorno de mezcla Gaussiana paramétrica, LM logró un Error Absoluto Medio (MAE) y un error estándar más bajos en la estimación de parámetros en comparación con SM, apoyando las afirmaciones de consistencia teórica.
• Conjuntos de Datos de Imágenes (MNIST, CIFER-10, CelebA):
  • LM logró una Distancia de Fréchet Inception (FID) más baja y mejores puntuaciones de Log-Verosimilitud Negativa (NLL) en comparación con las líneas base de SM.
  • El rendimiento mejoró con el rango de la aproximación de la Hessiana ($r$), alcanzando su punto máximo alrededor de $r=20-30$.
  • Relación Velocidad-Calidad: LM logró un FID más bajo para un tiempo de ejecución comparable y requirió menos pasos de muestreo inverso para alcanzar una alta calidad visual. Por ejemplo, en MNIST, LM produjo dígitos reconocibles tras solo 20 iteraciones inversas, mientras que las muestras de SM permanecían borrosas.

Significancia y Reivindicaciones
El artículo afirma que la Coincidencia de Verosimilitud ofrece una alternativa estadísticamente más rigurosa a la Coincidencia de Puntuación al apuntar directamente a la verosimilitud de los datos a través de un marco de cuasi-verosimilitud. Al incorporar información de segundo orden (la Hessiana), el método captura densidades de transición más ricas y mejora la eficiencia del muestreo, particularmente en regímenes de pocos pasos y en variedades de datos complejas y curvas. Los autores enfatizan que, si bien LM incurre en costos computacionales más altos por paso debido al modelado de la Hessiana, la reducción en el número de pasos de muestreo requeridos y la mejora en la calidad de generación (FID y NLL) validan su efectividad para entrenar modelos generativos de alta fidelidad. El trabajo cierra la brecha entre la eficiencia estadística de MLE y el éxito práctico de los modelos de difusión.