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
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
ML For Scientists: L13
By carol cuesta
