Accelerated Forward Models of the Universe
Going beyond the Training Set
[Video Credit: N-body simulation Francisco Villaescusa-Navarro]
Carolina Cuesta-Lazaro
Flatiron Institute
NYU
\mathcal{10^6}
Forward Model
Evaluations
DESI
Forward Modeling All The Way To Baryons
Adapted from arXiv:1804.03097
Symmetries
Connected to Underlying Physics
Hydro sims
Nbody + Empirical
Halo Occupation Distribution (HOD)
EFT bias expansion
Matter Density
Galaxy Distribution
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Carolina Cuesta-Lazaro @ Learning the Universe 2026
[Video credit: Francisco Villaescusa-Navarro]
Gas density
Gas temperature
Subgrid model 1
Subgrid model 2
Subgrid model 3
Subgrid model 4
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Accelerated Forward Models that can beyond their training set?
- Interpolate across parameter space
Carolina Cuesta-Lazaro @ Learning the Universe 2026
p_\theta(\delta_\mathrm{Simulations}|\mathcal{C}, \mathcal{A})
Generate New Samples
Field Level Likelihood!
Simulated
Emulated
p_\theta(\mathcal{C}, \mathcal{A}|\delta_\mathrm{Obs})
["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]
p(x|\theta) = \int p(x|z, \theta) p(z|\theta) \, dz
Initial Conditions
Simulator
Simulator
GANS
Deep Belief Networks
2006
VAEs
Normalising Flows
BigGAN
Diffusion Models
2014
2017
2019
2022
A folk music band of anthropomorphic autumn leaves playing bluegrass instruments
Contrastive Learning
2023
Meanwhile, on Earth...
2026
"Write a C compiler"
AGI?
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Goal: Estimate unknown p(x1) from samples
x_0 \sim \rho_0
Base
Target
T: \Omega \to \Omega
Transport Map
x_1 \sim \rho_1 \quad \text{via} \quad T(x_0) = x_1
x_1 \sim \rho_1
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Autoregressive in Frequency
Accelerated Forward Models that can beyond their training set?
- Interpolate across parameter space
Carolina Cuesta-Lazaro @ Learning the Universe 2026
1) Train a local mapping on small volume simulations
p_\theta(\delta_\mathrm{Simulations}|\delta_\mathrm{PM}, \mathcal{C}, \mathcal{A})
ICs Matched Cheap Gravity solver
PARTICLE MESH
GALAXIES/GAS
2) Run inference over larger volumes
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Yao Zhang
GOTHAM: From Approximate Matter to Dark Matter Haloes
Shivam Pandey
["Teaching Dark Matter simulations to speak the halo language" Pandey, Lanusse, Modi, Wandelt arXiv:2409.11401]
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Yao Zhang
Shivam Pandey
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Carolina Cuesta-Lazaro @ Learning the Universe 2026
5000 CPU hours
~5 mins 1 GPU
1 \, \mathrm{Gpc}/h
150k CPU hours
~15 GPU hours
~1 hour on 16 GPUs
3 \, \mathrm{Gpc}/h
Forward Modeling All The Way To Baryons
Adapted from arXiv:1804.03097
Symmetries
Connected to Underlying Physics
Hydro sims
Nbody + Empirical
Halo Occupation Distribution (HOD)
EFT bias expansion
Matter Density
Galaxy Distribution
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Amanda Lue
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Diffusion Models for Galaxies
25 \, \mathrm{Mpc}/h
210 \, \mathrm{Mpc}/h
p(\delta_g|\delta_m, \mathcal{C}, \mathcal{A})
Galaxies
N-body
N-body
Galaxies
p(\delta_g|\delta_m, \mathcal{C}, \mathcal{A})
Amanda Lue
Supernovae Feedback
N-body
p(\delta_g|\delta_m, \mathcal{C}, \mathcal{A})
Galaxies
p(\delta_g|\delta_m, \mathcal{C}, \mathcal{A})
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Diffusion Models for Galaxies
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Testing the models at the Field Level
25 \, \mathrm{Mpc}/h
4 TNG parameters
25 \, \mathrm{Mpc}/h
28 TNG parameters
50 \, \mathrm{Mpc}/h
35 TNG parameters
Joint constraints on Galaxies may be a must
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Max Lee
BIND
["BIND (Baryonic INpainting with Deep learning): A Field-level Emulator for Galaxy Groups and Clusters" Lee, Genel, Haiman, Bryan, Lovell, Hadzhiyska
arxiv:2609.10709v1]
Max Lee
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Max Lee
Fully Self Consistent Samples!
(Jointly constraining Galaxies and Baryons)
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Niall Jeffrey
Alternative: Macrocanonical Generators
\min_{\theta} \; \frac{1}{2} \left\| \mathbb{E}_{\varepsilon \sim p_0}\!\left[ \phi\big(F_\theta(\varepsilon)\big) \right] - \phi^* \right\|^2
Summary Statistic (spatially averaged)
Wavelets
["Macrocanonical Generator Networks: a data-scarce solution for cosmology AI surrogate simulations"
Jeffrey, Wandelt (in prep)]Carolina Cuesta-Lazaro @ Learning the Universe 2026
Nicolas Chartier
Correcting clustering with displacement potentials
["Point set clustering correction with a displacement potential"
Chartier, Bairagi, Bartlett, Ho, Wandelt (in prep)]\mathbf{x}' = \mathbf{x} + \nabla\phi(\mathbf{x})
\begin{aligned}
\Delta'_{\boldsymbol{k}} \approx\; & \Delta_{\boldsymbol{k}} + k^2 \phi_{\boldsymbol{k}}
+ \sum_{\boldsymbol{q}} (\boldsymbol{k}\cdot\boldsymbol{q})\,\phi_{\boldsymbol{q}}
\Bigg( \Delta_{\boldsymbol{q}-\boldsymbol{k}} + \frac{1}{2}\,\boldsymbol{k}\cdot(\boldsymbol{k}-\boldsymbol{q})\,\phi_{\boldsymbol{k}-\boldsymbol{q}}
+ \frac{1}{2}\sum_{\boldsymbol{q}'} (\boldsymbol{k}\cdot\boldsymbol{q}')\,\phi_{\boldsymbol{q}'}\,
\Delta_{\boldsymbol{q}+\boldsymbol{q}'-\boldsymbol{k}} \Bigg)
\end{aligned}
\Delta_{\mathbf{k}} \quad \text{and} \quad \Delta'_{\mathbf{k}}
Before Correction
After Correction
Accelerated Forward Models that can beyond their training set?
- Interpolate across parameter space
Carolina Cuesta-Lazaro @ Learning the Universe 2026
- Scale up in volume
- Interpolating across subgrid implementations
Can we learn a general and continuous representation of Baryonic feedback?
Gas
Galaxies
p(
, z_\mathrm{baryons})
Dark Matter
Baryonic fields
Marginalize over a broader set of subgrid physics
Interpolate between simulators
Mingshau Liu
(Ming)
Constrain z via multi-wavelength observations
["Continuous Representations of Baryonic Feedback for Robust Inference from Multiple Simulation Suites"
Liu, Cuesta-Lazaro
NeurIPs ML4PS 2025]Carolina Cuesta-Lazaro @ Learning the Universe 2026
Trained on:
TNG, SIMBA, Astrid, EAGLE
z = f(x)
Encoder
z_\mathrm{baryons}
1) Encoder
Gas
Galaxies
p(
p(
, z_\mathrm{baryons})
Dark Matter
Baryonic fields
2) Probabilistic Decoder
Carolina Cuesta-Lazaro @ Learning the Universe 2026
p(
, z_\mathrm{baryons})
Dark Matter
Baryonic fields
\mathcal{O}(10)
(Test suite)
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Gas Density
Temperature
Astrid
EAGLE
\alpha = 0
\alpha = 0.25
\alpha = 0.5
\alpha = 0.75
\alpha = 1
Interpolating over Simulations
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Generalizing to unseen simulations: Magneticum
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Accelerated Forward Models that can beyond their training set?
- Interpolate across parameter space
Carolina Cuesta-Lazaro @ Learning the Universe 2026
- Scale up in volume
- Interpolating across subgrid implementations
- Replacing add hoc subgrid models
Learned Subgrid Models
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Leena Iwamoto
Olga Borodina
Black Hole powered jets regulate star formation
But jets interact with the turbulent interstellar medium!
Text
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Leena Iwamoto
Olga Borodina
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Leena Iwamoto
Olga Borodina
A Matryoshka of Scales
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Fast emulation of haloes from Particle Mesh matter +
HOD Galaxies
Carolina Cuesta-Lazaro @ Learning the Universe 2026
1 \, \mathrm{Gpc} / h
2 \, \mathrm{Gpc} / h
3 \, \mathrm{Gpc} / h
Quijote
AbacusSummit
MilleniumTNG
L-Galaxies
Carolina Cuesta-Lazaro @ Learning the Universe 2026
The challenge of AFM in production: Shivam's pain
["Joint cosmological parameter inference and initial condition reconstruction with Stochastic Interpolants" Cuesta-Lazaro, Bayer, Albergo et al NeurIPs 2024 ML for the Physical Sciences]
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Shortcut: Accelerated Backward Models
1) Likelihood not necessarily Gaussian
2) Forward model no need differentiable
3) Amortized
Marginalizing over ICs
Fixing ICs
HMC: Marginalizing over ICs
Carolina Cuesta-Lazaro @ Learning the Universe 2026
True
Reconstructed
\delta_\mathrm{Obs}
\delta_\mathrm{ICs}
p(\delta_\mathrm{ICs}, \theta|\delta_\mathrm{Obs})
Can We Avoid Simulating Large Volumes?
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Chris Lovell
We know the Large Scale Behaviour
p(\delta_\mathrm{ICs}, \theta|\delta_\mathrm{Obs})
p(\delta_\mathrm{ICs}, \theta|\delta_\mathrm{Obs}, \delta_\mathrm{L})
Large Scale Reconstruction
1024^3
1 \, (\mathrm{Gpc}/h)^3
True
\delta_\mathrm{Galaxies}
\delta_\mathrm{ICs}
Reconstructed
["Reconstructing Cosmological Initial Conditions from Late-Time Structure with Convolutional Neural Networks"
Shallue, Eisenstein 2022]["Initial conditions from galaxies: machine-learning subgrid correction to standard reconstruction"
Parker, Bayer, Seljak 2025]Carolina Cuesta-Lazaro @ Learning the Universe 2026
PT
Power Spectrum
Cross Correlation
Peculiar Velocities
True
Reconstructed
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Fast emulation of haloes from Particle Mesh matter +
HOD Galaxies
Fast emulation of Physical Models of Galaxies
Carol's wishlist for AFM next year (sign up after the talk)
- Interpolate across parameter space
Carolina Cuesta-Lazaro @ Learning the Universe 2026
- Scale up in volume
- Interpolating across subgrid implementations
- Replacing add hoc subgrid models
Models that don't need to be retrained for changes in the background cosmology (w0wa)
Models that take advantage of PT on large scales to only do inference on small scales
Forward models that connect to the physics of galaxy formation and incorporate additional observables to jointly constrain it
Machine learned subgrid models that run in Arepo
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Shivam Pandey
CHARM: From Approximate Matter to Dark Matter Haloes
Particle Mesh for Gravity
p(\mathrm{Baryons}|\mathrm{DM}, \mathcal{C}, \mathcal{A})
Gas Properties
Density
Temperature
["BaryonBridge: Interpolants models for fast hydrodynamical simulations" Horowitz, Cuesta-Lazaro, Yehia ML4Astro workshop 2025]Probabilistic
Local
Carolina Cuesta-Lazaro @ Learning the Universe 2026
Simulated
Emulated
Lyman Alpha Flux
LyA Forest
AFM - LTU meeting - 2026
By carol cuesta
