Package initialization

Probabilistic Modelling with Turing.jl

David Widmann (@devmotion )

Uppsala University, Sweden

Julia User Group Munich, 5 July 2021

About me

• 👨‍🎓 PhD student at the Department of Information Technology and the Centre for Interdisciplinary Mathematics (CIM) in Uppsala

  • BSc and MSc in Mathematics (TU Munich)

  • Human medicine (LMU and TU Munich)

• 👨‍🔬 Research topic: "Uncertainty-aware deep learning"

  • I.e., statistics, probability theory, machine learning, and computer science

• 💻 Julia programming, e.g.,

  • SciML, in particular DelayDiffEq.jl

  • Turing ecosystem

SIR model

Classical mathematical model of infectious diseases (e.g., influenza or COVID-19):

$$\begin{array}{rl}\frac{\mathrm{d}S}{\mathrm{d}t}(t)& =-\beta S(t)I(t)/N\\ \frac{\mathrm{d}I}{\mathrm{d}t}(t)& =\beta S(t)I(t)/N-\gamma I(t)\\ \frac{\mathrm{d}R}{\mathrm{d}t}(t)& =\gamma I(t)\end{array}$$

with infection rate $\beta >0$, recovery/death rate $\gamma >0$, and constant total population $N=S(t)+I(t)+R(t)=\mathrm{c}\mathrm{o}\mathrm{n}\mathrm{s}\mathrm{t}$.

$S(0)$:

990

$I(0)$:

10

$R(0)$:

0

$\beta$:

0.5

$\gamma$:

0.25

Reproduction number: $R_0=\beta /\gamma \approx$ 2.0

Inference problem

Model can be used to make predictions or reason about possible interventions

But...

We do not know parameters such as

• transmission rate $\beta$

• recovery rate $\gamma$

• initial number of infected individuals $I(0)$

We might just know

• population size $N$

• initial number of removed individuals $R(0)=0$

• noisy (i.e., slightly incorrect) number of newly infected individuals for some days

Probabilistic modelling

Idea: model uncertainty of unknown parameters with probability distributions

Bayesian approach

Model uncertainty with conditional probability

$$p_{\mathrm{\Theta }|Y}(\theta |y)$$

of unknown parameters $\theta$ given observations $y$.

E.g.,

• $\theta$: unknown parameters $\beta$, $\gamma$, and $I(0)$ that we want to infer

• $y$: number of newly infected individuals

Bayes' theorem

Conditional probability $p_{\mathrm{\Theta }|Y}(\theta |y)$ can be calculated as

$$p_{\mathrm{\Theta }|Y}(\theta |y)=\frac{p_{Y|\mathrm{\Theta }}(y|\theta )p_{\mathrm{\Theta }}(\theta )}{p_Y(y)}$$

Workflow

• Choose model $p_{Y|\mathrm{\Theta }}(y|\theta )$

  • e.g., describes how daily cases of newly infected individuals depend on parameters

• Choose prior $p_{\mathrm{\Theta }}(\theta )$

  • should incorporate initial beliefs and knowledge about $\theta$

  • e.g., $\beta$ and $\gamma$ are positive and $I(0)$ is a natural number

• Compute

  $$p_{\mathrm{\Theta }|Y}(\theta |y)=\frac{p_{Y|\mathrm{\Theta }}(y|\theta )p_{\mathrm{\Theta }}(\theta )}{p_Y(y)}=\frac{p_{Y|\mathrm{\Theta }}(y|\theta )p_{\mathrm{\Theta }}(\theta )}{\int p_{Y|\mathrm{\Theta }}(y|{\theta }^{\prime })p_{\mathrm{\Theta }}({\theta }^{\prime })\mathrm{d}{\theta }^{\prime }}$$

Example: Coin flip

Coin with unknown probability of heads

• $k$-times head in $n$ coinflips (observation)

• unknown probability $p$ of heads (parameter)

Likelihood

$$p_{Y|\mathrm{\Theta }}(k|p)=\mathrm{Binom}(k;n,p)$$

$n$:

10

$p$:

0.5

Prior

We choose a Beta distribution as prior:

$$p_{\mathrm{\Theta }}(p)=\mathrm{Beta}(p;\alpha ,\beta )$$

$\alpha$:

1

$\beta$:

1

Posterior

For these choices of likelihood and prior, Bayes' theorem yields

$$p_{\mathrm{\Theta }|Y}(p|k)=\mathrm{B}\mathrm{e}\mathrm{t}\mathrm{a}(\alpha +k,\beta +n-k)$$

true $p$:

0.5

$n$:

50

$\alpha$:

1

$\beta$:

1

Discrete approximation

Idea: Approximate $p_{\mathrm{\Theta }|Y}(\theta |y)$ with a weighted mixture of point measures

$$p_{\mathrm{\Theta }|Y}(\cdot |y)\approx \sum _i{w}_i{\delta }_{{\theta }_i}(\cdot )$$

where $w_i>0$ and $\sum _i{w}_i=1$.

This implies that

$$\mathbb{E}(\mathrm{\Theta }|Y=y)\approx \sum _i{w}_i{\theta }_i$$

and more generally

$$\mathbb{E}(\varphi (\mathrm{\Theta })|Y=y)\approx \sum _i{w}_i\varphi ({\theta }_i)$$

number of samples:

