Increasing the Precision of Surrogate Models for Weak Lensing Mass Maps with Flow Matching

Resumen Técnico: Aumento de la Precisión de los Modelos Sustitutos para Mapas de Masa de Lente Gravitacional Débil con Flow Matching

Enunciado del Problema
La lente gravitacional débil proporciona una sonda directa de la estructura a gran escala del cosmos, codificando información no gaussiana en mapas de convergencia (masa) que es inaccesible para las estadísticas de dos puntos tradicionales. El modelado directo preciso de estos mapas es esencial para la inferencia cosmológica, dependiendo típicamente de simulaciones de $N$-cuerpos computacionalmente costosas. Aunque se han propuesto emuladores de aprendizaje profundo, en particular las Redes Generativas Antagónicas (GAN), para acelerar este proceso, los métodos existentes sufren limitaciones significativas. Los sustitutos basados en GAN anteriores a menudo no logran coincidir con las estadísticas resumidas con suficiente precisión (típicamente errores del 5–10%), producen mapas sobre-suavizados, luchan por capturar la distribución correcta de los conjuntos de mapas generados (fallando en reproducir la distribución probabilística completa en lugar de solo el "modo") y exhiben inestabilidad en el entrenamiento. Además, la generación condicional a través de parámetros cosmológicos continuos (por ejemplo, densidad de materia $\Omega_m$ y amplitud de agrupamiento $\sigma_8$) sigue siendo un desafío debido a las dificultades de aprendizaje heterogéneas a través del espacio de etiquetas.

Metodología
Los autores proponen un marco de Flow Matching Condicional de Etiquetas Residuales (RLCFM) para generar mapas de convergencia de alta fidelidad condicionados a parámetros cosmológicos. El enfoque se basa en Flujos Normalizantes Continuos (CNF) y Flow Matching (FM), que aprenden un campo vectorial dependiente del tiempo para transportar una distribución previa a una distribución de datos objetivo.

Los componentes metodológicos clave incluyen:

1. Flow Matching Condicional de Etiquetas Residuales: En lugar de entrenar un flujo condicional de etiquetas estándar, los autores introducen una reparametrización en un espacio residual. Restan la media específica de la etiqueta ($\mu_0(y)$) de los mapas objetivo antes del entrenamiento. Esto elimina desplazamientos de media triviales dependientes de la etiqueta, permitiendo que la red se concentre en aprender la estructura residual transferible. Este diseño tiene como objetivo estabilizar la interpolación a través de parámetros cosmológicos.
2. Prior Gaussiana Específica de Etiqueta: El modelo emplea una prior gaussiana específica de etiqueta $p(x|y)$ impulsada por datos en lugar de una única prior compartida. La media y la varianza de la prior se derivan de las estadísticas de la distribución condicional $q(x|y)$ en $t=0$, transportando esta prior específica a la distribución objetivo en $t=1$.
3. Objetivo de Entrenamiento: La red se entrena para minimizar la pérdida de Flow Matching, regresando el campo vectorial dependiente del tiempo que transporta la prior específica de la etiqueta a los mapas de convergencia objetivo. El entrenamiento se realiza en un conjunto de datos de 57 modelos cosmológicos ($\Omega_m, \sigma_8$) dentro de un marco $\Lambda$CDM plano, que comprende 684.000 mapas de $128 \times 128$ píxeles.
4. Preprocesamiento de Datos: Los mapas de convergencia se transforman mediante un logaritmo (aproximando un campo lognormal) y se estandarizan. El modelo aprende a generar el campo residual $r = \tilde{x} - \mu_0(y)$, que luego se reconstruye en la inferencia añadiendo de nuevo la media específica de la etiqueta.

Contribuciones Clave

• Arquitectura Novel: Introducción de la arquitectura RLCFM, que combina priores específicas de etiquetas con entrenamiento en espacio residual para abordar la inestabilidad y los desafíos de interpolación del modelado generativo condicional en cosmología.
• Fidelidad Mejorada: El método logra una precisión significativamente mayor que los puntos de referencia anteriores de GAN. Reduce los errores típicos en estadísticas básicas (distribuciones de píxeles/picos, espectro de potencia) de $<10\%$ a $<1\%$, y en estadísticas de orden superior (bispectro, funcionales de Minkowski) de $<20\%$ a $<5\%$.
• Precisión a Nivel de Distribución: A diferencia de las GAN que a menudo ajustan el "modo" de los datos, el modelo RLCFM reproduce con éxito la distribución completa de estadísticas resumidas generadas a partir de condiciones iniciales aleatorias, coincidiendo con la variabilidad de las simulaciones de $N$-cuerpos.
• Interpolación Robusta: El modelo demuestra un rendimiento estable a través del espacio de parámetros cosmológicos, incluyendo en cosmologías de prueba retenidas y en los bordes de la cuadrícula, con errores que aumentan solo moderadamente en los límites en comparación con el centro.

Resultados
Los autores evaluaron el modelo RLCFM frente a simulaciones de $N$-cuerpos utilizando una suite completa de métricas:

• Calidad Visual: Los mapas generados son visualmente indistinguibles de las referencias de $N$-cuerpos, con una evolución morfológica consistente a medida que varían $\Omega_m$ y $\sigma_8$.
• Estadísticas Resumidas:
  • Distribuciones de Píxeles y Picos: Las distancias de Wasserstein-1 normalizadas son típicamente $<0.01$.
  • Espectro de Potencia ($C_\ell$): Las diferencias fraccionarias son típicamente $<1\%$.
  • Bispectro ($B_\ell$): Las diferencias fraccionarias son típicamente $<5\%$.
  • Funcionales de Minkowski ($V_0, V_1, V_2$): Las diferencias fraccionarias son generalmente inferiores al 3%.
  • Matrices de Correlación: El contraste de la norma de Frobenius entre las matrices de correlación de RLCFM y $N$-cuerpos es típicamente $<3\%$.
  • Similitud Estructural (MS-SSIM): La significancia de la diferencia en las puntuaciones de similitud estructural es baja, lo que indica que el modelo no sufre de colapso de modos y preserva el realismo estructural multiescala.
• Eficiencia: El modelo se entrenó en aproximadamente 7 horas (65.000 pasos) en una sola GPU NVIDIA A100, requiriendo significativamente menos épocas para alcanzar un rendimiento estable en comparación con las líneas base de GAN.

Significado y Afirmaciones
El artículo afirma que el flow matching condicional, cuando se aumenta con una formulación en espacio residual y una prior impulsada por datos, proporciona un marco potente y eficiente para generar mapas de masa cosmológicos. El significado principal radica en la usabilidad práctica del emulador: captura la señal cosmológica con alta fidelidad a través de estadísticas de bajo y alto orden, al tiempo que soporta múltiples formas de análisis de datos. Los autores afirman que este trabajo acerca al campo a un emulador de mapas de masa práctico capaz de apoyar el modelado directo para futuras sondeos de gran campo (por ejemplo, Euclid, LSST) sin el costo computacional de ejecutar simulaciones completas de $N$-cuerpos para cada punto de parámetro cosmológico. El método se presenta como un paso hacia pipelines de inferencia cosmológica completamente generativos y basados en simulaciones.