Bayesian deep learning for cosmology and gravitational waves¶

Tom Charnock

Institut d'Astrophysique de Paris

Slides available at presentations.charnock.fr/2020/03/15/Paris/PCCP/BayesianDeepLearning

What is a neural network?¶

What is Bayesian statistics?¶

Why there is a problem with traditionally trained neural networks¶

What is a neural network formulated as a Bayesian model?¶

How can Bayesian neural networks be implemented?¶

What is a neural network?¶

A neural network is just a parameterised function¶

$$\Large\mathbb{NN}(\boldsymbol{\omega},\boldsymbol{\alpha}) : {\bf x} \to {\bf y}$$

Some examples of possible ${\bf x}$ and ${\bf y}$¶

Mapping inputs to ouputs: Classification, regression, etc.¶

Learning latent representations¶

Predicting posterior distribution of targets¶

Bayesian statistics¶

Statistical models¶

Aleatoric uncertainty¶

Uncertainty due to the true underlying statistical model¶

Epistemic uncertainty¶

Uncertainty which arises due to the lack of knowledge about the true statistical model¶

Prior belief¶

Quantifying what we don't know...¶

Characterising the epistemic uncertainty
(although it is not the epistemic uncertainty)

Creating a statistical model¶

Joint probability of data and parameters¶

The likelihood and the posterior¶

Likelihood-free (simulation based) inference¶

Bayes theorem¶

Epistemic uncertainty in the values of the parameter can be reduced under the evidence of observed data¶

$$\Large\mathcal{P}(\theta|\mathcal{D},\psi)=\frac{\mathcal{L}(\mathcal{D}|\theta,\psi)\mathscr{p}(\theta|\psi)}{\mathscr{p}(\mathcal{D}|\psi)}$$

Neural networks formulated as Bayesian models¶

How likely is any predicted output?¶

$$\mathcal{P}(\mathscr{y}|{\bf x})=\int_\Omega\int_\mathcal{A}d\omega d\alpha\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha)\mathscr{p}(\omega|\alpha)\mathscr{p}(\alpha)$$

A neural network is deterministic¶

How likely is any predicted output?¶

$$\mathcal{P}(\mathscr{y}|{\bf x})=\int_\Omega\int_\mathcal{A}d\omega d\alpha\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha)\mathscr{p}(\omega|\alpha)\mathscr{p}(\alpha)$$

Why neural networks do not work...¶

Choose a neural network¶

$$\begin{align*} \mathcal{P}(\mathscr{y}|{\bf x})&=\int_\Omega\int_\mathcal{A}d\omega d\alpha\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha)\mathscr{p}(\omega|\alpha)\mathscr{p}(\alpha)\\ &=\int_\Omega\int_\mathcal{A}d\omega d\alpha\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha)\mathscr{p}(\omega|\alpha)\delta(\alpha-\alpha^*)\\ &=\int_\Omega d\omega\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha^*)\mathscr{p}(\omega|\alpha^*)\\ \end{align*}$$

Maximum likelihood estimates of the weights¶

$$\omega_\textrm{MLE}=\underset{\omega\in\Omega}{\textrm{argmax}}\,\mathcal{L}(\mathcal{D}|\omega,\alpha^*)$$

Deterministic neural network¶

$$\begin{align*} \mathcal{P}(\mathscr{y}|{\bf x})&=\int_\Omega d\omega\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha^*)\mathscr{p}(\omega|\alpha^*)\\ &=\int_\Omega d\omega\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha^*)\delta(\omega-\omega_\textrm{MLE})\\ &=\mathcal{L}(\mathscr{y}|{\bf x},\omega_\textrm{MLE},\alpha^*)\\ &=1 \end{align*}$$

No characterisation of the epistemic uncertainty¶

Neural density estimators¶

Mixture density networks¶

Normalising flows¶

But...¶

