Learning Objectives & Program

$+ p(x) =$

• Morning: Introduction to Neural Networks and Understanding Uncertainties
  • What are Neural Networks and how to code them
  • Combining Neural Networks and Probabilities to perform Bayesian Inference
  Hands-on tutorial: Estimating Masses of Galaxy Clusters in JAX/Flax/Tensorflow Probability


• Afternoon: Deep Generative Modeling and Application to Inverse Problems
  • How to solve inverse problems from a Bayesian point of view
  • How to model high dimensional distributions with Deep Generative Models
  Hands-on tutorials:
  • Building your own Normalizing Flow in JAX/Flax/Tensorflow Probability
  • Deconvolving Galaxy Images from the HSC survey with Deep Generative Priors

What is a neural network?

Simplest architecture: Multilayer Perceptron (MLP)

• Series of Dense a.k.a Fully connected layers: $$ h = \sigma(W x + b)$$ where:
  • $\sigma$ is the activation function (e.g. ReLU, Sigmoid, etc.)
  • $W$ is a multiplicative weight matrix
  • $b$ is an additive bias parameter


• This defines a parametric non-linear function $f_\theta(x)$


• MLPs are universal function approximators
  Nota bene: only asymptotically true!

How do you use it to approximate functions?

• Assume a loss function that should be small for good approximations on a training set of data points $(x_i, y_i)$ $$ \mathcal{L} = \sum_{i} ( y_i - f_\theta(x_i))^2 $$


• Optimize the parameters $\theta$ to minimize the loss function by gradient descent $$ \theta_{t+1} = \theta_t - \eta \nabla_\theta \mathcal{L} $$

Different neural architectures for different types of data

Performance can be improved for particular types of data by making use of inductive biases

$$ h = \sigma(W \ast x + b)$$

• Image -> Convolutional Networks
  • Convolutional layers are translation invariant
  • Pixel topology is incorporated by construction


• Time Series -> Recurrent Networks
  • Temporal information and variable length is incorporated by construction


• Graph-stuctured data -> Graph Neural Networks
  • Graph topology is incorporated by construction
  • Can be seen as generalisation of convolutional networks

Our case study: Dynamical Mass Measurement for Galaxy Clusters

Figures and data from Ho et al. 2019

Our goal: train a neural network to estimate the cluster mass

What we will feed the model:

• Richness
• Velocity Dispersion
• Information about member galaxies:
  • radial distribution
  • stellar mass distribution
  • LOS velocity distribution

Training data from MultiDark Planck 2 N-body simulation (Klypin et al. 2016) with 261287 clusters.

JAX: NumPy + Autograd + XLA

• JAX uses the NumPy API
  => You can copy/paste existing code, works pretty much out of the box
  
   
• JAX is a successor of autograd
  => You can transform any function to get forward or backward automatic derivatives (grad, jacobians, hessians, etc)
  
   
• JAX uses XLA as a backend
  => Same framework as used by TensorFlow (supports CPU, GPU, TPU execution)

import jax.numpy as np

m = np.eye(10) # Some matrix

def my_func(x):
	return m.dot(x).sum() 

x = np.linspace(0,1,10)
y = my_func(x)
						

from jax import grad

df_dx = grad(my_func)
y = df_dx(x)

Now you try it!

def loss_fn(params, x, y):
  mse = (model.apply(params, x) - y)**2
  return jnp.mean(mse)

Raise your hand if you manage to get a cluster mass prediction!

First attempt with an MSE loss

class MLP(nn.Module):
@nn.compact
def __call__(self, x):
	x = nn.relu(nn.Dense(128)(x))
	x = nn.relu(nn.Dense(128)(x))
	x = nn.tanh(nn.Dense(64)(x))
	x = nn.Dense(1)(x)
	return x

def loss_fn(params, x, y):
	prediction = model.apply(params, x)
	return jnp.mean( (prediction - y)**2 )

• The prediction appears to be biased... What is going on?

Let's try to understand the neural network output by looking at the loss function

• Using an $\ell_2$ loss learns the mean of the $p(x | y)$


• Using an $\ell_1$ loss learns the median of $p(x|y)$


• In general, training a neural network for regression doesn't achieve de mode of the distribution.
  
  
  Check this blogpost and this notebook to learn how to do that:

Let's take a step back and think as Bayesians

We will need a new tool to go probabilistic!

TensorFlow Probability

import jax
import jax.numpy as jnp
from tensorflow_probability.substrates import jax as tfp
tfd = tfp.distributions

### Create a mixture of two scalar Gaussians:
gm = tfd.MixtureSameFamily(
	mixture_distribution=tfd.Categorical(
					probs=[0.3, 0.7]),
	components_distribution=tfd.Normal(
					loc=[-1., 1],      
					scale=[0.1, 0.5]))

