Strong Gravitational Lensing Posterior Sampling in Pixel-Space Using Diffusion Models and Recurrent Inference Machines

Resumen Técnico: Muestreo de la Posterior de Lente Gravitacional Fuerte en el Espacio de Píxeles utilizando Modelos de Difusión y Máquinas de Inferencia Recurrente

Planteamiento del Problema
Modelar lentes gravitacionales fuertes de galaxia-galaxia para inferir la distribución de brillo de una galaxia fuente y la distribución de masa de una galaxia en primer plano es un problema inverso computacionalmente exigente. Este desafío se ve exacerbado en el régimen de observaciones de alta resolución y alta relación señal-ruido (SNR), donde los supuestos paramétricos de baja dimensión (por ejemplo, leyes de potencia elípticas simples) son insuficientes para capturar la complejidad de los datos. Para modelar observaciones hasta el nivel del ruido, se requieren representaciones pixeladas de alta dimensión tanto de la fuente como de la masa en primer plano. Sin embargo, representar la lente y la fuente como imágenes pixeladas hace que el modelo de propagación (forward model) sea no lineal con respecto a la masa en primer plano, creando un problema inverso de alta dimensión y no lineal caracterizado por distribuciones posteriores complejas y multimodales, así como severas degeneraciones de parámetros. Los algoritmos de muestreo tradicionales y los enfoques estándar de aprendizaje automático suelen tener dificultades para muestrear estas distribuciones de manera efectiva.

Metodología: El Marco DiRIM
Los autores proponen la Máquina de Inferencia Recurrente de Difusión (DiRIM), un método bayesiano diseñado para generar muestras conjuntas de la posterior de la luminosidad de la galaxia fuente ($s$) y la distribución de masa en primer plano (representada como un mapa de convergencia, $\kappa$) condicionadas a las observaciones ($y$). El marco integra dos paradiglos distintos de aprendizaje automático:

1. Modelado Generativo Basado en Difusión: El método utiliza modelos de difusión basados en puntuación (score-based) para muestrear de distribuciones de probabilidad complejas. El proceso implica una ecuación diferencial estocástica (SDE) hacia adelante que añade ruido gradualmente a los datos, y una SDE inversa que reconstruye los datos a partir del ruido. Para muestrear de la posterior $p(x|y)$, la SDE inversa requiere la puntuación (score) de la posterior $\nabla_x \log p_t(x|y)$.
2. Máquinas de Inferencia Recurrente (RIMs): Las RIMs se emplean para resolver el problema inverso mediante el aprendizaje de un algoritmo de inferencia iterativo basado en datos. En lugar de producir un único estimador de Máxima Posterior (MAP), el marco de la RIM se adapta aquí para refinar iterativamente la estimación de la puntuación de la posterior.

La Arquitectura DiRIM:

• Entrada: La red neuronal $g_\theta$ (una arquitectura U-Net) toma como entrada los parámetros del modelo con ruido en el tiempo $t$ ($x_t$), la observación $y$, el paso de tiempo $t$, los gradientes de la verosimilitud ($\nabla L$) y los residuales de la imagen lente normalizados ($\delta y / \sigma_y$).
• Refinamiento Iterativo: A diferencia de los modelos de difusión estándar que realizan un único paso de eliminación de ruido por cada paso de tiempo, DiRIM utiliza el marco de la RIM para actualizar iterativamente su estimación de la fuente y el mapa de convergencia denotados ($\hat{x}_0$) dentro de cada paso de tiempo de difusión.
• Estimación de la Puntuación (Score): La red es entrenada para minimizar una pérdida de coincidencia de puntuación de eliminación de ruido condicional (conditional denoising score-matching loss). La salida final de las iteraciones de la RIM ($\hat{x}_0^{(M)}$) se utiliza para calcular la puntuación de la posterior requerida para la SDE inversa:
  $$\nabla_{x_t} \log p_t(x_t|y) = \frac{\hat{x}_0^{(M)} - x_t}{\sigma^2(t)}$$
