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