Implicit and Explicit Simulation-Based Bayesian Inference for Cosmology
François Lanusse
slides at eiffl.github.io/talks/Prague2023
Full-Field Inference
- Instead of analytically evaluating the likelihood of
sub-optimal summary statistics, let us build a model of the full observables.
⟹ The simulator becomes the physical model. - Each component of the full probabilistic model of the observables is now tractable, but at the cost of a large number of latent variables.
Benefits of a forward modeling approach
- Fully exploits the information content of the data (aka "full field inference").
- Easy to incorporate systematic effects.
- Easy to combine multiple cosmological probes by joint simulations.
(Porqueres et al. 2021)
For now, let's ignore the elephant in the room:
Do we have reliable enough models for the full complexity of the data?
Do we have reliable enough models for the full complexity of the data?
...so why is this not mainstream?
The Challenge of Full-Field Inference
- If we are interested in the likelihood p(x|θ) we need to
marginalize over stochastic latent variables:
p(x|θ)=∫p(x,z|θ)dz=∫p(x|z,θ)p(z|θ)dz
- If x is the lensing angular power spectrum, this marginalization can be approximated analytically:
p(Cℓ|θ)≃N(Cℓ;Cthℓ,Σ)
- If x is e.g. a full shear map, this marginalization is analytically intractable.
⟹ We need to find a tractable way to compute the marginal likelihood p(x|θ).
How to perform inference over forward simulation models?
- Implicit Inference: Treat the simulator as a black-box with only the ability to sample from the joint distribution
(x,θ)∼p(x,θ)a.k.a.
- Simulation-Based Inference (SBI)
- Likelihood-free inference (LFI)
- Approximate Bayesian Computation (ABC)
- Explicit Inference: Treat the simulator as a probabilistic model and perform inference over the joint posterior
p(θ,z|x)∝p(x|z,θ)p(z,θ)p(θ)a.k.a.
- Bayesian Hierarchical Modeling (BHM)
⟹ For a given simulation model, both methods should converge to the same posterior!
Implicit Inference
The land of Neural Density Estimation
Black-box Simulators Define Implicit Distributions
- A black-box simulator defines p(x|θ) as an implicit distribution, you can sample from it but you cannot evaluate it.
- Key Idea: Use a parametric distribution model Pφ to approximate the implicit distribution P.
True P
Samples xi∼P
Model Pφ
- Once trained, you can typically sample from Pφ and/or evaluate the likelihood pφ(x|θ).
Conditional Density Estimation with Neural Networks
- I assume a forward model of the observations: p(x)=p(x|θ) p(θ) All I ask is the ability to sample from the model, to obtain D={xi,θi}i∈N
- I am going to assume qϕ(θ|x) a parametric conditional density
- Optimize the parameters ϕ of qϕ according to minϕ∑i−logqϕ(θi|xi) In the limit of large number of samples and sufficient flexibility qϕ∗(θ|x)≈p(θ|x)
⟹ One can asymptotically recover the posterior by
optimizing a parametric estimator over
the Bayesian joint distribution
the Bayesian joint distribution
⟹ One can asymptotically recover the posterior by
optimizing a Deep Neural Network over
a simulated training set.
a simulated training set.
A variety of algorithms
Lueckmann, Boelts, Greenberg, Gonçalves, Macke (2021) 
A few important points:
- Amortized inference methods, which estimate p(θ|x), can greatly speed up posterior estimation once trained.
- Sequential Neural Posterior/Likelihood Estimation methods can actively sample simulations needed to refine the inference.
My Practical Recipe to Apply Neural Density Estimation
A two-steps approach to Likelihood-Free Inference
- Step I Automatically learn a optimal low-dimensional summary statistic y=fφ(x) typically y will have the same dimensionality as θ.
- Step II: Use Neural Density Estimation in low dimension to either:
- build an estimate pϕ of the likelihood function p(y | θ) (Neural Likelihood Estimation)
- build an estimate pϕ of the posterior distribution p(θ | y) (Neural Posterior Estimation)
Automated Summary Statistics Extraction
- Introduce a parametric function fφ to reduce the dimensionality of the data while preserving information.
Information-based loss functions
- Summary statistics y is sufficient for θ if, and only if, I(Y;Θ)=I(X;Θ)⇔p(θ|x)=p(θ|y)
- Variational Mutual Information Maximization
L = Ex,θ[logqϕ(θ|fφ(x))]≤I(Y;Θ)
Jeffrey, Alsing, Lanusse (2021)
Example of application: Likelihood-Free parameter inference with DES SV
Jeffrey, Alsing, Lanusse (2021)
Suite of N-body + raytracing simulations: D
Towards Optimal Full-Field Implicit Inference
by Neural Summarisation and Density Estimation
Work led by Denise Lanzieri and Justine Zeghal
An easy-to-use validation testbed: log-normal lensing simulations
DifferentiableUniverseInitiative/sbi_lens
JAX-based log-normal lensing simulation package
- 10x10 deg2 maps at LSST Y10 quality, conditioning the log-normal shift parameter on (Ωm,σ8,w0)
- Infer full-field posterior on cosmology:
- explicitly using an Hamiltonian-Monte Carlo (NUTS) sampler
- implicitly using a learned summary statistics and conditional density estimation.
but explicit inference yields intermediate data products
simulated observed data
posterior samples κ=f(z,θ) with z∼p(z,θ|x)
and not all neural compression techniques are equivalent
- Most papers applying neural techniques for inference have used sub-optimal compression techniques, e.g. Mean Square Error L=||fφ(x)−θ||22 ⟹ This is ok, contours will simply be larger than they could be.
- A lot of papers are still relying on assuming a proxy Gaussian likelihoods, i.e. estimating a mean and covariance from
simulations.
⟹ This is dangerous, can lead to biased contours.
Explicit Inference
Where the JAX things are
Simulators as Hierarchical Bayesian Models
- If we have access to all latent variables z of the simulator, then the joint log likelihood p(x|z,θ) is explicit.
- We need to infer the joint posterior p(θ,z|x) before marginalization to
yield p(θ|x)=∫p(θ,z|x)dz.
⟹ Extremely difficult problem as z is typically very high-dimensional. - Necessitates inference strategies with access to gradients of the likelihood. dlogp(x|z,θ)dθ;dlogp(x|z,θ)dz For instance: Maximum A Posterior estimation, Hamiltonian Monte-Carlo, Variational Inference.
⟹ The only hope for explicit cosmological inference is to have fully-differentiable cosmological simulations!
the hammer behind the Deep Learning revolution: Automatic Differentation
- Automatic differentiation allows you to compute analytic derivatives of arbitraty expressions:
If I form the expression y=a∗x+b, it is separated in fundamental ops: y=u+bu=a∗x then gradients can be obtained by the chain rule: ∂y∂x=∂y∂u∂u∂x=1×a=a - This is a fundamental tool in Machine Learning, and autodiff frameworks include TensorFlow and PyTorch.
Enters JAX: NumPy + Autograd + GPU
- JAX follows the NumPy api!
import jax.numpy as np - Arbitrary order derivatives
- Accelerated execution on GPU and TPU
jax-cosmo: Finally a differentiable cosmology library, and it's in JAX!
Campagne, Lanusse, Zuntz et al. (2023)
import jax.numpy as np
import jax_cosmo as jc
# Defining a Cosmology
cosmo = jc.Planck15()
# Define a redshift distribution with smail_nz(a, b, z0)
nz = jc.redshift.smail_nz(1., 2., 1.)
# Build a lensing tracer with a single redshift bin
probe = probes.WeakLensing([nz])
# Compute angular Cls for some ell
ell = np.logspace(0.1,3)
cls = angular_cl(cosmo_jax, ell, [probe])
Current main features
- Weak Lensing and Number counts probes
- Eisenstein & Hu (1998) power spectrum + halofit
- Angular Cℓ under Limber approximation
⟹ 3x2pt DES Y1 capable
let's compute a Fisher matrix
F=−Ep(x|θ)[Hθ(logp(x|θ))]
import jax
import jax.numpy as np
import jax_cosmo as jc
# .... define probes, and load a data vector
def gaussian_likelihood( theta ):
# Build the cosmology for given parameters
cosmo = jc.Planck15(Omega_c=theta[0], sigma8=theta[1])
# Compute mean and covariance
mu, cov = jc.angular_cl.gaussian_cl_covariance_and_mean(cosmo,
ell, probes)
# returns likelihood of data under model
return jc.likelihood.gaussian_likelihood(data, mu, cov)
# Fisher matrix in just one line:
F = - jax.hessian(gaussian_likelihood)(theta)
- No derivatives were harmed by finite differences in the computation of this Fisher!
- Only a small additional compute time compared to one forward evaluation of the model
Inference becomes fast and scalable
- Current cosmological MCMC chains take days, and typically require access to large computer clusters.
- Gradients of the log posterior are required for modern efficient and scalable inference techniques:
- Variational Inference
- Hamiltonian Monte-Carlo
- In jax-cosmo, we can trivially obtain exact gradients:
def log_posterior( theta ): return gaussian_likelihood( theta ) + log_prior(theta) score = jax.grad(log_posterior)(theta) - On a DES Y1 analysis, we find convergence in 70,000 samples with vanilla HMC, 140,000 with Metropolis-Hastings
DES Y1 posterior, jax-cosmo HMC vs Cobaya MH
(credit: Joe Zuntz)
the Fast Particle-Mesh scheme for N-body simulations
The idea: approximate gravitational forces by estimating densities on a grid.- The numerical scheme:
- Estimate the density of particles on a mesh
=> compute gravitational forces by FFT - Interpolate forces at particle positions
- Update particle velocity and positions, and iterate
- Estimate the density of particles on a mesh
- Fast and simple, at the cost of approximating short range interactions.
⟹ Only a series of FFTs and interpolations.
introducing FlowPM: Particle-Mesh Simulations in TensorFlow
Modi, Lanusse, Seljak (2020)
import tensorflow as tf
import flowpm
# Defines integration steps
stages = np.linspace(0.1, 1.0, 10, endpoint=True)
initial_conds = flowpm.linear_field(32, # size of the cube
100, # Physical size
ipklin, # Initial powerspectrum
batch_size=16)
# Sample particles and displace them by LPT
state = flowpm.lpt_init(initial_conds, a0=0.1)
# Evolve particles down to z=0
final_state = flowpm.nbody(state, stages, 32)
# Retrieve final density field
final_field = flowpm.cic_paint(tf.zeros_like(initial_conditions),
final_state[0])- Seamless interfacing with deep learning components
Mesh FlowPM: distributed, GPU-accelerated, and automatically differentiable simulations
- We developed a Mesh TensorFlow implementation that can scale on GPU clusters (horovod+NCCL).
- For a 20483 simulation:
- Distributed on 256 NVIDIA V100 GPUs
- Don't hesitate to reach out if you have a use case for model parallelism!
![]()
- Now developing the next generation of these tools in JAX
Hybrid physical/neural differential equations
Lanzieri, Lanusse, Starck (2022)
- x and v define the position and the velocity of the particles
- ϕPM is the gravitational potential in the mesh
→ We can use this parametrisation to complement the physical ODE with neural networks.
Fθ(x,a)=3Ωm2∇[ϕPM(x)∗F−1(1+fθ(a,|k|))]
Correction integrated as a Fourier-based isotropic filter fθ → incorporates translation and rotation symmetries
Projections of final density field
Camels simulations
![]()
PM simulations
![]()
PM+NN correction
![]()
Results
Forward Models in Cosmology
Linear Field
Final Dark Matter
Dark Matter Halos
Galaxies
⟶
N-body simulations
⟶
Group Finding
algorithms
Group Finding
algorithms
⟶
HOD models
HOD models
Differentiable sampling from Halo Occupation Distributions
Horowitz, Hahn, Lanusse, Modi, Ferraro (2022)
https://github.com/DifferentiableUniverseInitiative/DHOD
- Sampling from a discrete random distribution is classically not a differentiable operation
- Relaxed and reparameterized HOD sampling
- Relaxed Bernoulli distributions for centrals
Ncen=11+exp(−(log(p1−p)+ϵ)/τ) with ϵ∼Logistic(0,1). where τ is a temperature parameter - Relaxed Binomial distribution for satelittes
Nsat∼Binomial(N,p=⟨Nsat⟩N)
- Relaxed Bernoulli distributions for centrals
But where are your full-field constraints with Hierarchical Bayesian Inference?
towards reasonable resolution differentiable lensing simulations
Lanzieri, Lanusse, et al. (2023)
5x5 sq deg^2 convergence field, with sources at z=0.9
FlowPM simulation with 128^3 particles
Conclusion
Conclusion
Methodology for inference over simulators
- A change of paradigm from analytic likelihoods to simulators as physical model.
- State of the art Machine Learning models enable Likelihood-Free Inference over black-box simulators.
- Progress in differentiable simulators and inference methodology paves the way to full inference over probabilistic model.
- Ultimately, promises optimal exploitation of survey data, although the "information gap" against analytic likelihoods in realistic settings remains uncertain.
Thank you!
Implicit and Explicit Simulation-Based Bayesian Inference for Cosmology François Lanusse slides at eiffl.github.io/talks/Prague2023