Fast Sampling of Cosmological Initial Conditions with Gaussian Neural Posterior Estimation

Resumen Técnico: Muestreo Rápido de Condiciones Iniciales Cosmológicas con Estimación Neuronal de la Posterior Gaussiana

Planteamiento del Problema
La reconstrucción del campo de densidad de materia primordial (condiciones iniciales cosmológicas, o ICs) a partir de observaciones tardías de la estructura a gran escala (LSS) es un desafío fundamental en cosmología. Aunque el Universo temprano está bien descrito por fluctuaciones gaussianas lineales, la red cósmica actual es altamente no lineal y no gaussiana. Revertir esta evolución es un problema mal planteado: el mapeo de tiempos tempranos a tiempos tardíos es no inyectivo (diferentes ICs pueden dar lugar al mismo campo tardío) y la formación de estructuras en las escalas de los halos es caótica, lo que amplifica el ruido de medición.

Los enfoques probabilísticos tradicionales, como el Bayesian Origin Reconstruction from Galaxies (BORG), utilizan el muestreo de Monte Carlo de Hamiltoniano (HMC) para muestrear la distribución posterior. Sin embargo, estos métodos requieren simuladores diferenciables (que a menudo dependen de esquemas de aproximación) y son computacionalmente prohibitivos, requiriendo meses de tiempo de CPU para converger en espacios de parámetros de alta dimensión ($\gtrsim 10^6$ parámetros). Por el contrario, muchos enfoques de aprendizaje automático en este dominio dependen de estimadores de punto (por ejemplo, U-Nets) que proporcionan un único "mejor ajuste" de reconstrucción sin cuantificar la incertidumbre, o de modelos generativos (por ejemplo, difusión basada en puntuación/score-based) que siguen siendo computacionalmente costosos incluso en el tiempo de inferencia.

Metodología
Los autores proponen un marco de Inferencia Basada en Simulación (SBI) denominado Estimación Neuronal de la Posterior Gaussiana (NPE). La innovación central consiste en modelar la distribución posterior de las condiciones iniciales, $p(\mathbf{z}|\mathbf{x})$, como una distribución Gaussiana en el campo inferido $\mathbf{z}$, con una matriz de covarianza diagonal en el espacio de Fourier.

La metodología procede de la siguiente manera:

1. Formulación Estadística:

   • Prior: Se asume que el prior $p(\mathbf{z})$ es un campo aleatorio gaussiano definido por el espectro de potencia de materia lineal en un alto desplazamiento al rojo ($z=127$).
   • Likelihood (Verosimilitud): La verosimilitud $p(\mathbf{x}|\mathbf{z})$ se modela como una gaussiana con una media dependiente de la observación $\hat{\mathbf{z}}_\theta(\mathbf{x})$ y una matriz de precisión $\mathbf{Q}^L_\theta$. Crucialmente, $\mathbf{Q}^L_\theta$ está restringida a ser diagonal en el espacio de Fourier (y aproximadamente invariante ante rotaciones), reflejando la invariancia de traslación de la física.
   • Posterior: La combinación de un prior gaussiano y una verosimilitud gaussiana produce una posterior gaussiana. La matriz de precisión de la posterior es la suma de la matriz de precisión del prior y la de la verosimilitud ($\mathbf{Q}_\theta = \mathbf{Q}_P + \mathbf{Q}^L_\theta$).

2. Arquitectura Neuronal:

   • Estimador MAP: La media de la posterior, $\hat{\boldsymbol{\mu}}_\theta(\mathbf{x})$, se modela utilizando una arquitectura U-Net 3D. Para mejorar el rendimiento, los autores introducen una capa de activación personalizada que separa la reconstrucción en dos escalas: un factor de escala entrenable $\alpha_\theta(\mathbf{k})$ gestiona la evolución lineal a gran escala, mientras que la U-Net se enfoca en reconstruir los residuos no lineales a pequeña escala.
   • Matriz de Precisión: Se aprenden los elementos diagonales de la matriz de precisión de la verosimilitud, $D^L_\theta(k)$. Los autores probaron dos parametrizaciones: (1) tratar los valores como parámetros libres para cada modo de Fourier, y (2) imponer simetría de rotación mediante una red neuronal densa que toma el número de onda $k$ como entrada. El enfoque de la matriz "libre" mostró una estabilidad ligeramente mejor.

