= model(data...) # instantiate model on the data
m = vi(m, vi_alg) # perform VI on `m` using the VI method `vi_alg`, which returns a `VariationalPosterior` q
Variational inference (VI) in Turing.jl
In this post we’ll have a look at what’s know as variational inference (VI), a family of approximate Bayesian inference methods, and how to use it in Turing.jl as an alternative to other approaches such as MCMC. In particular, we will focus on one of the more standard VI methods called Automatic Differentation Variational Inference (ADVI).
Here we will focus on how to use VI in Turing and not much on the theory underlying VI. If you are interested in understanding the mathematics you can checkout our write-up or any other resource online (there a lot of great ones).
Using VI in Turing.jl is very straight forward. If model
denotes a definition of a Turing.Model
, performing VI is as simple as
Thus it’s no more work than standard MCMC sampling in Turing.
To get a bit more into what we can do with vi
, we’ll first have a look at a simple example and then we’ll reproduce the tutorial on Bayesian linear regression using VI instead of MCMC. Finally we’ll look at some of the different parameters of vi
and how you for example can use your own custom variational family.
We first import the packages to be used:
using Random
using Turing
using Turing: Variational
using StatsPlots, Measures
Random.seed!(42);
Simple example: Normal-Gamma conjugate model
The Normal-(Inverse)Gamma conjugate model is defined by the following generative process
\[\begin{align} s &\sim \mathrm{InverseGamma}(2, 3) \\ m &\sim \mathcal{N}(0, s) \\ x_i &\overset{\text{i.i.d.}}{=} \mathcal{N}(m, s), \quad i = 1, \dots, n \end{align}\]
Recall that conjugate refers to the fact that we can obtain a closed-form expression for the posterior. Of course one wouldn’t use something like variational inference for a conjugate model, but it’s useful as a simple demonstration as we can compare the result to the true posterior.
First we generate some synthetic data, define the Turing.Model
and instantiate the model on the data:
# generate data
= randn(2000); x
@model function model(x)
~ InverseGamma(2, 3)
s ~ Normal(0.0, sqrt(s))
m for i in 1:length(x)
~ Normal(m, sqrt(s))
x[i] end
end;
# Instantiate model
= model(x); m
Now we’ll produce some samples from the posterior using a MCMC method, which in constrast to VI is guaranteed to converge to the exact posterior (as the number of samples go to infinity).
We’ll produce 10 000 samples with 200 steps used for adaptation and a target acceptance rate of 0.65
If you don’t understand what “adaptation” or “target acceptance rate” refers to, all you really need to know is that NUTS
is known to be one of the most accurate and efficient samplers (when applicable) while requiring little to no hand-tuning to work well.
setprogress!(false)
= sample(m, NUTS(), 10_000); samples_nuts
┌ Info: Found initial step size
└ ϵ = 0.0125
Now let’s try VI. The most important function you need to now about to do VI in Turing is vi
:
@doc(Variational.vi)
vi(model, alg::VariationalInference)
vi(model, alg::VariationalInference, q::VariationalPosterior)
vi(model, alg::VariationalInference, getq::Function, θ::AbstractArray)
Constructs the variational posterior from the model
and performs the optimization following the configuration of the given VariationalInference
instance.
Arguments
model
:Turing.Model
orFunction
z ↦ log p(x, z) wherex
denotes the observationsalg
: the VI algorithm usedq
: aVariationalPosterior
for which it is assumed a specialized implementation of the variational objective used exists.getq
: function taking parametersθ
as input and returns aVariationalPosterior
θ
: only required ifgetq
is used, in which case it is the initial parameters for the variational posterior
Additionally, you can pass
- an initial variational posterior
q
, for which we assume there exists a implementation ofupdate(::typeof(q), θ::AbstractVector)
returning an updated posteriorq
with parametersθ
. - a function mapping \(\theta \mapsto q_{\theta}\) (denoted above
getq
) together with initial parametersθ
. This provides more flexibility in the types of variational families that we can use, and can sometimes be slightly more convenient for quick and rough work.
By default, i.e. when calling vi(m, advi)
, Turing use a mean-field approximation with a multivariate normal as the base-distribution. Mean-field refers to the fact that we assume all the latent variables to be independent. This the “standard” ADVI approach; see Automatic Differentiation Variational Inference (2016) for more. In Turing, one can obtain such a mean-field approximation by calling Variational.meanfield(model)
for which there exists an internal implementation for update
:
@doc(Variational.meanfield)
meanfield([rng, ]model::Model)
Creates a mean-field approximation with multivariate normal as underlying distribution.
Currently the only implementation of VariationalInference
available is ADVI
, which is very convenient and applicable as long as your Model
is differentiable with respect to the variational parameters, that is, the parameters of your variational distribution, e.g. mean and variance in the mean-field approximation.
@doc(Variational.ADVI)
struct ADVI{AD} <: AdvancedVI.VariationalInference{AD}
Automatic Differentiation Variational Inference (ADVI) with automatic differentiation backend AD
.
Fields
samples_per_step::Int64
: Number of samples used to estimate the ELBO in each optimization step.max_iters::Int64
: Maximum number of gradient steps.adtype::Any
: AD backend used for automatic differentiation.
To perform VI on the model m
using 10 samples for gradient estimation and taking 1000 gradient steps is then as simple as:
# ADVI
= ADVI(10, 1000)
advi = vi(m, advi); q
┌ Info: [ADVI] Should only be seen once: optimizer created for θ
└ objectid(θ) = 0x543dcd7376e4e21e
Unfortunately, for such a small problem Turing’s new NUTS
sampler is so efficient now that it’s not that much more efficient to use ADVI. So, so very unfortunate…
With that being said, this is not the case in general. For very complex models we’ll later find that ADVI
produces very reasonable results in a much shorter time than NUTS
.
And one significant advantage of using vi
is that we can sample from the resulting q
with ease. In fact, the result of the vi
call is a TransformedDistribution
from Bijectors.jl, and it implements the Distributions.jl interface for a Distribution
:
q isa MultivariateDistribution
true
This means that we can call rand
to sample from the variational posterior q
histogram(rand(q, 1_000)[1, :])
and logpdf
to compute the log-probability
logpdf(q, rand(q))
5.049886481894164
Let’s check the first and second moments of the data to see how our approximation compares to the point-estimates form the data:
var(x), mean(x)
(1.0149828942385808, -0.03802639687322374)
mean(rand(q, 1000); dims=2)...,) (
(1.0244278299554765, -0.032889543434005124)
That’s pretty close! But we’re Bayesian so we’re not interested in just matching the mean. Let’s instead look the actual density q
.
For that we need samples:
= rand(q, 10000);
samples size(samples)
(2, 10000)
= histogram(
p1 1, :]; bins=100, normed=true, alpha=0.2, color=:blue, label="", ylabel="density"
samples[
)density!(samples[1, :]; label="s (ADVI)", color=:blue, linewidth=2)
density!(samples_nuts, :s; label="s (NUTS)", color=:green, linewidth=2)
vline!([var(x)]; label="s (data)", color=:black)
vline!([mean(samples[1, :])]; color=:blue, label="")
= histogram(
p2 2, :]; bins=100, normed=true, alpha=0.2, color=:blue, label="", ylabel="density"
samples[
)density!(samples[2, :]; label="m (ADVI)", color=:blue, linewidth=2)
density!(samples_nuts, :m; label="m (NUTS)", color=:green, linewidth=2)
vline!([mean(x)]; color=:black, label="m (data)")
vline!([mean(samples[2, :])]; color=:blue, label="")
plot(p1, p2; layout=(2, 1), size=(900, 500), legend=true)
For this particular Model
, we can in fact obtain the posterior of the latent variables in closed form. This allows us to compare both NUTS
and ADVI
to the true posterior \(p(s, m \mid x_1, \ldots, x_n)\).
The code below is just work to get the marginals \(p(s \mid x_1, \ldots, x_n)\) and \(p(m \mid x_1, \ldots, x_n)\). Feel free to skip it.
# closed form computation of the Normal-inverse-gamma posterior
# based on "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
function posterior(μ₀::Real, κ₀::Real, α₀::Real, β₀::Real, x::AbstractVector{<:Real})
# Compute summary statistics
= length(x)
n = mean(x)
x̄ = sum(xi -> (xi - x̄)^2, x)
sum_of_squares
# Compute parameters of the posterior
= κ₀ + n
κₙ = (κ₀ * μ₀ + n * x̄) / κₙ
μₙ = α₀ + n / 2
αₙ = β₀ + (sum_of_squares + n * κ₀ / κₙ * (x̄ - μ₀)^2) / 2
βₙ
return μₙ, κₙ, αₙ, βₙ
end
= posterior(0.0, 1.0, 2.0, 3.0, x)
μₙ, κₙ, αₙ, βₙ
# marginal distribution of σ²
# cf. Eq. (90) in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
= InverseGamma(αₙ, βₙ)
p_σ² = z -> pdf(p_σ², z)
p_σ²_pdf
# marginal of μ
# Eq. (91) in "Conjugate Bayesian analysis of the Gaussian distribution" by Murphy
= μₙ + sqrt(βₙ / (αₙ * κₙ)) * TDist(2 * αₙ)
p_μ = z -> pdf(p_μ, z)
p_μ_pdf
# posterior plots
= plot()
p1 histogram!(samples[1, :]; bins=100, normed=true, alpha=0.2, color=:blue, label="")
density!(samples[1, :]; label="s (ADVI)", color=:blue)
density!(samples_nuts, :s; label="s (NUTS)", color=:green)
vline!([mean(samples[1, :])]; linewidth=1.5, color=:blue, label="")
plot!(range(0.75, 1.35; length=1_001), p_σ²_pdf; label="s (posterior)", color=:red)
vline!([var(x)]; label="s (data)", linewidth=1.5, color=:black, alpha=0.7)
xlims!(0.75, 1.35)
= plot()
p2 histogram!(samples[2, :]; bins=100, normed=true, alpha=0.2, color=:blue, label="")
density!(samples[2, :]; label="m (ADVI)", color=:blue)
density!(samples_nuts, :m; label="m (NUTS)", color=:green)
vline!([mean(samples[2, :])]; linewidth=1.5, color=:blue, label="")
plot!(range(-0.25, 0.25; length=1_001), p_μ_pdf; label="m (posterior)", color=:red)
vline!([mean(x)]; label="m (data)", linewidth=1.5, color=:black, alpha=0.7)
xlims!(-0.25, 0.25)
plot(p1, p2; layout=(2, 1), size=(900, 500))
Bayesian linear regression example using ADVI
This is simply a duplication of the tutorial on Bayesian linear regression (much of the code is directly lifted), but now with the addition of an approximate posterior obtained using ADVI
.
As we’ll see, there is really no additional work required to apply variational inference to a more complex Model
.
Random.seed!(1);
using FillArrays
using RDatasets
using LinearAlgebra
# Import the "Default" dataset.
= RDatasets.dataset("datasets", "mtcars");
data
# Show the first six rows of the dataset.
first(data, 6)
Row | Model | MPG | Cyl | Disp | HP | DRat | WT | QSec | VS | AM | Gear | Carb |
---|---|---|---|---|---|---|---|---|---|---|---|---|
String31 | Float64 | Int64 | Float64 | Int64 | Float64 | Float64 | Float64 | Int64 | Int64 | Int64 | Int64 | |
1 | Mazda RX4 | 21.0 | 6 | 160.0 | 110 | 3.9 | 2.62 | 16.46 | 0 | 1 | 4 | 4 |
2 | Mazda RX4 Wag | 21.0 | 6 | 160.0 | 110 | 3.9 | 2.875 | 17.02 | 0 | 1 | 4 | 4 |
3 | Datsun 710 | 22.8 | 4 | 108.0 | 93 | 3.85 | 2.32 | 18.61 | 1 | 1 | 4 | 1 |
4 | Hornet 4 Drive | 21.4 | 6 | 258.0 | 110 | 3.08 | 3.215 | 19.44 | 1 | 0 | 3 | 1 |
5 | Hornet Sportabout | 18.7 | 8 | 360.0 | 175 | 3.15 | 3.44 | 17.02 | 0 | 0 | 3 | 2 |
6 | Valiant | 18.1 | 6 | 225.0 | 105 | 2.76 | 3.46 | 20.22 | 1 | 0 | 3 | 1 |
# Function to split samples.
function split_data(df, at=0.70)
= size(df, 1)
r = Int(round(r * at))
index = df[1:index, :]
train = df[(index + 1):end, :]
test return train, test
end
# A handy helper function to rescale our dataset.
function standardize(x)
return (x .- mean(x; dims=1)) ./ std(x; dims=1)
end
function standardize(x, orig)
return (x .- mean(orig; dims=1)) ./ std(orig; dims=1)
end
# Another helper function to unstandardize our datasets.
function unstandardize(x, orig)
return x .* std(orig; dims=1) .+ mean(orig; dims=1)
end
function unstandardize(x, mean_train, std_train)
return x .* std_train .+ mean_train
end
unstandardize (generic function with 2 methods)
# Remove the model column.
select!(data, Not(:Model))
# Split our dataset 70%/30% into training/test sets.
= split_data(data, 0.7)
train, test = copy(train)
train_unstandardized
# Standardize both datasets.
= standardize(Matrix(train))
std_train = standardize(Matrix(test), Matrix(train))
std_test
# Save dataframe versions of our dataset.
= DataFrame(std_train, names(data))
train_cut = DataFrame(std_test, names(data))
test_cut
# Create our labels. These are the values we are trying to predict.
= train_cut[:, :MPG]
train_label = test_cut[:, :MPG]
test_label
# Get the list of columns to keep.
= filter(x -> !in(x, ["MPG"]), names(data))
remove_names
# Filter the test and train sets.
= Matrix(train_cut[:, remove_names]);
train = Matrix(test_cut[:, remove_names]); test
# Bayesian linear regression.
@model function linear_regression(x, y, n_obs, n_vars, ::Type{T}=Vector{Float64}) where {T}
# Set variance prior.
~ truncated(Normal(0, 100); lower=0)
σ²
# Set intercept prior.
~ Normal(0, 3)
intercept
# Set the priors on our coefficients.
~ MvNormal(Zeros(n_vars), 10.0 * I)
coefficients
# Calculate all the mu terms.
= intercept .+ x * coefficients
mu return y ~ MvNormal(mu, σ² * I)
end;
= size(train)
n_obs, n_vars = linear_regression(train, train_label, n_obs, n_vars); m
Performing VI
First we define the initial variational distribution, or, equivalently, the family of distributions to consider. We’re going to use the same mean-field approximation as Turing will use by default when we call vi(m, advi)
, which we obtain by calling Variational.meanfield
. This returns a TransformedDistribution
with a TuringDiagMvNormal
as the underlying distribution and the transformation mapping from the reals to the domain of the latent variables.
= Variational.meanfield(m)
q0 typeof(q0)
MultivariateTransformed{TuringDiagMvNormal{Vector{Float64}, Vector{Float64}}, Stacked{Vector{Any}, Vector{UnitRange{Int64}}}} (alias for Bijectors.TransformedDistribution{DistributionsAD.TuringDiagMvNormal{Array{Float64, 1}, Array{Float64, 1}}, Stacked{Array{Any, 1}, Array{UnitRange{Int64}, 1}}, ArrayLikeVariate{1}})
= ADVI(10, 10_000) advi
ADVI{AutoForwardDiff{nothing, Nothing}}(10, 10000, AutoForwardDiff())
Turing also provides a couple of different optimizers:
TruncatedADAGrad
(default)DecayedADAGrad
as these are well-suited for problems with high-variance stochastic objectives, which is usually what the ELBO ends up being at different times in our optimization process.
With that being said, thanks to Requires.jl, if we add a using Flux
prior to using Turing
we can also make use of all the optimizers in Flux
, e.g. ADAM
, without any additional changes to your code! For example:
using Flux, Turing
using Turing.Variational
vi(m, advi; optimizer=Flux.ADAM())
just works.
For this problem we’ll use the DecayedADAGrad
from Turing:
= Variational.DecayedADAGrad(1e-2, 1.1, 0.9) opt
AdvancedVI.DecayedADAGrad(0.01, 1.1, 0.9, IdDict{Any, Any}())
= vi(m, advi, q0; optimizer=opt)
q typeof(q)
MultivariateTransformed{TuringDiagMvNormal{Vector{Float64}, Vector{Float64}}, Stacked{Vector{Any}, Vector{UnitRange{Int64}}}} (alias for Bijectors.TransformedDistribution{DistributionsAD.TuringDiagMvNormal{Array{Float64, 1}, Array{Float64, 1}}, Stacked{Array{Any, 1}, Array{UnitRange{Int64}, 1}}, ArrayLikeVariate{1}})
Note: as mentioned before, we internally define a update(q::TransformedDistribution{<:TuringDiagMvNormal}, θ::AbstractVector)
method which takes in the current variational approximation q
together with new parameters z
and returns the new variational approximation. This is required so that we can actually update the Distribution
object after each optimization step.
Alternatively, we can instead provide the mapping \(\theta \mapsto q_{\theta}\) directly together with initial parameters using the signature vi(m, advi, getq, θ_init)
as mentioned earlier. We’ll see an explicit example of this later on!
To compute statistics for our approximation we need samples:
= rand(q, 10_000); z
Now we can for example look at the average
= vec(mean(z; dims=2)) avg
12-element Vector{Float64}:
0.2247834146410775
-0.013829593516254388
0.3784414549269388
-0.1041372990835297
-0.07426907586649932
0.6099641438734019
0.004764607373310724
0.06714781450354046
-0.0805586847942715
0.1272379673416078
0.1837727654429606
-0.6038295824282297
The vector has the same ordering as the model, e.g. in this case σ²
has index 1
, intercept
has index 2
and coefficients
has indices 3:12
. If you forget or you might want to do something programmatically with the result, you can obtain the sym → indices
mapping as follows:
= bijector(m, Val(true));
_, sym2range sym2range
(intercept = UnitRange{Int64}[2:2], σ² = UnitRange{Int64}[1:1], coefficients = UnitRange{Int64}[3:12])
For example, we can check the sample distribution and mean value of σ²
:
histogram(z[1, :])
union(sym2range[:σ²]...)] avg[
1-element Vector{Float64}:
0.2247834146410775
union(sym2range[:intercept]...)] avg[
1-element Vector{Float64}:
-0.013829593516254388
union(sym2range[:coefficients]...)] avg[
10-element Vector{Float64}:
0.3784414549269388
-0.1041372990835297
-0.07426907586649932
0.6099641438734019
0.004764607373310724
0.06714781450354046
-0.0805586847942715
0.1272379673416078
0.1837727654429606
-0.6038295824282297
Note: as you can see, this is slightly awkward to work with at the moment. We’ll soon add a better way of dealing with this.
With a bit of work (this will be much easier in the future), we can also visualize the approximate marginals of the different variables, similar to plot(chain)
:
function plot_variational_marginals(z, sym2range)
= []
ps
for (i, sym) in enumerate(keys(sym2range))
= union(sym2range[sym]...) # <= array of ranges
indices if sum(length.(indices)) > 1
= 1
offset for r in indices
= density(
p :];
z[r, ="$(sym)[$offset]",
title=10,
titlefontsize="",
label="Density",
ylabel=1.5mm,
margin
)push!(ps, p)
+= 1
offset end
else
= density(
p first(indices), :];
z[="$(sym)",
title=10,
titlefontsize="",
label="Density",
ylabel=1.5mm,
margin
)push!(ps, p)
end
end
return plot(ps...; layout=(length(ps), 1), size=(500, 2000), margin=4.0mm)
end
plot_variational_marginals (generic function with 1 method)
plot_variational_marginals(z, sym2range)
And let’s compare this to using the NUTS
sampler:
= sample(m, NUTS(), 10_000); chain
┌ Info: Found initial step size
└ ϵ = 0.24003906249999996
plot(chain; margin=12.00mm)