Variational inference

Fast approximate Bayesian inference for real-time applications

Estimated time: 1.5-2 hours | Advanced session | Requires: MCMC only (can skip stochastic sessions)

Introduction

Do this session if you need faster inference than MCMC (real-time forecasting, many datasets) or want to understand modern ML-influenced approaches to Bayesian inference.

So far in this course, we’ve used Markov Chain Monte Carlo (MCMC) for Bayesian inference. MCMC is asymptotically exact and provides reliable uncertainty quantification. But it has a key limitation: it’s slow.

For a single dataset with a simple model, MCMC is usually affordable: the model in this session fits in seconds. The problem is what happens next:

Real-time forecasting: Update predictions as new data arrives every hour
Many datasets: Analyse hundreds of outbreaks from different locations
High-dimensional models: Age-stratified models with 50+ parameters
Rapid prototyping: Quickly explore many model variants

This is where variational inference (VI) comes in. VI trades some accuracy for speed improvements - the gains depend on problem complexity, from no useful saving at all on models as small as this session’s to orders of magnitude for large neural networks or models with many latent variables.

Objectives

By the end of this session, you will be able to:

Understand the variational inference framework and ELBO optimisation
Implement mean-field and full-rank Gaussian approximations
Compare VI and MCMC trade-offs (speed vs accuracy)
Choose appropriate inference methods for different scenarios
Apply VI to real-time forecasting workflows

Setup

using Distributions
using DataFrames
using CSV
using Plots
using StatsPlots
using Turing
using Turing.Variational: vi, q_meanfield_gaussian, q_fullrank_gaussian
using ADTypes: AutoForwardDiff
using FlexiChains ## chain objects and diagnostics
using Random
using LinearAlgebra
using DrWatson
using Statistics: mean, std, cor
using MFIIDD ## the course package, which holds the SIR implementation

Random.seed!(1234)

TaskLocalRNG()

Coming from Stan?

The closest equivalent is Stan’s ADVI (e.g. vb() in rstan, $variational() in cmdstanr, or CmdStanModel.variational() in cmdstanpy). The interface is similar: define a model and call the VI method instead of the sampling method. The same trade-offs apply — faster than full MCMC but approximate.

Load Tristan da Cunha data

flu_tdc = CSV.read(datadir("flu_tdc_1971.csv"), DataFrame)

scatter(
    flu_tdc.time, flu_tdc.obs,
    xlabel = "Time (days)", ylabel = "Daily incidence",
    label = "Observed cases", markersize = 4, color = :red,
    title = "Tristan da Cunha 1971 Influenza Outbreak",
    legend = :topright
)

Simple SIR model

For clarity, we’ll use a simple SIR model.

These two functions live in src/sir_model.jl and arrived with using MFIIDD above.

function sir_ode!(du, u, p, t)
    S, I, R = u
    R_0, D_inf = p
    N = S + I + R

    β = R_0 / D_inf
    ν = 1 / D_inf

    du[1] = -β * S * I / N          # dS/dt
    du[2] = β * S * I / N - ν * I   # dI/dt
    du[3] = ν * I                   # dR/dt
end

function simulate_sir(θ, init_state, times)
    u0 = [init_state[:S], init_state[:I], init_state[:R]]
    prob = remake(
        SIR_BASE_PROBLEM;
        u0 = u0,
        p = [θ[:R_0], θ[:D_inf]],
        tspan = (times[1], times[end]),
    )
    sol = solve(prob, Tsit5(), saveat = times)

    df = DataFrame(
        time = sol.t,
        S = [u[1] for u in sol.u],
        I = [u[2] for u in sol.u],
        R = [u[3] for u in sol.u],
    )

    # Daily incidence: new infections, the flow out of S
    df.Inc = [0.0; -diff(df.S)]

    return df
end

## `times` starts at 0.0, the left edge of the first observation day, because the
## incidence over that day is S(0) - S(1) and there is no way to form it without a
## point at 0. `traj.Inc[i + 1]` is then the incidence over the day the i-th
## observation counts
times = [0.0; flu_tdc.time]

## The island as the ship lands, at the start of day 1. Of the five islanders it
## brought back, three had been ill during the eight day voyage and are past their
## infectious period, so they start recovered; the two who fell ill on landing
## start infectious. N = 284, the island population
init_state = Dict(:S => 279.0, :I => 2.0, :R => 3.0)

Dict{Symbol, Float64} with 3 entries:
  :I => 2.0
  :R => 3.0
  :S => 279.0

Turing model for inference

# The priors, named once. The model draws from them and the transform below
# inverts them, and naming them here keeps the two from disagreeing.
const PRIORS = (
    R_0 = Uniform(1.0, 20.0),
    D_inf = Uniform(1.0, 10.0),
    ρ = Uniform(0.1, 1.0),
)

