Variance Reduction for Non-Log-Concave Sampling with Applications to Inverse Problems

Resumen Técnico: Reducción de la Varianza para el Muestreo No Log-Cóncavo con Aplicaciones a Problemas Inversos

Planteamiento del Problema
El muestreo de distribuciones de alta dimensión y no log-cóncavas con densidades no normalizadas es un desafío fundamental en el aprendizaje automático, particularmente cuando los gradientes exactos de la función potencial no están disponibles. En tales escenarios, los gradientes deben aproximarse mediante estimaciones estocásticas, las cuales suelen exhibir una alta varianza bajo un presupuesto computacional fijo por iteración. Si bien las técnicas de reducción de la varianza (por ejemplo, SGD con momentum, STORM, PAGE) han mejorado la convergencia en la optimización no convexa, su aplicación al muestreo de distribuciones no log-cóncavas ha sido ampliamente inexplorada. Además, los métodos existentes de muestreo posterior para problemas inversos utilizando modelos generativos basados en puntuación (SGM) operan típicamente en un régimen de lote completo (full-batch), lo que limita su aplicabilidad a conjuntos de datos a gran escala comunes en imágenes 3D y dinámicas.

Metodología
Los autores proponen un marco unificado que adapta las técnicas de reducción de la varianza de la optimización no convexa al dominio del muestreo. Introducen dos estimadores unificados:

1. Difusión de Langevin con Reducción de Varianza de Lipschitz Medio (ML-VRLD): Este estimador unifica SGD con momentum y su contraparte estocástica. Asume la continuidad de Lipschitz de los gradientes exactos (Suposición 1) y proporciona un tratamiento unificado donde el estimador es un promedio ponderado de un gradiente de mini-lote y una actualización basada en momentum.
2. Difusión de Langevin con Reducción de Varianza de Lipschitz de Muestra (SL-VRLD): Este estimador unifica PAGE y STORM. Asume la continuidad de Lipschitz de los gradientes estocásticos en el sentido del error cuadrático medio (Suposición 4). Incorpora un término de corrección basado en la diferencia entre los gradientes estocásticos actuales y pasados, lo que mejora la precisión de la estimación bajo la condición de Lipschitz de muestra.

Ambos métodos están diseños para lograr una convergencia de Información de Fisher (FI) relativa $\epsilon$ con $O(1)$ computaciones de gradiente estocástico por iteración, contrastando con el requerimiento de $O(Ld\sigma^2/\epsilon)$ de la Dinámica de Langevin de Gradiente Estocástico (SGLD) estándar.

El marco se extiende a problemas inversos (por ejemplo, reconstrucción de MRI y CT) integrando estos estimadores en un proceso de muestreo de Monte Carlo de Langevin (LMC) con priors SGM. La formulación del muestreo posterior utiliza una secuencia de distribuciones objetivo intermedias con esquemas de recocido (annealing) y suavizado para guiar las iteraciones desde un prior suavizado hacia el verdadero posterior. Se establecen garantías teóricas bajo suposiciones de error de estimación de SGM acotado.

Contribuciones Clave

• Análisis Unificado: El artículo proporciona el primer análisis unificado de los estimadores de reducción de varianza (específicamente ML-VRLD y SL-VRLD) para el muestreo de distribuciones no log-cóncavas.
• Mejora de la Complejidad:
  • Tanto ML-VRLD como SL-VRLD logran una convergencia de FI relativa $\epsilon$ con $O(1)$ computaciones de gradiente por iteración, mejorando significativamente la complejidad por iteración de SGLD.
  • SL-VRLD reduce aún más la complejidad computacional total de $O(L^3d^3\sigma^2/\epsilon^3)$ (SGLD) a $O(L^2d^2\sigma^2/\epsilon^3)$.
• Garantías de Convergencia: Los autores establecen tasas de convergencia no asintóticas en distancia de variación total (TV) al cuadrado bajo una suposición de desigualdad de Poincaré y prueban la convergencia débil a la distribución objetivo para ambos algoritmos.
• Extensión a Problemas Inversos: El marco se extiende al muestreo posterior con priors SGM, estableciendo las mismas garantías de convergencia bajo un error de estimación de SGM acotado. Esto aborda la limitación de los métodos de muestreo posterior de lote completo existentes.

Resultados

• Experimentos Sintéticos: En modelos de mezcla Gaussiana (GMM) y problemas inversos de detección comprimida, los métodos propuestos demuestran una convergencia más rápida hacia valores de FI relativa más bajos en comparación con SGLD. SGLD con un tamaño de lote fijo converge a un nivel de FI suboptimal, mientras que ML-VRLD y SL-VRLD alcanzan la precisión objetivo con el mismo tamaño de lote fijo. SL-VRLD alcanza consistentemente la distribución objetivo más rápido que ML-VRLD.
• Validación Estadística: En una tarea de muestreo posterior con una verdad de campo conocida, SGLD produjo muestras con una varianza significativamente inflada. ML-VRLD y SL-VRLD redujeron significativamente esta varianza, siendo SL-VRLD el que logró la coincidencia más cercana con la varianza del posterior de la verdad de campo.
• Aplicaciones de Imagen:
  • Reconstrucción de MRI: En el conjunto de datos fastMRI con factores de aceleración de 8 y 12, ML-VRLD y SL-VRLD produjeron detalles anatómicos más nítidos y menos artefactos en comparación con las líneas base (incluyendo DPS, PnP-DM y SGLD). SL-VRLD alcanzó las puntuaciones de PSNR y SSIM más altas.
  • Reconstrucción de CT de Ángulo Escaso: En el conjunto de datos 2DeteCT con 180 y 360 vistas de proyección, los métodos propuestos superaron a las líneas base en la preservación de detalles estructurales finos, demostrando el beneficio de la reducción de la varianza en regímenes de alto ruido y submuestreo.

Significancia
El artículo afirma proporcionar las primeras garantías de convergencia no asintótica para el muestreo no log-cóncavo utilizando estimadores de reducción de varianza con $O(1)$ computaciones de gradiente por iteración. Al extender estas técnicas a problemas inversos con priors SGM, este trabajo permite el muestreo posterior principista para tareas de imagen a gran escala donde los métodos de lote completo son computacionalmente prohibitivos. Los resultados empíricos validan que la reducción de la varianza mejora consistentemente la calidad de la muestra y la precisión estadística bajo restricciones fijas de memoria y computación, ofreciendo una solución escalable para la inferencia bayesiana de alta dimensión.