Performance Tips

Julia is a dynamic programming language that allows for high performance computing. However, Julia features optional typing can lead to slow performance. 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 completely 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:

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: 362 samples with 1 evaluation per sample.
 Range (min … max):  11.741 ms … 262.848 ms  ┊ GC (min … max): 0.00% … 94.87%
 Time  (median):     12.357 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   13.819 ms ±  13.364 ms  ┊ GC (mean ± σ):  7.57% ±  7.89%

   ▁▅▅█▆▇ ▁                                                     
  ▄████████▇▇▅▄▃▂▂▅▄▄▃▄▅▄▃▄▂▂▂▂▁▁▁▁▂▁▁▁▂▃▃▃▃▃▃▃▄▃▃▄▂▄▃▃▁▁▄▂▁▁▂ ▃
  11.7 ms         Histogram: frequency by time         17.1 ms <

 Memory estimate: 22.92 MiB, allocs estimate: 37.

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: 420 samples with 1 evaluation per sample.
 Range (min … max):  11.038 ms …  15.037 ms  ┊ GC (min … max): 0.00% … 9.72%
 Time  (median):     11.733 ms               ┊ GC (median):    0.00%
 Time  (mean ± σ):   11.928 ms ± 625.086 μs  ┊ GC (mean ± σ):  1.45% ± 3.57%

         ▁▆▇▄▆▇▅█▄▂▂▁▄ ▂ ▂                                      
  ▃▃▃▄▄▄▆███████████████▆█▄▅▄▄▃▃▃▃▃▁▃▁▁▃▁▃▄▃▃▃▃▄▃▆▁▃▅▃▄▃▃▄▃▃▃▃ ▄
  11 ms           Histogram: frequency by time         13.7 ms <

 Memory estimate: 15.28 MiB, allocs estimate: 21.

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.239 μs … 408.487 μs  ┊ GC (min … max): 0.00% … 92.55%
 Time  (median):     2.821 μs               ┊ GC (median):    0.00%
 Time  (mean ± σ):   3.198 μs ±   9.767 μs  ┊ GC (mean ± σ):  6.94% ±  2.26%

   ▁▄▆▇▇▇▆▅▅▆▇██▇▇▆▅▄▃▂▁▁   ▁         ▁▁▁   ▁▂▂▁▁▁▂▁▁▁▂▃▂▂▁ ▁ ▃
  ▅██████████████████████████▇█▇██▇█▇█████▇██████████████████ █
  2.24 μs      Histogram: log(frequency) by time      4.93 μ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. Therefore, 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