Joint inference of weak lensing convergence map and cosmology with diffusion models

Resumen Técnico: Inferencia Conjunta de Mapas de Convergencia de Lente Débil y Cosmología con Modelos de Difusión

Planteamiento del Problema
Los análisis actuales de lentes gravitacionales débiles suelen depender de estadísticas resumidas (por ejemplo, funciones de correlación de dos puntos), las cuales descartan la información no gaussiana presente en la distribución de materia, particularmente en escalas pequeñas y no lineales. Si bien los enfoques de campo completo ofrecen una forma fundamentada de capturar toda la información estadística, los métodos existentes enfrentan una dicotomía:

1. Inferencia Marginal: Los métodos basados en redes neuronales infieren parámetros cosmológicos ($\theta$) directamente de los datos, pero marginalizan sobre el campo de convergencia subyacente ($\kappa$), fallando en la reconstrucción del mapa de masa.
2. Reconstrucción de Campos: Los métodos como el mapeo de masa reconstruyen la convergencia $\kappa$ a partir de la cizalladura ruidosa ($\gamma_{obs}$), pero tratan los parámetros cosmológicos como fijos, en lugar de inferirlos conjuntamente.
3. Inferencia Explícita (MCMC): Los enfoques como BORG realizan una inferencia conjunta ($p(\theta, \kappa | \gamma_{obs})$) utilizando Cadenas de Markov Monte Carlo (MCMC). Sin embargo, estos requieren un modelo de propagación (forward model) diferenciable que vincule las condiciones iniciales con los observables, lo que limita la complejidad de la física que se puede emplear, y sufren de altos costos computacionales en espacios de alta dimensión.

Los autores buscan desarrollar un método para la inferencia conjunta que reconstruya simultáneamente el campo de convergencia e infiera los parámetros cosmológicos sin requerir un modelo de propagación diferenciable o un muestreo MCMC iterativo durante la inferencia.

Metodología: JADE
Los autores presentan JADE (Joint Architecture for fielD and cosmological parameter Estimation), un marco de inferencia implícita basado en modelos de difusión condicionales.

• Inferencia Implícita mediante Modelos de Difusión: En lugar de aprender una función de verosimilitud, JADE aprende la distribución posterior conjunta $p(\theta, \kappa | \gamma_{obs})$ directamente de simulaciones. El método utiliza Flow Matching, un marco generativo donde una red neuronal aprende un campo de velocidad para transportar una distribución base (ruido gaussiano) hacia la distribución objetivo. El modelo está entrenado para predecir la muestra "limpia" ($x_1 = [\theta, \kappa]$) a partir de un interpolante ruidoso ($x_t$) dado el campo de cizalladura observado.
• Arquitectura Transformer: El núcleo de JADE es un Vision Transformer en el espacio de píxeles adaptado para datos multimodales.
  • Secuencia de Tokens Unificada: El campo de convergencia $\kappa$ se descompone en parches, y los parámetros cosmológicos $\theta$ se proyectan en embeddings de la misma dimensión. Estos se concatenan en una única secuencia de tokens.
  • Condicionamiento: El campo de cizalladura ruidoso observado $\gamma_{obs}$ también es parcheado y proyectado, y luego se concatena a la secuencia. El mecanismo de atención del transformer permite que los tokens objetivo (campo y parámetros) atiendan al input de condicionamiento (cizalladura) en todas las posiciones espaciales.
  • Salida: La red emite estimaciones de la densidad de la convergencia y de los parámetros cosmológicos de forma simultánea y denotada.
• Inferencia Amortizada: Una vez entrenado, el modelo puede generar muestras de la posterior para cualquier nueva observación integrando una Ecuación Diferencial Ordinaria (ODE) condicionada al $\gamma_{obs}$ específico. Este proceso es "amortizado", lo que significa que no se requiere reentrenamiento ni optimización por observación (como MCMC).

Configuración de la Simulación
El método fue validado en datos de lentes débiles simulados:

• Cosmología: Un modelo $w$CDM con seis parámetros ($\Omega_c, \Omega_b, \sigma_8, h_0, n_s, w_0$).
• Generación de Datos: Los mapas de convergencia se generaron utilizando un modelo log-normal 2D (un modelo fenomenológico para el crecimiento no lineal) en lugar de simulaciones completas de N-cuerpos, para permitir la generación rápida del gran conjunto de entrenamiento (100,000 muestras).
• Observables: Campos de cizalladura simulados con profundidad LSST Año 10 (Y10) con ruido de forma, cubriendo un parche de $5 \times 5$ deg$^2$ con $128 \times 128$ píxeles.

