Generative Solutions for Cosmic Problems
[Video Credit: N-body simulation Francisco Villaescusa-Navarro]
Carolina Cuesta-Lazaro
IAIFI Fellow, MIT / Center for Astrophysics
1-Dimensional
Machine Learning
Secondary anisotropies
Galaxy formation
Intrinsic alignments
DESI, DESI-II, Spec-S5
Euclid / LSST
Simons Observatory
CMB-S4
Ligo
Einstein
The era of Big Data Cosmology
xAstrophysics
5-Dimensional
w_0, w_a, f\sigma_8, \Omega_m, \sum m_\nu
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
The cost of cosmological simulations
AbacusSummit
330 billion particles in 2 Gpc/h volume
\mathcal{O}(100)
60 trillion particles
~ 8TBs per simulation
15M CPU hours
(TNG50 ~100M cpu hours)
\mathcal{O}(10^4)
ML Requirements
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Simulation efficient methods
Can we store a simulation inside a neural network?
Prompting simulators
1
Fast Emulators + Likelihood Models
Continuous Fields and Compression
2
Controllable
Simulators
3
x_1
x_2
Model
p_\theta(x)
Training Samples
x_\mathrm{train}
Generative Models 1o1
Evaluate probabilities
Low p(x)
High p(x)
Generate Novel Samples
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
z \sim p(z)
x \sim p(x)
x = f(z)
p(x) = p(z = f^{-1}(x)) \left| \det J_{f^{-1}}(x) \right|
Probability mass conserved locally
Inference a la gradient descent
z = f^{-1}(x)
1) Tractable
2) f maximally expressive
Loss = Maximize likelihood training data
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
\frac{dx_t}{dt} = v^\phi_t(x_t)
x_1 = x_0 + \int_0^1 v^\phi_t(x_t) dt
\frac{d p(x_t)}{dt} = - \nabla \left( v^\phi_t(x_t) p(x_t) \right)
In continuous time
Continuity Equation
Loss requires solving an ODE!
Diffusion, Flow matching, Interpolants... All ways to avoid this at training time
[Image Credit: "Understanding Deep Learning" Simon J.D. Prince]
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
\frac{dx_t}{dt} = v_\theta(t,x_t)
= (1-t)x_0 + t x_1
v(t,x_t) = \mathbb{E}\left[\partial_t I | x_t=x \right]
Can we regress the velocity field directly?
Turned maximum likelihood into a regression problem!
I(t,x_0,x_1) = x_t = \alpha_t x_0 + \beta_t x_1
\mathcal{L} = \int_0^1 \mathbb{E}_{x_0,x_1} \left[v_\theta(t, x_t) - \partial_t I(t,x_0,x_1) \right] ^2 dt
Interpolant
+ \gamma_t z
Stochastic Interpolant
Expectation over all possible paths that go through xt
["Stochastic Interpolants: A Unifying framework for flows and diffusion" Albergo et al arXiv:2303.08797]
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Diffusion Models
Reverse diffusion: Denoise previous step
Forward diffusion: Add Gaussian noise (fixed)
Prompt
A person half Yoda half Gandalf
Denoising = Regression
Fixed base distribution:
Gaussian
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
["A point cloud approach to generative modeling for galaxy surveys at the field level"
Cuesta-Lazaro and Mishra-Sharma
ICML AI4Astro 2023, arXiv:2311.17141]
Base Distribution
Target Distribution
- Sample
- Evaluate
Long range correlations
Huge pointclouds (20M)
Homogeneity and isotropy
Siddharth Mishra-Sharma
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Fixed Initial Conditions / Varying Cosmology
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
p(\theta|x) = \frac{p(x|\theta)p(\theta)}{p(x)}
Diffusion model
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
CNN
Diffusion
Increasing Noise
p(\sigma_8|\delta_m)
p(\sigma_8|\delta_m + 0.01 \epsilon)
p(\sigma_8|\delta_m + 0.02 \epsilon)
["Diffusion-HMC: Parameter Inference with Diffusion Model driven Hamiltonian Monte Carlo" Mudur, Cuesta-Lazaro and Finkbeiner NeurIPs 2023 ML for the physical sciences, arXiv:2405.05255]
Nayantara Mudur
CNN
Diffusion
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
["A Cosmic-Scale Benchmark for Symmetry-Preserving Data Processing"
Balla, Mishra-Sharma, Cuesta-Lazaro et al
NeurIPs 2024 NuerReps arXiv:2410.20516]
E(3) Equivariant architectures
Benchmark models
["Geometric and Physical Quantities Improve E(3) Equivariant Message Passing" Brandstetter et al arXiv:2110.02905]
Symmetry-preserving ML
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
The bitter lesson by Rich Sutton
The biggest lesson that can be read from 70 years of AI research is that general methods that leverage computation are ultimately the most effective, and by a large margin. [...]
methods that continue to scale with increased computation even as the available computation becomes very great. [...]
We want AI agents that can discover like we can, not which contain what we have discovered.
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
["A Cosmic-Scale Benchmark for Symmetry-Preserving Data Processing"
Balla, Mishra-Sharma, Cuesta-Lazaro et al
NeurIPs 2024 NuerReps arXiv:2410.20516]
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Memory scaling point clouds and voxels
h_0
h_1
h_5
h_4
h_2
h_3
h_6
e_{01}
e_{12}
Graph
\mathcal{O}(10^6)
\mathcal{O}(10^8)
Nodes
Edges
3D Mesh
\mathcal{O}(10^{10})
Voxels
Both data representations scale badly with increasing resolution
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Representing Continuous Fields
(x,y,z,t)
\Psi
\dot \Psi
Continuous in space and time
x500 Compression!
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Reconstructing dark matter back in time
Stochastic Interpolants
NF
p(\delta_\mathrm{ICs}, \theta|\delta_\mathrm{Obs}) =
p(\delta_\mathrm{ICs}|\delta_\mathrm{Obs})
p(\theta|\delta_\mathrm{ICs},\delta_\mathrm{Obs})
["Joint cosmological parameter inference and initial condition reconstruction with Stochastic Interpolants" Cuesta-Lazaro, Bayer, Albergo et al NeurIPs 2024 ML for the Physical Sciences]
Adrian Bayer
Mount Fuji?
p(x_1|x_0)
x_0
?
x_1
s
["Probabilistic Forecasting with Stochastic Interpolants and Foellmer Processes" Chen et al arXiv:2403.10648 (Figure adapted from arXiv:2407.21097)]
Generative SDE
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Guided simulations with fuzzy constraints
Simulate what you need
(and sometimes what you want)
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
1. Generative models are more than fast emulators: robust field-level likelihood models
2. Continuous fields can be used to represent cosmological fields
3. Cosmological field level inference can be made efficient with generative models
Conclusions
Can we make them more simulation efficient?
Compression + efficient data format
Carolina Cuesta-Lazaro IAIFI/MIT @ Princeton 2024
Can generally make simulators more controllable!
PrincetonML2024
By carol cuesta
