I. Importance weight approximation

Also, the importance weight approximation in step 3 gets easier. We can solve it out, instead of estimating it.

Actually, the nasty part \(\mathbb P(y|\textbf X,\theta)\) can be calculated by the following integration

\[ \mathbb P(y|\textbf X,\theta)=\int \mathbb P(y|\textbf X,\theta,\beta)p(\beta)d\beta. \]

Where \(p(\beta)\) is the prior of \(\beta\): \[ \beta|\sigma_\beta\sim N(0,\sigma^2_b\sigma^2). \]

For the likelihood term, we have \[ \hat\beta\sim N\left(\beta,(\textbf X^T\textbf X)^{-1}\sigma^2\right), \] where \[ \hat\beta=(\textbf X^T\textbf X)^{-1}\textbf X^TY. \]

Since the prior and likelihood are both Gaussian, the posterior distribution of \(\beta\) should also be Gaussian. Thus

\[ \begin{aligned} \int \mathbb P(y|\textbf X,\theta,\beta)p(\beta)d\beta &=(2\pi)^{-p}\sigma^{-2}|\textbf X^T\textbf X|^{1/2}\sigma_{\beta}^{-1} \int \exp\left\{-\frac1{2\sigma^2}\left({(\beta-\hat\beta)^T(\textbf X^T\textbf X)(\beta-\hat\beta)+\beta^T \beta/\sigma_\beta^2}\right)\right\}d\beta\\ &=\frac{C_0C_1}{C_2}\int \exp\left\{-\frac1{2\sigma^2}(\beta-\hat\beta_{ridge})^T\left(\textbf X^T\textbf X+\frac{\textbf I}{\sigma_b^2}\right)(\beta-\hat\beta_{ridge})\right\}d\beta\\ &=\frac{C_0C_1}{C_2C_3}, \end{aligned} \]

where \(\hat\beta=(\textbf X^T\textbf X)^{-1}\textbf X^T Y\), \(\hat\beta_{ridge}=(\textbf X^T\textbf X+\textbf I/\sigma^2_b)^{-1}\textbf X^T Y\), \(C_0=(2\pi)^{-p}\sigma^{-2}|\textbf X^T\textbf X|^{1/2}\sigma_b^{-1}\), \(C_1=\exp\left\{-\frac1{2\sigma^2}\hat\beta^T(\textbf X^T\textbf X)\hat\beta\right\}\), \(C_2=\exp\left\{-\frac1{2\sigma^2}\hat\beta_{ridge}^T(\textbf X^T\textbf X+\textbf I/\sigma_b^2)\hat\beta_{ridge}\right\}\), and \(C_3=(2\pi)^{-p/2}\sigma^{-1}\left|\textbf X^T\textbf X+\textbf I/\sigma_b^2\right|^{-1/2}\).

Finally we have

\[ \mathbb P(y|\textbf X,\theta)=\frac{|\textbf X^T\textbf X|^{1/2}\cdot|\textbf X^T\textbf X+\textbf I/\sigma_b^2|^{1/2}}{(2\pi)^{p/2}\sigma}\exp\left\{-\frac1{2\sigma^2}Y^T\textbf X\left((\textbf X^T\textbf X)^{-1}-(\textbf X^T\textbf X+\textbf I/\sigma^2_b)^{-1}\right)\textbf X^TY\right\}. \]

III. Simulation

Algorithm code

#' VI.ridge: use variational inference in ridge regression
#' @param X: variables in linear model
#' @param Y: response in linear model
#' @param N: the number of importance sampling
#' @return the estimated beta, result of ridge regression

