Neural Compression and Neural Density Estimation for Cosmological Inference

Justine Zeghal

justine.zeghal@umontreal.ca

Bayesian Deep Learning for Cosmology and Time Domain Astrophysics 3rd ed, Paris, May 22

Lambda Cold Dark Matter (   CDM)

\Lambda

The simplest model that best describes our observations is

\Lambda \text{CDM}

Relying only on a few parameters:

\Omega_c,\: \Omega_b,\:\Omega_\Lambda,\: h_0, \: n_s, \sigma_8\:.

Suggesting: ordinary matter, cold dark matter (CDM), and dark energy Λ as an explanation of the accelerated expansion.

Goal: determine the value of those parameters based on our observations.

Credit: ESA

How to constrain cosmological parameters?

x

For which we have an analytical likelihood function.

This likelihood function connects our compressed observations to the cosmological parameters.

p(t|\theta)
t = f(x)

Bayes theorem:

\underbrace{p(\theta|t=t_0)}_{\text{posterior}}
\underbrace{p(\theta)}_{\text{prior}}
\underbrace{p(t = t_0|\theta)}_{\text{likelihood}}
\propto

Planck Collaboration et al. (2020)

We need to update our inference methods

The traditional way of constraining cosmological parameters misses information.

This results in constraints on cosmological parameters that are not precise.

Credit: Natalia Porqueres

DES Y3 Results (with SBI).

Jeffrey et al. (2024)

\underbrace{p(\theta|x=x_0)}_{\text{posterior}}
\underbrace{p(\theta)}_{\text{prior}}
\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}
\propto

Bayes theorem:

We can build a simulator to map the cosmological parameters to the data.

Prediction

Inference

Full-field inference: extracting all cosmological information

x
\theta

 Simulator

\underbrace{p(x = x_0|\theta)}_{\text{likelihood}}

Full-field inference: extracting all cosmological information

Depending on the simulator’s nature we can either perform

  • Explicit inference 
  • Implicit inference 
x
\theta
z
f

 Simulator

Full-field inference: extracting all cosmological information

  • Explicit inference 

Explicit joint likelihood

 

p(x| \theta, z)
z
x

Initial conditions of the Universe

Large Scale Structure

p(\theta, z \: | \: x) \propto
p(z\:|\:\theta) p(\theta)

Needs an explicit simulator to sample the joint posterior through MCMC:

p(x| \theta, z)

We need to sample in extremely

high-dimension

→ gradient-based sampling schemes.

x
\theta
z
f
\sigma^2
\mathcal{N}

Credit:  Chirag Modi, François Lanusse, Mustafa Mustafa, Uroŝ Seljak

Depending on the simulator’s nature we can either perform

  • Explicit inference 
  • Implicit inference 

Full-field inference: extracting all cosmological information

x
\theta

 Simulator

z
f

  • Implicit inference 

It does not matter if the simulator is explicit or implicit because all we need are simulations

(\theta_i, x_i)_{i=1...N}

This approach typically involve 2 steps:

2) Implicit inference on these summary statistics to approximate the posterior.

1) compression of the high dimensional data into summary statistics. Without loosing cosmological information!

Summary statistics

t = f_{\varphi}(x)
p_{\Phi}(\theta | f_{\varphi}(x))

Full-field inference: extracting all cosmological information

x
\theta

 Simulator

z
f

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Neural Posterior Estimation with Differentiable Simulators

ICML 2022 Workshop on Machine Learning for Astrophysics

 

Justine Zeghal, François Lanusse, Alexandre Boucaud,

Benjamin Remy and Eric Aubourg

Implicit Inference

1) Draw N parameters 

\theta_i \sim p(\theta)

2) Draw N simulations 

x_i \sim p(x\:|\:\theta = \theta_i)

3) Train a neural density estimator on                                           to approximate the quantity of interest

(\theta_i, x_i)_{i=1...N}

4) Approximate the posterior from the learned quantity

\text{Sample } x \sim p(x)
p_{\phi}(x)

Algorithm

Normalizing Flows

p_x(x)
p_z(z)
f^{-1}_1
f^{-1}_2
f_1
f_2

Normalizing Flows

p_x(x)
p_z(z)
f_1

Normalizing Flows

p_x(x)
p_z(z)
f_1

Normalizing Flows

p_x(x)
p_z(z)
f_1
f_2

Normalizing Flows

p_x(x)
p_z(z)
f^{-1}_2
f_1
f_2

Normalizing Flows

p_x(x)
p_z(z)
f^{-1}_2
f_1
f_2
f^{-1}_1

Normalizing Flows

p_x(x)
p_z(z)
\log p_z(f^{-1}(x))
+ \log \displaystyle\left\lvert det \frac{\partial f^{-1}(x)}{\partial x}\right\rvert

