Notes on Computing the FLASH Objective

Indirect method

Recall the FLASH model: \[ Y = LF' + E \]

When updating loading \(l_k\), we are optimizing over \(g_{l_k}\) and \(q_{l_k}\). \(g_{l_k} \in \mathcal{G}\) is the prior on the elements of the \(k\)th column of the loadings matrix: \[ l_{1k}, \ldots, l_{nk} \sim^{iid} g_{l_k} \] \(q_{l_k}\) is an arbitrary distribution which enters the problem via the variational approach. For convenience, I drop the subscripts in the following.

The part of the objective that depends on \(g\) and \(q\) is \[ F(g, q) := E_q \left[ -\frac{1}{2} \sum_i (A_i l_i^2 - 2 B_i l_i) \right] + E_q \log \frac{g(\mathbf{l})}{q(\mathbf{l})} \] with \[ A_i = \sum_j \tau_{ij} Ef^2_j \text{ and } B_i = \sum_j \tau_{ij} R_{ij} Ef_j, \] (\(R\) is the matrix of residuals (excluding factor \(k\)) and \(Ef_j\) and \(Ef^2_j\) are the expected values of \(f_{jk}\) and \(f_{jk}^2\) with respect to the distribution \(q_{f_k}\) fitted during the factor update.)

As Lemma 2 in the paper shows (see Appendix A.2), this expression is optimized by setting \(s_j^2 = A_j\) and \(x_j = B_j s_j^2\), and then solving the EBNM problem, where the EBNM model is: \[ \mathbf{x} = \mathbf{\theta} + \mathbf{e},\ \theta_1, \ldots, \theta_n \sim^{iid} g,\ e_j \sim N(0, s_j^2) \]

Solving the EBNM problem gives \[\hat{g} = {\arg \max}_g\ p(x \mid g) \] and \[ \hat{q} = p(\theta \mid x, \hat{g}) \]

Finally, to update the overall objective, we need to compute \(E_q \log \frac{g(\mathbf{l})}{q(\mathbf{l})}\). FLASH uses a clever trick, noticing that \[ E_{\hat{q}} \log \frac{\hat{g}(\mathbf{l})}{\hat{q}(\mathbf{l})} = F(\hat{g}, \hat{q}) + \frac{1}{2} \sum_j \left[ \log 2\pi s_j^2 + (1/s_j^2) E_{\hat{q}} (x_j - \theta_j)^2 \right] \] (See Appendix A.4.)

Direct method

When using ebnm_pn, however, it seems possible to compute \(E_q \log \frac{g(\mathbf{l})}{q(\mathbf{l})}\) directly. Since the elements \(l_1, \ldots, l_n\) are i.i.d. from \(g\) (by the FLASH model) and the posterior distributions are mutually independent (by the EBNM model), \[ E_q \log \frac{g(\mathbf{l})}{q(\mathbf{l})} = \sum_j E_{q_j} \log \frac{g(l_j)}{q(l_j)} \]

I drop the subscripts \(j\). Write \[ g \sim \pi_0 \delta_0 + (1 - \pi_0) N(0, 1/a) \] and \[ q \sim \tilde{\pi}_0 \delta_0 + (1 - \tilde{\pi}_0) N(\tilde{\mu}, \tilde{\sigma}^2) \] (I parametrize the normals differently to follow the code more closely.)

Then \[\begin{aligned} E_q \log \frac{g(l)}{q(l)} &= \tilde{\pi_0} \log \frac{\pi_0}{\tilde{\pi}_0} + \int (1 - \tilde{\pi}_0) \text{dnorm}(x; \tilde{\mu}, \tilde{\sigma}^2) \log \frac{(1 - \pi_0)\text{dnorm}(x; 0, 1/a)} {(1 - \tilde{\pi}_0)\text{dnorm}(x; \tilde{\mu}, \tilde{\sigma}^2)}\ dx \\ &= \tilde{\pi_0} \log \frac{\pi_0}{\tilde{\pi}_0} + (1 - \tilde{\pi_0}) \log \frac{1 - \pi_0}{1 - \tilde{\pi}_0} \\ &\ + \int (1 - \tilde{\pi}_0) \text{dnorm}(x; \tilde{\mu}, \tilde{\sigma}^2) \log \left( \sqrt{a \tilde{\sigma}^2} \exp \left( -\frac{ax^2}{2} + \frac{(x - \tilde{\mu})^2}{2 \tilde{\sigma}^2} \right) \right) \ dx \\ &= \tilde{\pi_0} \log \frac{\pi_0}{\tilde{\pi}_0} + (1 - \tilde{\pi_0}) \log \frac{1 - \pi_0}{1 - \tilde{\pi}_0} + \frac{1 - \tilde{\pi_0}}{2} \log (a \tilde{\sigma}^2) \\ &\ - \frac{(1 - \tilde{\pi}_0)a}{2} E_{N(x; \tilde{\mu}, \tilde{\sigma}^2)} x^2 + \frac{1 - \tilde{\pi}_0}{2 \tilde{\sigma}^2} E_{N(x; \tilde{\mu}, \tilde{\sigma}^2)} (x - \tilde{\mu})^2 \\ &= \tilde{\pi_0} \log \frac{\pi_0}{\tilde{\pi}_0} + (1 - \tilde{\pi_0}) \log \frac{1 - \pi_0}{1 - \tilde{\pi}_0} + \frac{1 - \tilde{\pi_0}}{2} \log (a \tilde{\sigma}^2) - \frac{(1 - \tilde{\pi}_0)a}{2} (\tilde{\mu}^2 + \tilde{\sigma}^2) + \frac{1 - \tilde{\pi}_0}{2} \end{aligned}\]

So we should be able to calculate \(E_q \log \frac{g(l)}{q(l)}\) as follows:

calc_KL <- function(x, s, g) {
  pi0 <- g$pi0
  w <- 1 - g$pi0
  a <- g$a
  
  wpost <- ebnm:::wpost_normal(x, s, w, a) # 1 - \tilde{\pi}_0
  pi0post <- 1 - wpost[wpost < 1] # \tilde{\pi}_0
  pmean_cond <- ebnm:::pmean_cond_normal(x, s, a) # \tilde{\mu}
  pvar_cond <- ebnm:::pvar_cond_normal(s, a) # \tilde{\sigma}^2

  KLa <- pi0post * log(pi0 / pi0post)
  KLb <- wpost * log(w / wpost)
  KLc <- (wpost / 2) * log(a * pvar_cond)
  KLd <- -(wpost / 2) * a * (pvar_cond + pmean_cond^2)
  KLe <- (wpost / 2)
  sum(KLa) + sum(KLb) + sum(KLc) + sum(KLd) + sum(KLe)
}

Alternate objective

So maybe we could instead optimize (over \(g\)) calc_KL(x, s, g) - 0.5 * sum((1 / s^2) * (Etheta2(x, s, g) - 2 * x * Etheta(x, s, g))? If the quantity calculated by calc_KL is indeed correct, then this is simply another way to write the objective \(F(g, q)\) (with the constraint that \(q\) is the posterior induced by \(g\)).

alt_obj <- function(x, s, pi0, a) {
  w <- 1 - pi0
  
  wpost <- ebnm:::wpost_normal(x, s, w, a)
  pi0post <- 1 - wpost[wpost < 1]
  pmean_cond <- ebnm:::pmean_cond_normal(x, s, a)
  pvar_cond <- ebnm:::pvar_cond_normal(s, a)
  
  Etheta <- wpost * pmean_cond
  Etheta2 <- wpost * (pmean_cond^2 + pvar_cond)
  
  KLa <- pi0post * log(pi0 / pi0post)
  KLb <- wpost * log(w / wpost)
  KLc <- (wpost / 2) * log(a * pvar_cond)
  KLd <- -(wpost / 2) * a * (pvar_cond + pmean_cond^2)
  KLe <- (wpost / 2)
  KL <- sum(KLa) + sum(KLb) + sum(KLc) + sum(KLd) + sum(KLe)
  
  return(KL - 0.5 * sum(2 * pi * s^2 
                        + (1 / s^2) * (x^2 - 2 * x * Etheta + Etheta2)))
}

Session information

sessionInfo()

R version 3.4.3 (2017-11-30)
Platform: x86_64-apple-darwin15.6.0 (64-bit)
Running under: macOS Sierra 10.12.6

Matrix products: default
BLAS: /Library/Frameworks/R.framework/Versions/3.4/Resources/lib/libRblas.0.dylib
LAPACK: /Library/Frameworks/R.framework/Versions/3.4/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.0.1   Rcpp_0.12.17      digest_0.6.15    
 [4] rprojroot_1.3-2   R.methodsS3_1.7.1 backports_1.1.2  
 [7] git2r_0.21.0      magrittr_1.5      evaluate_0.10.1  
[10] stringi_1.1.6     whisker_0.3-2     R.oo_1.21.0      
[13] R.utils_2.6.0     rmarkdown_1.8     tools_3.4.3      
[16] stringr_1.3.0     yaml_2.1.17       compiler_3.4.3   
[19] htmltools_0.3.6   knitr_1.20