Bridging data-driven priors via the score function for posterior sampling -- Comparative review and experimental study

Resumen Técnico: Puentes entre Priores Basados en Datos mediante la Función de Puntuación para el Muestreo de la Posterior

Planteamiento del Problema
Las tareas de restauración de imágenes (IR), como la eliminación de ruido, el desenfoque, el inpainting y la superresolución, se formulan típicamente como problemas inversos lineales ($y = Ax + n$). Dentro de un marco bayesiano, la resolución de estos problemas mal definidos requiere definir una distribución a priori $p(x)$ para regularizar la solución. Mientras que los priores tradicionales basados en modelos (por ejemplo, la Variación Total) están bien comprendidos, a menudo no logran capturar las estructuras sofisticadas de las imágenes naturales. Por el contrario, los priores basados en datos aprendidos mediante aprendizaje profundo (por ejemplo, Redes Generativas Antagónicas, Autoencoders Variacionales, Flujos de Normalización) ofrecen una expresividad superior, pero plantean desafíos significativos para el muestreo de la posterior.

Un desafío principal en la IR bayesiana es la alta dimensionalidad del espacio de la imagen y la falta de suavidad de las funciones de regularización, lo que dificulta la eficiencia de los algoritmos de Monte Carlo por Cadenas de Markov (MCMC). Los métodos existentes suelen recurrir al Algoritmo de Langevin No Ajustado (ULA) para muestrear de la posterior. Sin embargo, el ULA puede tener dificultades con problemas severamente mal condicionados debido a las restricciones del tamaño de paso vinculadas a la constante de Lipschitz de la verosimilitud. Además, muchos de los avanzados priores basados en datos (como aquellos basados en denoisers o modelos generativos basados en puntuación) se definen implícitamente, lo que dificulta el muestreo directo sin un marco teórico unificado.

Metodología
El artículo propone un marco unificado para muestrear de distribuciones posteriores definidas por diversos priores basados en datos, siempre que puedan expresarse mediante su función de puntuación o score function ($s(x) = \nabla_x \log p(x)$). La metodología central se basa en el Muestreador de Gibbs Dividido con Langevin (LwSGS).

1. Aumento de Datos Asintóticamente Exacto (AXDA): El algoritmo introduce una variable auxiliar $z$ para descomponer la distribución posterior en una forma aumentada:
   $$\pi_\rho(x, z) \propto \exp\left[ -f(x, y) - g(z) - \frac{1}{2\rho^2}\|x - z\|^2 \right]$$
   donde $f(x,y)$ es el término de fidelidad de los datos y $g(z)$ es el potencial de regularización. A medida que el parámetro de acoplamiento $\rho \to 0$, la distribución marginal de $x$ converge a la posterior objetivo.

2. Muestreo de Gibbs Dividido (SGS): El algoritmo alterna entre el muestreo de $x$ y $z$:

   • Muestreo de $x$: Dada $z$, la distribución condicional $p(x|y, z)$ es Gaussiana (para problemas inversos gaussianos lineales), lo que permite un muestreo eficiente y en forma cerrada. Crucialmente, este paso desacopla el tamaño de paso de la constante de Lipschitz de la verosimilitud, abordando una limitación importante del ULA estándar.
   • Muestreo de $z$: Dada $x$, la distribución condicional $p(z|x)$ se muestrea utilizando un paso de Monte Carlo de Langevin (LMC). Este paso utiliza explícitamente la función de puntuación del prior:
     $$z^{(t+1)} = z^{(t)} + \gamma \left( s(z^{(t)}) + \frac{1}{\rho^2}(x^{(t)} - z^{(t)}) \right) + \sqrt{2\gamma}\epsilon^{(t)}$$
     donde $s(z) = \nabla \log p(z)$ es la función de puntuación del prior.

3. Representación Unificada de Priores: El artículo demuestra que cuatro clases distintas de priores basados en datos pueden unificarse bajo esta perspectiva de la función de puntuación:

   • Regularización por Denosing (RED): La función de puntuación se deriva directamente del residuo de eliminación de ruido ($s(x) \propto -(x - D(x))$).
   • Modelos Generativos Basados en Puntuación (SGM): La función de puntuación se aprende explícitamente mediante el ajuste de puntuación (score matching) o modelos de difusión probabilística de eliminación de ruido.
   • Flujos de Normalización (NF): La función de puntuación se calcula mediante diferenciación automática de la densidad logarítmica derivada de la fórmula de cambio de variables.
   • Regularizadores de Cresta Convexa (CRR): La función de puntuación corresponde a la salida de una CNN de una sola capa libre de sesgo.