3. Entrenamiento y Datos:

   • El modelo se entrena con 2,000 pares de campos de sobredensidad de materia oscura inicial ($z=127$) y final ($z=0$) de la suite de simulaciones de N-cuerpos Quijote.
   • El modelo hacia adelante (forward model) utiliza simulaciones de N-cuerpos completas y no diferenciables (Gadget-III), evitando la necesidad de aproximaciones diferenciables requeridas por HMC.
   • El entrenamiento minimiza el log-posterior negativo (Ec. 7) durante 30 épocas en una sola GPU Nvidia A100, tomando aproximadamente 1.5 horas.

Contribuciones Clave

• Muestreo Eficiente: Al asumir una posterior gaussiana con una matriz de precisión diagonal en el espacio de Fourier, los autores permiten la generación de miles de muestras de la posterior en segundos (aproximadamente 3 segundos para 1,000 muestras en una sola GPU), órdenes de magnitud más rápido que los métodos existentes de HMC o de difusión.
• Simuladores No Diferenciables: El método no requiere que el modelo hacia adelante sea diferenciable, lo que permite el uso de simulaciones de N-cuerpos de vanguardia y completas sin esquemas de aproximación.
• Ajuste Analítico para Muestreo Rápido: Los autores demuestran que los valores de la diagonal de la matriz de precisión aprendida $D_\theta(k)$ pueden ajustarse con precisión mediante una función analítica simple (una gaussiana más una constante) hasta $k \approx 0.2 \, h/\text{Mpc}$. Esto permite que cualquier estimador de punto (MAP) entrenado se convierta en un muestreador de la posterior rápido sin necesidad de reentrenar la matriz de precisión completa.
• Cuantificación de la Incertidumbre: A diferencia de los estimadores de punto estándar, este método proporciona una distribución posterior completa, permitiendo una cuantificación rigurosa de la incertidumbre para la reconstrucción del campo.

Resultos
Los autores validan el método utilizando una simulación fiduciaria no presente en el conjunto de entrenamiento:

• Consistencia Visual y Estadística: Las muestras generadas se asemejan visualmente a la verdad de campo (ground truth). Los estadísticos de resumen (espectro de potencia, función de transferencia, correlación cruzada) muestran una alta concordancia con las condiciones iniciales reales. El espectro de potencia coincide con la verdad dentro de un 1–2% hasta la escala de Nyquist ($k_{Nyq} \approx 0.4 \, h/\text{Mpc}$).
• Correlación Cruzada: La correlación cruzada entre las muestras generadas y la verdad permanece alta ($\gtrsim 70\%$) hasta $k \approx 0.3 \, h/\text{Mpc}$.
• Consistencia Bayesiana: Una prueba de consistencia dependiente de la escala confirma que las diferencias normalizadas entre los modos reales y la media de la posterior siguen una distribución normal estándar ($N(0,1)$) a través de varios bins de $k$. Esto verifica que las incertididades inferidas están bien calibradas.
• Robustez: El método mantiene su rendimiento incluso cuando se añade ruido gaussiano al campo observado final.

Significado y Pretensiones
El artículo sostiene que este enfoque proporciona una vía eficiente en términos de simulación para la SBI secuencial en cosmología. Al reducir la complejidad de la posterior a una forma gaussiana con una matriz de precisión diagonal, los autores logran un equilibrio entre la precisión física y la velocidad computacional que era previamente inalcanzable para la reconstrucción de campos de alta dimensión.

Los autores señalan modestamente que la aproximación gaussiana de Fourier-diagonal es suficiente para el rango de escala seleccionado ($k \lesssim 0.4 \, h/\text{Mpc}$), pero que podría requerir matrices de precisión más sofisticadas (por ejemplo, basadas en convoluciones o wavelets) para escalas más pequeñas o cuando los efectos observacionales rompan la invariancia de traslación. Enfatizan que, si bien el trabajo actual asume parámetros cosmológicos fijos y observaciones de materia oscura idealizadas, el marco está diseñado para ser extendido a cosmologías variables y sondeos galácticos realistas en trabajos futuros. La principal significancia radica en permitir la generación de miles de realizaciones constreñidas en segundos, allanando el camino para análisis estadísticos robustos de la LSS que anteriormente eran computacionalmente inviables.