500

The approximation can be constructed with e.g.

• importance sampling

• sequential Monte Carlo (SMC)

• Markov Chain Monte Carlo (MCMC)

Standalone samplers

One main design of the Turing ecosystem is to have many smaller packages that

• can be used independently of Turing

• are supported by Turing

For MCMC, currently we have e.g.

• AbstractMCMC.jl

• AdvancedHMC.jl

• AdvancedMH.jl

• AdvancedPS.jl

• EllipticalSliceSamplig.jl

• MCMCChains.jl

Additionally, Turing supports also DynamicHMC.jl.

AdvancedMH.jl

AdvancedMH.jl contains an implementation of the Metropolis-Hastings algorithm.

We have to define the unnormalized log density, i.e., the sum of the log-likelihood and the log prior, as a function of the unknown parameter.

The model is fully specified by the log density function.

We use a Metropolis-Hastings algorithm with a random walk proposal.

AbstractMCMC.jl

AbstractMCMC.jl defines an interface for MCMC algorithms.

The interface is supposed to make it easy to implement MCMC algorithms. The default definitions provide users with

• progress bars,

• support for user-provided callbacks,

• support for thinning and discarding initial samples,

• support for sampling with a custom stopping criterion,

• support for sampling multiple chains, serially or in parallel with multiple threads or multiple processes,

• an iterator and a transducer for sampling Markov chains.

The main user-facing API is sample.

Some common keyword arguments:

• discard_initial (default: 0): number of initial samples that are discarded

• thinning (default: 1): factor by which samples are thinned

• chain_type (default: Any): type of the returned chain

• callback (default: nothing): callback that is called after every sampling step

• progress (default: AbstractMCMC.PROGRESS[] which is true initially): toggles progress logging

AbstractMCMC.jl allows you to sample multiple chains with the following three algorithms:

• MCMCSerial(): sample in serial (no parallelization)

• MCMCThreads(): parallel sampling with multiple threads

• MCMCDistributed(): parallel sampling with multiple processes

The iterator interface allows us to e.g. plot the samples in every step:

MCMCChains.jl

MCMCChains.jl contains tools for analyzing and visualizing MCMC chains.

Posterior predictive check

It is often useful to generate simulated data using samples from the posterior distribution and compare them to the original data.

(see, e.g., "Statistical Rethinking" (section 3.3.2 - model checking) and the documentation)

Probabilistic programming

Source: Lecture slides

Developing probabilistic models and inference algorithms is a time-consuming and error-prone process.

• Probabilistic model written as a computer program

• Automatic inference (integral part of the programming language)

Advantages:

• Fast development of models

• Expressive models

• Widely applicable inference algorithms

Turing

Turing is a probabilistic programming language (PPL).

General design

• Probabilistic models are implemented as a Julia function

• One may use any Julia code inside of the model

• Random variables and observations are declared with the ~ operator:

  @model function mymodel(x, y)
      ...
      # random variable `a` with prior distribution `dist_a`
      a ~ dist_a
  
      ...
  
      # observation `y` with data distribution `dist_y`
      y ~ dist_y
      ...
  end

• PPL is implemented in DynamicPPL.jl, including @model macro

• Turing.jl integrates and reexports different parts of the ecosystem such as the PPL, inference algorithms, and tools for automatic differentiation

Coinflip model

A possible Turing model:

coinflip is a Julia function: it creates a model of type DynamicPPL.Model that stores

• the name of the model,

• the generative function of the model (i.e., the function above that declares how to run the model),

• and the arguments of the model and their default values.

We can look up the documentation of coinflip:

And we can inspect how the @model macro rewrites the function definition:

Since coinflip is a regular Julia function, we can also extend it. E.g., we can allow users to specify the results of the coin flips as a vector of true (heads) and false (tail):

Inference

We choose the same parameters as above: $n=$ 50, $k=$ 25, $\alpha =$ 1, $\beta =$ 1

If you call the model without arguments, it is executed and all variables are sampled from the prior.

Algorithm: Metropolis-HastingsHamiltonian Monte CarloNo-UTurn SamplerParticle Gibbs

Analysis

Posterior predictive check

Similar to above, we can check if the posterior predictions fit the observed data.

In more complicated models it can be tedious to rewrite how to sample the observations explicitly. Fortunately, we can use the following feature of Turing:

If the left hand side of a ~ expression is missing, it is sampled even if it is an argument to the model.

generated_quantities can be used to compute/sample an array of the return values of a model for each set of samples in the chain:

Automatic differentiation

Documentation

Algorithms such as Hamiltonian Monte Carlo (HMC) or no-U-turn sampler (NUTS) require the gradient of the log density function as well.

Turing computes the gradient automatically with automatic differentiation (AD). Different backends and algorithms are supported:

• ForwardDiff.jl: forward-mode AD, the default backend

• Tracker.jl: reverse-mode AD

• ReverseDiff.jl: reverse-mode AD, has to be loaded explicitly (optional cache for some models)

• Zygote.jl: reverse-mode AD, has to be loaded explicitly

If not set explicitly in the sampler, Turing uses the currently active global AD backend. It can be set with

• setadbackend(:forwarddiff),

• setadbackend(:tracker),

• setadbackend(:reversediff), or

• setadbackend(:zygote).

Alternatively, the backend can be set explicitly in the sampler: