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 into 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_i E_{q_i} \log \frac{g(l_i)}{q(l_i)} \]

I drop the subscripts \(i\). Write \[ g \sim \pi_0 \delta_0 + (1 - \pi_0) N(0, 1/a) \] and \[ q \sim (1 - \tilde{w}) \delta_0 + \tilde{w} N(\tilde{\mu}, \tilde{\sigma}^2) \] (I use different parametrizations to make the derivation cleaner and to follow the code more closely.)

Then \[\begin{aligned} E_q \log \frac{g(l)}{q(l)} &= (1 - \tilde{w}) \log \frac{\pi_0}{1 - \tilde{w}} + \int \tilde{w}\ \text{dnorm}(x; \tilde{\mu}, \tilde{\sigma}^2) \log \frac{(1 - \pi_0)\text{dnorm}(x; 0, 1/a)} {\tilde{w}\ \text{dnorm}(x; \tilde{\mu}, \tilde{\sigma}^2)}\ dx \\ &= (1 - \tilde{w}) \log \frac{\pi_0}{1 - \tilde{w}} + \tilde{w} \log \frac{1 - \pi_0}{\tilde{w}} \\ &\ + \int \tilde{w}\ \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 \end{aligned}\]

The last integral is equal to \[\begin{aligned} \frac{\tilde{w}}{2} \log (a \tilde{\sigma}^2) &- \frac{\tilde{w} a}{2} E_{N(x; \tilde{\mu}, \tilde{\sigma}^2)} x^2 + \frac{\tilde{w}}{2 \tilde{\sigma}^2} E_{N(x; \tilde{\mu}, \tilde{\sigma}^2)} (x - \tilde{\mu})^2 \\ &= \frac{\tilde{w}}{2} \log (a \tilde{\sigma}^2) - \frac{\tilde{w} a}{2} (\tilde{\mu}^2 + \tilde{\sigma}^2) + \frac{\tilde{w}}{2} \end{aligned}\]

Thus \[ E_q \log \frac{g(l)}{q(l)} = (1 - \tilde{w}) \log \frac{\pi_0}{1 - \tilde{w}} + \tilde{w} \log \frac{1 - \pi_0}{\tilde{w}} + \frac{\tilde{w}}{2} \left( \log (a \tilde{\sigma}^2) - a(\tilde{\mu}^2 + \tilde{\sigma}^2) + 1 \right) \]