Change of Variable Formula:

f^{-1}_1
f^{-1}_2
f_1
f_2
\log p_x(x)
=

Normalizing Flows

p_x(x)
p_z(z)
\log p_z(f^{-1}(x))
+ \log \displaystyle\left\lvert det \frac{\partial f^{-1}(x)}{\partial x}\right\rvert

Change of Variable Formula:

f^{-1}_1
f^{-1}_2
f_1
f_2
\log p_x(x)
=

Normalizing Flows

p_x(x)
p_z(z)
f^{-1}_1
f^{-1}_2
f_1
f_2

We need to learn the mapping

to approximate the complex distribution.

\begin{array}{ll} D_{KL}(p_x(x)||p_x^{\phi}(x)) &=\mathbb{E}_{p_x(x)}\left[ \log\left(p_x(x)\right) \right] - \mathbb{E}_{p_x(x)}\left[ \log\left(p_x^{\phi}(x)\right) \right] \end{array}
\begin{array}{ll} \implies Loss = - \mathbb{E}_{x \sim p_x(x)}\left[ \log\left(p_x^{\phi}(x)\right) \right]\\ \end{array}

From simulations only!

A lot of simulations..

Truth

Approximation

With a few simulations it's hard to approximate the posterior distribution.

→ we need more simulations

BUT if we have a few simulations

and the gradients

 

(also know as the score)

\nabla_{\theta} \log p(\theta | x)

then it's possible to have an idea of the shape of the distribution.

How gradients can help reduce the number of simulations?

How to train NFs with gradients?

How to train NFs with gradients?

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]

How to train NFs with gradients?

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]

How to train NFs with gradients?

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]

How to train NFs with gradients?

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(\theta, z |x) - \nabla_{\theta} \log p^{\phi}(\theta |x)\parallel_2^2 \right]

How to train NFs with gradients?

Problem: the gradient of current NFs lack expressivity

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(\theta, z |x) - \nabla_{\theta} \log p^{\phi}(\theta |x)\parallel_2^2 \right]

How to train NFs with gradients?

Problem: the gradient of current NFs lack expressivity

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(\theta, z |x) - \nabla_{\theta} \log p^{\phi}(\theta |x)\parallel_2^2 \right]

How to train NFs with gradients?

Problem: the gradient of current NFs lack expressivity

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(\theta, z |x) - \nabla_{\theta} \log p^{\phi}(\theta |x)\parallel_2^2 \right]

How to train NFs with gradients?

Problem: the gradient of current NFs lack expressivity

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(\theta, z |x) - \nabla_{\theta} \log p^{\phi}(\theta |x)\parallel_2^2 \right]

How to train NFs with gradients?

Problem: the gradient of current NFs lack expressivity

But to train the NF, we want to use both simulations and the gradients from the simulator:

Normalizing flows are trained by minimizing the negative log likelihood:

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(\theta | x)\right) \right]
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(\theta, z |x) - \nabla_{\theta} \log p^{\phi}(\theta |x)\parallel_2^2 \right]

Benchmark Metric

A metric

We use the Classifier 2-Sample Tests (C2ST) metric.

  • C2ST=0.5 (i.e “Impossible to differentiate 👍🏼”)
  • C2ST=1(i.e “Too easy to differentiate 👎🏻”)

distribution 1

distribution 2

Requirement: the true distributions is needed.

Results on a toy model

→ On a toy Lotka Volterra model, the gradients helps to constrain the distribution shape.

Results on a toy model

Without gradients

With gradients

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Simulation-Based Inference Benchmark for LSST Weak Lensing Cosmology

Justine Zeghal, Denise Lanzieri, François Lanusse, Alexandre Boucaud, Gilles Louppe, Eric Aubourg, Adrian E. Bayer

and The LSST Dark Energy Science Collaboration (LSST DESC)

  • do gradients help implicit inference methods?

In the case of weak lensing full-field analysis,

  • which inference method requires the fewest simulations?

x
\theta
z
f
\sigma^2
\mathcal{N}

We developed a fast and differentiable (JAX) log-normal mass maps simulator.

For our benchmark: a Differentiable Mass Maps Simulator

DifferentiableUniverseInitiative/​sbi-lens

Benchmark metric

Explicit inference theoretically and asymptotically converges to the truth.

Explicit inference and implicit inference yield comparable constraints.

C2ST = 0.6!

To use the C2ST we need the true posterior distribution.

→ We use the explicit full-field posterior.

Why?

  • Do gradients help implicit inference methods?             

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(x | \theta)\right) \right]
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(x, z |\theta) -{\nabla_{\theta} \log p^{\phi}(x|\theta )}\parallel_2^2 \right]

