An Unconditional Representation of the Conditional Score in Infinite-Dimensional Linear Inverse Problems

Resumen Técnico: Una Representación Incondicional de la Puntuación Condicional en Problemas de Inversión Lineal de Dimensión Infinita

Enunciado del Problema
Los modelos de difusión basados en puntuación (SDM) se han convertido en una herramienta poderosa para el muestreo de distribuciones posteriores en problemas de inversión bayesianos, ofreciendo priores flexibles que superan a los regularizadores tradicionales diseñados a mano. Sin embargo, los métodos existentes para el muestreo posterior enfrentan un compromiso crítico entre la eficiencia computacional y la escalabilidad.

1. Métodos basados en corrección (por ejemplo, agregar gradientes de verosimilitud a la puntuación o imponer consistencia con los datos) requieren evaluaciones repetidas del operador directo $A$ durante el proceso de muestreo. Esto crea un cuello de botella computacional prohibitivo para problemas a gran escala y de alta dimensión.
2. Los métodos de entrenamiento condicional directo aprenden una función de puntuación condicionada a las mediciones ($s(x, t; \mu_y)$). Aunque esto evita las evaluaciones directas durante el muestreo, aumenta drásticamente la dimensionalidad de entrada de la red neuronal (de $x$ a $(x, y)$), requiriendo datos de entrenamiento extensos y arquitecturas complejas para aproximar con precisión la puntuación condicional, particularmente en configuraciones de dimensión infinita.

El artículo aborda el desafío de realizar un muestreo posterior eficiente y escalable en problemas de inversión lineal de dimensión infinita sin incurrir en el alto costo en línea de las evaluaciones del modelo directo ni en la alta complejidad de muestra del entrenamiento condicional directo.

Metodología: Representación Incondicional de la Puntuación Condicional (UCoS)
Los autores proponen UCoS, un marco que desplaza la carga computacional del condicionamiento desde la fase de muestreo en línea hacia una fase de entrenamiento fuera de línea. La idea central es que, para problemas de inversión lineal con ruido gaussiano, la función de puntuación condicional puede derivarse exactamente a partir de una función de puntuación incondicional de un proceso de difusión modificado dependiente de la tarea mediante transformaciones afines.

• Fundamento Teórico:

  • El método opera en un espacio de Hilbert separable $\mathcal{H}$, asegurando la invarianza de la discretización.
  • En lugar de entrenar una red para aproximar $s(x, t; \mu_y)$, los autores definen un nuevo proceso de difusión $\tilde{X}_t$ cuya distribución es una mezcla del prior y un término gaussiano dependiente del operador directo $A$.
  • Demuestran (Teorema 3.7) que la puntuación condicional verdadera $s(x, t; \mu_y)$ está relacionada con la puntuación incondicional $\tilde{s}$ de este nuevo proceso mediante la identidad:
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    donde $r$ es una versión transformada de la puntuación incondicional, $\lambda(t)$ es un escalar dependiente del tiempo, y $\xi_t(x, y)$ es una transformación afín que involucra la medición $y$ y el operador adjunto $A^*$.
  • Crucialmente, el término $A^*\Gamma^{-1}y$ (donde $\Gamma$ es la covarianza del ruido) puede computarse una vez antes de que comience el muestreo. El proceso de muestreo subsiguiente requiere únicamente la evaluación de la red entrenada $r_\theta$ y sumas vectoriales simples, eliminando la necesidad de cualquier aplicación adicional de $A$ o $A^*$ durante los pasos de difusión inversa.

• Implementación:

  • Fase fuera de línea: Una red neuronal (específicamente un Operador de Red de Fourier, FNO) se entrena para aproximar la función $r$. El entrenamiento utiliza un objetivo de ajuste de puntuación de eliminación de ruido en el proceso modificado $\tilde{X}_t$. Esta fase requiere acceso a $A$ y $A^*$, pero no a datos de medición específicos $y$.
  • Fase en línea: Las muestras posteriores se generan resolviendo la EDE de tiempo inverso utilizando la puntuación transformada. El proceso es "libre de matrices", requiriendo únicamente el desplazamiento precomputado $A^*\Gamma^{-1}y$ y la evaluación de la red.

