ML for Scientists: L13

 

Carol(ina) Cuesta-Lazaro

This week in AI4Science

https://rikiyat.github.io/blog/weight-space-programming/

Why should I care?

Decision making

Decision making in science

Is the current Standard Model ruled out by data?

Mass density

Vacuum Energy Density

0

CMB

Supernovae

\Omega_m
\Omega_\Lambda

Observation

Ground truth

Prediction

Uncertainty

Is it safe to drive there?

p(\mathrm{theory} | \mathrm{data})

Interpretable Simulators

\mathrm{data}
\mathrm{theory}

Uncertainties are everywhere

x^i_1
x^i_2

Noise in features

+ correlations

Noise in finite data realization

\{x^1, x^2,...,x^N \}
\theta
p(\theta|x)
\phi

Uncertain parameters

Limited model architecture

Imperfect optimization

Ensembling / Bayesian NNs

\theta

Forward Model

Observable

x
\color{darkgray}{\Omega_m}, \color{darkgreen}{w_0, w_a},\color{purple}{f_\mathrm{NL}}\, ...

Dark matter

Dark energy

Inflation

Predict

Infer

Parameters

Inverse mapping

\color{darkgray}{\sigma}, \color{darkgreen}{v}, ...

Fault line stress

Plate velocity

\color{white}{p(\theta|x)} \color{black}{=} \frac{\color{white}{p(x|\theta)} \color{white}{p(\theta)}}{\color{white}{p(x)}}

Likelihood

Posterior

Prior

Evidence

p(\theta|x)
p(x|\theta)
p(\theta)
p(x)

Markov Chain Monte Carlo MCMC

Hamiltonian Monte  Carlo HMC

Variational Inference VI

p(x) = \int p(x|\theta) p(\theta) d\theta

If can evaluate posterior (up to normalization), but not sample

Intractable

Unknown likelihoods

Amortized inference

Scaling high-dimensional

Marginalization nuisance

 

["Polychord: nested sampling for cosmology" Handley et al]

["Fluctuation without dissipation: Microcanonical Langevin Monte Carlo" 
Robnik and Seljak]

The price of sampling

Higher Effective Sample Size (ESS) = less correlated samples

Number of Simulator Calls

Known likelihood

Differentiable simulators

The simulator samples the likelihood

x_\mathrm{final}, y_\mathrm{final} \sim \mathrm{Simulator}(f, \theta) = p(x_\mathrm{final}, y_\mathrm{final} \mid f, \theta)
p(x_\mathrm{final}, y_\mathrm{final}|f, \theta) = \int dz p(x_\mathrm{final}, y_\mathrm{final},z|f, \theta)

z: All possible trajectories

Maximize the likelihood of the training samples

\hat \phi = \argmax \left[ \log p_\phi (x_\mathrm{train}|\theta_\mathrm{train}) \right]

Model

p_\phi(x)

Training Samples

x_\mathrm{train}

Neural Likelihood Estimation NLE

x_\mathrm{final}, y_\mathrm{final} \sim p(x_\mathrm{final}, y_\mathrm{final} \mid f, \theta) p(f, \theta)

NLE

No implicit prior

Not amortized

Goodness-of-fit

Scaling with dimensionality of x

p(x|\theta)

Implicit marginalization

\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(x_i \mid \theta_i )

Neural Posterior Estimation NPE

D_{\text{KL}}(p(\theta \mid x) \parallel q_\phi(\theta \mid x)) = \mathbb{E}_{\theta \sim p(\theta \mid x)} \log \frac{p(\theta \mid x)}{q_\phi(\theta \mid x)}

Loss   Approximate variational posterior, q, to true posterior, p

q_\phi(\theta \mid x)
p(\theta \mid x)
Image Credit: "Bayesian inference; How we are able to chase the Posterior" Ritchie Vink

KL Divergence