# Evaluate probability
gm.log_prob(1.0)

Let's build a Conditional Density Estimator in JAX/Flax/TFP

class MDN(nn.Module):
	num_components: int

	@nn.compact
	def __call__(self, x):
		x = nn.relu(nn.Dense(128)(x))
		x = nn.relu(nn.Dense(128)(x))
		x = nn.tanh(nn.Dense(64)(x))

		# Instead of regressing directly the value of the mass, the network
		# will now try to estimate the parameters of a mass distribution.
		categorical_logits = nn.Dense(self.num_components)(x)
		loc = nn.Dense(self.num_components)(x)
		scale = 1e-3 + nn.softplus(nn.Dense(self.num_components)(x))

		dist =tfd.MixtureSameFamily(
			mixture_distribution=tfd.Categorical(logits=categorical_logits),
			components_distribution=tfd.Normal(loc=loc, scale=scale))
		
		# To make it understand the batch dimension
		dist =  tfd.Independent(dist)
		
		# Returns a distribution !
		return dist

Now you try it!

def loss_fn(params, x, y):
	q = model.apply(params, x)
	return jnp.mean( - q.log_prob(y[:,0]) )

Raise your hand if you manage to get a cluster mass prediction with this new model!

The limits of traditional cosmological inference

HSC cosmic shear power spectrum

HSC Y1 constraints on $(S_8, \Omega_m)$

(Hikage et al. 2018)

• Measure the ellipticity $\epsilon = \epsilon_i + \gamma$ of all galaxies
  $\Longrightarrow$ Noisy tracer of the weak lensing shear $\gamma$


• Compute summary statistics based on 2pt functions,
  e.g. the power spectrum


• Run an MCMC to recover a posterior on model parameters, using an analytic likelihood $$ p(\theta | x ) \propto \underbrace{p(x | \theta)}_{\mathrm{likelihood}} \ \underbrace{p(\theta)}_{\mathrm{prior}}$$

Full-Field 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.

...so why is this not mainstream?

The Challenge of Simulation-Based Inference

$$ p(x|\theta) = \int p(x, z | \theta) dz = \int p(x | z, \theta) p(z | \theta) dz $$ Where $z$ are stochastic latent variables of the simulator.

$\Longrightarrow$ This marginal likelihood is intractable!

Black-box Simulators Define Implicit Distributions

• A black-box simulator defines $p(x | \theta)$ as an implicit distribution, you can sample from it but you cannot evaluate it.
• Key Idea: Use a parametric distribution model $\mathbb{P}_\varphi$ to approximate the implicit distribution $\mathbb{P}$.

True $\mathbb{P}$

Samples $x_i \sim \mathbb{P}$

Model $\mathbb{P}_\varphi$

Deep Learning Approaches to Implicit Inference

Automated Summary Statistics Extraction

• Introduce a parametric function $f_\varphi$ to reduce the dimensionality of the data while preserving information.

Information-based loss functions

• 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 | f_\varphi(x)) ] \leq I(Y; \Theta) $$
  Jeffrey, Alsing, Lanusse (2021)

Unrolling the Probabilistic Learning Recipe

• I assume a forward model of the observations: \begin{equation} p( x ) = p(x | \theta) \ p(\theta) \nonumber \end{equation} All I ask is the ability to sample from the model, to obtain $\mathcal{D} = \{x_i, \theta_i \}_{i\in \mathbb{N}}$


• I am going to assume $q_\phi(\theta | x)$ a parametric conditional density


• Optimize the parameters $\phi$ of $q_{\phi}$ according to \begin{equation} \min\limits_{\phi} \sum\limits_{i} - \log q_{\phi}(\theta_i | x_i) \nonumber \end{equation} In the limit of large number of samples and sufficient flexibility \begin{equation} \boxed{q_{\phi^\ast}(\theta | x) \approx p(\theta | x)} \nonumber \end{equation}

Illustration on 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)$


• Infer full-field posterior on cosmology:
  • explicitly using an Hamiltonian-Monte Carlo (NUTS) sampler
  • implicitly using a learned summary statistics and conditional density estimation.

Example of application: Constraining Dark Matter Substructures

Brehmer, Mishra-Sharma, Hermans, Louppe, Cranmer (2019)

Example of application: Infering Microlensing Event Parameters

Zhang, Bloom, Gaudi, Lanusse, Lam, Lu (2021)

Example of application: Likelihood-Free parameter inference with DES SV

Jeffrey, Alsing, Lanusse (2021)

Suite of N-body + raytracing simulations: $\mathcal{D}$

A motivating example: Image Deconvolution

Hubble Space Telescope

some deep neural network

Simulated Ground-Based Telescope

Example: GalaxyGAN (Schawinski et al. 2017)

