Probabilistic Programming for
Inverse Problems in Physical Sciences

Atılım Güneş Baydin, Lukas Heinrich, Wahid Bhimji,
Lei Shao, Saeid Naderiparizi, Andreas Munk,Jialin Liu, Bradley Gram-Hansen, Gilles Louppe,Lawrence Meadows, Philip Torr, Victor Lee, Prabhat,Kyle Cranmer, Frank Wood
Stanford SLAC AI Seminar
25 Sep 2020

Simulation and physical sciences
Computational models and simulation are key to scientific advance at all scales

Particle physics

Drug discovery

Nuclear physics

Weather

Material design

Climate scienceCosmology2Introducing a new way to use existing simulatorsProbabilistic
programming

Simulation

3

Simulators
Parameters

Outputs (data)

Simulator

4

Simulators
Parameters

Outputs (data)

Simulator

Prediction:● Simulate forward evolution of the system● Generate samples of output

5

Simulators
Parameters

Outputs (data)

Simulator

Prediction:● Simulate forward evolution of the system● Generate samples of output

6

Simulators
Parameters

Outputs (data)

Simulator

Prediction:● Simulate forward evolution of the system● Generate samples of output

WE NEED THE INVERSE!
7

Simulators
Parameters

Outputs (data)

Simulator

Prediction:● Simulate forward evolution of the system● Generate samples of output

Inference:● Find parameters that can produce (explain) observed data
● Inverse problem
8● Often a manual process

Simulators
Parameters

Inferredparameters

Gene network

Outputs (data)

Simulator

Observed data

Gene expression

9

Simulators
Parameters

Inferredparameters

Earthquake
location &characteristics

Outputs (data)

Simulator

Observed data

Seismometerreadings

10

Simulators
Parameters

Inferredparameters

Event analyses &new particle
discoveries

Outputs (data)

Simulator

Observed data

Particle detector
readings

11

Inverting a simulatorParameters

Inferredparameters

Outputs (data)

Simulator

Observed data

Probabilistic programming is a machine learning framework allowing us to● write programs that define probabilistic models● run automated Bayesian inference of parametersconditioned on observed outputs (data)
Pyro

EdwardStan
12

Inverting a simulator
Parameters

Inferredparameters

Outputs (data)

Simulator

Observed data

Probabilistic programming is a machine learning framework allowing us to● write programs that define probabilistic models● Has been
limited totoy and small-scaleproblems● run automatedBayesianinferenceof parameters● Normally requires one to implement a probabilistic
conditionedon observed outputs (data)
model from scratch in the chosen language/system
Pyro

EdwardStan
13

Inverting a simulator
Parameters

Inferredparameters

Outputs (data)

Simulator

Observed data

Key idea:
Many simulators are stochastic and they define probabilistic models
by sampling random numbers

14

Inverting a simulator
Parameters

Outputs (data)
Code

Inferredparameters

Simulator

Observed data

Key idea:
Many simulators are stochastic and they define probabilistic models
by sampling random numbers
Simulators are probabilistic programs!
15

Inverting a simulator
Parameters

Outputs (data)
Code

Inferredparameters

Simulator

Observed data

Key idea:
Many simulators are stochastic and they define probabilistic models
by sampling random numbers
Simulators are probabilistic programs!
We “just” need an infrastructure to execute them as such

16

A new probabilistic programming system forexisting simulators (in any language)
based on PyTorchProbabilistic
execution
Parameters

Outputs (data)
Code

Inferredparameters

Simulator

Observed data18

Probabilistic
execution
ParametersOutputs (data)
Code

Inferredparameters●●

Simulator

Observed data

Run forward & catch all random choices (“hijack” all calls to RNG)
Record an execution trace: a record of all parameters, random choices, outputs

19

Probabilistic
execution
Parameters

Outputs (data)
Code

Inferredparameters●●

Simulator

Observed data

Run forward & catch all random choices (“hijack” all calls to RNG)
Record an execution trace: a record of all parameters, random choices, outputs

Probabilistic Programming eXecution protocol
C++, C#, Dart, Go, Java, JavaScript, Lua, Python, Rust and others
20Probabilistic
execution
Parameters

Outputs (data)
CodeInferredparameters●●

Simulator

Observed data

Run forward & catch all random choices (“hijack” all calls to RNG)
Record an execution trace: a record of all parameters, random choices, outputs