Contribuciones Clave

1. Identidad Exacta en Dimensiones Infinitas: El artículo establece una identidad teórica rigurosa (Teorema 3.7) que conecta la puntuación condicional con una puntuación incondicional dependiente de la tarea mediante transformaciones afines en espacios de Hilbert de dimensión infinita. Esto generaliza el principio más allá de discretizaciones específicas de dimensión finita.
2. Análisis de Convergencia: Los autores proporcionan un límite de error cuantitativo (Teorema 4.1) para el proceso generativo, caracterizando explícitamente las contribuciones del error provenientes de la pérdida de entrenamiento, la discretización numérica y la inicialización. Esto justifica la convergencia del método hacia la medida posterior verdadera.
3. Escalabilidad y Eficiencia: Al desacoplar el condicionamiento del bucle de muestreo, UCoS logra la eficiencia computacional de los métodos incondicionales (sin evaluaciones directas durante el muestreo) mientras mantiene la precisión de los métodos condicionales. Evita la maldición de la dimensionalidad asociada con el entrenamiento de redes en el espacio conjunto $(x, y)$.

Resultados Experimentales
Los autores validan UCoS en tres problemas de inversión: relleno de imágenes (inpainting), tomografía computarizada (CT) de vistas escasas y desenfoque de imágenes.

• Relleno de imágenes (Prior Gaussiano): En un escenario donde la posterior es gaussiana y las puntuaciones pueden computarse en forma cerrada, UCoS se desempeña de manera comparable al método condicional exacto y a otras líneas base, confirmando la identidad teórica.
• Imágenes de CT y Desenfoque (Priors Aprendidos): Utilizando los conjuntos de datos LIDC-IDRI y CelebA respectivamente, UCoS se compara frente a SDE ALD, Muestreo Posterior de Difusión (DPS), métodos basados en proyección (Proj) y entrenamiento Condicional directo.
  • Precisión: UCoS produce muestras con sesgo y desviación estándar comparables o mejores que los métodos incondicionales y significativamente mejores que el método Condicional al utilizar arquitecturas de baja complejidad (por ejemplo, 32 nodos por capa).
  • Eficiencia: UCoS iguala la velocidad de muestreo del método Condicional (ambos evitan evaluaciones directas) pero supera significativamente a los métodos basados en corrección (DPS, Proj, SDE ALD), los cuales requieren pasadas directas repetidas.
  • Robustez: El rendimiento del método Condicional se degrada notablemente con menos parámetros de red debido a la complejidad aumentada de aproximar la distribución conjunta, mientras que UCoS mantiene alta calidad con arquitecturas más simples.

Significado y Afirmaciones
El artículo afirma que UCoS ofrece una solución de "lo mejor de ambos mundos": elimina el cuello de botella computacional de las evaluaciones del modelo directo durante el muestreo (una limitación de los métodos basados en corrección) mientras evita la alta complejidad de muestra y las demandas arquitectónicas del entrenamiento condicional directo.

Los autores enfatizan que su enfoque es invariante a la discretización, lo que lo hace adecuado para problemas de inversión de dimensión infinita donde las garantías de rendimiento a menudo fallan en aproximaciones de dimensión finita de alta dimensión. Reconocen que el método es dependiente de la tarea, requiriendo reentrenamiento si el operador directo $A$ cambia, y que la correspondencia exacta depende de la linealidad del modelo directo y de la gaussianidad de la verosimilitud (permitiendo el paso de "completar el cuadrado"). El trabajo se presenta como una prueba de concepto para una nueva clase de muestreadores basados en difusión que desplazan la complejidad del condicionamiento a la fase de entrenamiento sin error de aproximación en la formulación teórica.