Linear inverse problems

Deconvolution

Inpainting

Denoising

A Bayesian view of the problem

• $p(y | x)$ is the data likelihood, which contains the physics
  


• $p(x)$ is our prior knowledge on the solution.

Classical examples of signal priors

Sparse
$$ \log p(x) = \parallel \mathbf{W} x \parallel_1 $$

Gaussian $$ \log p(x) = x^t \mathbf{\Sigma^{-1}} x $$

Total Variation $$ \log p(x) = \parallel \nabla x \parallel_1 $$

But what about this?

Why isn't it easy?

• The curse of dimensionality put all points far apart in high dimension

Distance between pairs of points drawn from a Gaussian distribution.

• Classical methods for estimating probability densities, i.e. Kernel Density Estimation (KDE) start to fail in high dimension because of all the gaps

The Variational Auto-Encoder

The famous Variational Auto-Encoder

Normalizing Flows

Still a latent variable model, except the mapping $f_\theta$ is made to be bijective.

Dinh et al. 2016

Normalizing Flows

• Assumes a bijective mapping between data space $x$ and latent space $z$ with prior $p(z)$: $$ z = f_{\theta} ( x ) \qquad \mbox{and} \qquad x = f^{-1}_{\theta}(z)$$
• Admits an explicit marginal likelihood: $$ \log p_\theta(x) = \log p(z) + \log \left| \frac{\partial f_\theta}{\partial x} \right|(x) $$

$\Longrightarrow$ The challenge is in designing mappings $f_\theta$ that are both: easy to invert, easy to compute the jacobian of.

One example of NF: RealNVP

Jacobian of an affine coupling layer

In a standard affine RealNVP, one Flow layer is defined as: $$ \begin{matrix} y_{1\ldots d} &=& x_{1 \ldots d} \\ y_{d+1\ldots N} &=& x_{d+1 \ldots N} ⊙ \sigma_\theta(x_{1 \ldots d}) + \mu_\theta(x_{1 \ldots d}) \end{matrix} $$ where $\sigma_\theta$ and $\mu_\theta$ are unconstrained neural networks.

We will call this layer an affine coupling.

A Physicist's approach to the deconvolution problem: let's build a model

Probabilistic model

$$ x \sim ? $$

$$ x \sim \mathcal{N}(z, \Sigma) \quad z \sim ? $$
latent $z$ is a denoised galaxy image

$$ x \sim \mathcal{N}( \mathbf{P} z, \Sigma) \quad z \sim ?$$
latent $z$ is a super-resolved and denoised galaxy image

$$ x \sim \mathcal{N}( \mathbf{P} (\Pi \ast z), \Sigma) \quad z \sim ? $$
latent $z$ is a deconvolved, super-resolved, and denoised galaxy image

$$ x \sim \mathcal{N}( \mathbf{P} (\Pi \ast g_\theta(z)), \Sigma) \quad z \sim \mathcal{N}(0, \mathbf{I}) $$
latent $z$ is a Gaussian sample
$\theta$ are parameters of the model

Bayesian Inference a.k.a. Uncertainty Quantification

The Bayesian view of the problem: $$ p(z | x ) \propto p_\theta(x | z, \Sigma, \mathbf{\Pi}) p(z)$$ where:

• $p( z | x )$ is the posterior
• $p( x | z )$ is the data likelihood, contains the physics
• $p( z )$ is the prior

Data
$x_n$

Truth
$x_0$

Posterior samples
$g_\theta(z)$

$\mathbf{P} (\Pi \ast g_\theta(z))$

Median

Data residuals
$x_n - \mathbf{P} (\Pi \ast g_\theta(z))$

Standard Deviation

Now you try it!

Raise your hand when you reach the Maximum Likelihood Solution i.e. end of Step III.

How to solve the posterior inference problem by Variational Inference?

• The core idea of Variational Inference is to assume a model distribution $q_\phi$ and we fit it to the unknown posterior.
• How do we fit a distribution to another distribution?
  We minize the KL divergence $$ \mathbb{D}_{KL}\left( q_\phi(z) \parallel p(z | x) \right) = \mathbb{E}_{z \sim q_{\phi}} \left[ \log \frac{q_\phi(z)}{p(z | x)} \right] $$

  $$ = \mathbb{E}_{z \sim q_{\phi}} \left[ \log q(z) \right] - \mathbb{E}_{z \sim q_{\phi}} \left[ \log p(z | x) \right] $$

  $$ = \mathbb{E}_{z \sim q_{\phi}} \left[ \log q(z) \right] - \mathbb{E}_{z \sim q_{\phi}} \left[ \log p(x | z) + \log p(z) - \log p(x)\right] $$

  $$ = \underbrace{\mathbb{D}_{KL}\left( q_\phi(z) \parallel p(z) \right) - \mathbb{E}_{z \sim q_{\phi}} \left[ \log p(x | z) \right]}_{\mbox{Evidence Lower Bound}} + \log p(x) $$