@model function sir_model(times, init_state, n_obs)
    # Priors
    R_0 ~ PRIORS.R_0
    D_inf ~ PRIORS.D_inf
    ρ ~ PRIORS.ρ

    # Simulate
    θ = Dict(:R_0 => R_0, :D_inf => D_inf)
    traj = simulate_sir(θ, init_state, times)

    # Likelihood
    lambdas = [max(ρ * traj.Inc[i + 1], 1e-10) for i in 1:n_obs]
    obs ~ arraydist(Poisson.(lambdas))

    return traj
end

model = sir_model(times, init_state, length(flu_tdc.obs)) | (; obs = flu_tdc.obs)

DynamicPPL.Model{typeof(sir_model), (:times, :init_state, :n_obs), (), (), Tuple{Vector{Float64}, Dict{Symbol, Float64}, Int64}, Tuple{}, DynamicPPL.CondFixContext{DynamicPPL.Condition, VarNamedTuple{(:obs,), Tuple{Vector{Int64}}}, DynamicPPL.DefaultContext}, false}(sir_model, (times = [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0  …  50.0, 51.0, 52.0, 53.0, 54.0, 55.0, 56.0, 57.0, 58.0, 59.0], init_state = Dict(:I => 2.0, :R => 3.0, :S => 279.0), n_obs = 59), NamedTuple(), CondFixContext{DynamicPPL.Condition}(VarNamedTuple(obs = [0, 1, 0, 10, 6, 32, 47, 37, 29, 11, 13, 8, 2, 2, 2, 0, 2, 1, 2, 2, 5, 2, 3, 2, 3, 5, 7, 8, 5, 7, 4, 12, 7, 5, 3, 4, 1, 0, 5, 2, 2, 0, 6, 1, 1, 2, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 1],), DynamicPPL.DefaultContext()))

Reminder: arraydist and | conditioning

obs ~ arraydist(Poisson.(lambdas)) declares a vector of independent Poisson observations, one per time point. The | operator, applied when creating the model instance below, conditions the model on the actual data — see the MCMC session for a full explanation of this pattern.

In your own work, you should run prior and posterior predictive checks when fitting models. We skip them here to focus on comparing VI with MCMC.

Baseline: MCMC with NUTS

Let’s start with full MCMC to establish a baseline:

println("Running MCMC with NUTS...")
## A short run first, so that the timing below measures sampling and not Julia
## compiling the sampler for the first time
sample(model, NUTS(0.65), 100, progress=false)
t_mcmc = @elapsed chain_mcmc = sample(model, NUTS(0.65), 500, progress=false)
println("MCMC time: $(round(t_mcmc, digits=1)) seconds")

Running MCMC with NUTS...
┌ Info: Found initial step size
└   ϵ = 0.2
┌ Info: Found initial step size
└   ϵ = 0.05
MCMC time: 0.4 seconds

Code explainer

@elapsed - returns the time in seconds to execute the expression.
NUTS(0.65) - sets the target acceptance rate during warmup. 0.65 is a good default; use 0.8-0.9 for difficult posteriors.

# Plot MCMC posteriors
plot(chain_mcmc[[:R_0, :D_inf, :ρ]], size=(900, 600))

# Summary
summarystats(chain_mcmc)

╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────╮
│   iter    collapsed                                                          │
│   chain   collapsed                                                          │
│ ↓ stat  = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95]          │
│                                                                              │
│ Parameters (3) ── AbstractPPL.VarName                                        │
│  Float64  R_0, D_inf, ρ                                                      │
│                                                                              │
│ Extras (14)                                                                  │
│  Float64  n_steps, is_accept, acceptance_rate, log_density,                  │
│           hamiltonian_energy, hamiltonian_energy_error,                      │
│           max_hamiltonian_energy_error, tree_depth, numerical_error,         │
│           step_size, nom_step_size, logprior, loglikelihood, logjoint        │
│                                                                              │
│ Summary                                                                      │
│   param    mean     std    mcse  ess_bulk  ess_tail    rhat      q5  …       │
│     R_0  1.8431  0.1721  0.0124  190.6839  196.0334  1.0027  1.6242  …       │
│   D_inf  3.2148  0.4001  0.0297  184.1912  239.2827  1.0036  2.6557  …       │
│       ρ  0.9904  0.0097  0.0006  191.8459  202.8282  0.9998  0.9713  …       │
╰──────────────────────────────────────────────────────────────────────────────╯

Note the time printed above and compare it with the variational runs below. Each timing cell runs the method once before measuring it, so the figure is the cost of sampling and not the cost of Julia compiling the sampler for the first time, which on a model this small would otherwise dominate everything.

Approach 1: mean-field variational inference

What is variational inference?

MCMC works by drawing samples from the posterior distribution. This gives us exact results (in the limit), but requires many sequential steps as the sampler explores parameter space.

Variational inference takes a completely different approach: instead of sampling from the true posterior $p(\theta | \text{data})$ , VI finds a simpler approximation $q(\theta)$ that’s “close” to the true posterior.

The key idea is to turn inference into an optimisation problem:

Choose a family of simple distributions $q(\theta)$ (e.g., Gaussians)
Optimise parameters of $q$ to make it as close as possible to the true posterior
Use gradient descent instead of MCMC sampling

Why should this ever be faster? Optimisation can use parallel gradient computations and converges in a fixed number of steps, whereas MCMC must run sequentially until convergence, which can take many thousands of iterations. Whether that pays off depends on the model, and on the small one in this session it does not.

The ELBO objective

How do we measure “closeness” between $q(\theta)$ and the true posterior? VI uses the Kullback-Leibler (KL) divergence, which measures how much information is lost when we use $q$ to approximate the true posterior.

Directly minimising KL divergence is intractable (we can’t compute the true posterior!), but we can instead maximise an equivalent quantity called the Evidence Lower Bound (ELBO):

$\text{ELBO}(q) = \mathbb{E}_{q}[\log p(\text{data}, \theta)] - \mathbb{E}_{q}[\log q(\theta)]$

The first term rewards $q$ for placing probability mass where the joint distribution $p(\text{data}, \theta)$ is high. The second term (entropy of $q$ ) prevents $q$ from collapsing to a point estimate. Together, they push $q$ to approximate the posterior well.

Maximising ELBO is equivalent to minimising $\text{KL}(q || p)$ , and importantly, we can estimate the ELBO using samples from $q$ - no need to evaluate the intractable posterior directly.

Mathematical detail

The derivation of the ELBO from KL divergence involves Jensen’s inequality and properties of logarithms. If you are not familiar with these concepts, you can safely accept the ELBO formula above on faith and skip ahead to the practical section below — the key intuition is simply that a higher ELBO means a better approximation.

Mean-field approximation

The simplest VI uses a mean-field approximation, which assumes all parameters are independent:

$q(\theta) = q(\theta_1) \times q(\theta_2) \times \cdots \times q(\theta_n)$

For our SIR model with three parameters: $q(R_0, D_{\text{inf}}, \rho) = \mathcal{N}(\mu_1, \sigma_1^2) \times \mathcal{N}(\mu_2, \sigma_2^2) \times \mathcal{N}(\mu_3, \sigma_3^2)$

This factorisation means we only need to learn 6 numbers (3 means + 3 variances) instead of the full joint distribution. The optimisation becomes tractable, but we lose the ability to capture correlations between parameters.

Note: The actual optimisation happens in unconstrained space (using log transforms for positive parameters, logit for bounded parameters), then samples are transformed back.

Passing q_meanfield_gaussian to vi tells Turing to construct an initial mean-field variational distribution: one independent Gaussian per model parameter, in unconstrained space. All means start at zero and all standard deviations at one. The vi function then adjusts these means and standard deviations to maximise the ELBO.

println("\nRunning Mean-Field VI...")
vi(model, q_meanfield_gaussian, 100; adtype=AutoForwardDiff(), show_progress=false)  ## compile first
t_vi_mf = @elapsed result_mf = vi(model, q_meanfield_gaussian, 500; adtype=AutoForwardDiff(), show_progress=false)
q_mf = result_mf.q
stats_mf = result_mf.info
println("Mean-field VI time: $(round(t_vi_mf, digits=1)) seconds")
println("Final ELBO: ", round(stats_mf[end].elbo, digits=2))