VI.ridge <- function (X, Y, sigma.b, N = 50) { 
    prior <- function (sigma2) {
        return (1 / sigma2)
    }   
    
    p <- ncol(X)
    n <- nrow(X)
    
    # intercept
    beta.hat <- numeric(p + 1)
    beta.hat[1] <- mean(Y)
    
    # center the columns of X and Y
    Y <- Y - mean(Y)
    X <- t(t(X) - colMeans(X))
    
    # preprocess to make it faster
    XtX <- t(X) %*% X 
    Xy  <- t(X) %*% Y
    
    Z.exp <- t(Xy) %*% (solve(XtX) - solve(XtX + diag(p) / (sigma.b ^ 2))) %*% Xy
    # Z \propto 1/sigma * exp(-Z.exp / 2sigma^2)
    
    sigma2 <- runif (N, min = 0, max = 10)
    weight <- prior(sigma2)
    s      <- numeric(p)
    
    mu     <- matrix(0, nrow = N, ncol = p) # each row is a sample
    
    for (i in 1: N) {
        s2 <- sigma2[i] / (diag(XtX) + 1 / (sigma.b ^ 2))
        converge <- FALSE
        iter <- 0
        while (!converge && iter < 50) {
            converge = TRUE
            iter <- iter + 1
            
            for (k in 1: p) {
                record <- mu[i, k]
              
                mu[i, k] <- s2[k] / sigma2[i] * 
                            (Xy[k] - mu[i, ] %*% XtX[, k]  + XtX[k, k] * mu[i, k])
                
                if (abs(mu[i, k] - record) > 1e-4)
                    converge <- FALSE
            }
        }
        
        Z <- exp(-Z.exp / (2 * sigma2[i])) / sqrt(sigma2[i])
        weight[i] <- Z * weight[i]
    }
    
    # take the weighted average of the sampled posterior mean
    weight <- weight / sum(weight)
    beta.hat[2: (p + 1)] <- t(mu) %*% weight
    
    result <- list()
    result$coef <- beta.hat
    result$weight <- weight
    result$mu <- mu
    
    return (result)
}
#VI.coef <- VI.ridge(X, Y, sqrt(1 / lambda))

Generate data with \(\textbf X\in\mathbb R^{50\times30}\).

set.seed(1)

# generate data
sigma   <- 3
sigma.b <- 0.5
p       <- 200
n       <- 300

true.beta <- rnorm (p, 0, sd = sigma.b * sigma)
X <- matrix (rnorm (n * p), nrow = n, ncol = p)
Y <- rnorm (n, X %*% true.beta + 1, sigma)

Compared with ridge regression, where the tuning parameter in ridge regression is chosen by cross-validation:

library(lasso2)
library(MASS)
library(glmnet)

cvfit <- cv.glmnet(X, Y, alpha = 0)
lambda <- cvfit$lambda.1se
# remove the intercept
ridge.coef <- coef(cvfit, s = "lambda.1se")[2: (p + 1)]

VI.result <- VI.ridge(X, Y, sqrt(1 / lambda))
# remove the intercept
VI.coef   <- VI.result$coef[2: (p + 1)]

plot (VI.coef, ridge.coef, main = 'VI v.s. ridge regression',
      xlab = 'VI coefficient', ylab = 'ridge coefficient')
abline(a = 0, b = 1)

Past versions of unnamed-chunk-3-1.png

The VI result and the ridge regression result are close to each other. As ridge regression finds the exact posterior estimate, VI returns a reasonable approximation.

Session information

sessionInfo()

R version 3.6.0 (2019-04-26)
Platform: x86_64-w64-mingw32/x64 (64-bit)
Running under: Windows 10 x64 (build 17134)

Matrix products: default

locale:
[1] LC_COLLATE=English_United States.1252 
[2] LC_CTYPE=English_United States.1252   
[3] LC_MONETARY=English_United States.1252
[4] LC_NUMERIC=C                          
[5] LC_TIME=English_United States.1252    

attached base packages:
[1] stats     graphics  grDevices utils     datasets  methods   base     

other attached packages:
[1] glmnet_2.0-18 foreach_1.4.4 Matrix_1.2-17 MASS_7.3-51.4 lasso2_1.2-20

loaded via a namespace (and not attached):
 [1] Rcpp_1.0.1       knitr_1.23       whisker_0.3-2    magrittr_1.5    
 [5] workflowr_1.4.0  lattice_0.20-38  stringr_1.4.0    tools_3.6.0     
 [9] grid_3.6.0       xfun_0.7         git2r_0.25.2     htmltools_0.3.6 
[13] iterators_1.0.10 yaml_2.2.0       rprojroot_1.3-2  digest_0.6.19   
[17] fs_1.3.1         codetools_0.2-16 glue_1.3.1       evaluate_0.14   
[21] rmarkdown_1.13   stringi_1.4.3    compiler_3.6.0   backports_1.1.4