Training the NF with simulations and gradients:

Loss = 

  • Do gradients help implicit inference methods? 

- \mathbb{E}_{p(x, \theta)}\left[ \log\left(p^{\phi}(x | \theta)\right) \right]

Training the NF with simulations and gradients:

Loss = 

  • Do gradients help implicit inference methods? 

(\theta_i, x_i)_{i=1...N}
\nabla_{\theta} \log \hat{p}(x |\theta)
+ \: \lambda \: \displaystyle \mathbb{E}\left[ \parallel \nabla_{\theta} \log p(x, z |\theta) -{\nabla_{\theta} \log p^{\phi}(x|\theta )}\parallel_2^2 \right]
\nabla_{\theta}\log p(x|\theta)
\nabla_{\theta}\log p(x,z|\theta)

(from the simulator)

→ For this particular problem, the gradients from the simulator are too noisy to help.

  • Do gradients help implicit inference methods?     

→ Implicit inference requires 1500 simulations.

→ In the case of perfect gradients it does not significantly help.

→ Simple distribution all the simulations seems to help locate the posterior distribution.

  • do gradients help implicit inference methods?

In the case of weak lensing full-field analysis,

  • which inference method requires the fewest simulations?

x
\theta
z
f
\sigma^2
\mathcal{N}

→  No, it does not help to reduce the number of simulations because the gradients of the simulator are too noisy.

→  Even with marginal gradients the gain is not significant.

→  For now, we now that implicit inference requires 1500 simulations.

  • which inference method requires the fewest simulations?

What about explicit inference?

  • which inference method requires the fewest simulations?

What about explicit inference?

→ Explicit inference requires

simulations.

10^5

  • which inference method requires the fewest simulations?

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Optimal Neural Summarisation for Full-Field Weak Lensing Cosmological Implicit Inference

Denise Lanzieri*, Justine Zeghal*, T. Lucas Makinen, François Lanusse, Alexandre Boucaud and Jean-Luc Starck

​* equal contibutions

x
t = f_{\varphi}(x)
\text{A statistic } t \text{ is said to be sufficient for the parameters } \theta \text{ if }
p(\theta \: | \: x) = p(\theta \: | \: t) \: \text{ with } \: t=f(x)

How to extract all the information?

It is only a matter of the loss function used to train the compressor.

Definition: Sufficient Statistic

Two main compression schemes

\hat{\theta} = f_\varphi(x)

Regression Losses

Two main compression schemes

Text

\hat{\theta} = f_\varphi(x)

Regression Losses

Two main compression schemes

Which learns a moment of the posterior distribution.

Mean Squared Error (MSE) loss:

\mathcal{L}_{\text{MSE}} = \mathbb{E}_{p(x,\theta)} \left[\parallel \theta - f_\varphi(x) \parallel ^2 \right]

Which learns a moment of the posterior distribution.

\hat{\theta} = f_\varphi(x)

Regression Losses

Two main compression schemes

Mean Squared Error (MSE) loss:

\mathcal{L}_{\text{MSE}} = \mathbb{E}_{p(x,\theta)} \left[\parallel \theta - f_\varphi(x) \parallel ^2 \right]

→ Approximate the mean of the posterior.

\hat{\theta} = f_\varphi(x)

Regression Losses

Two main compression schemes

Which learns a moment of the posterior distribution.

Mean Squared Error (MSE) loss:

\mathcal{L}_{\text{MSE}} = \mathbb{E}_{p(x,\theta)} \left[\parallel \theta - f_\varphi(x) \parallel ^2 \right]

→ Approximate the mean of the posterior.

\hat{\theta} = f_\varphi(x)

Regression Losses

Mean Absolute Error (MAE) loss:

\mathcal{L}_{\text{MAE}} = \mathbb{E}_{p(x,\theta)} \left[| \theta - f_\varphi(x) |\right]

Two main compression schemes

Which learns a moment of the posterior distribution.

Mean Squared Error (MSE) loss:

\mathcal{L}_{\text{MSE}} = \mathbb{E}_{p(x,\theta)} \left[\parallel \theta - f_\varphi(x) \parallel ^2 \right]

→ Approximate the mean of the posterior.

\hat{\theta} = f_\varphi(x)

Regression Losses

Mean Absolute Error (MAE) loss:

\mathcal{L}_{\text{MAE}} = \mathbb{E}_{p(x,\theta)} \left[| \theta - f_\varphi(x) |\right]

→ Approximate the median of the posterior.

Two main compression schemes

Which learns a moment of the posterior distribution.

Regression Losses

Two main compression schemes

\mu_1 = \mu_2

Regression Losses

Two main compression schemes

\mu_1 = \mu_2

Regression Losses

Two main compression schemes

\mu_1 = \mu_2