Uniquely label each choice
at runtime by “addresses” ofstack frames

21

Probabilistic
execution
Parameters

Outputs (data)
Code

Inferredparameters

Simulator

Observed data●●

Run forward & catch all random choices (“hijack” all calls to RNG)
Record an execution trace: a record of all parameters, random choices, outputs●○ Conditioning: compare simulated output and observed data
Approximate the distribution of parameters that can produce (explain)observed data, using inference engines like Markov-chain Monte Carlo (MCMC)

22

Probabilistic
execution
Parameters

Outputs (data)
CodeInferredparameters●●●

Simulator

Observed data

Simulators = giant probability models so inference is hard andRun forward & catch all random choices (“hijack” all calls to RNG)
computationally costly
Record an execution trace: a record of all parameters, random choices, outputs● Need to run simulator up to millions of times● Simulator execution and MCMC inference are sequential● MCMC has “burn-in period” and autocorrelation○ Conditioning: compare simulated output and observed data
Approximate the distribution of parameters that can produce (explain)observed data, using inference engines like Markov-chain Monte Carlo (MCMC)

23

Probabilistic
execution
Parameters

Outputs (data)
CodeInferredparameters●●●Simulator

Observed data

Simulators = giant probability models so inference is hard andRun forward & catch all random choices (“hijack” all calls to RNG)
computationally costly
Record an execution trace: a record of all parameters, random choices, outputs● Need to run simulator up to millions of times● Simulator execution and MCMC inference are sequential● MCMC has “burn-in period” and autocorrelation○ Conditioning: compare simulated output and observed data
Approximate the distribution of parameters that can produce (explain)
But we can amortize the cost of inference using deep learningobserved data, using inference engines like Markov-chain Monte Carlo (MCMC)

24Distributed
data generation

Training data(execution traces)

Sim.

Trace

NNSim.

TraceNN⋮Sim.

Distributed
training⋮HPC

Trace

NN

Trained NN⋮
HPC

Training (recording simulator behavior)● Deep recurrent neural network learns all random choices in simulator● Dynamic NN: grows with simulator complexity○ Layers get created as we learn more of the simulator○ 100s of millions of parameters in particle physics simulation● Costly, but amortized: we need to train only once per given model

Distributed
data generation

Training data(execution traces)

Sim.

TraceNNSim.

Trace

NN

⋮Sim.

Trained NN

Observed data

Distributedtraining
Trained NN
⋮⋮HPC

Trace

NN

HPC

Distributed inference

Traces reproducing
observed data

NN

Sim.

Trace

NN

Sim.

Trace

NN

⋮⋮
Sim.

Inferredparameters
⋮

HPC

Trace

Inference (controlling simulator behavior)● Trained deep NN makes intelligent choices given data observation● Embarrassingly parallel distributed inference● No “burn in period”● No autocorrelation: every sample is independent

Trained NN

Observed data

Distributed inference

Traces reproducing
observed data

NN

Sim.

Trace

NN

Sim.

Trace

NN⋮⋮Sim.

Inferredparameters⋮HPCTrace

Inference (controlling simulator behavior)● Trained deep NN makes intelligent choices given data observation● Embarrassingly parallel distributed inference● No “burn in period”● No autocorrelation: every sample is independent

Trained NN

Observed data

Inferredparameters

Probabilistic programming with simulators
https://github.com/pyprob/pyprob●

Probabilistic programming system for simulators and HPC, based on PyTorchDistributed training and inference, eﬀicient support for multi-TB distribution files
Optimized memory usage, parallel trace processing and combination
https://github.com/pyprob/ppx●Probabilistic Programming eXecution protocol
Simulator and inference/NN executed in separate processes and machines across network
Using Google flatbuﬀers to support C++, C#, Dart, Go, Java, JavaScript, Lua, Python, Rust and others
Probabilistic programming analogue to Open Neural Network Exchange (ONNX) for deep learning
Pyprob_cpp, RNG front end for C++ simulators https://github.com/pyprob/pyprob_cpp
Containerized workflow
Develop locally, deploy to HPC on many nodes for experiments

etalumis → | ← simulate
Atılım GüneşBaydin

LukasHeinrich

Wahid
BhimjiLei
Shao

SaeidNaderiparizi

AndreasMunk

JialinLiu

Bradley
Gram-Hansen

