Variational Bayes

Settings

$$\mathbf{y} = X_{n\times p}\mathbf{\beta}_{p\times 1} + \mathbf{\epsilon} \quad \mathbf{\epsilon} \sim N_{n}(0, \sigma^2 I)$$ $$\beta_{j} \sim \begin{cases} N(0, \sigma^2 \sigma_{\beta}^2) & \gamma_{j} = 1\\ 0 & \gamma_{j} = 0 \end{cases} \quad \quad P(\gamma_{j} = 1) = \pi$$

Variational approach

$$\begin{align*} \color{red}{\log p(\mathbf{y}|X, \sigma^2)} &= \log \int p(\mathbf{y}, \mathbf{\beta}| X, \sigma^2) d\mathbf{\beta} \\ &= \int q(\mathbf{\beta}) \log \frac{p(\mathbf{y}, \mathbf{\beta}|X, \sigma^2)}{q(\mathbf{\beta})} d\mathbf{\beta} + \int q(\mathbf{\beta}) \log \frac{q(\mathbf{\beta})}{p(\mathbf{\beta}|\mathbf{y}, X, \sigma^2)} d \mathbf{\beta} \\ &\color{red}{= \mathcal{L}(q) + D_{KL}(q||p(\mathbf{\beta}|\mathbf{y}, X, \sigma^2))} \end{align*}$$

Assume the full factorized variational approximation $$q(\mathbf{\beta}) = \prod_{j=1}^{p} q_{j}(\beta_{j})$$

We want to max $F(q)$, min $D_{KL}(q||p(\beta|\mathbf{y}, \mathbf{x}, \sigma^2))$.

