Regression

In this tutorial, we will illustrate how to use the AdaptEllipticalSliceSampler.jl package, to perform Bayesian computation using the adaptive generalized elliptical slice sampler^[1] in regression settings. We consider four settings in this tutorial: (1) Bayesian linear regression, (2) Bayesian generalized linear regression (3) high-dimensional linear sparse regression^{[2] [3]}, and (4) generalized ReLU regression.

Download this tutorial as a Jupyter notebook

Bayesian Linear Regression

In this section, we consider using AGESS in the context of standard Bayesian linear regression. Letting $\mathbf{Y} \in \mathbb{R}^N$ be the response variable and $\mathbf{X} \in \mathbb{R}^{N \times D}$ be the set of covariates, we can specify our model as follows:

\[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 start by first generating data from our model.

import Random
import LogExpFunctions
using AdaptEllipticalSliceSampler
using Distributions
using Plots
using LinearAlgebra
using Turing, MCMCChains, StatsPlots

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)

([1.1840703709557707, -1.6437370448488724, -1.6181942515709322, -0.6108573527050607, 0.4212909718278306, 0.3366640790635556, -0.6178538133826769, -1.9858204398026311, -1.5941439357192193, 1.0309437423699197], [1.171685221802266 -1.177965631527714 … 0.029681358521803834 1.230714630625736; -0.30729971233631087 -0.5501257332792802 … 0.9497156839835412 -0.7604970817954496; … ; 0.9685819653962602 -2.142791626046223 … -0.7609766068757015 0.8187625833900045; 2.1012750843330763 0.08399913120329579 … -0.12142493814515284 -1.8514310605297297], [4.2509485214568254, 3.702067516660768, -3.653119910060239, -0.28567345139499334, -2.210986356964136, -0.8846381981280158, 5.936056534923045, 3.724379258679478, -1.4427757741863716, 2.033800483057324  …  5.339047977921442, 4.924010342402756, 1.1051788074805757, -7.155263299094693, 2.1151175613179456, 1.4371440995778353, 1.4916887518714412, -1.7523530126411566, 5.845578252046485, 2.7464899232331477])

We can see that we generated the data from our model under $\sigma = 0.5$ ($N = 1000$, $D = 10$). As of version 0.2.0 of AdaptEllipticalSliceSampler.jl, the package is now integrated into the AbstractMCMC.jl framework, allowing us to utilize the Turing ecosystem. Thus, we have two options: (1) we can write a function that calculates the log posterior density or (2) we can specify a Turing.jl model and directly use that to conduct Bayesian computation. While one can easily write a function for evaluating the log_posterior density (while making sure to transform $\sigma^2$ into an unconstrained space), let's look at how to use AGESS using a Turing.jl model. Luckily, Turing.jl (and the AdaptEllipticalSliceSampler.jl package) will automatically transform constrained parameters into an unconstrained space for sampling, then transform it back to the original space before returning the chain back to the user. Thus, in cases where we can use Turing.jl, we can not worry about remembering Jacobians!

@model function linear_regression(X::AbstractMatrix{Y}, y::AbstractVector{Y}) where {Y<:AbstractFloat}
    N, D = size(X)
    # Make sure dimensions conform
    @assert length(y) == N

    # Start with priors
    β ~ MvNormal(zeros(D), I)
    σ² ~ InverseGamma(1.0, 1.0)

    # Specify Likelihood
    y ~ MvNormal(X * β,  σ² * I)

end;

model = linear_regression(X, y)