• Finally, we need a flexible enough model distribution, why not a Normalizing Flow :-) ?

Let's set the stage: Gravitational lensing

Galaxy shapes as estimators for gravitational shear

$$ e = \gamma + e_i \qquad \mbox{ with } \qquad e_i \sim \mathcal{N}(0, I)$$

• We are trying the measure the ellipticity $e$ of galaxies as an estimator for the gravitational shear $\gamma$

Gravitational Lensing as an Inverse Problem

Shear $\gamma$

Convergence $\kappa$

First realization: The score is all you need!

• Whether you are looking for the MAP or sampling with HMC or MALA, you only need access to the score of the posterior: $$\frac{\color{orange} d \color{orange}\log \color{orange}p\color{orange}(\color{orange}x \color{orange}|\color{orange} y\color{orange})}{\color{orange} d \color{orange}x}$$
  • Gradient descent: $x_{t+1} = x_t + \tau \nabla_x \log p(x_t | y) $
  • Langevin algorithm: $x_{t+1} = x_t + \tau \nabla_x \log p(x_t | y) + \sqrt{2\tau} n_t$

• The score of the full posterior is simply: $$\nabla_x \log p(x |y) = \underbrace{\nabla_x \log p(y |x)}_{\mbox{known explicitly}} \quad + \quad \underbrace{\nabla_x \log p(x)}_{\mbox{known implicitly}}$$ $\Longrightarrow$ "all" we have to do is model/learn the score of the prior.

Neural Score Estimation by Denoising Score Matching (Vincent 2011)

• Denoising Score Matching: An optimal Gaussian denoiser learns the score of a given distribution.
  • If $x \sim \mathbb{P}$ is corrupted by additional Gaussian noise $u \in \mathcal{N}(0, \sigma^2)$ to yield $$x^\prime = x + u$$
  • Let's consider a denoiser $r_\theta$ trained under an $\ell_2$ loss: $$\mathcal{L}=\parallel x - r_\theta(x^\prime, \sigma) \parallel_2^2$$
  • The optimal denoiser $r_{\theta^\star}$ verifies: $$\boxed{\boldsymbol{r}_{\theta^\star}(\boldsymbol{x}', \sigma) = \boldsymbol{x}' + \sigma^2 \nabla_{\boldsymbol{x}} \log p_{\sigma^2}(\boldsymbol{x}')}$$

$\boldsymbol{x}'$

$\boldsymbol{x}$

$\boldsymbol{x}'- \boldsymbol{r}^\star(\boldsymbol{x}', \sigma)$

$\boldsymbol{r}^\star(\boldsymbol{x}', \sigma)$

Example of one chain during annealing

Validating Posterior Sampling under a Gaussian prior

Some details: We don't actually know the marginal posterior score!

• We know the following quantities:
  • Annealed likelihood (analytically): $p_\sigma(y | x) = \mathcal{N}(y; \mathbf{A} x, \mathbf{\Sigma} + \sigma^2 \mathbf{I})$
  • Annealed prior score (by score matching): $\nabla_x \log p_\sigma(x)$
• But, unfortunately: $\boxed{p_\sigma(x|y) \neq p_\sigma(y|x) \ p_\sigma(x)}$ $\Longrightarrow$ We don't know the marginal posterior score!

• We cannot use the reverse SDE/ODE of diffusion models to sample from the posterior. $$\mathrm{d} x = [f(x, t) - g^2(t) \underbrace{\nabla_x \log p_t(x|y)}_{\mbox{unknown}} ] \mathrm{d}t + g(t) \mathrm{d} w$$

Proposed sampling strategy

• Even if not equivalent to the marginal posterior score, $\nabla_x \log p_{\sigma^2}(y | x) + \nabla_x \log p_{\sigma^2}(x)$ still has good properties:
  • Tends to an isotropic Gaussian distribution for large $\sigma$
  • Corresponds to the target posterior for $\sigma=0$


• If we simulate this SDE sufficiently slowly (i.e. timescale of change of $\sigma$ is much larger than the timescale of the SDE) we can expect to sample from the target posterior.

Illustration on $\kappa$-TNG simulations

Remy, Lanusse, et al. (2022)

True convergence map

Traditional Kaiser-Squires

Wiener Filter

Posterior Mean (ours)

Posterior samples

Reconstruction of the HST/ACS COSMOS field

• COSMOS shear data from Schrabback et al. 2010
• Prior learned from $\kappa$-TNG simulation from Osato et al. 2021.

Massey et al. (2007)

Remy et al. (2022) Posterior mean

Remy et al. (2022) Posterior samples