Implicit and Explicit Simulation-Based Inference for Cosmology

the Rubin Observatory Legacy Survey of Space and Time

1000 images each night, 15 TB/night for 10 years

18,000 square degrees, observed once every few days

Tens of billions of objects, each one observed $\sim1000$ times

We need to rethink all stages of data analysis for modern surveys

Bosch et al. 2017

Jeffrey, Lanusse, et al. 2020

(Hikage et al. 2018)

Cheng et al. 2020

Galaxies are no longer blobs.
Signals are no longer Gaussian.
Cosmological likelihoods are no longer tractable.

$\Longrightarrow$ This is the end of the analytic era...

Cosmological Simulation-Based Inference

Instead of trying to analytically evaluate the likelihood of sub-optimal summary statistics, let us build a forward model of the full observables.
$\Longrightarrow$ The simulator becomes the physical model.

Each component of the model 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)

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)

$\Longrightarrow$ For a given simulation model, both methods should converge to the same posterior!

Conventional Recipe for Full-Field Implicit Inference...

A two-steps approach to Implicit Inference

Automatically learn an optimal low-dimensional summary statistic $y = f_\varphi(x)$
Use Neural Density Estimation to either:
- build an estimate $p_\phi$ of the likelihood function $p(y \ | \ \theta)$ (Neural Likelihood Estimation)
- build an estimate $p_\phi$ of the posterior distribution $p(\theta \ | \ y)$ (Neural Posterior Estimation)

An easy-to-use experimentation testbed: log-normal lensing simulations

DifferentiableUniverseInitiative/sbi_lens
JAX-based log-normal lensing simulation package

10x10 deg $^2$ maps at LSST Y10 quality, conditioning the log-normal shift parameter on $(\Omega_m, \sigma_8, w_0)$

Provides explicit posterior by Hamiltonian-Monte Carlo (NUTS) sampling in reasonable time

Can be used to sample a practically infinite number of maps for implicit inference

Information Point of View on Neural Summarisation

Learning Sufficient Statistics

Summary statistics $y$ is sufficient for $\theta$ if $I(Y; \Theta) = I(X; \Theta) \Leftrightarrow p(\theta | x ) = p(\theta | y)$
Variational Mutual Information Maximization $\mathcal{L} \ = \ \mathbb{E}_{x, \theta} [ \log q_\phi(\theta | y=f_\varphi(x)) ] \leq I(Y; \Theta)$ (Barber & Agakov variational lower bound)
Jeffrey, Alsing, Lanusse (2021)

Can you train any Normalizing Flow by Score Matching?

Let us assume an analytic distribution with known score, and train $p_\varphi(x)$ by Score Matching: $\mathcal{L} = \mathbb{E}_{x \sim p} \left[ \left| \nabla_x \log p(x) - \nabla_x \log p_\varphi(x) \right|^2 \right]$

True distribution and score

Trained RealNVP and model score

Why? The score of a RealNVP with affine coupling are trivial.