Regression Losses

Two main compression schemes

\mu_1 = \mu_2

Regression Losses

\neq

The mean is not guaranteed to be a sufficient statistic.

Two main compression schemes

Mutual information maximization

Two main compression schemes

p(\theta \: | \: x) = p(\theta \: | \: t(x)) \: \Leftrightarrow I(\theta, x) = I (\theta, t(x))

Mutual information maximization

By definition:

Two main compression schemes

I(\theta,t)
H(\theta)
H(t)
H(\theta|t)
H(t|\theta)
I(\theta,t) = H(t) - H(t|\theta)
p(\theta \: | \: x) = p(\theta \: | \: t(x)) \: \Leftrightarrow I(\theta, x) = I (\theta, t(x))

Mutual information maximization

By definition:

Two main compression schemes

I(\theta,t)
H(\theta)
H(t)
H(\theta|t)
H(t|\theta)
I(\theta,t) = H(t) - H(t|\theta)
p(\theta \: | \: x) = p(\theta \: | \: t(x)) \: \Leftrightarrow I(\theta, x) = I (\theta, t(x))

Mutual information maximization

By definition:

Two main compression schemes

I(\theta,x)
I(\theta,t) = H(t) - H(t|\theta)
p(\theta \: | \: x) = p(\theta \: | \: t(x)) \: \Leftrightarrow I(\theta, x) = I (\theta, t(x))

Mutual information maximization

By definition:

Two main compression schemes

I(\theta,t) = H(t) - H(t|\theta)
p(\theta \: | \: x) = p(\theta \: | \: t(x)) \: \Leftrightarrow I(\theta, x) = I (\theta, t(x))

Mutual information maximization

By definition:

Two main compression schemes

I(\theta,t) = H(t) - H(t|\theta)
p(\theta \: | \: x) = p(\theta \: | \: t(x)) \: \Leftrightarrow I(\theta, x) = I (\theta, t(x))

Mutual information maximization

By definition:

Two main compression schemes

I(\theta,t) = H(t) - H(t|\theta)

→ should build sufficient statistics according to the definition.

p(\theta \: | \: x) = p(\theta \: | \: t(x)) \: \Leftrightarrow I(\theta, x) = I (\theta, t(x))

Mutual information maximization

By definition:

Two main compression schemes

For our benchmark

Log-normal LSST Y10 like

differentiable

simulator

t = f_{\varphi}(x)
DifferentiableUniverseInitiative/​sbi-lens

1. We compress using one of the losses.

Benchmark procedure:

2. We compare their extraction power by comparing their posteriors.

For this, we use implicit inference, which is fixed for all the compression strategies.

p(\theta \: | \: x) = p(\theta \: | \: t) \: \text{ with } \: t=f(x)

Numerical results

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Outline

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Correcting Model Misspecification with Conditional Optimal Transport

Justine Zeghal, Benjamin Remy, Laurence Perreault-Levasseur, Yashar Hezaveh

Preliminary results*

What happens when the simulation model differs from the true physical model?

With full-field inference, we are now only relying on simulations, and we work at the pixel level.

We cannot escape this, as there may be physics that we do not understand or cannot model computationally.

A way to correct this bias is to learn a mapping  to transform one

simulation into another

x_1 = \phi_1(x_0)

and we would like it to be the optimal transport mapping in the sense that        is minimally transformed to match its PM counterpart.

x_0

OT Flow Matching enables to learn an OT mapping between two random distributions.

f^{-1}_1
f^{-1}_2
f_1
f_2
\text{of } dx = v_t(x)dt

Need to learn discrete transformations

f_1 \text{ and } f_2.

Need to learn a continuous  transformation

f_t \text{ solution }

Credit: https://mlg.eng.cam.ac.uk/blog/2024/01/20/flow-matching.html

Credit: Tong et al., 2023

Credit: Albergo et al., 2023

Flow Matching

Optimal Transport Flow Matching

Optimal Transport Flow Matching

Optimal Transport Flow Matching

Optimal Transport Flow Matching

Optimal Transport Flow Matching

Preliminary Results

Conclusion

Which full-field inference methods require the fewest simulations?

How to build sufficient statistics?

Can we perform implicit inference with fewer simulations?

How to deal with model misspecification?

Gradients can be beneficial, depending on your simulation model. 

Explicit inference requires 100 times more simulations than implicit inference.

Mutual Information Maximization

We can learn an optimal transport mapping.

\theta
f
\mathcal{N}

 Simulator

Summary statistics

t = f_{\varphi}(x)
x
p_{\Phi}(\theta | f_{\varphi}(x))

Thank you for your attention!

Talk BDL3

By Justine Zgh

Talk BDL3

  • 188

More from Justine Zgh