Resumen Técnico: Estimación del Score Posterior Exacto para la Resolución de Problemas Inversos Lineales

Planteamiento del Problema

Los problemas inversos lineales consisten en recuperar una señal desconocida $x_0$ a partir de mediciones ruidosas $y = Ax_0 + \eta$, donde $A$ es un operador de propagación conocido y $\eta$ es el ruido de observación. Estos problemas son omnipresentes en el procesamiento de imágenes (por ejemplo, inpainting, superresolución, desenfoque) y suelen ser mal planteados, lo que significa que múltiples señales son consistentes con la misma observación. El objetivo es muestrear de la distribución posterior $p(x_0|y)$ en lugar de encontrar un único estimador puntual, para capturar la incertidumbre y equilibrar la fidelidad de los datos con la plausibilidad de la prioridad (prior).

Si bien los modelos de difusión y basados en flujo proporcionan potentes prioridades de datos al aprender a revertir la corrupción gaussiana, estos proporcionan intrínsecamente el score incondicional $\nabla_{x_t} \log p(x_t)$. Resolver problemas inversos requiere el score posterior $\nabla_{x_t} \log p(x_t|y)$. Los enfoques existentes se dividen en dos categorías con limitaciones significativas:

1. Métodos sin entrenamiento (por ejemplo, DPS, DDNM, $\Pi$GDM) mantienen fijo un denoiser preentrenado y aplican correcciones aproximadas de correspondencia de mediciones (por ejemplo, pasos de gradiente, proyecciones) en cada paso inverso. Estos introducen sesgo, lo que a menudo conduce al sobresuavizado, estructuras alucinadas o una incertidumbre mal calibrada.
2. Métodos basados en entrenamiento (por ejemplo, Palette, InDI) entrenan un nuevo modelo condicional desde cero o mediante destilación. Aunque son expresivos, a menudo abandonan la estructura de denoising del prior, requiriendo que la red aprenda la dependencia del operador de extremo a extremo, lo que impide un inicio rápido (warm-starting) eficiente desde checkpoints preentrenados y aumenta los costes de entrenamiento.

Metodología: Exact Posterior Score (EPS)

El artículo deriva una expresión de forma cerrada para el score posterior exacto bajo problemas inversos gaussianos lineales e interpolantes gaussianos generales. La idea central es que el muestreo posterior sigue siendo un problema de denoising pero con una entrada consciente de la medición y una covarianza de ruido anisotrópica dependiente del operador.

Derivación Teórica

Los autores demuestran que el score posterior en el tiempo $t$ puede expresarse como:
$$\nabla_{x_t} \log p(x_t|y) = \frac{1}{\beta_t^2} \left( \alpha_t D_{\Sigma^\star(t)}(\mu^\star(x_t, y, t)) - x_t \right)$$
donde:

• $D_{\Sigma^\star(t)}$ es un denoiser anisotrópico que estima $E[x_0 | x_0 + \xi = \mu^\star]$ con $\xi \sim \mathcal{N}(0, \Sigma^\star(t))$.
• $\mu^\star(x_t, y, t)$ es el pivote posterior, una fusión ponderada por precisión del estado ruidoso actual $x_t$ y la medición $y$.
• $\Sigma^\star(t)$ es la covarianza posterior, derivada de los parámetros del problema inverso ($A, \sigma_y$) y el esquema de difusión ($\alpha_t, \beta_t$).

Crucialmente, $\mu^\star$ actúa como un estadístico suficiente para $(x_t, y)$, y $\Sigma^\star$ captura la incertidumbre anisotrópica: las direcciones observadas por $A$ se vuelven más ciertas, mientras que las direcciones no observadas (en el nul espacio de $A$) permanecen inciertas.

El Objetivo de Entrenamiento EPS

Los autores definen el Exact Posterior Score (EPS) como un objetivo de entrenamiento que aprovecha esta identidad:

1. Transformación de Entrada: En lugar de entrenar sobre la entrada ruidosa estándar $x_t$, la red se entrena sobre el pivote $\mu^\star(x_t, y, t)$.
2. Ruido Anisotrópico: Los objetivos de entrenamiento apuntan al denoising de $\mu^\star$ bajo la covariancia anisotrópica específica $\Sigma^\star(t)$.
3. Simulación Eficiente: Una contribución teórica clave (Proposición 3) muestra que la corrupción anisotrópica requerida es inducida naturalmente por la construcción de forma cerrada de $\mu^\star$. Por lo tanto, se puede entrenar utilizando el muestreo de ruido isotrópico estándar ($x_t = \alpha_t x_0 + \beta_t \epsilon$) y calcular $\mu^\star$ analíticamente, evitando la necesidad de muestrear directamente el ruido anisotrópico.
4. Arquitectura: La arquitectura de la red es idéntica a los denoisers preentrenados estándar, con los canales de entrada expandidos para concatenar el pivote $\mu^\star$ y la observación bruta $y$. Esto permite que EPS sea entrenado desde cero o, más eficientemente, ajustado (fine-tuned) a partir de un denoiser incondicional preentrenado.