$$p(\mathbf{y},\mathbf{\beta}|X, \sigma^2, \sigma_{\beta}^2, \pi) = \left( \frac{1}{\sqrt{2\pi \sigma^2}} \right)^{n} \exp \left( -\frac{\|\mathbf{y} - X\mathbf{\beta}\|^2}{2\sigma^2} \right) \prod_{j=1}^{p}(\pi N(\beta_{j};0,\sigma^2\sigma_{\beta}^2) + (1-\pi) \delta_{0}(\beta_{j}))$$ $$\begin{align*} \mathcal{L}(q) &= \int q(\mathbf{\beta}) \log \frac{p(\mathbf{y}, \mathbf{\beta}|X, \sigma^2, \sigma_{\beta}^2, \pi)}{q(\mathbf{\beta})} d\mathbf{\beta} \\ &= \int q(\mathbf{\beta}) \log \frac{p(\mathbf{y}|\mathbf{\beta}, X, \sigma^2) p(\mathbf{\beta}|\sigma_{\beta}^2, \pi)}{q(\mathbf{\beta})} d\mathbf{\beta} \\ &= \mathbb{E}_{q}\left[ \log N_{n}(\mathbf{y} - X\mathbf{\beta}; 0, \sigma^2I) \right] + \sum_{j=1}^{p} \mathbb{E}_{q}\left[ \log \frac{p_{j}(\beta_{j})}{q_{j}(\beta_{j})} \right] \quad p_{j}(\beta_{j}) = \text{prior of }\beta_{j} \\ &= -\frac{n}{2}\log 2\pi \sigma^2 - \frac{\|\mathbf{y} - X \mathbb{E}_{q}(\mathbf{\beta})\|^2}{2\sigma^2} - \frac{\sum_{j=1}^{p}\mathbf{x}_{j}^{T}\mathbf{x}_{j} \mathbb{V}ar_{q}(\beta_{j}) }{2\sigma^2} - \sum_{j=1}^{p} D_{KL}(q_{j}||p_{j}) \end{align*}$$ The variational Bayes algorithm iteratively updates each approximate marginal distribution $q_{j}$ using $$\color{red}{\log q_{j}(\beta_{j}) \leftarrow \mathbb{E}_{q_{j'}, j' \neq j}\left[ \log p(\mathbf{y}, \mathbf{\beta})\right] + C}$$ $$\begin{align*} \mathbb{E}_{q_{j'}, j'\neq j} \left[\log p(\mathbf{y}, \mathbf{\beta}) \right] &= \mathbb{E}_{q_{j'}, j'\neq j} \left[ -\frac{\|\mathbf{y} - X\mathbf{\beta}\|^2}{2\sigma^2} + \sum_{j'=1}^{p} \log p_{j'}(\beta_{j'}) \right] + C \\ &= -\frac{\|\mathbf{y} - X_{-j} \mathbb{E}_{q}(\mathbf{\beta_{-j}}) - \mathbf{x}_{j}\beta_{j} \|^2}{2\sigma^2} + \log p_{j}(\beta_{j}) + C \\ &= -\frac{\|\mathbf{y}_{res,-j} - \mathbf{x}_{j}\beta_{j} \|^2}{2\sigma^2} + \log p_{j}(\beta_{j}) + C \end{align*}$$

Exponentiating the result $$q_{j}(\beta_{j}) \propto \exp\left( -\frac{\|\mathbf{y}_{res,-j} - \mathbf{x}_{j}\beta_{j} \|^2}{2\sigma^2} \right)p_{j}(\beta_{j})$$ which is a conditional posterior distribution for $\beta_{j}$. $$\beta_{j} \sim_{q_{j}} \begin{cases} N(\tilde{\mu}_{j}, \tilde{\sigma}_{j}^{2}) & \text{w.p } \alpha_{j} \\ 0 & \text{w.p } 1-\alpha_{j} \end{cases}$$

$$\tilde{\sigma}_{j}^{2} = \frac{\sigma^2\sigma_{\beta}^2}{\mathbf{x}_{j}^{T}\mathbf{x}_{j}\sigma_{\beta}^2 + 1} = \left( \frac{\mathbf{x}_{j}^{T}\mathbf{x}_{j}}{\sigma^2} + \frac{1}{\sigma^2\sigma_{\beta}^2} \right)^{-1}$$ $$\tilde{\mu}_{j} = \frac{\mathbf{x}_{j}^{T}\mathbf{y}_{res,-j}}{\mathbf{x}_{j}^{T}\mathbf{x}_{j} + 1/\sigma_{\beta}^2} = \frac{\tilde{\sigma}_{j}^{2}}{\sigma^2} \mathbf{x}_{j}^{T}\mathbf{y}_{res,-j}$$ $$\frac{\alpha_{j}}{1-\alpha_{j}} = \frac{\pi}{1-\pi} \frac{\tilde{\sigma}_{j}}{\sigma_{\beta}\sigma} \exp(\frac{\tilde{\mu}_{j}^{2}}{2\tilde{\sigma}_{j}^{2}}) \quad \quad \alpha_{j} \propto \pi N(\frac{\mathbf{x}_{j}^{T}\mathbf{y}_{res,-j}}{\mathbf{x}_{j}^{T}\mathbf{x}_{j}}; 0, \frac{\sigma^2}{\mathbf{x}_{j}^{T}\mathbf{x}_{j}} + \sigma^2\sigma_{\beta}^2 )$$

Session information

sessionInfo()

R version 3.5.1 (2018-07-02)
Platform: x86_64-apple-darwin15.6.0 (64-bit)
Running under: macOS High Sierra 10.13.6

Matrix products: default
BLAS: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRblas.0.dylib
LAPACK: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRlapack.dylib

locale:
[1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8

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

loaded via a namespace (and not attached):
 [1] workflowr_1.1.1   Rcpp_0.12.18      digest_0.6.15    
 [4] rprojroot_1.3-2   R.methodsS3_1.7.1 backports_1.1.2  
 [7] git2r_0.23.0      magrittr_1.5      evaluate_0.11    
[10] stringi_1.2.4     whisker_0.3-2     R.oo_1.22.0      
[13] R.utils_2.6.0     rmarkdown_1.10    tools_3.5.1      
[16] stringr_1.3.1     yaml_2.2.0        compiler_3.5.1   
[19] htmltools_0.3.6   knitr_1.20