\mathcal{L}(\phi) = -\mathbb{E}_{\theta \sim p(\theta \mid x)} \left[ \log q_\phi(\theta \mid x) \right]
\mathcal{L}(\phi) = -\int p(\theta \mid x) \log q_\phi(\theta \mid x) \, d\theta
= -\int \frac{p(x \mid \theta) p(\theta)}{p(x)} \log q_\phi(\theta \mid x) \, d\theta.
\propto -\int p(x \mid \theta) p(\theta) \log q_\phi(\theta \mid x) \, d\theta.
D_{\text{KL}}(p(\theta \mid x) \parallel q_\phi(\theta \mid x)) = \mathbb{E}_{\theta \sim p(\theta \mid x)} \log \frac{p(\theta \mid x)}{q_\phi(\theta \mid x)}

Need samples from true posterior

Run simulator

p(x,\theta)
\phi^* = \argmin_{\phi}

Minimize KL

\mathbb{E}_{p(\theta, x)}
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid x_i)
\mathcal{L}(\phi) = -\mathbb{E}_{(\theta, x) \sim p(\theta, x)} \left[ \log q_\phi(\theta \mid x) \right]
(\theta_i, x_i) \sim p(\theta, x)

Amortized Inference!

Run simulator

Neural Posterior Estimation NPE

Neural Compression

x
s = F_\eta(x)

High-Dimensional

Low-Dimensional

p(\theta|x) = p(\theta|s)

s is sufficient iif

Neural Compression: MI

I(s(x), \theta)

Maximise

Mutual Information

I(\theta, s(x)) = D_{\text{KL}}(p(\theta, s(x)) \parallel p(\theta)p(s(x)))
\theta, s(x) \, \, \mathrm{independent} \rightarrow p(\theta, s(x)) = p(\theta)p(s(x))
s(x)
\theta
I(s(x), \theta)
\theta, s(x)
I(\theta, s(x)) = \mathbb{E}_{p(\theta, s(x))} \left[ \log \frac{p(\theta, s(x))}{p(\theta)p(s(x))} \right]
= \mathbb{E}_{p(\theta, s(x))} \left[ \log \frac{p(\theta \mid s(x))p(s(x))}{p(\theta) p(s(x))} \right]
p(\theta, s(x)) = p(\theta \mid s(x)) p(s(x))

Need true posterior!

I(\theta, s(x)) \approx \mathbb{E}_{p(\theta, s(x))} \left[ \log \frac{q_\phi(\theta \mid s(x))}{p(\theta)} \right]
\mathcal{L}(\phi, \eta) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid s_\eta(x_i))

NLE

No implicit prior

Not amortized

Goodness-of-fit

Scaling with dimensionality of x

p(x|\theta)

NPE

p(\theta|x)

Amortized

Scales well to high dimensional x

Goodness-of-fit? 

Robustness?

Fixed prior

Implicit marginalization

\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid x_i)
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(x_i \mid \theta_i )

Implicit marginalization

Do we actually need Density Estimation?

Just use binary classifiers!

x, \theta \sim p(x,\theta)
x, \theta \sim p(x)p(\theta)
y = 1
y = 0
\theta
x

Binary cross-entropy

Sample from simulator

Mix-up

Likelihood-to-evidence ratio

r(x,\theta) = \frac{p(x\mid \theta)}{p(x)}

Likelihood-to-evidence ratio

p(x,\theta|y) = p(x,\theta) \, \, \mathrm{if} \, \, y=1
p(x,\theta|y) = p(x)p(\theta) \, \, \mathrm{if} \, \, y=0
p(y=1 \mid x,\theta) = \frac{p(x,\theta|y=1)p(y)}{p(x,\theta)}
\left(p(x,\theta \mid y=0) + p(x,\theta \mid y=1) \right)p(y)
p(y=1 \mid x,\theta) = \frac{p(x,\theta)}{p(x)p(\theta) + p(x,\theta)}
= \frac{p(x,\theta)}{p(x)p(\theta)}
= \frac{r(x,\theta)}{r(x,\theta) + 1}

NLE

No implicit prior

Not amortized

Goodness-of-fit

Scaling with dimensionality of x

p(x|\theta)

NPE

p(\theta|x)

NRE

\frac{p(x|\theta)}{p(x)}

Amortized

Scales well to high dimensional x

Implicit marginalization

\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid x_i)
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(x_i \mid \theta_i )

