An Algorithm Based on Variational Methods

workflowr

Ignored files:
    Ignored:    .RData
    Ignored:    .Rhistory
    Ignored:    .Rproj.user/
    Ignored:    analysis/.RData
    Ignored:    analysis/.Rhistory

It turns out that we can apply the derivation of variational algorithm in SuSiE paper (appendix B) to SuSiE-mixture model. Before we start, we will modify the model assumption a little bit, so that the derivation can be applied.

• Model

\[ \begin{aligned} y &= \textbf X\boldsymbol\beta_0+\sum_{l=1}^L\textbf X\boldsymbol\beta_{l}+\epsilon,\epsilon\sim N(0,\sigma^2 I)\\ \boldsymbol\beta_0&\sim N(0,\sigma_b^2 I),\boldsymbol\beta_{l}=\boldsymbol\gamma_{l}\beta_{l},\\ \boldsymbol\gamma_l&\sim Mult(1,\pi),\beta_{l}\sim N(0,\sigma_{l}^2). \end{aligned} \]

The difference is, we assume \(\boldsymbol\beta_1\sim N(0,\sigma_b^2 I)\) instead of \(\boldsymbol\beta_1\sim N(0,\sigma^2\sigma_b^2 I)\). We will explain the reason later.

• Additive effects model

It derives a variational inference algorithm based on the additive effects model

\[ \boldsymbol y=\sum_{l=1}^L \boldsymbol\mu_l+\textbf e,\textbf e\sim N(0,\sigma^2I),\boldsymbol\mu_l\sim g_l. \]

The \(\boldsymbol\mu_l\)’s are variables, and there is no specific assumption on their distributions \(g_l\). Let \(q_l\) be the functions that approximate the posterior distribution of \(\boldsymbol\mu_l\), then the ELBO is given as follows

\[ F(q,g,\sigma^2;y)=-\frac n2\log(2\pi\sigma^2)-\frac 1{2\sigma^2}\mathbb E_q\left[\|y-\sum_{l=1}^L\mu_l\|_2^2\right] +\sum_{l=1}^L\mathbb E_{q_l}\left[\log\frac{g_l(\mu_l)}{q_l(\mu_l)}\right]. \]

Let

\[ ERSS(y,\bar\mu,\bar{\mu^2})=\mathbb E_q\|y-\sum_{l=1}^L\mu_l\|_2^2. \]

Then if \(g_l\) and \(q_l\) are irrelevant to \(\sigma^2\), then we can maximize ELBO w.r.t \(\sigma^2\) and get

\[ \hat\sigma^2=ERSS(y,\bar\mu,\bar{\mu^2})/n. \]

In SuSiE-mixture model, we introduce an additional term \(\textbf X\boldsymbol \beta_0\). We can easily plug that into the additive effects model. Let \(\boldsymbol\mu_{0}=\textbf X\boldsymbol \beta_1\), and \(y=\sum_{l=0}^L\mu_l+e\). By assuming the distribution to be \(N(0,XX^T\sigma^2_b)\) instead of \(N(0,XX^T\sigma^2_b\sigma^2)\), the distribution \(g_0\) will be irrelevant to \(\sigma^2\), so that the formula to estimate \(\sigma^2\) will remain the same as above.

• Iterative Bayesian stepwise selection (IBSS) algorithm

Based on the additive effects model, specifically for SuSiE, it turns out only the first and second moments of \(\beta\) are required to optimize over \(q_l\) and \(g_l\). We can derive the IBSS algorithm as follows

• while not converge do
  • \(\bar r\leftarrow y - X\sum_{l=1}^L \bar\beta_l\)
  • for l in 1,…, L do
    • \(\bar r_l \leftarrow \bar r+X\beta_l\)
    • \(\sigma_{0l}^2\leftarrow \arg\max_{\sigma_0^2}SER(\bar r_l;\sigma_0^2,\sigma^2)\)
    • \((\alpha_l,\mu_{1l},\sigma_{1l})\leftarrow SER(X,\bar r_l;\sigma^2,\sigma_{0l}^2)\)
    • \(\bar \beta_l\leftarrow \alpha_l \circ \mu_{1l}\)
    • \(\bar {\beta^2_l} \leftarrow \alpha_l\circ (\sigma_{1l}^2+\mu_{1l}^2)\)
  • \(\sigma^2\leftarrow ERSS(y,\bar \beta,\bar {\beta^2})/n\).
• return \(\sigma^2, \boldsymbol{\sigma}_0^2, \boldsymbol\alpha,\boldsymbol{\beta}_1,\boldsymbol{\mu}_1\)

For SuSiE-mixture model, we take the addition term \(X\beta_1\) into account, and update \(\beta_1\) in each step. Then the algorithm will be

• while not converge do
  • \(\bar r\leftarrow y - X\sum_{l=0}^L \bar\beta_l\)
  • // update \(\beta_0\)
  • \(\bar r_0 \leftarrow \bar r+X\beta_0\)
  • \(\sigma_b^2\leftarrow \arg\max_{\sigma_b^2}ELBO_{ridge}(\bar r_0;\sigma_b^2,\sigma^2)\)
  • \((\bar \beta_{0},\bar{\beta^2_0})\leftarrow Ridge(\bar r_0,X,\sigma^2,\sigma_b^2)\)
  • for l in 1,…, L do
    • \(\bar r_l \leftarrow \bar r+X\beta_l\)
    • \(\sigma_{0l}^2\leftarrow \arg\max_{\sigma_0^2}l_{SER}(\bar r_l;\sigma_0^2,\sigma^2)\)
    • \((\alpha_l,\mu_{1l},\sigma_{1l})\leftarrow SER(X,\bar r_l;\sigma^2,\sigma_{0l}^2)\)
    • \(\bar \beta_l\leftarrow \alpha_l \circ \mu_{1l}\)
    • \(\bar {\beta^2_l} \leftarrow \alpha_l\circ (\sigma_{1l}^2+\mu_{1l}^2)\)
  • \(\sigma^2\leftarrow ERSS(y,\bar \beta,\bar {\beta^2})/n\).
• return \(\sigma^2, \sigma_b^2,\boldsymbol{\sigma}_0^2, \boldsymbol\alpha,\boldsymbol{\sigma}_1,\boldsymbol{\mu}_1\)

See experiment2 for the simulation for this algorithm.