$$\begin{align*} \mathcal{P}(\mathcal{P}(\mathscr{y}|{\bf x})|{\bf x})&=\int_\Omega d\omega\,\mathcal{L}(\mathcal{P}(\mathscr{y}|{\bf x})|{\bf x},\omega,\alpha^*)\mathscr{p}(\omega|\alpha^*)\\ &=\int_\Omega d\omega\,\mathcal{L}(\mathcal{P}(\mathscr{y}|{\bf x})|{\bf x},\omega,\alpha^*)\delta(\omega-\omega_\textrm{MLE})\\ &=\mathcal{L}(\mathcal{P}(\mathscr{y}|{\bf x})|{\bf x},\omega_\textrm{MLE},\alpha^*)\\ &=1 \end{align*}$$

Bayesian neural networks¶

$$\mathcal{P}(\mathscr{y}|{\bf x})=\int_\Omega\int_\mathcal{A}d\omega d\alpha\,\mathcal{L}(\mathscr{y}|{\bf x},\omega,\alpha)\mathscr{p}(\omega|\alpha)\mathscr{p}(\alpha)$$

Networks not used in a Bayesian context, but working with distributions¶

Model $\mathscr{p}(\omega|\alpha)$ with some distribution¶

Maybe model $\mathscr{p}(\alpha)$ with some distribution (although probably not)¶

Networks used in a Bayesian context¶

Almost intractable posterior of network parameters¶

$$\mathcal{P}(\omega,\alpha|\mathcal{D})$$

Almost certainly intractable evidence¶

$${\mathscr{p}(\mathcal{D})=\int_\Omega\int_\mathcal{A}d\omega d\alpha\mathcal{L}(\mathcal{D}|\omega,\alpha)\mathscr{p}(\omega|\alpha)\mathscr{p}(\alpha)}$$

Practical implementations¶

Numerically approximate methods using the exact distribution
Markov chain Monte Carlo methods

Markov chains¶

Stationary distribution, $\pi(\omega)\approx\mathcal{P}(\omega|\mathcal{D},\alpha^*)$,
ensured by detailed balance

$$\pi(\omega^i)t(\omega^{i+1}|\omega^i)=\pi(\omega^{i+1})t(\omega^i|\omega^{i+1})$$

Dropping the evidence¶

$$\begin{align*} \frac{t(\omega^{i+1}|\omega^i)}{t(\omega^i|\omega^{i+1})}&=\frac{\mathcal{P}(\omega^{i+1}|\mathcal{D},\alpha^*)}{\mathcal{P}(\omega^i|\mathcal{D},\alpha^*)}\\\\ &=\frac{\mathcal{L}(\mathcal{D}|\omega^{i+1},\alpha^*)\mathscr{p}(\omega^{i+1}|\alpha^*)/\mathscr{p}(\mathcal{D}|\alpha^*)}{\mathcal{L}(\mathcal{D}|\omega^i,\alpha^*)\mathscr{p}(\omega^i|\alpha^*)/\mathscr{p}(\mathcal{D}|\alpha^*)}\\\\ &=\frac{\mathcal{L}(\mathcal{D}|\omega^{i+1},\alpha^*)\mathscr{p}(\omega^{i+1}|\alpha^*)}{\mathcal{L}(\mathcal{D}|\omega^i,\alpha^*)\mathscr{p}(\omega^i|\alpha^*)}\\\\ &=\frac{\varrho(\omega^{i+1})}{\varrho(\omega^i)} \end{align*}$$

Metropolis Rosenbluth-Hastings¶

Problems with Metropolis-Hastings¶

High rejection rate¶

Extremely expensive in many dimensions (random walk)¶

Hamiltonian Monte Carlo¶

Mass matrix describes correlations - these are not known a priori¶

Quasi-Newonian HMC¶

Use the approximate Hessian, ${\bf B}$, for additional geometric information¶