DynamicPPL.Model{typeof(Main.linear_regression), (:X, :y), (), (), Tuple{Matrix{Float64}, Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext, false}(Main.linear_regression, (X = [1.171685221802266 -1.177965631527714 … 0.029681358521803834 1.230714630625736; -0.30729971233631087 -0.5501257332792802 … 0.9497156839835412 -0.7604970817954496; … ; 0.9685819653962602 -2.142791626046223 … -0.7609766068757015 0.8187625833900045; 2.1012750843330763 0.08399913120329579 … -0.12142493814515284 -1.8514310605297297], y = [4.2509485214568254, 3.702067516660768, -3.653119910060239, -0.28567345139499334, -2.210986356964136, -0.8846381981280158, 5.936056534923045, 3.724379258679478, -1.4427757741863716, 2.033800483057324  …  5.339047977921442, 4.924010342402756, 1.1051788074805757, -7.155263299094693, 2.1151175613179456, 1.4371440995778353, 1.4916887518714412, -1.7523530126411566, 5.845578252046485, 2.7464899232331477]), NamedTuple(), DynamicPPL.DefaultContext())

Now that we have constructed the model, we can use AGESS to conduct Bayesian computation by simply calling two functions: AGESSSampler() and sample().

n_MCMC = 10_000
sampler = AGESSSampler(model, n_MCMC)
results = sample(model, sampler, 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     = 1.49 seconds
Compute duration  = 1.49 seconds
parameters        = β[1], β[2], β[3], β[4], β[5], β[6], β[7], β[8], β[9], β[10], σ²
internals         = lp

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

We will start by discarding the initial part of the chain due to burn-in and will assess the results using describe().

# Discard first 2500 iterations due to burn-in
results = results[2501:end,:,:]
describe(results)

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

Iterations        = 2501:1:10000
Number of chains  = 1
Samples per chain = 7500
Wall duration     = 1.49 seconds
Compute duration  = 1.49 seconds
parameters        = β[1], β[2], β[3], β[4], β[5], β[6], β[7], β[8], β[9], β[10], σ²
internals         = lp

Summary Statistics

  parameters      mean       std      mcse    ess_bulk    ess_tail      rhat   ⋯
      Symbol   Float64   Float64   Float64     Float64     Float64   Float64   ⋯

        β[1]    1.1570    0.0163    0.0002   5014.8909   3138.4541    1.0005   ⋯
        β[2]   -1.6451    0.0162    0.0002   4895.1858   2854.1499    1.0006   ⋯
        β[3]   -1.6209    0.0162    0.0002   5432.4816   3154.7118    1.0004   ⋯
        β[4]   -0.5915    0.0154    0.0002   4884.7873   3107.5498    0.9999   ⋯
        β[5]    0.4329    0.0159    0.0002   4970.9530   3097.4966    0.9999   ⋯
        β[6]    0.3331    0.0155    0.0002   5273.8825   3225.9524    1.0001   ⋯
        β[7]   -0.5978    0.0156    0.0002   5218.5250   2950.9723    1.0043   ⋯
        β[8]   -1.9606    0.0162    0.0002   5046.8622   2717.6087    1.0005   ⋯
        β[9]   -1.5953    0.0162    0.0002   5237.6577   3366.3075    1.0005   ⋯
       β[10]    1.0340    0.0167    0.0002   5293.3710   3440.4684    1.0003   ⋯
          σ²    0.2576    0.0116    0.0002   4656.1634   3101.5824    1.0002   ⋯

                                                                1 column omitted

Quantiles

  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

        β[1]    1.1246    1.1461    1.1572    1.1681    1.1885
        β[2]   -1.6768   -1.6562   -1.6453   -1.6342   -1.6135
        β[3]   -1.6528   -1.6318   -1.6210   -1.6100   -1.5890
        β[4]   -0.6214   -0.6018   -0.5920   -0.5813   -0.5605
        β[5]    0.4016    0.4224    0.4330    0.4437    0.4645
        β[6]    0.3026    0.3229    0.3330    0.3437    0.3631
        β[7]   -0.6281   -0.6082   -0.5977   -0.5876   -0.5666
        β[8]   -1.9929   -1.9719   -1.9605   -1.9494   -1.9289
        β[9]   -1.6272   -1.6062   -1.5954   -1.5841   -1.5638
       β[10]    1.0011    1.0227    1.0341    1.0451    1.0672
          σ²    0.2359    0.2496    0.2572    0.2652    0.2812

Notice how the parameter names that we specified also follow in the subsequent evaluations of the chain–-this holds for calling plot(results), ess(results), autocorplot(results), etc. instead of using the default plotting function, let's visualize only the positive coefficients along with the true value (due to the dimension).

# Plot trace plot
plot(results[:, findall(β .> 0),:])
# Plot true values, one per coefficient's own subplot
true_vals = β[findall(β .> 0)]
hline!(reshape(true_vals, 1, :), subplot = reshape(1:2:2*length(true_vals), 1, :),
       color = :red, linestyle = :dash)

Similarly, we can view the trace plot of $\sigma^2$.

# Plot trace plot, don't forget to transform the transformed variables back
plot(results[:σ²], label = false)
# Plot true value
hline!([0.25], color = :red, label = false)

Lastly, we can plot the log posterior density at every iteration of the Markov chain to potentially detect convergence issues.

plot(results[:lp], legend = false)

High-Dimensional Sparse Linear Regression

In many modern applications, the number of covariates may be larger than the sample size ($N < D$), often requiring regularization in order to achieve good predictive properties. Here we consider the sparse regression setting, where sparsity is induced via continuous shrinkage through global-local shrinkage priors. Here, we consider the horseshoe prior^{[2] [3]}. Specifically, we consider the following model:

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

\[\beta_j \sim \mathcal{N}(0, \sigma^2\tau^2\lambda_j^2) \;\;\;\;\; p(\sigma^2) \propto \frac{1}{\sigma^2},\]

\[\tau \sim C^+(0,1) \;\;\;\;\; \lambda_j \sim C^+(0,1),\]

for $i = 1, \dots, N$ and $j = 1, \dots, D$, where $C^+$ denotes the half-Cauchy distribution. We will first start by generating data, where the design matrix has correlated covariates. We will consider the case where we have 25 observations and 25 covariates. Here we will generate a dataset where the probability of $\beta_j = 0$ is $0.9$ ($1 \le j \le 25$).

function gen_data_AR1(N::T, P::T; sparsity::Y = 0.8, ρ::Y = 0.2,
                      σ_sq::Y = 1.0) where {Y<:AbstractFloat, T<:Integer}
    Σ = ones(P, P)
    for i in 1:P
      for j in 1:P
        Σ[i,j] = ρ^(abs(i - j))
      end
    end
    Σ[diagind(Σ)] .= 1
    X = zeros(N, P)
    X .= rand(MultivariateNormal(zeros(P), Σ), N)'
    β = zeros(P)
    for i in 1:P
        if rand(Bernoulli(1 - sparsity)) == 1
          β[i] = (rand() * 3 + 1) * (-1)^i
        end
    end
    if sum(β) == 0
      β[1] = (rand() * 3 + 1) * (-1)^1
    end

    Y_obs = rand(MultivariateNormal(X * β, σ_sq * diagm(ones(N))))

    return X, Y_obs, β
end

Random.seed!(123)

N = 50
D = 25

X, y, β = gen_data_AR1(N, D, ρ = 0.7, sparsity = 0.9, σ_sq = 1.0)

([0.8082879284649668 -0.2355185009078168 … 2.1952434858659013 1.4455194595375302; 0.4694770042412599 -0.7617789966031957 … -0.8139537433814701 0.32963692676718603; … ; 0.8997726927030626 -0.13588404285163935 … -1.0202408659575293 -0.7975464687328946; -0.13619889472272229 -0.48904002253757134 … -1.0667748660476324 -0.24415843476965626], [-7.087267835848337, 5.205687034665962, -2.9744791427205506, 2.721384400176828, -9.042270062149427, -5.781077205943315, 5.9846966163908295, 3.54759428164667, 6.227456840612067, 2.4912291860230695  …  -11.072039215600881, 4.196282397376866, 3.2177888069892715, 3.9834053074485922, -9.839587460552332, -2.686279615079939, 0.2743214758522353, 0.6333020085530112, 9.189151700095069, 7.141090576424618], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 3.7980509144925407  …  2.9162735141224414, 0.0, 0.0, 0.0, 0.0, -1.176280495800773, 0.0, 0.0, 0.0, 0.0])

We visualize the correlation structure of the design matrix.

heatmap(cor(X))

We can also view the values of $\boldsymbol{\beta}$.

scatter(β,legend = false)

Specification of the Model

While we could directly specify a function evaluating the log posterior density (remembering to log-transform $\tau$ and $\lambda$), we will instead simply specify the model using the Turing.jl framework.

@model function HS_regression(X::AbstractMatrix{Y}, y::AbstractVector{Y}) where {Y<:AbstractFloat}
    N,D = size(X)
    @assert length(y) == N

    # Define Half-Cauchy distribution
    half_cauchy = truncated(Cauchy(0, 1); lower=0)

    # Priors
    τ ~ half_cauchy                 # Global Shrinkage param
    λ ~ filldist(half_cauchy, D)    # Local Shrinkage param
    log_σ² ~ Flat()                 # Flat prior to create Jeffrey's prior
    σ² = exp(log_σ²)
    β ~ MvNormal(zeros(D), σ² .* Diagonal((λ .* τ).^2))

    # Likelihood
    y ~ MvNormal(X * β, σ² * I)
end

model = HS_regression(X, y)

DynamicPPL.Model{typeof(Main.HS_regression), (:X, :y), (), (), Tuple{Matrix{Float64}, Vector{Float64}}, Tuple{}, DynamicPPL.DefaultContext, false}(Main.HS_regression, (X = [0.8082879284649668 -0.2355185009078168 … 2.1952434858659013 1.4455194595375302; 0.4694770042412599 -0.7617789966031957 … -0.8139537433814701 0.32963692676718603; … ; 0.8997726927030626 -0.13588404285163935 … -1.0202408659575293 -0.7975464687328946; -0.13619889472272229 -0.48904002253757134 … -1.0667748660476324 -0.24415843476965626], y = [-7.087267835848337, 5.205687034665962, -2.9744791427205506, 2.721384400176828, -9.042270062149427, -5.781077205943315, 5.9846966163908295, 3.54759428164667, 6.227456840612067, 2.4912291860230695  …  -11.072039215600881, 4.196282397376866, 3.2177888069892715, 3.9834053074485922, -9.839587460552332, -2.686279615079939, 0.2743214758522353, 0.6333020085530112, 9.189151700095069, 7.141090576424618]), NamedTuple(), DynamicPPL.DefaultContext())

Running AGESS

Similarly to before, once we have constructed the model, we can sample by simply calling two functions.

# Specify the number of MCMC iterations
n_MCMC = 100_000

# Run AGESS
sampler = AGESSSampler(model, n_MCMC)
results = sample(model, sampler, n_MCMC)

Chains MCMC chain (100000×53×1 Array{Float64, 3}):

Iterations        = 1:1:100000
Number of chains  = 1
Samples per chain = 100000
Wall duration     = 7.3 seconds
Compute duration  = 7.3 seconds
parameters        = τ, λ[1], λ[2], λ[3], λ[4], λ[5], λ[6], λ[7], λ[8], λ[9], λ[10], λ[11], λ[12], λ[13], λ[14], λ[15], λ[16], λ[17], λ[18], λ[19], λ[20], λ[21], λ[22], λ[23], λ[24], λ[25], log_σ², β[1], β[2], β[3], β[4], β[5], β[6], β[7], β[8], β[9], β[10], β[11], β[12], β[13], β[14], β[15], β[16], β[17], β[18], β[19], β[20], β[21], β[22], β[23], β[24], β[25]
internals         = lp

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

After running AGESS, we can view the trace plots of the non-zero coefficients.

# Discard initial 25% of chain due to burn-in
results = results[25_001:end,:,:]

# Get β parameters
β_samps = get(results, :β)
p = plot()
for i in findall(β .!= 0)
    p = plot!(β_samps.β[i][1:10:end], legend = false)
end

p
hline!(β[findall(β .!= 0)], line = :dash, color =:black)

We can also view the trace plots of the coefficients that are equal to zero.

p = plot()
for i in findall(β .== 0)
    p = plot!(β_samps.β[i][1:10:end], legend = false)
end

p
hline!([0.0], line = :dash, color =:black)

Next we will view the trace plot of $\sigma^2$.

plot(exp.(results[:log_σ²][1:10:end]), legend = false)
hline!([1], line = :dash, color =:black)

Lastly, we can plot the log pdf of the posterior at every iteration of the Markov chain to potentially detect convergence issues.

plot(results[:lp], legend = false)

Generalized ReLU Regression

As discussed in the main manuscript^[1], AGESS can handle non-differentiable target functions. In this section of the tutorial, we consider a generalized regression setting where the posterior distribution is not differentiable everywhere. Therefore, alternative samplers such as HMC^[4] are not suitable for these types of target distributions. Let $\mathbf{Y} \in \mathbb{R}^N$ be the response variable and $\mathbf{X} \in \mathbb{R}^{N \times D}$ be the covariates of interest. Consider the following model, inspired by density discontinuity modeling^[5]:

\[Y_i \sim Bernoulli(\Phi(\mu_i)) \;\;\;\;\;\Phi(z) = \frac{e^z}{1 + e^z} \;\;\;\;\; \mu_i = \max(0, \mathbf{x}_i'\boldsymbol{\beta}),\]

for $i = 1, \dots, N$, where $\boldsymbol{\beta} \sim \mathcal{N}_{D}(\mathbf{0}, \mathbf{I})$. We will first start by specifying a function to generate data from our model. We will consider the simple case of when $D = 2$, allowing for easy visualization of the target distribution.

function generate_data(N::T, P::T, ν::Y = 6.0) where {Y<:AbstractFloat, T<:Integer}
    β = randn(P) * (2 * log(P))^(1.0 / 4)
    β .*= sqrt(rand(Gamma(ν/2, 2/ ν)))
    x::Matrix{typeof(ν)} = randn(N, P) .+ randn() * 0.5
    μ = zeros(N)
    y::Vector{typeof(N)} = zeros(Int64, N)

    for i in 1:N
        μ[i] = max(0, dot(x[i,:], β))
        y[i] = rand(Binomial(1, LogExpFunctions.logistic(μ[i])), 1)[1]
    end


    return β, x, μ, y
end

Random.seed!(123)

D = 2
N = 1000
β, x, μ, y = generate_data(N, D)

([-0.3611372440159916, -0.8183513546797343], [0.25957106302208094 -0.9556035799933374; 0.20180178045009275 -0.4422196885740821; … ; 0.9358808315839681 -0.06709585879776; -0.42646210250593763 0.010337450394495056], [0.6882787058982559, 0.28901294238138364, 0.2554305477921972, 0.9893159580995416, 0.0, 1.19171914741668, 0.0, 1.2432768875517346, 0.0, 0.27356355703349333  …  0.32521446983604785, 0.0, 0.0, 0.0, 0.4929140191342167, 0.9870706892189209, 0.0, 1.409099581139658, 0.0, 0.14555168184199005], [1, 1, 1, 1, 1, 1, 1, 1, 1, 0  …  1, 1, 0, 1, 0, 1, 0, 1, 1, 0])

Specification of the Log Posterior Density

Now that we have generated the data, we can construct a function which evaluates the posterior log pdf. As in the previous sections, it is important that we specify a function that is efficiently implemented. We will first start by constructing a function that evaluates the log likelihood, and then a function that evaluates the log posterior density.

function log_likelihood(β::AbstractVector{Y}, x::Matrix{Y},
                        y::Vector{T}) where {Y <:AbstractFloat, T<:Integer}
    log_lik::Float64 = 0.0
    z::Float64 = 0.0
    for i in eachindex(y)
        @views z = dot(x[i,:], β)
        if z < 0.0
            z = 0.0
        end
        log_lik -= log1p(exp(-(sign(y[i] - 0.5) * z)))
    end

    return log_lik
end

function log_posterior_ReLU(β::AbstractVector{Y}, x::Matrix{Y},
                            y::Vector{T}) where {Y <:AbstractFloat, T<:Integer}
    log_lik::Float64 = log_likelihood(β, x, y) - 0.5 * dot(β, β)
    return log_lik
end

log_posterior_ReLU (generic function with 1 method)

Running AGESS

Now that we have specified a function to efficiently evaluate the log posterior density, we can use the AGESS function to generate samples from the posterior distribution.

# Specify the dimension of the target distribution
P = D
# Specify the number of MCMC iterations
n_MCMC = 10_000

# Let's specify the param_names this time
param_names = [string("β_",i) for i in 1:D]
# Run AGESS
results = AGESS(β -> log_posterior_ReLU(β, x, y), n_MCMC, P; param_names = param_names)

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

Iterations        = 1:1:10000
Number of chains  = 1
Samples per chain = 10000
Wall duration     = 1.28 seconds
Compute duration  = 1.28 seconds
parameters        = β_1, β_2
internals         = lp

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

Discard the initial 2500 samples to burnin

results = results[2501:end,:,:]
describe(results)

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

Iterations        = 2501:1:10000
Number of chains  = 1
Samples per chain = 7500
Wall duration     = 1.28 seconds
Compute duration  = 1.28 seconds
parameters        = β_1, β_2
internals         = lp

Summary Statistics

  parameters      mean       std      mcse    ess_bulk    ess_tail      rhat   ⋯
      Symbol   Float64   Float64   Float64     Float64     Float64   Float64   ⋯

         β_1   -0.3665    0.0946    0.0012   6078.8333   4057.4212    1.0001   ⋯
         β_2   -0.8262    0.1063    0.0013   6274.1341   4111.5716    0.9999   ⋯

                                                                1 column omitted

Quantiles

  parameters      2.5%     25.0%     50.0%     75.0%     97.5%
      Symbol   Float64   Float64   Float64   Float64   Float64

         β_1   -0.5564   -0.4274   -0.3656   -0.3052   -0.1828
         β_2   -1.0420   -0.8947   -0.8252   -0.7533   -0.6242

Using the output of AGESS, we can visualize the samples from the target distribution.

scatter(results[:β_1], results[:β_2], alpha = 0.1,
        legend = false)
scatter!([β[1]], [β[2]], color = "red")

Additionally, we can plot the log pdf of the posterior at every iteration of the Markov chain to potentially detect convergence issues.

plot(results[:lp], legend = false)

Conclusion

Adaptive generalized elliptical slice sampling^[1] can be utilized in many generalized regression settings, including target distributions that are non-differentiable. While limited by computational resources in this tutorial, AGESS can be used to sample from relatively high-dimensional target distributions. Comparisons between AGESS and alternative samplers can be found in the main manuscript^[1].

# Free the sampled chains before the next tutorial's page is built.
results = β_samps = nothing
GC.gc()

This page was generated using Literate.jl.