$\left\{ \begin{array}{ll} x_{1:d} &= z_{1:d} \\ x_{d+1:D} &= a(z_{1:d}) \times z_{d+1:D} + b(z_{1:d}) \end{array} \right.$

Smooth Normalizing Flows

Adapted from Kohler et al. 2021

Key idea I: Use $\mathcal{C}^\infty$ coupling layers $\left\{ \begin{array}{ll} x_{1:d} &= z_{1:d} \\ x_{d+1:D} &= \sigma_{(a,b,c)}(z_{d+1:D}) \end{array} \right.$ with $\sigma(x) := c\cdot x + \frac{1-c}{1+\exp({-\rho(x)})}$ , $\rho(x) := a \cdot \left(\log\left(\frac{x}{1-x}\right) + b\right)$ and $a$ , $b$ , $c$ learned using a neural network.

$\Longrightarrow$ Retains non-trivial and expressive gradients. but is generally non-analytically invertible.

Key Idea II: Use a numerical inverse for $\sigma^{-1}(x)$ and implicit function theorem to compute its gradient.

$\Longrightarrow$ Computes inverse function and its gradients by solving a root-finding problem.

Neural Posterior Estimation with Differentiable Simulators

Zeghal, et al. (2022)

$\mathcal{L} = \mathbb{E}_{x,\theta,z \sim p} \left[ - \log p_\varphi(\theta| x) + \lambda \left| \nabla_\theta \log p(\theta, z | x) - \nabla_\theta \log p_\varphi(\theta|x) \right|^2 \right]$

Illustration on Lotka-Volterra

If only things were that easy...

Kiessling et al. (2015)

Tenneti et al. (2015)

Tidal interactions with local gravitational potential can lead to coherent intrinsic galaxy alignments which mimics gravitational lensing.

Very complicated effect in details, no single analytic model for all galaxy types
$\Longrightarrow$ study requires expensive hydrodynamical simulations

Let's take a step back: The 3D pose estimation problem

(credit: Murphi et al. 2021)

$\Longrightarrow$ We want to model a distribution

$p(\mathbf{R} | y)$ where

$\mathbf{R} \in \mathrm{SO}(3)$ , the Lie group of 3D rotations.

What you need to know about diffusion models in $\mathbb{R}^n$

Song et al. (2021)

The SDE defines a marginal distribution $p_t(x)$ as the convolution of the target distribution $p(x)$ with a noise kernel $p_{t|s}(\cdot | x_s)$ : $p_t(x) = \int p(x_s) p_{t|s}(x | x_s) d x_s$
For a given forward SDE that evolves $p(x)$ to $p_T(x)$ , there exists a reverse SDE that evolves $p_T(x)$ back into $p(x)$ . It involves having access to the marginal score $\nabla_x \log_t p(x)$ .
You can sample by solving the associated ODE (aka probability flow ODE): $\mathrm{d} \mathbf{x} = [\mathbf{f}(\mathbf{x}, t) - g(t)^2 \nabla \log p_t(\mathbf{x})] \mathrm{d} t$

Going back to the heat equation

The heat kernel is the solution of the heat diffusion equation and corresponds to the
transition density of Brownian motion

$p_{t | s}(\cdot | x_s)$ for

$t>s$ .

On $\mathbb{R}^n$ the heat kernel is a Gaussian distribution $\mathcal N(0, t \mathbf{I})$
Knowing the solution of the heat equation allows us to easily sample the marginal distribution $p_t(x) = \int p_s( x_s ) p_{t|s}(x | x_s) d x_s$

$\Longrightarrow$ On a closed manifold like

$\mathrm{SO}(3)$ , the heat kernel is not a Gaussian distribution anymore!

Solution of the heat equation on SO(3)

On

$\mathrm{SO}(3)$ the heat kernel can be expressed as (Nikolayev & Savyolov, 1970):

$f_\epsilon(\omega) = \sum_{\ell=0}^{\infty} (2 \ell +1) \exp(- l (l+1) \epsilon^2) \frac{\sin((\ell + 1/2) \omega)}{\sin(\omega/2)}$ where

$\omega \in \left( -\pi, \pi \right]$ is the rotation angle of an axis-angle representation of

$\mathrm{SO}(3)$ ,

$\epsilon$ is a concentration parameter.

$\Longrightarrow$ Can be robustly approximated by truncation or closed form expressions (Matthies et al. 1988).

Isotropic Gaussian Distribution on SO(3)

The isotropic Gaussian distribution on SO(3) is defined as: $\mathcal{IG}_{\mathrm{SO(3)}}(\mathbf{x}; \mathbf{\mu}, \epsilon) = f_\epsilon(\arccos( 2^{-1} \mathrm{tr}(\mathbf{\mu}^T \mathbf{x}) - 1 ) )$ where $\mathbf{x}$ and $\mathbf{\mu}$ are rotation matrices and $f_\epsilon$ is the density function with variance $\epsilon$ .
Like the Gaussian distribution, it is closed under convolution, corresponds to the solution of the heat diffusion process on SO(3).

$\Longrightarrow$ Sampling from the marginal distribution at time

$t$ becomes very easy:

$x \sim p(x), \mathbf{u} \sim \mathcal{IG}_{\mathrm{SO(3)}}(\mathbf{Id}, \epsilon(t)), \mathbf{x}^\prime = \mathbf{u} \mathbf{x}$

Results on toy distributions

Jagvaral, Lanusse, Mandelbaum, AAAI 2024

Unconditional distribution modeling.
Conditional distribution modeling $p(\mathbf{R} | y)$ simply by making the score network conditional $s_\theta( \mathbf{R}, \epsilon, y)$

Hybrid physical/neural differential equations

Lanzieri, Lanusse, Starck (2022)

$\left\{ \begin{array}{ll} \frac{d \color{#6699CC}{\mathbf{x}} }{d a} & = \frac{1}{a^3 E(a)} \color{#6699CC}{\mathbf{v}} \\ \frac{d \color{#6699CC}{\mathbf{v}}}{d a} & = \frac{1}{a^2 E(a)} F_\theta( \color{#6699CC}{\mathbf{x}} , a), \\ F_\theta( \color{#6699CC}{\mathbf{x}}, a) &= \frac{3 \Omega_m}{2} \nabla \left[ \color{#669900}{\phi_{PM}} (\color{#6699CC}{\mathbf{x}}) \right] \end{array} \right.$

$\mathbf{x}$ and $\mathbf{v}$ define the position and the velocity of the particles
$\phi_{PM}$ is the gravitational potential in the mesh

$\to$ We can use this parametrisation to complement the physical ODE with neural networks.

$F_\theta(\mathbf{x}, a) = \frac{3 \Omega_m}{2} \nabla \left[ \phi_{PM} (\mathbf{x}) \ast \mathcal{F}^{-1} (1 + \color{#996699}{f_\theta(a,|\mathbf{k}|)}) \right]$

Correction integrated as a Fourier-based isotropic filter

$f_{\theta}$

$\to$ incorporates translation and rotation symmetries

Defining Custom Distributed Ops in JAX

 from jax.experimental import mesh_utils
from jax.sharding import PositionalSharding

# Create a Sharding object to distribute a value across devices:
sharding = PositionalSharding(mesh_utils.create_device_mesh((8,)))

# Create an array of random values:
x = jax.random.normal(jax.random.PRNGKey(0), (8192, 8192))
# and use jax.device_put to distribute it across devices:
y = jax.device_put(x, sharding.reshape(4, 2))
jax.debug.visualize_array_sharding(y)

---

┌──────────┬──────────┐
│  TPU 0   │  TPU 1   │
├──────────┼──────────┤
│  TPU 2   │  TPU 3   │
├──────────┼──────────┤
│  TPU 6   │  TPU 7   │
├──────────┼──────────┤
│  TPU 4   │  TPU 5   │
└──────────┴──────────┘

Distributed arrays and automatic parallelization

 from jaxlib.hlo_helpers import custom_call

def _rms_norm_fwd_cuda_lowering(ctx, x, weight, eps):
	[...]
	opaque = gpu_ops.create_rms_norm_descriptor([...])

	out = custom_call(
		b"rms_forward_affine_mixed_dtype",
		result_types=[
			ir.RankedTensorType.get(x_shape, w_type.element_type),
			ir.RankedTensorType.get((n1,), iv_element_type),
		],
		operands=[x, weight],
		backend_config=opaque,
		operand_layouts=default_layouts(x_shape, w_shape),
		result_layouts=default_layouts(x_shape, (n1,)),
	).results
	return out

_rms_norm_fwd_p = core.Primitive("rms_norm_fwd")
mlir.register_lowering(
	_rms_norm_fwd_p,
	_rms_norm_fwd_cuda_lowering,
	platform="gpu",
)

Custom operations for GPUs with C++ and CUDA

def partition(mesh, arg_shapes, result_shape):
	result_shardings = jax.tree.map(lambda x: x.sharding, result_shape)
	arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
	return mesh, fft, \
		supported_sharding(arg_shardings[0], arg_shapes[0]), \
		(supported_sharding(arg_shardings[0], arg_shapes[0]),)

def infer_sharding_from_operands(mesh, arg_shapes, result_shape):
	arg_shardings = jax.tree.map(lambda x: x.sharding, arg_shapes)
	return supported_sharding(arg_shardings[0], arg_shapes[0])

@custom_partitioning
def my_fft(x):
	return fft(x)

my_fft.def_partition(
	infer_sharding_from_operands=infer_sharding_from_operands,
	partition=partition)

jax.experimental.custom_partitioning

JAX cuDecomp interface led by Wassim Kabalan (IN2P3/APC)

JAX>=v0.4.1 defines sharded tensors and a "computation follows data" philosophy.

jaxlib provides a helper to define custom CUDA lowering

Recent API allows us to define custom partitioning schemes compatible with a primitive

Implicit and Explicit Simulation-Based Inference for Cosmology

François Lanusse