Running Mean-Field VI...
[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.
[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.
Mean-field VI time: 1.2 seconds
Final ELBO: -391.97

Code explainer

q_meanfield_gaussian - a constructor passed to vi that builds the initial mean-field Gaussian variational family (independent Normals per parameter, standard normals in unconstrained space).
vi(model, family, n_iters; ...) - runs variational inference, optimising the ELBO for n_iters steps. Returns a VIResult whose .q field holds the optimised distribution and .info holds the training statistics.
adtype=AutoForwardDiff() - specifies ForwardDiff.jl for computing gradients. Forward-mode AD is efficient when the number of parameters is small (like here).
stats_mf[end].elbo - the final ELBO value. Higher (less negative) is better, indicating a closer approximation to the true posterior.

Monitoring ELBO convergence

A critical diagnostic for VI is checking that the ELBO has converged — i.e., that the optimisation has run for long enough. We can plot the ELBO over iterations and look for a plateau:

elbo_vals_mf = [s.elbo for s in stats_mf]
plot(elbo_vals_mf, xlabel="Iteration", ylabel="ELBO",
     title="ELBO convergence (mean-field)", label="ELBO",
     linewidth=2, color=:blue)

The ELBO should increase (become less negative) and eventually plateau. If the curve is still climbing at the end, you need more iterations. Typical counts range from a few hundred to a few thousand, depending on model complexity — for our simple 3-parameter model, 500 iterations is usually sufficient, but more complex models may need 1000–5000 or more.

The mean-field run above took about as long as MCMC did. On a model this small there is nothing to save. The question worth asking here is whether it is accurate.

Before any of that can be compared with MCMC, the draws have to be put back on the scale of the parameters.

# `vi` optimises in the unconstrained space Turing transforms the model into, so
# `rand(q, n)` returns numbers on that scale: an R_0 of about -3 rather than 1.8. Each
# parameter goes back through the inverse of the bijector belonging to its prior.
const Bijectors = Turing.Bijectors

"Map draws from a variational approximation back to the parameter scale."
function to_constrained(samples)
    columns = map(enumerate(keys(PRIORS))) do (i, name)
        back = Bijectors.inverse(Bijectors.bijector(PRIORS[name]))
        name => [back(samples[i, j]) for j in axes(samples, 2)]
    end
    return DataFrame(columns)
end

Variational draws are not on the parameter scale

Forgetting this step is quiet rather than loud. rand(q, n) returns a perfectly ordinary matrix of numbers, and nothing complains if you label its rows R_0, D_inf and ρ and plot them against an MCMC posterior. The result is a comparison between the real posterior and a set of numbers that are not the parameters at all: an R_0 of about -3.1 for a quantity that cannot go below 1, and a ρ of about +4.6 for one that cannot go above it.

# Sample from variational approximation and put the draws on the parameter scale
samples_vi_mf = rand(q_mf, 10_000)
vi_mf_df = to_constrained(samples_vi_mf)

# Plot comparison
p1 = histogram(vec(chain_mcmc[:R_0]), bins=30, alpha=0.5, normalize=:pdf,
               label="MCMC", xlabel="R_0", ylabel="Density")
histogram!(p1, vi_mf_df.R_0, bins=30, alpha=0.5, normalize=:pdf,
           label="VI (mean-field)")

p2 = histogram(vec(chain_mcmc[:D_inf]), bins=30, alpha=0.5, normalize=:pdf,
               label="MCMC", xlabel="D_inf", ylabel="Density")
histogram!(p2, vi_mf_df.D_inf, bins=30, alpha=0.5, normalize=:pdf,
           label="VI (mean-field)")

plot(p1, p2, layout=(1, 2), size=(900, 400),
     title=["R_0 Comparison" "D_inf Comparison"])

The independence assumption problem

Mean-field assumes parameters are independent. But in epidemic models, $R_0$ and $D_{\text{inf}}$ are strongly correlated in the posterior, and here that correlation is positive: the cell below prints it at about 0.95.

Why does this correlation arise? Consider the SIR model dynamics:

The transmission rate $\beta = R_0 / D_{\text{inf}}$ determines how fast susceptibles become infected
If we increase $R_0$ while keeping $\beta$ roughly constant (to match observed epidemic speed), we must also increase $D_{\text{inf}}$
Conversely, a shorter infectious period goes with a lower $R_0$ , since $\beta$ is their ratio and the data pin $\beta$ down

In other words, there are multiple combinations of ( $R_0$ , $D_{\text{inf}}$ ) that produce similar epidemic trajectories. The data constrains their ratio more tightly than their individual values. This creates a ridge of high posterior probability running diagonally through parameter space - exactly the kind of correlation mean-field VI cannot capture.

The sign follows from that argument. Holding $\beta = R_0 / D_{\text{inf}}$ roughly fixed, a larger $R_0$ has to be paid for with a larger $D_{\text{inf}}$ , so the two move together and the correlation comes out positive. The cell below measures it rather than asserting it, which is worth doing whenever you are about to claim a direction.

# Check correlation in MCMC samples
println("MCMC correlation (R_0, D_inf): ",
        round(cor(vec(chain_mcmc[:R_0]), vec(chain_mcmc[:D_inf])), digits=3))

# Mean-field VI forces correlation to 0
println("VI mean-field correlation (R_0, D_inf): ",
        round(cor(vi_mf_df.R_0, vi_mf_df.D_inf), digits=3))

MCMC correlation (R_0, D_inf): 0.949
VI mean-field correlation (R_0, D_inf): 0.004

What we see is that the mean-field approximation forces this correlation to zero, which can lead to:

Overconfident marginals: Each parameter appears more certain than it really is
Implausible joint samples: Combinations that the true posterior considers unlikely
Poor predictive performance: Forward simulations may not match observed data well

Approach 2: full-rank Gaussian VI

Capturing correlations

The solution to the independence problem is to use a full covariance matrix instead of a diagonal one:

$q(\theta) = \mathcal{N}(\mu, \Sigma)$

where $\Sigma$ is a full $n \times n$ covariance matrix (not just diagonal elements).

For our 3-parameter model, this means:

Mean-field: 6 parameters (3 means + 3 variances)
Full-rank: 9 parameters (3 means + 6 unique covariance entries, since $\Sigma$ is symmetric)

The additional parameters allow the approximation to capture correlations. The elliptical contours of a full Gaussian can be tilted and stretched to match the ridge-like structure of the true posterior.

The trade-off is that more parameters means more to optimise, and for high-dimensional problems the covariance matrix grows as $O(n^2)$ , which can become prohibitive; in those cases low-rank approximations or structured covariances are used. With three parameters that cost is invisible: both families run the same 500 iterations, and the extra covariance entries are a handful of numbers. Expect the two timings below to come out close enough that the difference is not worth reading.

println("\nRunning Full-Rank VI...")
vi(model, q_fullrank_gaussian, 100; adtype=AutoForwardDiff(), show_progress=false)  ## compile first
t_vi_full = @elapsed result_full = vi(model, q_fullrank_gaussian, 500; adtype=AutoForwardDiff(), show_progress=false)
q_full = result_full.q
stats_full = result_full.info
println("Full-rank VI time: $(round(t_vi_full, digits=1)) seconds")
println("Final ELBO: ", round(stats_full[end].elbo, digits=2))

Running Full-Rank VI...
[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.
[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.
Full-rank VI time: 0.3 seconds
Final ELBO: -399.77

Code explainer

q_fullrank_gaussian - constructor passed to vi that builds a full-rank Gaussian variational family. Captures correlations between parameters via a lower-triangular Cholesky factor for the covariance matrix.

# Sample from full-rank approximation, on the parameter scale again
samples_vi_full = rand(q_full, 10_000)
vi_full_df = to_constrained(samples_vi_full)

# Check correlation - should now be closer to MCMC
println("VI full-rank correlation (R_0, D_inf): ",
        round(cor(vi_full_df.R_0, vi_full_df.D_inf), digits=3))

VI full-rank correlation (R_0, D_inf): 0.964

The full-rank approximation recovers a correlation close to the one MCMC found, where mean-field was pinned at zero by construction. The correlation is not the only casualty, though. Mean-field reports the $R_0$ and $D_{\text{inf}}$ marginals about three times too narrow, while $\rho$ comes out slightly wider than MCMC rather than narrower. The narrowing falls on the pair whose correlation it discarded, which is the cost of assuming independence. Full-rank errs slightly wide on all three. Those narrow marginals make mean-field’s credible intervals untrustworthy in practice.

Plotting code

# Scatter plot comparison
p1 = scatter(vec(chain_mcmc[:R_0]), vec(chain_mcmc[:D_inf]),
             alpha=0.3, label="MCMC", markersize=2,
             xlabel="R_0", ylabel="D_inf",
             title="MCMC: Captures Correlation")

p2 = scatter(vi_mf_df.R_0, vi_mf_df.D_inf,
             alpha=0.3, label="VI (mean-field)", markersize=2,
             xlabel="R_0", ylabel="D_inf",
             title="Mean-Field VI: No Correlation")

p3 = scatter(vi_full_df.R_0, vi_full_df.D_inf,
             alpha=0.3, label="VI (full-rank)", markersize=2,
             xlabel="R_0", ylabel="D_inf",
             title="Full-Rank VI: Captures Correlation")

plot(p1, p2, p3, layout=(1, 3), size=(1200, 400))

Comparison: MCMC vs VI

Quantitative comparison

Let’s compare posterior means and standard deviations:

# Extract summaries
mcmc_summary = summarystats(chain_mcmc)
println("MCMC Summary:")
display(mcmc_summary)

println("\n" * "="^60)
println("Mean-Field VI Summary:")
println("  R_0: mean = ", round(mean(vi_mf_df.R_0), digits=2),
        ", std = ", round(std(vi_mf_df.R_0), digits=2))
println("  D_inf: mean = ", round(mean(vi_mf_df.D_inf), digits=2),
        ", std = ", round(std(vi_mf_df.D_inf), digits=2))
println("  ρ: mean = ", round(mean(vi_mf_df.ρ), digits=2),
        ", std = ", round(std(vi_mf_df.ρ), digits=4))

println("\n" * "="^60)
println("Full-Rank VI Summary:")
println("  R_0: mean = ", round(mean(vi_full_df.R_0), digits=2),
        ", std = ", round(std(vi_full_df.R_0), digits=2))
println("  D_inf: mean = ", round(mean(vi_full_df.D_inf), digits=2),
        ", std = ", round(std(vi_full_df.D_inf), digits=2))
println("  ρ: mean = ", round(mean(vi_full_df.ρ), digits=2),
        ", std = ", round(std(vi_full_df.ρ), digits=4))

MCMC Summary:

============================================================
Mean-Field VI Summary:
  R_0: mean = 1.85, std = 0.05
  D_inf: mean = 3.24, std = 0.14
  ρ: mean = 0.99, std = 0.0122

============================================================
Full-Rank VI Summary:
  R_0: mean = 1.93, std = 0.22
  D_inf: mean = 3.41, std = 0.51
  ρ: mean = 0.99, std = 0.0126

╭─FlexiSummary (9 statistics) ─────────────────────────────────────────────────╮
│   iter    collapsed                                                          │
│   chain   collapsed                                                          │
│ ↓ stat  = [mean, std, mcse, ess_bulk, ess_tail, rhat, q5, q50, q95]          │
│                                                                              │
│ Parameters (3) ── AbstractPPL.VarName                                        │
│  Float64  R_0, D_inf, ρ                                                      │
│                                                                              │
│ Extras (14)                                                                  │
│  Float64  n_steps, is_accept, acceptance_rate, log_density,                  │
│           hamiltonian_energy, hamiltonian_energy_error,                      │
│           max_hamiltonian_energy_error, tree_depth, numerical_error,         │
│           step_size, nom_step_size, logprior, loglikelihood, logjoint        │
│                                                                              │
│ Summary                                                                      │
│   param    mean     std    mcse  ess_bulk  ess_tail    rhat      q5  …       │
│     R_0  1.8431  0.1721  0.0124  190.6839  196.0334  1.0027  1.6242  …       │
│   D_inf  3.2148  0.4001  0.0297  184.1912  239.2827  1.0036  2.6557  …       │
│       ρ  0.9904  0.0097  0.0006  191.8459  202.8282  0.9998  0.9713  …       │
╰──────────────────────────────────────────────────────────────────────────────╯

Speed comparison

println("Speed comparison (from runs above):")
println("="^50)
println("MCMC (500 samples):      $(round(t_mcmc, digits=1)) seconds")
println("Mean-field VI (500 iter): $(round(t_vi_mf, digits=1)) seconds")
println("Full-rank VI (500 iter):  $(round(t_vi_full, digits=1)) seconds")
println("="^50)
println("MCMC time / mean-field VI time: $(round(t_mcmc / t_vi_mf, digits=1))")
println("MCMC time / full-rank VI time:  $(round(t_mcmc / t_vi_full, digits=1))")

Speed comparison (from runs above):
==================================================
MCMC (500 samples):      0.4 seconds
Mean-field VI (500 iter): 1.2 seconds
Full-rank VI (500 iter):  0.3 seconds
==================================================
MCMC time / mean-field VI time: 0.4
MCMC time / full-rank VI time:  1.3

Once compilation is out of the way, there is almost nothing between the three on this model, whatever machine renders the page. Variational inference earns its keep on problems large enough for a real per-fit gap to open up, and then especially when that gap is multiplied across many datasets. A three-parameter SIR gives it nothing to work with. Comparing the two here is the right way to check that VI is telling the truth about the posterior, and the wrong way to decide whether you need it.

Trade-offs summary

Method	Speed	Accuracy	Correlations	Use Case
MCMC	Slower, and the gap grows with model size	Exact (asymptotically)	Yes	Final analysis, publications
Mean-field VI	Fast	Approximate	No	Quick exploration, many datasets
Full-rank VI	Fast	Better	Yes	Real-time forecasting

On this three-parameter model the timings above separate the three methods hardly at all. Read the table as the shape of the trade-off on problems large enough for it to appear, rather than as a promise about seconds here.

When to use each method

Use MCMC when:

You need exact inference for publication
You have time (single dataset, offline analysis)
Posterior may be multimodal or complex
You need reliable uncertainty quantification
Problem is low-to-moderate dimensional (<50 parameters)

Use mean-field VI when:

You need extremely fast approximate inference
Parameters are roughly independent
You’re doing exploratory analysis or prototyping
You need to analyse many datasets quickly
Rough uncertainty estimates are sufficient

Use full-rank VI when:

Speed is essential and MCMC is too slow (real-time systems)
Parameters are correlated (otherwise mean-field is fine)
You’ve validated that VI and MCMC give similar results for your problem

When VI fails

Gaussian VI (both mean-field and full-rank) assumes the posterior is well-approximated by a single Gaussian. This assumption breaks down in several important cases:

Multimodal posteriors: If the posterior has multiple peaks (e.g., a model with label-switching symmetry), a single Gaussian will either collapse onto one mode or spread across modes, giving a poor approximation in either case.
Heavy-tailed posteriors: Gaussian tails decay exponentially, so VI will systematically underestimate the probability of extreme parameter values. This leads to overconfident uncertainty intervals.
Strong nonlinear correlations: Full-rank VI captures linear correlations via the covariance matrix, but banana-shaped or other nonlinear posterior structures cannot be represented by any Gaussian.

When you suspect these issues, MCMC remains the more reliable choice. A useful diagnostic: if VI and MCMC give substantially different results, trust MCMC — VI is likely struggling with the posterior geometry.

Practical workflow

Recommended strategy

Start with mean-field VI: Quick exploration, check if model runs
Run MCMC on subset: Validate on small dataset or fewer iterations
Compare VI and MCMC: Are they similar?
- If yes and speed is critical → VI is acceptable for operational use
- If time allows → use MCMC for final analysis (more reliable)
For real-time systems using VI: Periodic MCMC checks to validate

Why speed matters in practice

On a single small analysis the choice hardly matters, as the timings above show. The difference becomes critical in operational settings:

Real-time epidemic forecasting: During an outbreak, you might need to update estimates as new data arrives - daily or even hourly. Take a model large enough that MCMC needs five minutes a fit and VI under a minute, which is a far bigger gap than the one this session’s three-parameter model shows. Fitting 100 regions is then over 8 hours per update cycle against about an hour. The per-fit gap is what does the work; the hundred regions multiply whatever gap you have rather than creating one.

Model development and debugging: When developing a new model, you’ll run inference hundreds of times as you iterate on the specification, fix bugs, and tune priors. On a model where the gap is real, faster iteration cycles accelerate development considerably.

Sensitivity analyses: Checking how results change with different priors, data subsets, or model variants requires many inference runs. VI makes comprehensive sensitivity analysis practical.

Ensemble forecasting: Combining predictions from multiple models improves forecast accuracy. VI’s speed makes it practical to run many model variants.

The practical recommendation: validate VI against MCMC once (to ensure they give similar results for your problem), then use VI for operational workflows where speed matters.

Exercises

Exercise 1: ELBO convergence and iteration count

Run mean-field VI with different numbers of iterations (50, 200, 500, 2000) and observe how the ELBO converges and how the posterior approximation changes.

Hint

Use a loop over iteration counts, storing the final ELBO and posterior mean of $R_0$ for each run. Remember to put the draws back on the parameter scale with to_constrained before summarising them.

Solution

iter_counts = [50, 200, 500, 2000]
results = []

for n_iter in iter_counts
    res = vi(model, q_meanfield_gaussian, n_iter; adtype=AutoForwardDiff(), show_progress=false)
    # Back to the parameter scale before summarising, as in the main text
    draws = to_constrained(rand(res.q, 5_000))
    push!(results, (
        n_iter = n_iter,
        final_elbo = res.info[end].elbo,
        R0_mean = mean(draws.R_0),
        R0_std = std(draws.R_0),
        elbo_trace = [s.elbo for s in res.info]
    ))
end

# Plot ELBO traces
p = plot(xlabel="Iteration", ylabel="ELBO",
         title="ELBO convergence by iteration count", legend=:bottomright)
for r in results
    plot!(p, r.elbo_trace, label="$(r.n_iter) iters", linewidth=1.5)
end
p

[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.
[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.
[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.
[ Info: The capability of the supplied target `LogDensityProblem` LogDensityProblems.LogDensityOrder{1}() is >= `LogDensityProblems.LogDensityOrder{1}()`. To make use of this, the `adtype` argument for AdvancedVI must be one of `AutoReverseDiff`, `AutoMooncake`, or `AutoEnzyme` in reverse mode.

# Summarise results
println("Iterations | Final ELBO | R₀ mean ± std")
println("-"^50)
for r in results
    println("$(lpad(r.n_iter, 9)) | $(lpad(round(r.final_elbo, digits=1), 10)) | $(round(r.R0_mean, digits=2)) ± $(round(r.R0_std, digits=2))")
end

Iterations | Final ELBO | R₀ mean ± std
--------------------------------------------------
       50 |     -427.3 | 2.93 ± 0.12
      200 |     -413.0 | 2.09 ± 0.08
      500 |     -391.6 | 1.86 ± 0.06
     2000 |     -392.5 | 1.81 ± 0.05

You should see that with too few iterations (e.g., 50), the ELBO has not converged and the posterior means may be inaccurate. By around 500 iterations the ELBO typically plateaus for this model, and increasing to 2000 gives negligible improvement. This tells us 500 iterations is sufficient here.

Exercise 2: Comparing VI and MCMC posteriors

Run both mean-field VI and full-rank VI, and compare their posteriors against the MCMC baseline. Produce overlay histograms for all three parameters and assess where VI agrees with MCMC and where it diverges.

Hint

Use the MCMC chain chain_mcmc and VI distributions q_mf and q_full that we already fitted above. Sample 10,000 draws from each VI distribution and overlay the histograms.

Solution

# Draw samples (reuse existing VI fits from above), on the parameter scale
samp_mf = to_constrained(rand(q_mf, 10_000))
samp_full = to_constrained(rand(q_full, 10_000))

param_names = ["R_0", "D_inf", "ρ"]
plots_list = []

for (i, pname) in enumerate(param_names)
    p = histogram(vec(chain_mcmc[Symbol(pname)]), bins=40, alpha=0.4,
                  normalize=:pdf, label="MCMC", xlabel=pname, ylabel="Density")
    histogram!(p, samp_mf[!, pname], bins=40, alpha=0.4,
               normalize=:pdf, label="Mean-field VI")
    histogram!(p, samp_full[!, pname], bins=40, alpha=0.4,
               normalize=:pdf, label="Full-rank VI")
    push!(plots_list, p)
end

plot(plots_list..., layout=(1, 3), size=(1200, 400),
     plot_title="Posterior comparison: MCMC vs VI")

# Quantitative comparison
println("Parameter | MCMC mean (std) | MF-VI mean (std) | FR-VI mean (std)")
println("-"^75)
for (i, pname) in enumerate(param_names)
    mcmc_vals = vec(chain_mcmc[Symbol(pname)])
    ## standard deviations to two significant figures, so ρ's does not round to 0.01
    println("$(rpad(pname, 9)) | $(round(mean(mcmc_vals), digits=2)) ($(round(std(mcmc_vals), sigdigits=2))) | " *
            "$(round(mean(samp_mf[!, pname]), digits=2)) ($(round(std(samp_mf[!, pname]), sigdigits=2))) | " *
            "$(round(mean(samp_full[!, pname]), digits=2)) ($(round(std(samp_full[!, pname]), sigdigits=2)))")
end

Parameter | MCMC mean (std) | MF-VI mean (std) | FR-VI mean (std)
---------------------------------------------------------------------------
R_0       | 1.84 (0.17) | 1.85 (0.054) | 1.93 (0.21)
D_inf     | 3.21 (0.4) | 3.24 (0.14) | 3.41 (0.5)
ρ         | 0.99 (0.0097) | 0.99 (0.012) | 0.99 (0.013)

Look at where the three methods agree and disagree. Typical findings for this model:

Posterior means are usually similar across all three methods — VI tends to find the right location.
Posterior widths may differ: mean-field VI can be overconfident (narrower) because it ignores correlations.
Full-rank VI generally gives a closer match to MCMC than mean-field, particularly for correlated parameters like $R_0$ and $D_\text{inf}$ .

If you see substantial disagreement between full-rank VI and MCMC, it may indicate that 500 VI iterations was not enough, or that the posterior has features (e.g., skewness) that even full-rank Gaussian VI cannot capture.

Learning points

Variational inference turns Bayesian inference into optimisation: find a simple distribution $q(\theta)$ that approximates the true posterior
Mean-field VI assumes parameter independence - cheap per iteration, but misses correlations that are common in epidemic models
Full-rank VI captures correlations using a full covariance matrix - good balance between speed and accuracy
ELBO (Evidence Lower Bound) is the objective function - higher values mean better approximations
Speed vs accuracy trade-off: VI trades exactness for speed, and the saving grows with the size of the problem rather than showing up on a model like this one; use VI for exploration and real-time systems, MCMC for final publication-quality analysis

VI is not a replacement for MCMC - it’s a complementary tool that enables applications where MCMC is too slow!

Going further

The Gaussian families used here are the most common in practice, but the variational family need not be Gaussian. Normalising flows use sequences of invertible transformations to build far more flexible approximations, handling multimodal and skewed posteriors that mean-field and full-rank Gaussians cannot — see Bijectors.jl and NormalizingFlows.jl for Julia implementations. Everything else about the workflow stays the same: you are still optimising one approximation, to one posterior, for one dataset.

Relaxing that assumption — learning the posterior as a function of the data, so that a new dataset costs nothing — is a bigger step, and the subject of the next session.

Next session

In Neural posterior estimation, we take the idea of replacing sampling with optimisation one step further: instead of optimising an approximation for a single dataset, we learn the posterior as a function of the data, so that every subsequent fit is essentially free.

References

Kucukelbir et al. (2017). Automatic Differentiation Variational Inference. JMLR 18(14), 1-45.
Blei et al. (2017). Variational Inference: A Review for Statisticians. JASA 112(518), 859-877. (Excellent introduction to VI)
Rezende & Mohamed (2015). Variational Inference with Normalizing Flows. ICML.
Zhang et al. (2019). Advances in Variational Inference. TPAMI 41(8), 2008-2026.
Turing.jl documentation - Probabilistic programming in Julia
AdvancedVI.jl - Variational inference algorithms for Turing