Simulation-based inference
Boomers Quantified Uncertainty. We Simulate It
[Video Credit: N-body simulation Francisco Villaescusa-Navarro]
IAIFI Fellow
Carolina Cuesta-Lazaro
DESC
Why should I care?
Decision making
Decision making in science
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Is the current Standard Model ruled out by data?
Mass density
Vacuum Energy Density
0
0
CMB
Supernovae
Ωm
\Omega_m
ΩΛ
\Omega_\Lambda
Observation
Ground truth
Prediction
Uncertainty
Is it safe to drive there?
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Better data needs better models
p(theory∣data)
p(\mathrm{theory} | \mathrm{data})
Interpretable Simulators
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
data
\mathrm{data}
theory
\mathrm{theory}
Uncertainties are everywhere
x1i
x^i_1
x2i
x^i_2
Noise in features
+ correlations
Noise in finite data realization
{x1,x2,...,xN}
\{x^1, x^2,...,x^N \}
θ
\theta
p(θ∣x)
p(\theta|x)
ϕ
\phi
Uncertain parameters
Limited model architecture
Imperfect optimization
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Ensembling / Bayesian NNs
θ
\theta
Forward Model
Observable
x
x
Ωm,w0,wa,fNL...
\color{darkgray}{\Omega_m}, \color{darkgreen}{w_0, w_a},\color{purple}{f_\mathrm{NL}}\, ...
Dark matter
Dark energy
Inflation
Predict
Infer
Parameters
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Inverse mapping
σ,v,...
\color{darkgray}{\sigma}, \color{darkgreen}{v}, ...
Fault line stress
Plate velocity
p(θ∣x)=p(x)p(x∣θ)p(θ)
\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(θ∣x)
p(\theta|x)
p(x∣θ)
p(x|\theta)
p(θ)
p(\theta)
p(x)
p(x)
Markov Chain Monte Carlo MCMC
Hamiltonian Monte Carlo HMC
Variational Inference VI
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
p(x)=∫p(x∣θ)p(θ)dθ
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
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
["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
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
xfinal,yfinal∼Simulator(f,θ)=p(xfinal,yfinal∣f,θ)
x_\mathrm{final}, y_\mathrm{final} \sim \mathrm{Simulator}(f, \theta) = p(x_\mathrm{final}, y_\mathrm{final} \mid f, \theta)
p(xfinal,yfinal∣f,θ)=∫dzp(xfinal,yfinal,z∣f,θ)
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
ϕ^=argmax[logpϕ(xtrain∣θtrain)]
\hat \phi = \argmax \left[ \log p_\phi (x_\mathrm{train}|\theta_\mathrm{train}) \right]
Model
pϕ(x)
p_\phi(x)
Training Samples
xtrain
x_\mathrm{train}
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Neural Likelihood Estimation NLE
xfinal,yfinal∼p(xfinal,yfinal∣f,θ)p(f,θ)
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∣θ)
p(x|\theta)
Implicit marginalization
L(ϕ)≈−N1i=1∑Nlogqϕ(xi∣θi)
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(x_i \mid \theta_i )
Neural Posterior Estimation NPE
DKL(p(θ∣x)∥qϕ(θ∣x))=Eθ∼p(θ∣x)logqϕ(θ∣x)p(θ∣x)
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ϕ(θ∣x)
q_\phi(\theta \mid x)
p(θ∣x)
p(\theta \mid x)
Image Credit: "Bayesian inference; How we are able to chase the Posterior" Ritchie Vink
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
KL Divergence
L(ϕ)=−Eθ∼p(θ∣x)[logqϕ(θ∣x)]
\mathcal{L}(\phi) = -\mathbb{E}_{\theta \sim p(\theta \mid x)} \left[ \log q_\phi(\theta \mid x) \right]
L(ϕ)=−∫p(θ∣x)logqϕ(θ∣x)dθ
\mathcal{L}(\phi) = -\int p(\theta \mid x) \log q_\phi(\theta \mid x) \, d\theta
=−∫p(x)p(x∣θ)p(θ)logqϕ(θ∣x)dθ.
= -\int \frac{p(x \mid \theta) p(\theta)}{p(x)} \log q_\phi(\theta \mid x) \, d\theta.
∝−∫p(x∣θ)p(θ)logqϕ(θ∣x)dθ.
\propto -\int p(x \mid \theta) p(\theta) \log q_\phi(\theta \mid x) \, d\theta.
DKL(p(θ∣x)∥qϕ(θ∣x))=Eθ∼p(θ∣x)logqϕ(θ∣x)p(θ∣x)
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
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
p(x,θ)
p(x,\theta)
ϕ∗=ϕargmin
\phi^* = \argmin_{\phi}
Minimize KL
Ep(θ,x)
\mathbb{E}_{p(\theta, x)}
L(ϕ)≈−N1i=1∑Nlogqϕ(θi∣xi)
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid x_i)
L(ϕ)=−E(θ,x)∼p(θ,x)[logqϕ(θ∣x)]
\mathcal{L}(\phi) = -\mathbb{E}_{(\theta, x) \sim p(\theta, x)} \left[ \log q_\phi(\theta \mid x) \right]
(θi,xi)∼p(θ,x)
(\theta_i, x_i) \sim p(\theta, x)
Amortized Inference!
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Run simulator
Neural Posterior Estimation NPE
Neural Compression
x
x
s=Fη(x)
s = F_\eta(x)
High-Dimensional
Low-Dimensional
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
p(θ∣x)=p(θ∣s)
p(\theta|x) = p(\theta|s)
s is sufficient iif
Neural Compression: MI
I(s(x),θ)
I(s(x), \theta)
Maximise
Mutual Information
I(θ,s(x))=DKL(p(θ,s(x))∥p(θ)p(s(x)))
I(\theta, s(x)) = D_{\text{KL}}(p(\theta, s(x)) \parallel p(\theta)p(s(x)))
θ,s(x)independent→p(θ,s(x))=p(θ)p(s(x))
\theta, s(x) \, \, \mathrm{independent} \rightarrow p(\theta, s(x)) = p(\theta)p(s(x))
s(x)
s(x)
θ
\theta
I(s(x),θ)
I(s(x), \theta)
θ,s(x)
\theta, s(x)
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
I(θ,s(x))=Ep(θ,s(x))[logp(θ)p(s(x))p(θ,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]
=Ep(θ,s(x))[logp(θ)p(s(x))p(θ∣s(x))p(s(x))]
= \mathbb{E}_{p(\theta, s(x))} \left[ \log \frac{p(\theta \mid s(x))p(s(x))}{p(\theta) p(s(x))} \right]
p(θ,s(x))=p(θ∣s(x))p(s(x))
p(\theta, s(x)) = p(\theta \mid s(x)) p(s(x))
Need true posterior!
I(θ,s(x))≈Ep(θ,s(x))[logp(θ)qϕ(θ∣s(x))]
I(\theta, s(x)) \approx \mathbb{E}_{p(\theta, s(x))} \left[ \log \frac{q_\phi(\theta \mid s(x))}{p(\theta)} \right]
L(ϕ,η)≈−N1i=1∑Nlogqϕ(θi∣sη(xi))
\mathcal{L}(\phi, \eta) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid s_\eta(x_i))
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
NLE
No implicit prior
Not amortized
Goodness-of-fit
Scaling with dimensionality of x
p(x∣θ)
p(x|\theta)
NPE
p(θ∣x)
p(\theta|x)
Amortized
Scales well to high dimensional x
Goodness-of-fit
Fixed prior
Implicit marginalization
L(ϕ)≈−N1i=1∑Nlogqϕ(θi∣xi)
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid x_i)
L(ϕ)≈−N1i=1∑Nlogqϕ(xi∣θ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,θ∼p(x,θ)
x, \theta \sim p(x,\theta)
x,θ∼p(x)p(θ)
x, \theta \sim p(x)p(\theta)
y=1
y = 1
y=0
y = 0
θ
\theta
x
x
Binary cross-entropy
Sample from simulator
Mix-up
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Likelihood-to-evidence ratio
r(x,θ)=p(x)p(x∣θ)
r(x,\theta) = \frac{p(x\mid \theta)}{p(x)}
Likelihood-to-evidence ratio
p(x,θ∣y)=p(x,θ)ify=1
p(x,\theta|y) = p(x,\theta) \, \, \mathrm{if} \, \, y=1
p(x,θ∣y)=p(x)p(θ)ify=0
p(x,\theta|y) = p(x)p(\theta) \, \, \mathrm{if} \, \, y=0
p(y=1∣x,θ)=p(x,θ)p(x,θ∣y=1)p(y)
p(y=1 \mid x,\theta) = \frac{p(x,\theta|y=1)p(y)}{p(x,\theta)}
(p(x,θ∣y=0)+p(x,θ∣y=1))p(y)
\left(p(x,\theta \mid y=0) + p(x,\theta \mid y=1) \right)p(y)
p(y=1∣x,θ)=p(x)p(θ)+p(x,θ)p(x,θ)
p(y=1 \mid x,\theta) = \frac{p(x,\theta)}{p(x)p(\theta) + p(x,\theta)}
=p(x)p(θ)p(x,θ)
= \frac{p(x,\theta)}{p(x)p(\theta)}
=r(x,θ)+1r(x,θ)
= \frac{r(x,\theta)}{r(x,\theta) + 1}
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
NLE
No implicit prior
Not amortized
Goodness-of-fit
Scaling with dimensionality of x
p(x∣θ)
p(x|\theta)
NPE
p(θ∣x)
p(\theta|x)
NRE
p(x)p(x∣θ)
\frac{p(x|\theta)}{p(x)}
Amortized
Scales well to high dimensional x
Goodness-of-fit
Fixed prior
Implicit marginalization
L(ϕ)≈−N1i=1∑Nlogqϕ(θi∣xi)
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid x_i)
L(ϕ)≈−N1i=1∑Nlogqϕ(xi∣θ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
L=−N1i=1∑N[yilog(pi)+(1−yi)log(1−pi)]
\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
Density Estimation 101
Maximize the likelihood of the training samples
ϕ^=argmax[logpϕ(xtrain)]
\hat \phi = \argmax \left[ \log p_\phi (x_\mathrm{train}) \right]
x1
x_1
x2
x_2
Model
pϕ(x)
p_\phi(x)
Training Samples
xtrain
x_\mathrm{train}
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
x1
x_1
x2
x_2
Trained Model
pϕ(x)
p_\phi(x)
Evaluate probabilities
Low Probability
High Probability
Generate Novel Samples
Simulator
Simulator
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Normalizing flows
[Image Credit: "Understanding Deep Learning" Simon J.D. Prince]
z∼p(z)
z \sim p(z)
x∼p(x)
x \sim p(x)
x=f(z)
x = f(z)
p(x)=p(z=f−1(x))detJf−1(x)
p(x) = p(z = f^{-1}(x)) \left| \det J_{f^{-1}}(x) \right|
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Bijective
Sample
Evaluate probabilities
Probability mass conserved locally
z0∼p(z)
z_0 \sim p(z)
zk=fk(zk−1)
z_k = f_k(z_{k-1})
logp(x)=logp(z0)−k=1∑Klog∣detJfk(zk−1)∣
\log p(x) = \log p(z_0) - \sum_{k=1}^{K} \log \left| \det J_{f_k} (z_{k-1}) \right|
Image Credit: "Understanding Deep Learning" Simon J.D. Prince
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Masked Autoregressive Flows
p(x)=i∏p(xi∣x1:i−1)
p(x) = \prod_i{p(x_i \,|\, x_{1:i-1})}
p(xi∣x1:i−1)=N(xi∣μi,(expαi)2)
p(x_i \,|\, x_{1:i-1}) = \mathcal{N}(x_i \,|\,\mu_i, (\exp\alpha_i)^2)
μi,αi=fϕi(x1:i−1)
\mu_i, \alpha_i = f_{\phi_i}(x_{1:i-1})
Neural Network
xi=ziexp(αi)+μi
x_i = z_i \exp(\alpha_i) + \mu_i
zi=(xi−μi)exp(−αi)
z_i = (x_i - \mu_i) \exp(-\alpha_i)
Sample
Evaluate probabilities
NLE
No implicit prior
Not amortized
Goodness-of-fit
Scaling with dimensionality of x
p(x∣θ)
p(x|\theta)
NPE
p(θ∣x)
p(\theta|x)
NRE
p(x)p(x∣θ)
\frac{p(x|\theta)}{p(x)}
Amortized
Scales well to high dimensional x
Goodness-of-fit
Fixed prior
Implicit marginalization
L(ϕ)≈−N1i=1∑Nlogqϕ(θi∣xi)
\mathcal{L}(\phi) \approx -\frac{1}{N} \sum_{i=1}^N \log q_\phi(\theta_i \mid x_i)
L(ϕ)≈−N1i=1∑Nlogqϕ(xi∣θ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
L=−N1i=1∑N[yilog(pi)+(1−yi)log(1−pi)]
\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
How good is your posterior?
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Test log likelihood
["Benchmarking simulation-based inference"
Lueckmann et al
arXiv:2101.04653]
Ep(x,θ)logp^(θ∣x)
\mathbb{E}_{p(x,\theta)} \log \hat{p}(\theta \mid x)
Posterior predictive checks
Observed
Re-simulated posterior samples
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Classifier 2 Sample Test (C2ST)
R∼p(θ∣x)
\mathrm{R} \sim p(\theta|x)
F∼p^(θ∣x)
\mathrm{F} \sim \hat{p}(\theta|x)
Real or Fake?
Benchmarking SBI
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
["Benchmarking simulation-based inference"
Lueckmann et al
arXiv:2101.04653]
Classifier 2 Sample Test (C2ST)
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
["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
∫Θp^(θ∣x=x0)dθ=1−α
\int_\Theta \hat{p}(\theta \mid x = x_0) d\theta = 1 - \alpha
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
["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−α
1 - \alpha
probability
θ∗
\theta^*
H
\mathcal{H}
Empirical Coverage Probability (ECP)
ECP=Ep(x,θ)[1[θ∈Hp^(θ∣x)(1−α)]]
\mathrm{ECP} = \mathbb{E}_{p(x,\theta)} \left[ \mathbb{1} \left[ \theta \in \mathcal{H}_{\hat{p}(\theta|x)}(1-\alpha)\right] \right]
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
["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
ECP=Ep(x,θ)[1[θ∈H(p^,α)]]
\mathrm{ECP} = \mathbb{E}_{p(x,\theta)} \left[ \mathbb{1} \left[ \theta \in \mathcal{H}(\hat{p}, \alpha) \right] \right]
=Ep(θ)[1[θ∈H(p^,α)]]
= \mathbb{E}_{p(\theta)} \left[ \mathbb{1} \left[ \theta \in \mathcal{H}(\hat{p}, \alpha) \right] \right]
p^(θ∣x)=p(θ)
\hat{p}(\theta \mid x) = p(\theta)
=∫H(p^),αdθp(θ)
= \int_{\mathcal{H}(\hat{p}),\alpha} d\theta p(\theta)
Ep(x,θ)[logp^(θ∣x)−logp(θ)]
\mathbb{E}_{p(x,\theta)} \left[ \log \hat{p}(\theta \mid x) - \log p(\theta) \right]
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
=1−α
= 1 - \alpha
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
["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]
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
Model mispecification
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
["Investigating the Impact of Model Misspecification in Neural Simulation-based Inference" Cannon et al arXiv:2209.01845]
More misspecified
SBI Resources
Carolina Cuesta-Lazaro IAIFI/MIT - Simulation-Based Inference
"The frontier of simulation-based inference" Kyle Cranmer, Johann Brehmer, and Gilles Louppe
Github repos
Review
cuestalz@mit.edu
Book
"Probabilistic Machine Learning: Advanced Topics" Kevin P. Murphey
Simulation-based inference Boomers Quantified Uncertainty. We Simulate It [Video Credit: N-body simulation Francisco Villaescusa-Navarro] IAIFI Fellow Carolina Cuesta-Lazaro DESC
DESC-SBI
By carol cuesta
DESC-SBI
- 124
