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 three settings in this tutorial: (1) Bayesian linear regression, (2) high-dimensional linear sparse regression^{[2] [3]}, and (3) generalized ReLU regression.

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

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$). Next, we can write a function that calculates the log posterior density. When conducting inference on variables that have positive support (i.e., $\sigma^2 > 0$), it is often useful to transform the variables to remove these constraints. Thus, we will transform $\sigma^2$ using a log transformation to remove the positivity constraints (don't forget the Jacobian term when calculating the log pdfs). Additionally, we will concatenate all of our variables into one vector $\text{Param} = [\boldsymbol{\beta}, \log(\sigma^2)]$.

function log_posterior(Param::AbstractVector{Y}, X::AbstractMatrix{Y},
                       y::AbstractVector{Y}) where {Y<: AbstractFloat}
    P = length(Param)
    ## Normal Likelihood
    @views lpdf = -0.5 * (1 / exp(Param[P])) *  norm(X * Param[1:P-1] - y)^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

log_posterior (generic function with 1 method)

Now that we have specified a function to efficiently evaluate the log posterior density, we can simply use the AGESS function to draw samples from the posterior distribution. When calling the function AGESS, we will specify an anonymous function with $\text{Param}$ as the only input to the function log_posterior, using the global parameter values for $\mathbf{X}$ and $\mathbf{Y}$.

### Specify the dimension of the target distribution
P = D + 1
### Specify the number of MCMC iterations
n_MCMC = 10000

### Run AGESS
results = AGESS(Param -> log_posterior(Param, X, y), n_MCMC, P)

AdaptEllipticalSliceSampler.MCMC_output([0.0 0.0 … 0.0 0.0; 0.03241642672644927 -0.6719675758389738 … 0.009835452313085797 0.3969716430768995; … ; 1.1556682359812287 -1.6205079265492548 … 1.021595069893918 -1.412786403395197; 1.1744977885537387 -1.64746509437571 … 1.0111450774183266 -1.335370541595082], [-7992.841473690834, -1949.4722301854213, -1008.8557014165999, -710.5456631899677, -677.6096870648968, -598.8699641483735, -522.893676104241, -449.4479842703263, -385.8841850644391, -357.754382688629  …  168.83877222469155, 167.51442431889822, 167.33198224783146, 169.0453491271659, 169.44604798185804, 170.30868686667404, 168.782404930476, 168.4792338544405, 172.39195407768702, 172.92520073697844], AGESS_MCMC_params(Main.var"#2#3"(), [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0 0.0 … 0.0 0.0; 0.0 1.0 … 0.0 0.0; … ; 0.0 0.0 … 1.0 0.0; 0.0 0.0 … 0.0 1.0], true, 6.0, 0.5, 0.05, 0.5, 0.05, 10000, 11), [0.00022191954814292705 -6.995690300311635e-6 … -1.1945116136814807e-5 3.9945388194638975e-5; -6.995690300311635e-6 0.00027235770259251437 … -1.1198666452464954e-6 2.80173781929415e-5; … ; -1.1945116136814807e-5 -1.1198666452464954e-6 … 0.0002707816617225753 -3.432108283606644e-5; 3.9945388194638975e-5 2.80173781929415e-5 … -3.432108283606644e-5 0.0020525290409753624], [1.1568496621534863, -1.6442686375027107, -1.6202843878329538, -0.5923767470966812, 0.43235908559143427, 0.3345463659971098, -0.5975752904329056, -1.9606585534960252, -1.5947052188853121, 1.033629061680873, -1.3547585081830846])

We can visualize some of the trace plots of the $\beta$ parameters using the output of AGESS.

### Plot trace plot
plot(results.samps[2501:n_MCMC, findall(β .> 0)])
### Plot true values
hline!(β[findall(β .> 0)])

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

### Plot trace plot, don't forget to transform the transformed variables back
plot(exp.(results.samps[2501:n_MCMC, D+1]))
### Plot true value
hline!([0.25])

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

plot(results.l_pdf[1:n_MCMC], 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 Log Posterior Density

When using the AdaptEllipticalSliceSampler.jl, it is crucial that we construct a function that can efficiently evaluate the log posterior density. We will first start by constructing two functions that evaluate the priors on $\boldsymbol{\beta}$ and $\boldsymbol{\lambda}$. Since $\lambda_j$ has positive support, we will take a log transformation of the $\lambda_j$ parameters ($j = 1, \dots, D$). Similarly, we will take a log transformation of $\tau$ and $\sigma$.

function prior_β(β::AbstractVector{Y}, τ::Y, λ::AbstractVector{Y},
                 σ::Y) where {Y<:AbstractFloat}
  lpdf::Float64 = 0.0
  for i in eachindex(β)
      @views lpdf += logpdf(Normal(0.0, exp(σ) * exp(τ) * exp(λ[i])), β[i])
  end
  return lpdf
end

function prior_λ(λ::AbstractVector{Y}) where {Y<:AbstractFloat}
  lpdf::Float64 = 0.0
  cauchy_d = Cauchy(0.0,1.0)
  for i in eachindex(λ)
      @views lpdf += logpdf(cauchy_d, exp(λ[i])) + λ[i]
  end
  return lpdf
end

prior_λ (generic function with 1 method)

Using these two functions, we can specify a function evaluating the log pdf of the target distribution of interest. We will concatenate all the parameters into one vector $\text{Param} = [\boldsymbol{\beta}, \log(\boldsymbol{\lambda}), \log(\tau), \log(\sigma)]$.

function log_posterior_HD(Param::AbstractVector{Y}, X::AbstractMatrix{Y},
                          y::AbstractVector{Y}, D::T) where {Y<:AbstractFloat, T<:Integer}
  lpdf::Float64 = 0.0
  ## Likelihood
  for i in eachindex(y)
    @views lpdf += logpdf(Normal(dot(X[i,:], Param[1:D]), exp(Param[2*D + 2])), y[i])
  end
  ## Prior
  lpdf += prior_β(Param[1:D], Param[2*D + 1], Param[(D + 1):2*D], Param[2*D + 2]) +
            prior_λ(Param[(D + 1):2*D]) + logpdf(Cauchy(0.0,1.0), exp(Param[2*D + 1])) + Param[2*D + 1]

  return lpdf
end

log_posterior_HD (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. Similarly to before, we will specify an anonymous function with $\text{Param}$ as the only input to the function log_posterior_HD, using the global parameter values for $\mathbf{X}$, $\mathbf{y}$, and $D$.

### Specify the dimension of the target distribution
P = 2*D + 2
### Specify the number of MCMC iterations
n_MCMC = 100000

### Run AGESS
results = AGESS(Param -> log_posterior_HD(Param, X, Y, D), n_MCMC, P)

AdaptEllipticalSliceSampler.MCMC_output([0.0 0.0 … 0.0 0.0; -0.14657331235852855 -0.2379737601887019 … -0.14991626721740858 0.14107604637413743; … ; 0.011816106371171246 0.021863132565760113 … -2.366654365888479 0.14142146761127827; 0.02112571653968324 0.025629505150798995 … -2.4428576405523557 0.16655377809066302], [-751.5465616725934, -328.60523625175796, -225.21345255126687, -220.98260530706278, -215.0038232061213, -209.81974533009384, -197.91659100611065, -193.6689419604744, -195.31089481950812, -195.92363804412696  …  -123.72049910150594, -123.95608887805227, -124.062384795074, -124.70919853152887, -123.60016238540969, -124.45483604773469, -124.87229298229747, -124.82778012224077, -126.43832850593896, -126.40424686085807], AGESS_MCMC_params(Main.var"#6#7"(), [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0  …  0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0 0.0 … 0.0 0.0; 0.0 1.0 … 0.0 0.0; … ; 0.0 0.0 … 1.0 0.0; 0.0 0.0 … 0.0 1.0], true, 6.0, 0.5, 0.05, 0.5, 0.05, 100000, 52), [0.013729712715597853 -0.006042746981915986 … -0.019611175904040178 0.00012010974224106597; -0.006042746981915986 0.024239622256590685 … 0.009930792779268768 -0.0004097034581138251; … ; -0.019611175904040178 0.009930792779268768 … 0.23370200345429468 -0.015187770329205956; 0.00012010974224106597 -0.0004097034581138251 … -0.015187770329205956 0.01353193930452958], [0.039831661644393, 0.12313365054752985, 0.033500363434703306, -0.05233756062999257, -0.04123076028426733, 0.400940434974939, -0.02041908995541667, -0.05095787927076799, 0.03940521963876113, 3.726653264086524  …  -0.3533859980907335, -0.4122226193122465, -0.26316813413027457, 2.1322824921304173, -0.1560708558480274, 0.02735647338806655, -0.4493971393290821, -0.2929220519460389, -1.951494220641405, 0.08123531892235633])

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

plot(results.samps[25001:10:end, findall(β .!= 0)], legend = false)
hline!(β[findall(β .!= 0)], line = :dash, color =:black)

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

plot(results.samps[25001:10:end, findall(β .== 0)], legend = false)
hline!(β[findall(β .== 0)], line = :dash, color =:black)

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

plot(exp.(results.samps[25001:10:end, 2*D + 2]).^2, 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.l_pdf[1:n_MCMC], 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 = 10000

### Run AGESS
results = AGESS(β -> log_posterior_ReLU(β, x, y), n_MCMC, P)

AdaptEllipticalSliceSampler.MCMC_output([0.0 0.0; -0.43293867644599 -0.56146856398129; … ; -0.41432287341757945 -0.875709834806713; -0.4475153166272507 -0.8528377996605969], [-693.1471805599322, -656.5861256703402, -656.0790773486126, -656.0609970801869, -655.9648124084979, -653.2757624962807, -653.4833417050382, -653.6177917947325, -654.03441826953, -654.596153317717  …  -653.3271489664228, -654.6928882022297, -653.7194660989785, -653.2916964948257, -656.505107779531, -653.7911921096525, -653.8992785465533, -653.6110568728842, -653.4131039241543, -653.6362769634308], AGESS_MCMC_params(Main.var"#9#10"(), [0.0, 0.0], [1.0 0.0; 0.0 1.0], true, 6.0, 0.5, 0.05, 0.5, 0.05, 10000, 2), [0.008403108701537983 0.001079084269487994; 0.001079084269487994 0.011867066337535898], [-0.36588603281090815, -0.8340346501916386])

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

# Discard the initial 2500 samples to burnin
scatter(results.samps[2501:10000, 1], results.samps[2501:10000, 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.l_pdf[1:n_MCMC], 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].