No need variational distribution

No implicit prior

Implicit marginalization

Approximately normalised

\begin{split} L &= - \frac{1}{N} \sum_{i=1}^N \left[ y_i \log(p_i) \right. \\ &\left. + (1 - y_i) \log(1 - p_i) \right] \end{split}

Not amortized

Implicit marginalization

Goodness-of-fit? 

Robustness?

Fixed prior

smsharma/awesome-neural-sbi

How good is your posterior?

Test log likelihood

["Benchmarking simulation-based inference" 
Lueckmann et al

arXiv:2101.04653]

\mathbb{E}_{p(x,\theta)} \log \hat{p}(\theta \mid x)

Posterior predictive checks

Observed

Re-simulated posterior samples

Classifier 2 Sample Test (C2ST)

\mathrm{R} \sim p(\theta|x)
\mathrm{F} \sim \hat{p}(\theta|x)

Real or Fake?

Benchmarking SBI

["Benchmarking simulation-based inference" 
Lueckmann et al

arXiv:2101.04653]

Classifier 2 Sample Test (C2ST)

["A Trust Crisis In Simulation-Based Inference? Your Posterior Approximations Can Be Unfaithful" 
Hermans et al

arXiv:2110.06581]

Much better than overconfident!

Coverage: assessing uncertainties

\int_\Theta \hat{p}(\theta \mid x = x_0) d\theta = 1 - \alpha
["A Trust Crisis In Simulation-Based Inference? Your Posterior Approximations Can Be Unfaithful" 
Hermans et al

arXiv:2110.06581]

Credible region (CR)

Not unique

High Posterior Density region (HPD)

Smallest "volume"

True value in CR with

1 - \alpha

probability

\theta^*
\mathcal{H}

Empirical Coverage Probability (ECP)

\mathrm{ECP} = \mathbb{E}_{p(x,\theta)} \left[ \mathbb{1} \left[ \theta \in \mathcal{H}_{\hat{p}(\theta|x)}(1-\alpha)\right] \right]
["Investigating the Impact of Model Misspecification in Neural Simulation-based Inference" 
Cannon et al arXiv:2209.01845 ]

Underconfident

Overconfident

Calibrated doesn't mean informative!

Always look at information gain too

\mathrm{ECP} = \mathbb{E}_{p(x,\theta)} \left[ \mathbb{1} \left[ \theta \in \mathcal{H}(\hat{p}, \alpha) \right] \right]
= \mathbb{E}_{p(\theta)} \left[ \mathbb{1} \left[ \theta \in \mathcal{H}(\hat{p}, \alpha) \right] \right]
\hat{p}(\theta \mid x) = p(\theta)
= \int_{\mathcal{H}(\hat{p}),\alpha} d\theta p(\theta)
\mathbb{E}_{p(x,\theta)} \left[ \log \hat{p}(\theta \mid x) - \log p(\theta) \right]
= 1 - \alpha
["A Trust Crisis In Simulation-Based Inference? Your Posterior Approximations Can Be Unfaithful" 
Hermans et al

arXiv:2110.06581]

["Calibrating Neural Simulation-Based Inference with Differentiable Coverage Probability" 
Falkiewicz et al

arXiv:2310.13402]

["A Trust Crisis In Simulation-Based Inference? Your Posterior Approximations Can Be Unfaithful" 
Hermans et al

arXiv:2110.06581]

Model mispecification

["Investigating the Impact of Model Misspecification in Neural Simulation-based Inference" 
Cannon et al arXiv:2209.01845]

More misspecified

Sequential SBI

["Benchmarking simulation-based inference" 
Lueckmann et al

arXiv:2101.04653]

[Image credit: https://www.mackelab.org/delfi/]

Sequential SBI

Real life scaling: Gravitational lensing

["A Strong Gravitational Lens Is Worth a Thousand Dark Matter Halos: Inference on Small-Scale Structure Using Sequential Methods" 
Wagner-Carena et al arXiv:2404.14487]

Money plot - Nov 18th -> Discuss in class Nov 19th

Transformer from scratch - Nov 4th

Nov 5th discuss in class