• Datos de Entrenamiento: El modelo se entrena con observaciones simuladas realistas generadas a partir de simulaciones magnetohidrodinámicas cosmológicas (IllustrisTNG) y perfiles analíticos. Las imágenes de las fuentes se extraen del conjunto de datos SKIRT-TNG (procesado con transferencia radiativa de polvo), y los mapas de convergencia son derivados de halos de simulaciones suavizados o renderizados a partir de perfiles analíticos (Ley de Potencia Elíptica + multipolos + subhalos NFW).

Contribuciones Clave

• Primer Muestreo Conjunto de la Posterior en el Espacio de Píxeles: Este trabajo presenta el primer método capaz de muestrear la posterior conjunta de la fuente y la masa en primer plano como imágenes pixeladas para lentes gravitacionales fuertes, yendo más allá de los puntos estimativos o los ajustes paramétricos de baja dimensión.
• Integración de RIMs y Difusión: El artículo introduce una arquitectura novedosa donde una RIM se incrusta dentro de un modelo de difusión condicional para refinar recursivamente las puntuaciones de la posterior, abordando la no linealidad y la alta dimensionalidad del problema de las lentes.
• Modelado al Nivel del Ruido: El método demuestra la capacidad de modelar observaciones hasta el nivel del ruido, una capacidad que a menudo se pierde cuando se utilizan modelos paramétricos simplificados.

Resultados
Los autores entrenaron y evaluaron DiRIM en dos conjuntos de datos: uno utilizando mapas de convergencia de simulaciones cosmológicas y otro utilizando perfiles analíticos.

• Calibración: Los gráficos de Pruebas de Exactitud con Puntos Aleatorios (TARP) indican que las muestras de la posterior están casi perfectamente calibradas, lo que sugiere que el estimador es insesgado.
• Bondad de Ajuste: La distribución de los valores de $\chi^2$ para las imágenes reconstruidas coincide estrechamente con la distribución teórica de $\chi^2$, confirmando que los modelos capturan los datos hasta el nivel del ruido.
• Manejo de la Complejidad: El método modela con éxito sistemas con distribaciones de masa complejas, incluyendo múltiples deflectores principales y subhalos. En casos donde los modelos paramétricos analíticos fallaron al intentar ajustar los datos al nivel del ruido debido a la falta de flexibilidad, DiRIM tuvo éxito.
• Detección de Subhalos: Al aplicarse a mapas analíticos que contienen subhalos NFW, el método detectó los subhalos en más del 98% de las muestras de la posterior, logrando una significancia de detección de $>3\sigma$.
• Muestreo de Degeneración: El método muestrea eficazmente las direcciones degeneradas en el espacio de parámetros, como la correlación entre el tamaño de la fuente y la pendiente radial del mapa de convergencia.
• Robustez (OOD): El modelo demostró robustez ante datos fuera de la distribución (OOD), infiriendo con éxito la morfología de fuentes y mapas de convergencia que diferían significativamente de la distribución de entrenamiento (por ejemplo, galaxias con forma de "7" generadas por IA y mapas de convergencia sumados).
• Ablación: Eliminar el componente de la RIM (estableciendo las iteraciones $M=1$) resultó en una inferencia sesgada y residuales con estructura significativa, confirmando que el refinamiento iterativo de las puntuaciones es esencial para el éxito del método.

Significancia y Reivindicaciones
El artículo afirma que DiRIM constituye una "primera prueba de concepto" sobre la capacidad de los modelos generativos de difusión combinados con máquinas de inferencia recurrente para resolver problemas inversos de alta dimensión y no lineales en astrofísica. Los autores enfatizan que su método aprovecha representaciones de alta dimensión y flexibles para extraer información de los datos que es inaccesible para el modelado paramétrico tradicional. Posicionan este trabajo como un paso necesario hacia la realización del potencial científico de los próximos grandes sondeos (como LSST y Euclid), que se espera descubran más de 100,000 lentes gravitacionales fuertes. Los autores señalan modestamente que, aunque el método es efectivo, persisten pequeños sesgos en sistemas complejos (por ejemplo, deflectores múltiples), y que el trabajo futuro abordará estos puntos y explorará formulaciones en el espacio latente para reducir los costos computacionales.