Initial Conditions Reconstruction with Stochastic Interpolants
[Video Credit: N-body simulation Francisco Villaescusa-Navarro]
Carolina Cuesta-Lazaro
IAIFI Fellow, MIT / Center for Astrophysics
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
Initial Conditions
early Universe
Cosmological Parameters
theory
\mathit{O}(2048^3)
p(\delta_{\mathrm{ICs}}, \mathcal{\theta}|\delta_{\mathrm{Obs}})
?
Observed Density Field
today
\mathit{O}(10)
\mathit{O}(2048^3)
+
\Omega_m,
\sigma_8
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
Hamiltonian Monte Carlo
1) Likelihood is intractable for realistic scenarios, but can get samples from simulator
2) Forward model has to be differentiable
(and relatively fast)
3) Not amortized
\log P(\delta_\mathrm{Obs}|\delta_\mathrm{ICs}, \mathcal{C}) = -\frac{1}{2} \sum_i \left(\frac{\mathcal{S}^i[\delta_{\mathrm{ICs}}, \mathcal{C}] - \delta_{\mathrm{Obs}}^i}{\sigma} \right)^2
["Field-Level Inference with Microcanonical Langevin Monte Carlo" Bayer, Seljak, Modi arXiv:2307.09504]
["Bayesian physical reconstruction of initial conditions from large scale structure surveys" Jasche, Wandelt arXiv:1203.3639]
["Posterior Sampling of the Initial Conditions of the Universe from Non-linear Large Scale Structures using Score-Based Generative Models" Legin et al arXiv:2304.03788]
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
\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)
Continuity Equation
Diffusion, Flow matching, Interpolants...
[Image Credit: "Understanding Deep Learning" Simon J.D. Prince]
Data
Base
Continuous Time Normalizing Flows
Missing pieces:
Simulation-free loss
SDE formulation
Can we regress the velocity field?
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
= (1-t)x_0 + t x_1
b(t,x, x_0) = \mathbb{E}\left[\partial_t I + \partial_t \gamma(t) \sqrt{t} z| x_t = x \right]
Simulation-free!
I_t = \alpha_t x_0 + \beta_t x_1
Interpolant
+ \gamma_t W_t
Expectation over all possible paths that go through xt
["Stochastic Interpolants: A Unifying framework for flows and diffusion" Albergo et al arXiv:2303.08797]
Stochastic Interpolants
dX_t = b(t, X_t, x_0) dt + \sigma_t dW_t
Stochastic
\mathcal{L} = \int_0^1 \mathbb{E}_{p(x_0,x_1), p(z)} |\hat{b}_\phi(t,x,x_0) - \partial_t I(t,x_0,x_1) - \partial_t \gamma(t) z|^2 dt
z \sim N(0,1)
\gamma(0) = \gamma(1) = 0
W_t = \sqrt{t} z
Generative SDE
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
Generative SDE
dX_t = b_t(X_t, x_0) dt + \sigma_t dW_t
3D U-Net
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
True
Reconstructed
\delta_\mathrm{Obs}
\delta_\mathrm{ICs}
Initial Conditions
Finals
L = 400 h^{-1} \mathrm{Mpc}, N = 32^3
L = 100 h^{-1} \mathrm{Mpc}, N = 64^3
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
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})
(Marginalizing over parameters)
1) Likelihood not necessarily Gaussian
2) Forward model no need differentiable
3) Amortized
\mathrm{SI}:
\mathrm{SI}:
\mathrm{HMC}:
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
p(\delta_{\mathrm{ICs}}|\mathrm{Condition})
N(\delta > \delta_\mathrm{th}) = \sum_i \Theta(\delta_i - \delta_\mathrm{th})
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
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 @ BASP 2025
Representing Continuous Fields
(x,y,z,t)
\Psi
\dot \Psi
Continuous in space and time
x500 Compression?
Can we store a simulation inside a neural network?
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
Sampling entire consistent trajectories, rather than just Initial Conditions
Scaling to large observed volumes
\mathcal{O}(5026^3)
Training on small volume simulations
# To Do
p(\delta_{\mathrm{ICs}}, \mathcal{\theta}|\delta_{\mathrm{Obs}})
?
Controllable Simulations
Carolina Cuesta-Lazaro IAIFI/MIT @ BASP 2025
BASP2025
By carol cuesta