Resultos Clave

1. Reconstrucción Conjunta: El modelo generó con éxito muestras de la posterior tanto de $\kappa$ como de $\theta$ condicionadas a la cizalladura ruidosa. La inspección visual y las pruebas estadísticas confirmaron que los mapas de convergencia reconstruidos recuperan las estructuras reales mientras reflejan correctamente la incertidumbre residual en las escalas pequeñas.
2. Consistencia Estadística:
   • Espectros de Potencia: Los espectros de potencia auto y cruzados de las muestras de la posterior coincidieron con los de simulaciones independientes extraídas en las cosmologías muestreadas, con un error relativo de ~6.7% promediado sobre bins y escalas.
   • PDFs de un Punto: Las distribuciones recuperadas exhibieron la forma log-normal esperada, confirmando la captura de estadísticas no gaussianas.
   • Residuales: El estadístico de chi-cuadrado reducido ($\chi^2_{red} \approx 1.002$) confirmó que los residuales entre las observaciones y las muestras de la posterior estaban dominados por el ruido de forma inyectado, tal como se esperaba.
3. Restricciones Cosmológicas: Las posteriores cosmológicas marginales ($p(\theta | \gamma_{obs})$) derivadas de JADE mostraron una excelente concordancia con las cadenas MCMC tradicionales de NUTS (No-U-Turn Sampler) en términos de regiones de credibilidad y direcciones de degeneración.
4. Calibración:
   • Puntuación MIRA: La posterior conjunta logró una puntuación de calibración MIRA (Mass In Random Areas) de $0.635 \pm 0.017$ (donde 0.667 es óptimo), indicando una buena calibración con una ligera sobreconfianza. La posterior cosmológica marginal obtuvo un puntaje de $0.659 \pm 0.0139$.
   • TARP: Las pruebas de Exactitud con Puntos Aleatorios confirmaron que las regiones de credibilidad alcanzaron la cobertura frecuentista correcta a través del espacio de parámetros.
5. Eficiencia: Generar una única muestra de la posterior conjunta toma aproximadamente 0.2 segundos (256 pasos de ODE). Extraer 1,000 muestras toma ~3 minutos, comparado con el costo computacional significativamente mayor de ejecutar cadenas MCMC (que requirieron ~2.18 millones de evaluaciones de gradiente para el benchmark).

Significancia y Reivindicaciones
El artículo afirma establecer la inferencia implícita basada en difusión como una ruta viable hacia los análisis de cosmología de campo completo. La principal significancia radica en unificar la reconstrucción del mapa de masa y la estimación de parámetros cosmológicos dentro de un único marco no diferenciable.

Las principales reivindicaciones sobre las ventajas del método incluyen:

• No Requiere un Modelo de Propagación Diferenciable: A diferencia de la inferencia de nivel de campo basada en MCMC, JADE no requiere gradientes analíticos del modelo de propagación, lo que teóricamente permite el uso de simuladores arbitrariamente complejos (por ejemplo, códigos de N-cuerpos completos o hidrodinámicos) una vez disponibles los datos de entrenamiento.
• Amortización: El modelo se entrena una vez y se aplica a cualquier nueva observación instantáneamente, evitando la necesidad de muestreo por observación.
• Inferencia Conjunta: Es el primer método presentado que infiere conjuntamente la posterior completa $p(\theta, \kappa | \gamma_{obs})$ utilizando un único modelo de difusión condicional, en lugar de marginalizar sobre el campo o inferir los parámetros por separado.

Limitaciones y Trabajo Futuro
Los autores señalan que la demostración actual depende de simulaciones log-normales, que son más rápidas que los códigos de N-cuerpos pero menos completas físicamente. Escalar a simuladores realistas y costosos presenta desafíos respecto a la densidad del conjunto de entrenamiento en el espacio de parámetros cosmológicos. Además, como método amortizado, JADE debe reentrenarse si el modelo de datos (por ejemplo, máscara de encuesta, modelo de ruido) cambia, a diferencia de los métodos de inferencia explícita que pueden adaptarse en el momento de la inferencia. Los autores proponen como trabajo futuro extender esta arquitectura a la inferencia conjunta de las condiciones iniciales y los parámetros cosmológicos, así como explorar enfoques de difusión latente para manejar campos de mayor resolución de manera más eficiente.