$$\begin{align*} \dot{\boldsymbol{\omega}} &={\bf B}{\bf M}^{-1}\boldsymbol{p} \\ \dot{\boldsymbol{p}} &={\bf B}\nabla\ln\varrho(\boldsymbol{\omega}) \end{align*}$$

Rescales the mass matrix so that steps are taken in a useful direction¶

Sampling the exact posterior of neural network parameters is possible...¶

But it is hyper-expensive and hyper-high-dimensional¶

There are only a small handfull of actual applications (bit.ly/NeuralBiasModel)
and most others are proof of concept

Exact treatment of approximate distributions
Variational inference

Approximate the posterior distribution of weights with analytic approximations¶

Analytic (or automatic) parameterised distributions¶

Normalising flows¶

Kullback–Leibler (KL) divergence¶

Amount of information lost by approximating $\mathcal{P}(\omega|\mathcal{D})$ with $\mathcal{Q}(\omega)$¶

$$\mathbb{KL}(\mathcal{P}||\mathcal{Q})=\int_\Omega d\omega\, \mathcal{P}(\omega|\mathcal{D})\ln\frac{\mathcal{P}(\omega|\mathcal{D})}{\mathcal{Q}(\omega)}$$

How variational inference works¶

For variational inference we normally consider

$$\mathbb{KL}(\mathcal{Q}||\mathcal{P})=\int_\Omega d\omega\,\mathcal{Q}(\omega)\ln\frac{\mathcal{Q}(\omega)}{\mathcal{P}(\omega|\mathcal{D})}$$
since the expectation value of $\mathcal{P}(\omega|\mathcal{D})$ is intractable.

Minimising the reverse KL divergence is equivalent to expectation-maximisation.

Evidence Lower BOund (ELBO)¶

How do we deal with the intractable posterior $\mathcal{P}(\omega|\mathcal{D})$?

Note, this is NOT Bayesian¶

Mean field approximation¶

Variational Bayes¶

Model the posterior of weights with some parameterised distribution¶

Update the parameters with by maximising the ELBO.¶

Which distributions?¶

Neural network parameters are not identifiable...¶

• Shallow local maxima

• Highly degenerate

• Super multimodal

We need something to model this!¶

MCDropout (the worst usable choice?)¶

Weird sampling of combinations of neurons, no parameters learned¶

DropConnect¶

At least its actually sampling weights, no parameters learned¶

Variational Dropout¶

Error added to weights, not weight posterior, no parameters learned¶

Reparameterisation trick¶

Actual characterised posterior distribution for weights, only Gaussian¶

Flowed reparameterisation¶

Actual characterised posterior distribution for weights¶

Variational Inference is tractable¶

It can run in a reasonable time¶

Variational Bayes isn't yet very Bayesian...¶

People get confused between Bayesian and neural networks using distributions¶

We should use the posterior distribution of weights and the form of the distribution¶

Instead we choose a (poor) choice of representation of the distribution and fit that¶

Epistemic uncertainty is not reduced (and may be increased or biased)¶

What do we actually want?¶

Likelihood-free inference¶

Use a (well) trained network to find summaries¶

Generate simulations drawn from prior and summarise¶

Make a kernel (or neural) density estimate (or use ABC) of the summary-parameter pairs¶

A slice through at the summarised observation is the approximate posterior

Pros and cons¶

Easy to use - can stop worrying about Bayesian neural networks¶

Guaranteed posterior (of summaries)¶

Density estimator can be precomputed and evaluated quickly¶

There is a choice of density estimator¶

Depends on the number of simulations (not different than other methods)¶

So... what have we learned?¶

Neural networks can be formulated as statistical models¶

Epistemic uncertainties can (in principle) be reduced via Bayes¶

Bayesian neural networks mean different things to different people¶

Actual (as in sampled) neural networks are REALLY expensive¶

Variational approximation to weight posterior exist, aren't Bayesian¶

Likelihood-free inference could provide a non-Bayesian neural network Bayesian method with neural networks¶