GillesLouppeLawrence
Meadows

PhilTorr

VictorLee

Prabhat

KyleCranmer

FrankWood30

Baydin, Heinrich, Bhimji, Shao, Naderiparizi, Munk, Liu, Gram-Hansen, Louppe, Meadows, Torr, Lee, Prabhat, Cranmer, Wood.
NeurIPS 201931

Tau lepton decay
We study tau lepton decay using the state-of-the-art Sherpa simulator (C++)
Coupled to a fast approximate calorimeter simulation in C++

Baydin, Heinrich, Bhimji, Shao, Naderiparizi, Munk, Liu, Gram-Hansen, Louppe, Meadows, Torr, Lee, Prabhat, Cranmer, Wood.
NeurIPS 201932

Latent variables in Sherpa
We found Sherpa to contain at least 25k addresses (latent variables)
Note: the simulator defines an unlimited number of latents due to Turing-complete hostlanguage (C++) and presence of many sampling loops

...

33

Common trace types in Sherpa
Approximately 450 trace types encountered
Trace type: unique sequencing of addresses (with diﬀerent sampled values)

...

34

Inference results●●●Achieved MCMC (RMH) “ground truth”
First tractable Bayesian inference for LHC physics○ Posterior over full latent space (>25k latent variables)○ Autocorrelation typically around 105
Amortized inference (IC) closely matches MCMC (RMH)○ No autocorrelation, embarrassingly parallel○ MCMC: 115 hours, IC: 30 minutes

Gelman-Rubin convergence diagnostic

Autocorrelation

Trace log-probability

More physics events

Etalumis gives access to all latent variables: allows answering
any model-based question

Etalumis gives access to all latent variables: allows answering
any model-based question

Interpretability
Latent probabilistic structure of 10 most frequent trace types

39

Interpretability
Latent probabilistic structure of 10 most frequent trace types

40

Interpretability
Latent probabilistic structure of 10 most frequent trace types
px

py

pz

Rejection
sampling

Decay
channelCalorimeter

Rejection
sampling
41

Interpretability
Latent probabilistic structure of 25 most frequent trace types
px

py

pz

Rejection
sampling

Decay
channelCalorimeter

Rejection
sampling
42

Interpretability
Latent probabilistic structure of 100 most frequent trace types
px

pypz

Rejection
sampling

Decay
channelCalorimeter

Rejection
sampling
43

Interpretability
Latent probabilistic structure of 250 most frequent trace types
pxpypzRejection
sampling

Decay
channelCalorimeter

Rejection
sampling
44

What’s next?

Current and upcoming work
Autodiﬀ through PPX protocol
Learning simulator surrogates (approximate forward simulators)
Rejection sampling loops (weighting schemes)
Rare event simulation for compilation (“prior inflation”)
Batching of open-ended traces for NN training
Distributed training of dynamic networks○ Recently ran on 32k CPU cores on Cori (largest-scale PyTorch MPI)
● User features: posterior code highlighting, etc.
● Other simulators: astrophysics, epidemiology, computer vision
●
●●
●
●
●46

Probabilistic programming isfor the first time practicalfor large-scale real-world science models
This is just the beginning ...

Spacecraft collision prevention
Collaboration with ESA

Simulation of composite materials
Munk et al. 2019. arXiv:1910.11950

Simulation-based inference in health
Schroeder de Witt et al. 2020. arXiv:2005.07062
Gram-Hansen et al. 2019. arXiv:1905.12432

Neural Information Processing Systems (NeurIPS) workshopExpecting your papers at the intersection of machine learningand physical sciences!
Paper deadline: 2 Oct 2020; workshop: 11 Dec 2020
https://ml4physicalsciences.github.io/

Thank you for listening

Stanford SLAC AI Seminar
25 Sep 2020

Live demo
Jupyter notebook:https://github.com/gbaydin/mlhep2020/b
lob/master/notebooks/probprog-physicsexample.ipynb

50References

