Performance Tips

Julia is a dynamic programming language that allows for high performance computing. However, Julia's optional typing can lead to slow performance if left unspecified. Since the AGESS function essentially only requires the user to specify a function evaluating the log target distribution, it is paramount that the user specifies an efficient implementation of this function, as this function will constantly be called in the AGESS function. Here, we will give a quick example of 3 evaluations of the same target distribution; each leading to different computational costs. While a full guide to writing performance oriented code in Julia is out of the scope of this documentation, here are some useful resources:

Download this page as a Jupyter notebook

Use of Turing Models

As of version 0.2.0, users are able to specify a Turing.jl model instead of supplying a function that efficiently evaluates the log posterior density. Thus, we refer users to the Turing.jl documentation for tips on improving performance.

Performance Tips for Turing.jl

Linear Regression

Consider the standard model for linear regression (see the "Regression" tutorial for more details):

\[Y_i \sim \mathcal{N}(\mathbf{x}_i' \boldsymbol{\beta}, \sigma^2),\]

\[\boldsymbol{\beta} \sim \mathcal{N}(\mathbf{0}, \mathbf{I}),\]

\[\sigma^2 \sim \text{Inv-Gamma}(1,1).\]

Let's consider a simple implementation of this function, where we do not specify any types:

import Random
import LogExpFunctions
using BenchmarkTools
using AdaptEllipticalSliceSampler
using Distributions
using Plots
using LinearAlgebra

function lm_log_posterior_1(Param, X, y)
    P = length(Param)
    N = length(y)
    lpdf = logpdf(MvNormal(X * Param[1:(P-1)],  exp(Param[P]) * diagm(ones(N))), y)
    lpdf += logpdf(MvNormal(zeros(P-1),  diagm(ones(P-1))), Param[1:(P-1)])
    lpdf += logpdf(InverseGamma(1, 1), exp(Param[P])) + Param[P]

    return lpdf
end

lm_log_posterior_1 (generic function with 1 method)

Let's generate some synthetic data and benchmark how long it takes to run lm_log_posterior_1.

Random.seed!(123)

function generate_data(N::T, D::T) where {T<:Integer}
    β = randn(D) * (2 * log(D))^(1.0 / 4)
    x = randn(N, D)
    y = zeros(Float64, N)
    for i in 1:N
        y[i] = randn() * 0.5 + dot(x[i,:], β)
    end

    return β, x, y
end

# Generate data with 1000 observations and 10 covariates
D = 10
β, X, y = generate_data(1000, D)

# Benchmark function
Param = ones(D + 1)
@benchmark lm_log_posterior_1($Param, $X, $y)

BenchmarkTools.Trial: 402 samples with 1 evaluation per sample.
 Range (min … max):   7.863 ms … 326.590 ms  ┊ GC (min … max):  0.00% … 97.21%
 Time  (median):      8.101 ms               ┊ GC (median):     0.00%
 Time  (mean ± σ):   12.440 ms ±  31.044 ms  ┊ GC (mean ± σ):  30.70% ± 14.26%

  █ ▄▁   ▂                                                      
  █▇██▆▆▇█▇▅▁▅▁▁▁▁▄▁▁▁▁▁▁▁▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▅ ▆
  7.86 ms       Histogram: log(frequency) by time      46.9 ms <

 Memory estimate: 22.92 MiB, allocs estimate: 39.

Let's see if we can improve on this by pre-allocating some of our variables and specifying the type of variables.

function lm_log_posterior_2(Param::AbstractVector{Y}, X::AbstractMatrix{Y},
                            y::AbstractVector{Y}, μ::AbstractVector{Y},
                            μ_0::AbstractVector{Y}, Σ_I_N::AbstractMatrix{Y},
                            Σ_I_P::AbstractMatrix{Y}) where {Y<:AbstractFloat}
    P = length(Param)
    @views μ .= X * Param[1:(P-1)]
    lpdf = logpdf(MvNormal(μ, exp(Param[P]) * Σ_I_N), y)
    @views lpdf += logpdf(MvNormal(μ_0,  Σ_I_P), Param[1:(P-1)])
    lpdf += logpdf(InverseGamma(1, 1), exp(Param[P])) + Param[P]

    return lpdf
end

# Pre-allocate parameters
μ = zeros(1000)
Σ_I_N = diagm(ones(1000))
Σ_I_P = diagm(ones(D))
μ_0 = zeros(D)
@benchmark lm_log_posterior_2($Param, $X, $y, $μ, $μ_0, $Σ_I_N, $Σ_I_P)

BenchmarkTools.Trial: 508 samples with 1 evaluation per sample.
 Range (min … max):  7.483 ms … 312.213 ms  ┊ GC (min … max):  0.00% … 97.32%
 Time  (median):     7.751 ms               ┊ GC (median):     0.00%
 Time  (mean ± σ):   9.854 ms ±  19.135 ms  ┊ GC (mean ± σ):  19.55% ± 10.93%

  ▆█▅            ▆▆▃                                           
  ████▁▅▅▄▁▁▁▁▄▅▇████▆▅▄▄▇▅▆▁▄▁▁▁▄▁▅▁▁▁▅▇█▅▁▁▁▄▅▁▁▁▁▁▁▁▁▁▁▄▁▆ ▇
  7.48 ms      Histogram: log(frequency) by time      13.7 ms <

 Memory estimate: 15.28 MiB, allocs estimate: 23.

We can see that there was a modest improvement in performance. We can see that we are allocating less memory. However, it is slow to actually construct these multivariate distributions and evaluate the log pdf of these distributions. We can just perform the calculations ourselves and get significantly better performance. Since Julia is compiled just-in-time, we should feel free to use for-loops as we please! (Unlike R)

function lm_log_posterior_3(Param::AbstractVector{Y}, X::AbstractMatrix{Y},
                            y::AbstractVector{Y}, ph::AbstractVector{Y}) where {Y<: AbstractFloat}
    P = length(Param)
    # Normal Likelihood
    @views ph .= X * Param[1:P-1]
    ph .-= y
    lpdf = -0.5 * (1 / exp(Param[P])) *  norm(ph)^2 -
            (0.5 * length(y) * Param[P])

    # Priors
    # Std Normal prior on coefficients
    @views lpdf += -0.5 * norm(Param[1:P-1])^2

    # IG(1,1) prior on scale parameter (log-transformed)
    lpdf += -1 * Param[P]  -  (1 / exp(Param[P]))

    return lpdf
end

ph = zeros(1000)
@benchmark lm_log_posterior_3($Param, $X, $y, $ph)

BenchmarkTools.Trial: 10000 samples with 9 evaluations per sample.
 Range (min … max):  2.322 μs … 171.322 μs  ┊ GC (min … max):  0.00% … 91.51%
 Time  (median):     2.512 μs               ┊ GC (median):     0.00%
 Time  (mean ± σ):   2.926 μs ±   7.072 μs  ┊ GC (mean ± σ):  11.67% ±  4.75%

     ▂▇█▆▄▁                                                    
  ▂▃▆██████▆▄▃▃▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▂▁▂▁▁▂▂▂▂▂▂▂▂▂▂▂▂▁▂▂▂▂▂ ▃
  2.32 μs         Histogram: frequency by time        4.17 μs <

 Memory estimate: 7.88 KiB, allocs estimate: 3.

We can see that we have a 1000-fold speed-up by just efficiently evaluating the log posterior density. This directly translates into a similar magnitude increase in the effective sample size per second achieved by AGESS.

Key takeaway

It is paramount to write efficient functions that evaluate the log posterior density when using AGESS.

Tips:

Packages like JET.jl can help catch inefficiencies in coding.
Pre-allocate variables (especially for intermediate computations)
@views can help reduce allocating new arrays when doing computations on subarrays

Block Updates

Writing an efficient log_posterior is crucial, but for high-dimensional target distributions we still perform 1-d updates during the burn-in stage, as well as randomly throughout (as controlled by single_step_prop). However, if your model has structure – such as a hierarchical model where we have conditional independence between blocks of parameters and do not need to compute the entire posterior to calculate the conditional density – we can significantly reduce the computational burden of these 1-d updates by using AGESSSampler's blocks. Here, we will provide an example of fitting a hierarchical model using these blocks (using a direct specification of the log target density, and separately using the Turing.jl ecosystem).

Consider a simple hierarchical model:

\[Y_{ig} \sim \mathcal{N}(\theta_g, 1), \qquad \theta_g \sim \mathcal{N}(\mu, 1).\]

We can simulate data as follows.

G = 10
n_g = 200
μ_true = 1.0
θ_true = μ_true .+ 0.5 .* randn(G)
data = [θ_true[g] .+ randn(n_g) for g in 1:G]

10-element Vector{Vector{Float64}}:
 [0.9013990275798104, 1.7307604482841636, -0.623100376503377, -0.13626100150521725, 2.3776935605153544, 1.5178728186576, 1.477497849118865, 2.239571689283943, 2.3690365722971816, -0.301515826607999  …  -0.16035580728950938, 2.790784059742017, 1.2119798149388559, -0.28000800914663115, -0.09858246591906183, 0.4575823058461617, 0.9470118576737067, 2.629708964219551, 0.16697007806513375, 1.2056572080848516]
 [-0.2895687139352876, 1.93968277845019, 0.4273493223250098, 0.3949525468384877, 0.667677415236551, -1.0942290547841371, 1.4310211762876364, 0.09382217818706776, -1.3510997237165698, -0.35247207611343623  …  0.1487673759140855, -0.5421755704049742, -0.7766918853168839, -0.7704489581622468, -0.963660293797732, -0.86011779784858, -0.3971408475574375, 0.7838050640559421, 0.5088274741608093, -1.2528090209821592]
 [-1.2948081769800162, 0.48685354802263714, 0.4263755143914767, 1.1748975472591585, -0.21515557053532808, 0.09632389451171897, 0.5891319124903814, -0.9140941207875157, -2.2365878774921466, 0.7723878401557069  …  1.203064638212455, 0.3372363311075946, -0.28949900948721297, 2.205152766314879, 1.2923008944670742, -0.9321550977366384, 3.576398697273053, 0.5976545197076552, 3.1056782071418176, 2.095882758653833]
 [3.8503460127098585, 0.880844685454626, 2.206836650214337, 1.8423658651058457, 0.1480994792823982, 1.4282533496319747, 2.3356999479513045, 0.7212006933358586, 2.77986461609928, 0.8593652469360468  …  2.292634093679932, 1.870178497961332, 1.289676995596944, 0.830432613388464, 1.0516599690628596, 2.0231416238297486, 2.4203898466657527, 3.3546169274048694, -0.005859359795444163, 2.6048664594752697]
 [-0.3626684322755145, -0.029564961268190304, 1.9501633720055718, 1.894668583577091, 1.4916487639951155, -0.5900977287577189, 0.9259378132992652, -0.1142665852342668, 0.6819602093361499, 1.086703119350226  …  0.41161074894409344, 1.6633974312699018, 0.09283895994066371, 0.5331190565561723, 2.238614903236373, 1.3739724041372503, -0.7200596787708362, 1.0459714883497535, 2.030651815407884, 0.8654255534897668]
 [0.5977450963162596, 1.836370442550111, 0.8564207576815919, 0.62878251171116, 0.9894110385224424, -2.049685798630513, 1.4860898455231764, 1.397702139385692, 0.6980894470544547, -0.9701654854571562  …  0.36054537137208303, 2.1771153415602593, -0.877038136718554, 0.0006603858513627991, -0.32903274181027686, 0.002424045741097669, -0.45544093326951574, -0.03375981636478087, 0.004024223310775288, 0.8272286936165316]
 [1.6121439856714603, 0.1925965909535161, -0.18216793599663383, 1.1648629790888203, 0.1846768246912467, 0.5638746490321391, 1.5072326874105646, 1.1635806679422702, 0.4303583479287658, 1.2629145168394098  …  1.528717807119023, 0.6789926360165481, 0.49633866939861354, 2.903162254791212, 0.6343836837930593, 0.48140717865674937, 0.49559608141161016, -0.40952400890913754, 1.113153770450085, -0.08956967309535702]
 [1.6329406331461342, 0.6192483872070771, 1.393648370407991, 0.26488483352855774, 2.866527015819166, 2.7139137605229253, 1.7719598616368009, 0.10623532494421428, 1.9505562825735772, 1.2783080463418335  …  0.8773324574756023, 1.5555650573529873, -0.7076037103098283, 0.49460788786734355, 0.5651157330310004, 1.3334725562573277, 0.5932483637806119, 3.0505066844125164, 1.2268209974440247, 1.388093005580309]
 [1.6201888600644427, 1.1246082469322847, 2.097238244257407, 0.8014206027736686, 0.13645015350504164, 2.052482096021071, -0.26616252263691653, 1.9498428109012775, 0.1447729342709031, 0.5905778626689122  …  1.5933034871256606, 1.7528424662085331, -0.021936575230382283, 0.7627075928006294, 1.61725793763542, 2.1182990738203267, 0.8429952005380029, 2.0321106523025145, 2.849380085321948, 0.7736069933552765]
 [-0.04176972235630627, 2.2944855096113597, 0.7714979425121953, 1.9490810432306551, 0.03337941903078778, 0.19647146125090587, 1.19532101732616, 2.18717052487714, 1.4532337909560795, 1.8381652965528872  …  -0.19456791883100544, 0.6228428458276778, 0.9781768432987676, 0.950620942144146, 1.7883235076918789, -0.3322375720398192, 1.6774304369228714, 1.1864766021297555, 1.6512230030361175, 0.901715245362382]

We will first start by explicitly specifying the blocks and the log target density.

# Param layout: Param[1] = μ, Param[1+g] = θ_g for g in 1:G.
function hier_log_posterior(Param::AbstractVector{Y}, data) where {Y<:AbstractFloat}
    μ = Param[1]
    lpdf = -0.5 * μ^2
    for g in eachindex(data)
        θ_g = Param[1 + g]
        lpdf += -0.5 * (θ_g - μ)^2
        for y in data[g]
            lpdf += -0.5 * (y - θ_g)^2
        end
    end
    return lpdf
end

hier_log_posterior (generic function with 1 method)

The full hier_log_posterior above touches every group's data on every call, even when only θ_g for one group is changing. A block's conditional gets to see the same full parameter vector, but only needs to return the parts of the density that actually depend on its own block. Here, updating θ_g only needs group g's own data:

function group_conditional(g::Integer, Param::AbstractVector{Y}, data) where {Y<:AbstractFloat}
    μ = Param[1]
    θ_g = Param[1 + g]
    lpdf = -0.5 * (θ_g - μ)^2
    for y in data[g]
        lpdf += -0.5 * (y - θ_g)^2
    end
    return lpdf
end

group_conditional (generic function with 1 method)

For $G=10$ groups of 200 observations each, that's roughly a 10-fold reduction in the amount of data touched per call:

Param = vcat(μ_true, θ_true)
@benchmark hier_log_posterior($Param, $data)

BenchmarkTools.Trial: 10000 samples with 10 evaluations per sample.
 Range (min … max):  1.865 μs …  16.690 μs  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     1.869 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   1.891 μs ± 219.208 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

  █                                                           ▁
  █▇▅▄▄▃▄▃▃▃▁▁▄▄▄▁▃▁▁▁▁▄▁▁▄▁▁▃▁▁▁▁▁▁▁▁▁▁▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▇█ █
  1.87 μs      Histogram: log(frequency) by time      2.68 μs <

 Memory estimate: 0 bytes, allocs estimate: 0.

@benchmark group_conditional(3, $Param, $data)

BenchmarkTools.Trial: 10000 samples with 719 evaluations per sample.
 Range (min … max):  174.388 ns … 236.341 ns  ┊ GC (min … max): 0.00% … 0.00%
 Time  (median):     175.335 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   176.960 ns ±   4.093 ns  ┊ GC (mean ± σ):  0.00% ± 0.00%

     █▅   ▃  ▂                                    ▁▄▃▁          ▁
  ▇▁▃██▃▁▅█▄▁█▅▇█▇▆▅▄▄▁▄▄▄▁▃▁▃▄▁▄▃▁▁▄▃▁▃▃▄▁▃▃▁▁▁▁▄█████▇▇▆▆▆▆▆▆ █
  174 ns        Histogram: log(frequency) by time        189 ns <

 Memory estimate: 0 bytes, allocs estimate: 0.

Wiring this into AGESSSampler just means building one AGESSBlock per group with its conditional, plus a block for μ left on the default (full hier_log_posterior) path: μ's own prior doesn't touch any θ_g, but each θ_g's prior depends on μ, so μ's full conditional needs every group's θ_g and so isn't separable the same way.

blocks = [AGESSBlock([1 + g]; conditional = p -> group_conditional(g, p, data)) for g in 1:G]
push!(blocks, AGESSBlock([1]))

n_MCMC = 10_000
chain = AGESS(p -> hier_log_posterior(p, data), n_MCMC, 1 + G; blocks = blocks)

Chains MCMC chain (10000×12×1 Array{Float64, 3}):

Iterations        = 1:1:10000
Number of chains  = 1
Samples per chain = 10000
Wall duration     = 0.77 seconds
Compute duration  = 0.77 seconds
parameters        = param_1, param_2, param_3, param_4, param_5, param_6, param_7, param_8, param_9, param_10, param_11
internals         = lp

Use `describe(chains)` for summary statistics and quantiles.

For Turing.jl models we can also utilize AGESSSampler's blocks. Using the same setup, we can set up our Turing.jl model as follows:

using Turing

@model function local_model(data_g, μ)
    θ ~ Normal(μ, 1)
    data_g .~ Normal(θ, 1.0)
end

@model function full_model(data)
    G = length(data)
    μ ~ Normal(0.0, 1.0)
    θ ~ filldist(Normal(μ, 1.0), G)
    for g in 1:G
        for i in eachindex(data[g])
            data[g][i] ~ Normal(θ[g], 1.0)
        end
    end
end

# Note: We can also write the model using the local model, but we will get warnings for growable
# arrays. Note that these warnings do not affect the correctness of the sampling scheme.
# @model function full_model(data, G)
#    μ ~ Normal(0, 1)
#    θ = Vector{Float64}(undef, G)
#    for g in 1:G
#        θ[g] ~ to_submodel(local_model(data[g], μ))
#    end
# end

# μ at index 1; group g's θ at index 1 + g
blocks = [AGESSBlock([1 + g], ctx -> local_model(data[g], ctx[1]), [1]) for g in 1:G]
push!(blocks, AGESSBlock([1]))  # μ stays on the default (full log_posterior) path

agessB = AGESSSampler(full_model(data), n_MCMC; blocks = blocks)
chain_turing = sample(full_model(data), agessB, n_MCMC)

Chains MCMC chain (10000×12×1 Array{Float64, 3}):

Iterations        = 1:1:10000
Number of chains  = 1
Samples per chain = 10000
Wall duration     = 2.15 seconds
Compute duration  = 2.15 seconds
parameters        = μ, θ[1], θ[2], θ[3], θ[4], θ[5], θ[6], θ[7], θ[8], θ[9], θ[10]
internals         = lp

Use `describe(chains)` for summary statistics and quantiles.

Key takeaway

If your model has natural conditional independence structure where part of the likelihood does not depend on a block of parameters, AGESSSampler's blocks keyword can significantly reduce computational costs – particularly in high-dimensional settings. However, correctness depends on conditional actually being a valid restriction of log_posterior. A block left without a conditional always falls back to the full log_posterior, so when in doubt, leave it out rather than risk a subtly wrong conditional.

This page was generated using Literate.jl.