Contribuciones Clave

• Unificación de Priores: El artículo establece que los priores RED, SGM, NF y CRR, a pesar de sus diferentes orígenes y motivaciones, pueden integrarse en un solo marco de muestreo a través de sus funciones de puntuación.
• Versatilidad Algorítmica: Demuestra que LwSGS es un algoritmo genérico capaz de manejar cualquier log-prior diferenciable con una función de puntuación explícita o eficientemente computable, superando las limitaciones de los esquemas simples de ULA.
• Desacoplamiento de Restricciones: Al dividir el problema, el método permite que el paso de muestreo para el prior sea independiente del condicionamiento del término de fidelidad de los datos, mejorando potencialmente las propiedades de mezcla en problemas mal definidos.

Resultados Experimentales
Los autores evaluaron el marco LwSGS frente a sus contrapartes basadas en ULA (PnP-ULA, NF-ULA y los nuevos SGM-ULA y CRR-ULA) en dos tareas: inpainting de imágenes y superresolución de una sola imagen.

• Conjuntos de Datos: Los experimentos se realizaron con 30 imágenes RGB de los conjuntos de datos FFHQ e ImageNet.
• Métricas de Rendimiento: El estudio utilizó PSNR, SSIM, LPIPS (similitud perceptual) e Intervalo de Autocorrelación Integrada (IAT) para medir la calidad de la reconstrucción y la eficiencia del muestreo.
• Hallazgos:
  • Calidad de Reconstrucción: LwSGS produjo consistentemente reconstrucciones de alta calidad que coinciden estrechamente con los detalles de la verdad de campo (ground-truth). En tareas de inpainting, llenó eficazmente las regiones faltantes; en superresolución, preservó detalles intrincados.
  • Comparación con ULA: LwSGS generalmente superó a los métodos basados en ULA, particularmente cuando se utilizaron priores RED y SGM, produciendo detalles más nítidos y mayor fidelidad. Incluso con los priores NF y CRR, LwSGS mostró un rendimiento competitivo o superior.
  • Eficiencia: El algoritmo LwSGS demostró una convergencia más rápida, indicada por valores de IAT más bajos y funciones de autocorrelación de decaimiento más rápido en comparación con ULA.
  • Cuantificación de la Incertidumbre: Como método de muestreo, LwSGS permitió la estimación de la incertidumbre píxel a píxel (desviación estándar), mostrando una mayor incertidumbre en las regiones de inpainting, como era de esperar.
  • Aplicación en el Mundo Real: El marco se aplicó para restaurar muestras geológicas reales (imágenes desenfocadas). Usando un prior RED, el método recuperó con éxito detalles nítidos a partir de imágenes desenfocadas, demostrando utilidad práctica en las geociencias.

Significancia y Reivindicaciones
El artículo afirma que el marco LwSGS proporciona una alternativa versátil y sistemática al ULA para problemas inversos bayesianos que involucran priores basados en funciones de puntuación. Su significancia radica en:

1. Amplia Aplicabilidad: No se limita a modelos generativos específicos, sino que puede manejar cualquier prior donde la función de puntuación sea accesible, cerrando la breancia entre diversos paradigmas de regularización basados en datos.
2. Eficiencia Práctica: La estrategia de división mitiga eficazmente los problemas de condicionamiento inherentes al ULA, lo que conduce a una mejor mezcla y una convergencia más rápida en espacios de alta dimensión.
3. Utilidad en el Mundo Real: La restauración exitosa de muestras geológicas demuestra que el método es lo suficientemente robusto para flujos de trabajo operativos donde los protocolos de adquisición pueden ser subóptimos.

Los autores concluyen que, si bien el presente trabajo se centra en priores específicos, el marco unificado es lo suficientemente general como para manejar cualquier distribución posterior definida por una amplia clase de priores basados en funciones de puntuación, allanando el camino para la futura integración de diversos modelos de aprendizaje profundo en la inferencia bayesiana.