Atılım Güneş Baydin, Lukas Heinrich, Wahid Bhimji, Lei Shao, Saeid Naderiparizi, Andreas
Munk, Jialin Liu, Bradley Gram-Hansen, Gilles Louppe, Lawrence Meadows, Philip Torr,Victor Lee, Prabhat, Kyle Cranmer, Frank Wood. 2019. “Eﬀicient Probabilistic Inference in theQuest for Physics Beyond the Standard Model.” NeurIPS 2019Atılım Güneş Baydin, Lei Shao, Wahid Bhimji, Lukas Heinrich, Lawrence F. Meadows, JialinLiu, Andreas Munk, Saeid Naderiparizi, Bradley Gram-Hansen, Gilles Louppe, Mingfei Ma,Xiaohui Zhao, Philip Torr, Kyle Cranmer, Victor Lee, Prabhat, Frank Wood. 2019. “Etalumis:Bringing Probabilistic Programming to Scientific Simulators at Scale.” InternationalConference for High Performance Computing, Networking, Storage, and Analysis - SC1951

Extra slidesInference engines●●●

Markov chain Monte Carlo○ Probprog-specific:
priorproposal■ Lightweight
Metropolis–Hastings■ Random-walkposterior
Metropolis–Hastings○ Sequential○ Autocorrelation in samples○ “Burn in” period
Importance sampling
We sample in trace space:○ Propose from prioreach sample (trace) is one full execution of○ Use learned proposal
the model/simulator!
parameterized by observations○ No autocorrelation or burn in○ Each sample is independent (parallelizable)
53
Others: variational inference, Hamiltonian Monte Carlo, etc.

Inference engines●●●Markov chain Monte Carlo○ Probprog-specific:
priorproposal■ Lightweight
Metropolis–Hastings■ Random-walkposterior
Metropolis–Hastings○ Sequential○ Autocorrelation in samples○ “Burn in” period
Importance sampling
We sample in trace space:○ Propose from prioreach sample (trace) is one full execution of○ Use learned proposal
the model/simulator!
parameterized by observations○ No autocorrelation or burn in○ Each sample is independent (parallelizable)
54
Others: variational inference, Hamiltonian Monte Carlo, etc.

Inference engines●●●Markov chain Monte Carlo○ Probprog-specific:
priorproposal■ Lightweight
Metropolis–Hastings■ Random-walkposterior
Metropolis–Hastings○ Sequential○ Autocorrelation in samples○ “Burn in” period
Importance sampling
We sample in trace space:○ Propose from prioreach sample (trace) is one full execution of○ Use learned proposal
the model/simulator!
parameterized by observations○ No autocorrelation or burn in○ Each sample is independent (parallelizable)
55
Others: variational inference, Hamiltonian Monte Carlo, etc.

Inference engines●●●

Markov chain Monte Carlo○ Probprog-specific:
priorproposal■ LightweightMetropolis–Hastings■ Random-walkposterior
Metropolis–Hastings○ Sequential○ Autocorrelation in samples○ “Burn in” period
Importance sampling
We sample in trace space:○ Propose from prioreach sample (trace) is one full execution of○ Use learned proposal
the model/simulator!
parameterized by observations○ No autocorrelation or burn in○ Each sample is independent (parallelizable)
56
Others: variational inference, Hamiltonian Monte Carlo, etc.

Probabilistic programming languages (PPLs)● Anglican (Clojure)● Church (Scheme)● Edward, TensorFlow Probability (Python, TensorFlow)● Pyro (Python, PyTorch)● Figaro (Scala)● Infer.NET (C#)● LibBi (C++ template library)● PyMC3 (Python)● Stan (C++)● WebPPL (JavaScript)
For more, see http://probabilistic-programming.org

57

CalorimeterFor each particle in the final state coming from Sherpa:
1. Determine whether it interacts with the calorimeter at all(muons and neutrinos don't)
2. Calculate the total mean number and spatial distribution ofenergy depositions from the calorimeter shower(simulating combined eﬀect of secondary particles )3. Draw a number of actual depositions from the total mean
and then draw that number of energy depositions according
to the spatial distribution

Training objective and data for IC• Minimize• Using stochastic gradient descent with Adam
• Infinite stream of minibatches

sampled from the model
59

Gelman-Rubin and autocorrelation formulae

From Eric B. Ford (Penn State): Bayesian Computing for Astronomical Data Analysis
http://astrostatistics.psu.edu/RLectures/diagnosticsMCMC.pdf

60

Gelman-Rubin and autocorrelation formulae

From Eric B. Ford (Penn State): Bayesian Computing for Astronomical Data Analysis
http://astrostatistics.psu.edu/RLectures/diagnosticsMCMC.pdf

61