Inferencia

En la inferencia, EPS utiliza el mismo sampler (por ejemplo, Euler ODE) que el backbone subyacente. La única modificación es reemplazar la llamada al denoiser estándar $D_\theta(x_t, t)$ por $D_\theta(\mu^\star(x_t, y, t), y, t)$. Esto requiere ningún gradiente de verosimilitud, ninguna proyección y ningún bucle de optimización interno durante el muestreo.

Contribuciones Clave

1. Identidad de Exact Posterior Score: La derivación de un score posterior de forma cerrada para problemas inversos gaussianos lineales, mostrando que el muestreo posterior se reduce al denoising en un pivote desplazado bajo una covariancia anisotrópica.
2. Algoritmo EPS: Un objetivo de entrenamiento y un algoritmo de muestreo que preserva la estructura de entrada/salida de los preentrenamientos estándar. Permite el ajuste eficiente de checkpoints preentrenados y utiliza el sampler del backbone sin cambios durante la inferencia.
3. Eficiencia y Rendimiento: EPS logra resultados de vanguardia con significativamente menos evaluaciones del denoiser (NFE) que los muestreadores basados en gradientes. Converge en $\sim20$ pasos, mientras que los baselines como DPS o $\Pi$GDM a menudo requieren más de 100 pasos para acercarse a un rendimiento similar.
4. Límite de Alto Ruido: El método proporciona un estimador de un solo paso para la media posterior $E[x_0|y]$ en el límite de alto ruido, ofreciendo un estimador puntual MMSE-óptimo y rápido sin muestreo iterativo.

Resultados Experimentales

Los autores evalúan EPS en cinco problemas inversos lineales (inpainting aleatorio, inpainting de caja, superresolución 4$\times$, desenfoque gaussiano, desenfoque de movimiento) utilizando los conjuntos de datos FFHQ e ImageNet.

• Métricas: EPS supera tanto a los baselines sin entrenamiento (DPS, DAPS, DDNM, $\Pi$GDM) como a los basados en entrenamiento (Palette) en términos de fidelidad puntual (PSNR, SSIM), calidad perceptual (LPIPS, FID) y calibración de la distribución (CRPS, MMD).
• Eficiencia de Muestreo: EPS alcanza su meseta de rendimiento dentro de $\sim20$ NFE. En contraste, los baselines de muestreo continúan mejorando lentamente hasta las 100 NFE y a menudo no logran alcanzar la asíntota de EPS.
• Eficiencia de Entrenamiento: Cuando se ajusta a partir de un checkpoint preentrenado, EPS converge en una fracción de las iteraciones requeridas por otros métodos basados en entrenamiento (por ejemplo, Palette), debido a la preservación de la geometría de denoising preentrenada.
• Tiempo de Ejecución: EPS es el resolvedor de problemas inversos más rápido entre los métodos comparados, ejecutándose esencialmente a la misma velocidad de tiempo de reloj que el denoiser incondicional preentrenado (aproximadamente 1.006$\times$), ya que evita los pasos hacia atrás (backward passes) requeridos por los métodos basados en gradientes.
• Ablaciones: Los experimentos confirman que el pivote desplazado $\mu^\star$ es el factor crítico para el rendimiento, y que la geometría anisotrópica aprendida por EPS es esencial, ya que el pivoting zero-shot (usar un denoiser preentrenado sin ajuste fino) funciona deficientemente.

Significado y Reivindicaciones

El artículo afirma que EPS proporciona un enfoque principista y calibrado para el muestreo posterior en problemas inversos lineales. Al derivar la geometría posterior exacta, cierra la brecha entre la eficiencia de los métodos sin entrenamiento y la expresividad de los métodos basados en entrenamiento.

• Alineación Estructural: A diferencia de los métodos de entrenamiento condicional que obligan a la red a aprender un nuevo camino condicional desde cero, EPS adapta el prior preentrenado a la geometría anisotrópica específica inducida por la medición. Esto permite un ajuste fino rápido y el reuso de potentes priors incondicionales.
• Transparencia: El pivote y la covarianza de forma cerrada ofrecen un manejo transparente para analizar el sampler, separando explícitamente las direcciones medidas de las no observadas.
• Practicidad: El método no requiere bucles de muestreo complejos ni computaciones de gradiente durante la inferencia, lo que lo hace adecuado para aplicaciones de baja latencia.

Los autores reconocen limitaciones: la derivación exacta asume un operador de propagación lineal y ruido de observación gaussiano. Los operadores no lineales o el ruido no gaussiano requerirían una linealización local o un entrenamiento contra la verdadera verosimilitud, para lo cual la forma cerrada no se aplica directamente. Además, los backbones de difusión latente requieren cuidado, ya que el decodificador puede volver no lineales en el espacio latente los operadores lineales en el espacio de píxeles.