Empirical Bayes Matrix Factorization

Abstract

Matrix factorization methods, which include Factor analysis (FA) and Principal Components Analysis (PCA), are widely used for inferring and summarizing structure in multivariate data. Many such methods use a penalty or prior distribution to achieve sparse representations (“Sparse FA/PCA”), and a key question is how much sparsity to induce. Here we introduce a general Empirical Bayes approach to matrix factorization (EBMF), whose key feature is that it estimates the appropriate amount of sparsity by estimating prior distributions from the observed data. The approach is very flexible: it allows for a wide range of different prior families and allows that each component of the matrix factorization may exhibit a different amount of sparsity. The key to this flexibility is the use of a variational approximation, which we show effectively reduces fitting the EBMF model to solving a simpler problem, the so-called “normal means” problem. We demonstrate the benefits of EBMF with sparse priors through both numerical comparisons with competing methods and through analysis of data from the GTEx (Genotype Tissue Expression) project on genetic associations across 44 human tissues. In numerical comparisons EBMF often provides more accurate inferences than other methods. In the GTEx data, EBMF identifies interpretable structure that agrees with known relationships among human tissues. Software implementing our approach is available at https://github.com/stephenslab/flashr.

1. Introduction

Matrix factorization methods are widely used for inferring and summarizing structure in multivariate data. In brief, these methods represent an observed $n\times p$ data matrix $Y$ as:

$$Y=L^{T}F+E$$ (1.1)

where $L$ is an $K\times n$ matrix, $F$ is a $K\times p$ matrix, and $E$ is an $n\times p$ matrix of residuals (whose entries we assume to be normally distributed, although the methods we develop can be generalized to other settings; see Section 6.4). Here we adopt the notation and terminology of factor analysis, and refer to $L$ as the “loadings” and $F$ as the “factors”.

The model (1.1) has many potential applications. One range of applications arise from “matrix completion” problems (e.g. Fithian et al., 2018): methods that estimate $L$ and $F$ in (1.1) from partially observed $Y$ provide a natural and effective way to fill in the missing entries. Another wide range of applications come from the desire to summarize and understand the structure in a matrix Y : in (1.1) each row of $Y$ is approximated by a linear combination of underlying “factors” (rows of), which—ideally—have some intuitive or scientific interpretation. For example, suppose $Y_{ij}$ represents the rating of a user $i$ for a movie $j$. Each factor might represent a genre of movie (“comedy”, “drama”, “romance”, “horror” etc), and the ratings for a user $i$ could be written as a linear combination of these factors, with the weights (loadings) representing how much individual $i$ likes that genre. Or, suppose $Y_{ij}$ represents the expression of gene $j$ in sample $i$. Each factor might represent a module of co-regulated genes, and the data for sample $i$ could be written as a linear combination of these factors, with the loadings representing how active each module is in each sample. Many other examples could be given across many fields, including psychology (Ford et al., 1986), econometrics (Bai and Ng, 2008), natural language processing (Bouchard et al., 2015), population genetics (Engelhardt and Stephens, 2010), and functional genomics (Stein-O’Brien et al., 2018).

The simplest approaches to estimating $L$ and/or $F$ in (1.1) are based on maximum likelihood or least squares. For example, Principal Components Analysis (PCA)—or, more precisely, truncated Singular Value Decomposition (SVD)—can be interpreted as fitting (1.1) by least squares, assuming that columns of $L$ are orthogonal and columns of $F$ are orthonormal (Eckart and Young, 1936). And classical factor analysis (FA) corresponds to maximum likelihood estimation of $L$, assuming that the elements of $F$ are independent standard normal and allowing different residual variances for each column of $Y$ (Rubin and Thayer, 1982). While these simple methods remain widely used, in the last two decades researchers have focused considerable attention on obtaining more accurate and/or more interpretable estimates, either by imposing additional constraints (e.g. non-negativity; Lee and Seung, 1999) or by regularization using a penalty term (e.g. Jolliffe et al., 2003; Witten et al., 2009; Mazumder et al., 2010; Hastie et al., 2015; Fithian et al., 2018), or a prior distribution (e.g. Bishop, 1999; Attias, 1999; Ghahramani and Beal, 2000; West, 2003). In particular, many authors have noted the benefits of sparsity assumptions on $L$ and/or $F$— particularly in applications where interpretability of the estimates is desired—and there now exists a wide range of methods that attempt to induce sparsity in these models (e.g. Sabatti and James, 2005; Zou et al., 2006; Pournara and Wernisch, 2007; Carvalho et al., 2008; Witten et al., 2009; Engelhardt and Stephens, 2010; Knowles and Ghahramani, 2011; Bhattacharya and Dunson, 2011; Mayrink et al., 2013; Yang et al., 2014; Gao et al., 2016; Hore et al., 2016; Ročková and George, 2016; Srivastava et al., 2017; Kaufmann and Schumacher, 2017; Frühwirth-Schnatter and Lopes, 2018; Zhao et al., 2018). Many of these methods induce sparsity in the loadings only, although some induce sparsity in both loadings and factors.

In any statistical problem involving sparsity, a key question is how strong the sparsity should be. In penalty-based methods this is controlled by the strength and form of the penalty, whereas in Bayesian methods it is controlled by the prior distributions. In this paper we take an Empirical Bayes approach to this problem, exploiting variational approximation methods (Blei et al., 2017) to obtain simple algorithms that jointly estimate the prior distributions for both loadings and factors, as well as the loadings and factors themselves.

Both EB and variational methods have been previously used for this problem (Bishop, 1999; Lim and Teh, 2007; Raiko et al., 2007; Stegle et al., 2010). However, most of this previous work has used simple normal prior distributions that do not induce sparsity. Variational methods that use sparsity-inducing priors include Girolami (2001), which uses a Laplace prior on the factors (no prior on the loadings which are treated as free parameters), Hochreiter et al. (2010) which extends this to Laplace priors on both factors and loadings, with fixed values of the Laplace prior parameters; Titsias and Lázaro-Gredilla (2011), which uses a sparse “spike-and-slab” (point normal) prior on the loadings (with the same prior on all $K$ loadings) and a normal prior on the factors; and Hore et al. (2016) which uses a spike-and-slab prior on one mode in a tensor decomposition. (While this work was in review further examples appeared, including Argelaguet et al. 2018, which uses normal priors on the loadings and point-normal on the factors.)

Our primary contribution here is to develop and implement a more general EB approach to matrix factorization (EBMF). This general approach allows for a wide range of potential sparsity-inducing prior distributions on both the loadings and the factors within a single algorithmic framework. We accomplish this by showing that, when using variational methods, fitting EBMF with any prior family can be reduced to repeatedly solving a much simpler problem—the “empirical Bayes normal means” (EBNM) problem—with the same prior family. This feature makes it easy to implement methods for any desired prior family—one simply has to implement a method to solve the corresponding normal means problem, and then plug this into our algorithm. This approach can work for both parametric families (e.g. normal, point-normal, laplace, point-laplace) and non-parametric families, including the “adaptive shrinkage” priors (unimodal and scale mixtures of normals) from Stephens (2017). It is also possible to accommodate non-negative constraints on either $L$ and/or $F$ by using non-negative prior families. Even simple versions of our approach—e.g. using point-normal priors on both factors and loadings— provide more generality than most existing EBMF approaches and software.

A second contribution of our work is to highlight similarities and differences between EBMF and penalty-based methods for regularizing $L$ and/or $F$. Indeed, our algorithm for fitting EBMF has the same structure as commonly-used algorithms for penalty-based methods, with the prior distribution playing a role analogous to the penalty (see Remark 3 later). While the general correspondence between estimates from penalized methods and Bayesian posterior modes (MAP estimates) is well known, the connection here is different, because the EBMF approach is estimating a posterior mean, not a mode (indeed, with sparse priors the MAP estimates of $L$ and $F$ are not useful because they are trivially 0). A key difference between the EBMF approach and penalty-based methods is that the EBMF prior is estimated by solving an optimization problem, whereas in penalty-based methods the strength of the penalty is usually chosen by cross-validation. This difference makes it much easier for EBMF to allow for different levels of sparsity in every factor and every loading: in EBMF one simply uses a different prior for every factor and loading, whereas tuning a separate parameter for every factor and loading by CV becomes very cumbersome.

The final contribution is that we provide an $\text{R}$ software package, flash (Factors and Loadings by Adaptive SHrinkage), implementing our flexible EBMF framework. We demonstrate the utility of these methods through both numerical comparisons with competing methods and through a scientific application: analysis of data from the GTEx (Genotype Tissue Expression) project on genetic associations across 44 human tissues. In numerical comparisons flash often provides more accurate inferences than other methods, while remaining computationally tractable for moderate-sized matrices (millions of entries). In the GTEx data, flash highlight both effects that are shared across many tissues (“dense” factors) and effects that are specific to a small number of tissues (“sparse” factors). These sparse factors often highlight similarities between tissues that are known to be biologically related, providing external support for the reliability of the results.

2. A General Empirical Bayes Matrix Factorization Model

We define the $K$-factor Empirical Bayes Matrix Factorization (EBMF) model as follows:

$$Y=\sum _{k=1}^K {\mathit{l}}_k{\mathit{f}}_k^T+E$$ (2.1)

$$l_{k1},\dots ,l_{kn}{~}^{iid}{g}_{{\mathit{l}}_k},g_{{\mathit{l}}_k}\in {𝒢}_{\mathit{l}}$$ (2.2)

$${f_{k1},\dots ,f_{kp}~}^{iid}{g}_{{\mathit{f}}_k},g_{{\mathit{f}}_k}\in {𝒢}_{\mathit{f}}$$ (2.3)

$$E_{ij}~N\left(\mathrm{0,1}/{\tau }_{ij}\right)\text{ with }\mathit{\tau }≔\left({\tau }_{ij}\right)\in 𝒯.$$ (2.4)

Here $Y$ is the $n\times p$ observed data matrix, ${\mathit{l}}_k$ is an $n$-vector (the $k$th set of “loadings”), ${\mathit{f}}_k$ is a $p$-vector (the $k$th “factor”), ${𝒢}_l$ and ${𝒢}_f$ are pre-specified (possibly non-parametric) families of distributions, $g_{{\mathit{l}}_k}$ and $g_{{\mathit{f}}_k}$ are unknown “prior” distributions that are to be estimated, $E$ is an $n\times p$ matrix of independent error terms, and $\mathit{\tau }$ is an unknown $n\times p$ matrix of precisions $\left({\tau }_{ij}\right)$ which is assumed to lie in some space $𝒯$. (This allows structure to be imposed on $\mathit{\tau }$, such as constant precision, ${\tau }_{ij}=\tau$, or column-specific precisions, ${\tau }_{ij}={\tau }_j$, for example.) Our methods allow that some elements of $Y$ may be “missing”, and can estimate the missing values (Section 4.1).

The term “Empirical Bayes” in EBMF means we fit (2.1)-(2.4) by obtaining point estimates for the priors $g_{l_k},g_{f_k},(k=1,\dots ,K)$ and approximate the posterior distributions for the parameters ${\mathit{l}}_k,{\mathit{f}}_k$ given those point estimates. This contrasts with a “fully Bayes” approach that, instead of obtaining point estimates for $g_{{\mathit{l}}_k},g_{{\mathit{f}}_k}$, would integrate over uncertainty in the estimates. This would involve specifying prior distributions for $g_{{\mathit{l}}_k},g_{{\mathit{f}}_k}$ as well as (perhaps substantial) additional computation. The EB approach has the advantage of simplicity - both conceptually and computationally—while enjoying many of the benefits of a fully Bayes approach. In particular it allows for sharing of information across elements of each loading/factor. For example, if the data suggest that a particular factor, ${\mathit{f}}_k$, is sparse, then this will be reflected in a sparse estimate of $g_{{\mathit{f}}_k}$, and subsequently strong shrinkage of the smaller elements of $f_{k1},\dots ,f_{kp}$ towards 0. Conversely, when the data suggest a non-sparse factor then the prior will be dense and the shrinkage less strong. By allowing different prior distributions for each factor and each loading, the model has the flexibility to adapt to any combination of sparse and dense loadings and factors. However, to fully capitalize on this flexibility one needs suitably flexible prior families ${𝒢}_l$ and ${𝒢}_f$ capable of capturing both sparse and dense factors. A key feature of our work is it allows for very flexible prior families, including non-parametric families.

Some specific choices of the distributional families ${𝒢}_l$ and ${𝒢}_f$ correspond to models used in previous work. In particular, many previous papers have studied the case with normal priors, where ${𝒢}_l$ and ${𝒢}_f$ are both the family of zero-mean normal distributions (e.g. Bishop, 1999; Lim and Teh, 2007; Raiko et al., 2007; Nakajima and Sugiyama, 2011). This family is particularly simple, having a single hyper-parameter, the prior variance, to estimate for each factor. However, it does not induce sparsity on either $L$ or $F$; indeed, when the matrix $Y$ is fully observed, the estimates of $L$ and $F$ under a normal prior (when using a fully factored variational approximation) are simply scalings of the singular vectors from an SVD of $Y$ (Nakajima and Sugiyama, 2011; Nakajima et al., 2013). Our work here extends these previous approaches to a much wider range of prior families that do induce sparsity on $L$ and/or $F$.

We note that the EBMF model (2.1)-(2.4) differs in an important way from the sparse factor analysis (SFA) methods in Engelhardt and Stephens (2010), which use a type of Automatic Relevance Determination prior (e.g. Tipping, 2001; Wipf and Nagarajan, 2008) to induce sparsity on the loadings matrix. In particular, SFA estimates a separate hyperparameter for every element of the loadings matrix, with no sharing of information across elements of the same loading. In contrast, EBMF estimates a single shared prior distribution for elements of each loading, which, as noted above, allows for sharing of information across elements of each loading/factor.

3. Fitting the EBMF Model

To simplify exposition we begin with the case $K=1$ (“rank 1”); see Section 3.6 for the extension to general $K$. To simplify notation we assume the families ${𝒢}_l,{𝒢}_f$ are the same, so we can write ${𝒢}_{\mathit{l}}={𝒢}_{\mathit{f}}=𝒢$. To further lighten notation in the case $K=1$ we use $g_{\mathit{l}},g_{\mathit{f}},\mathit{l},\mathit{f}$ instead of $g_{{\mathit{l}}_1},g_{{\mathit{f}}_1},{\mathit{l}}_1,{\mathit{f}}_1$.

Fitting the EBMF model involves estimating all of $g_{\mathit{l}},g_{\mathit{f}},\mathit{l},\mathit{f},\mathit{\tau }$. A standard EB approach would be to do this in two steps:

• Estimate $g_{\mathit{l}},g_{\mathit{f}}$ and $\mathit{\tau }$, by maximizing the likelihood:

  $$L\left(g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)≔\iint p(Y\mid \mathit{l},\mathit{f},\mathit{\tau })g_{\mathit{l}}\left(d{l}_1\right)\dots g_{\mathit{l}}\left(d{l}_n\right)g_{\mathit{f}}\left(d{f}_1\right)\dots g_{\mathit{f}}\left(d{f}_p\right)$$ (3.1)

  over $g_{\mathit{l}},g_{\mathit{f}}\in 𝒢,\mathit{\tau }\in 𝒯$. (This optimum will typically not be unique because of identifiability issues; see Section 3.8.)

• Estimate $\mathit{l}$ and $\mathit{f}$ using their posterior distribution: $p\left(\mathit{l},\mathit{f}\mid Y,{\hat{g}}_{\mathit{l}},{\hat{g}}_{\mathit{f}},\stackrel{\mathit{ˆ}}{\mathit{\tau }}\right)$.

However, both these two steps are difficult, even for very simple choices of $𝒢$. Instead, following previous work (see Introduction for citations) we use variational approximations to approximate this approach. Although variational approximations are known to typically under-estimate uncertainty in posterior distributions, our focus here is on obtaining useful point estimates for $\mathit{l},\mathit{f}$; results shown later demonstrate that the variational approximation can perform well in this task.

3.1. The Variational Approximation

The variational approach—see Blei et al. (2017) for review—begins by writing the log of the likelihood (3.1) as:

$$l\left(g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)≔\text{log}L\left(g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$$ (3.2)

$$\begin{array}{rr}& \\ & =F\left(q,g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)+D_{KL}(q\parallel p)\end{array}$$ (3.3)

where

$$F\left(q,g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)=\int q(\mathit{l},\mathit{f})\text{log}\frac{p\left(Y,\mathit{l},\mathit{f}\mid g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)}{q(\mathit{l},\mathit{f})}d\mathit{l}d\mathit{f},$$ (3.4)

and

$$D_{KL}\left(q\right|\left|p\right)=-\int q(\mathit{l},\mathit{f})\text{log}\frac{p\left(\mathit{l},\mathit{f}\mid Y,g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)}{q(\mathit{l},\mathit{f})}d\mathit{l}d\mathit{f}$$ (3.5)

is the Kullback–Leibler divergence from $q$ to $p$. This identity holds for any distribution $q(\mathit{l},\mathit{f})$. Because $D_{KL}$ is non-negative, it follows that $F\left(q,g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$ is a lower bound for the $\text{log}$ likelihood:

$$l\left(g_l,g_{\mathit{f}},\mathit{\tau }\right)\ge F\left(q,g_l,g_{\mathit{f}},\mathit{\tau }\right)$$ (3.6)

with equality when $q(\mathit{l},\mathit{f})=p\left(\mathit{l},\mathit{f}\mid Y,g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$.

In other words,

$$l\left(g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)=\underset{q}{\text{max}} F\left(q,g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right),$$ (3.7)

where the maximization is over all possible distributions $q(\mathit{l},\mathit{f})$. Maximizing $l\left(g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$ can thus be viewed as maximizing $F$ over $q,g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }$. However, as noted above, this maximization is difficult. The variational approach simplifies the problem by maximizing $F$ but restricting the family of distributions for $q$. Specifically, the most common variational approach — and the one we consider here—restricts $q$ to the family $𝒬$ of distributions that “fully-factorize”:

$$𝒬=\left\{q:q(\mathit{l},\mathit{f})=\prod _{i=1}^n q_{l,i}\left(l_i\right)\prod _{j=1}^p q_{f,j}\left(f_j\right)\right\}.$$ (3.8)

The variational approach seeks to optimize $F$ over $q,g_l,g_{\mathit{f}},\mathit{\tau }$ with the constraint $q\in Q$. For $q\in Q$ we can write $q(\mathit{l},\mathit{f})=q_{\mathit{l}}\left(\mathit{l}\right)q_{\mathit{f}}\left(\mathit{f}\right)$ where $q_{\mathit{l}}\left(\mathit{l}\right)={\prod }_{i=1}^n q_{l,i}\left(l_i\right)$ and $q_{\mathit{f}}\left(\mathit{f}\right)=$ ${\prod }_{j=1}^p q_{f,j}\left(f_j\right)$, and we can consider the problem as maximizing $F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$.

3.2. Alternating Optimization

We optimize $F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$ by alternating between optimizing over variables related to l $\left[\left(q_l,g_{\mathit{l}}\right)\right]$, over variables related to $\mathit{f}\left[\left(q_{\mathit{f}},g_{\mathit{f}}\right)\right]$, and over $\mathit{\tau }$. Each of these steps is guaranteed to increase (or, more precisely, not decrease) $F$, and convergence can be assessed by (for example) stopping when these optimization steps yield a very small increase in $F$. Note that $F$ may be multi-modal, and there is no guarantee that the algorithm will converge to a global optimum. The approach is summarized in Algorithm 1

The key steps in Algorithm 1 are the maximizations in Steps 4–6.

Step 4, the update of $\mathit{\tau }$, involves computing the expected squared residuals:

$${\overline{R^2}}_{ij}≔{\text{E}}_{q_{\mathit{l}},q_{\mathit{f}}}\left[{\left(Y_{ij}-l_i{f}_j\right)}^2\right]$$ (3.9)

$$\begin{array}{rr}& \\ & ={\left[Y_{ij}-{\text{E}}_{q_{\mathit{l}}}\left(l_i\right){\text{E}}_{q_{\mathit{f}}}\left(f_j\right)\right]}^2-{\text{E}}_{q_{\mathit{l}}}{\left(l_i\right)}^2{\text{E}}_{q_{\mathit{f}}}{\left(f_j\right)}^2+{\text{E}}_{q_{\mathit{l}}}\left(l_i^2\right){\text{E}}_{q_{\mathit{f}}}\left(f_j^2\right).\end{array}$$ (3.10)

This is straightforward provided the first and second moments of $q_{\mathit{l}}$ and $q_{\mathit{f}}$ are available (see Appendix A.1 for details).

Steps 5 and 6 are essentially identical except for switching the role of $\mathit{l}$ and $\mathit{f}$. One of our key results is that each of these steps can be achieved by solving a simpler problem—the Empirical Bayes normal means (EBNM) problem. The next subsection (3.3) describes the EBNM problem, and the following subsection (3.4) details how this can be used to solve Steps 5 and 6.

3.3. The EBNM Problem

Suppose we have observations $\mathit{x}=\left(x_1,\dots ,x_n\right)$ of underlying quantities $\mathit{\theta }=\left({\theta }_1,\dots ,{\theta }_n\right)$, with independent Gaussian errors with known standard deviations $\mathit{s}=\left(s_1,\dots ,s_n\right)$. Suppose further that the elements of $\mathit{\theta }$ are assumed i.i.d. from some distribution, $g\in 𝒢$. That is,

$$\mathit{x}\mid \mathit{\theta }~N_n\left(\mathit{\theta },\text{diag}\left(s_1^2,\dots ,s_n^2\right)\right)$$ (3.11)

$$\begin{array}{rr}& \\ {\theta }_1,\dots ,{\theta }_n& {~}^{\text{iid }}g,g\in 𝒢\end{array}$$ (3.12)

where $N_n(\mathit{\mu },\Sigma )$ denotes the $n$-dimensional normal distribution with mean $\mathit{\mu }$ and covariance matrix $\Sigma$.

By solving the EBNM problem we mean fitting the model (3.11)-(3.12) by the following two-step procedure:

1. Estimate $g$ by maximum (marginal) likelihood:

   $$\hat{g}=\text{arg}\underset{g\in 𝒢}{\text{max}} \prod _j \int p\left(x_j\mid {\theta }_j,s_j\right)g\left(d{\theta }_j\right).$$ (3.13)
2. Compute the posterior distribution for $\mathit{\theta }$ given $\hat{g}$,

   $$p(\mathit{\theta }\mid \mathit{x},\mathit{s},\hat{g})\propto \prod _j \hat{g}\left({\theta }_j\right)p\left(x_j\mid {\theta }_j,s_j\right).$$ (3.14)

Later in this paper we will have need for the posterior first and second moments, so we define them here for convenience:

$${\bar{\theta }}_j≔E\left({\theta }_j\mid \mathit{x},\mathit{s},\hat{g}\right)$$ (3.15)

$${\overline{{\theta }^2}}_j≔E\left({\theta }_j^2\mid \mathit{x},\mathit{s},\hat{g}\right).$$ (3.16)

Formally, this procedure defines a mapping (which depends on the family 𝒢) from the known quantities $(\mathit{x},s)$, to $(\hat{g},p)$, where $\hat{g},p$ are given in (3.13) and (3.14). We use $\mathit{\text{EBNM}}$ to denote this mapping:

$$\mathit{\text{EBNM}}(\mathit{x},\mathit{s})=(\hat{g},p)$$ (3.17)

Remark 1 Solving the EBNM problem is central to all our algorithms, so it is worth some study. A key point is that the EBNM problem provides an attractive and flexible way to induce shrinkage and/or sparsity in estimates of $\theta$. For example, if $\theta$ is truly sparse, with many elements at or near 0, then the estimate $\hat{g}$ will typically have considerable mass near 0, and the posterior means (3.15) will be “shrunk” strongly toward 0 compared with the original observations. In this sense solving the EBNM problem can be thought of as a model-based analogue of thresholding-based methods, with the advantage that by estimating $g$ from the data the EBNM approach automatically adapts to provide an appropriate level of shrinkage. These ideas have been used in wavelet denoising (Clyde and George, 2000; Johnstone et al., 2004; Johnstone and Silverman, 2005a; Xing et al., 2016), and false discovery rate estimation (Thomas et al., 1985; Stephens, 2017) for example. Here we apply them to matrix factorization problems.

3.4. Connecting the EBMF and EBNM Problems

The EBNM problem is well studied, and can be solved reasonably easily for many choices of $𝒢$ (e.g. Johnstone and Silverman, 2005b; Koenker and Mizera, 2014a; Stephens, 2017). In Section 4 we give specific examples; for now our main point is that if one can solve the EBNM problem for a particular choice of $𝒢$ then it can be used to implement Steps 5 and 6 in Algorithm 1 for the corresponding EBMF problem. The following Proposition formalizes this for Step 5 of Algorithm 1; a similar proposition holds for Step 6 (see also Appendix A).

Proposition 2 Step 5 in Algorithm 1 is solved by solving an EBNM problem. Specifically

$$\text{arg}\underset{q_l,g_l}{\text{max}} F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)=\mathit{\text{EBNM}}\left(\stackrel{\mathit{ˆ}}{\mathit{l}}\left(Y,\stackrel{-}{\mathit{f}},\stackrel{-}{{\mathit{f}}^2},\mathit{\tau }\right),{\mathit{s}}_l\left(\stackrel{-}{{\mathit{f}}^2},\mathit{\tau }\right)\right)$$ (3.18)

where the functions $\stackrel{\mathit{ˆ}}{\mathit{l}}:R^{n\times p}\times R^p\times R^p\times R^{n\times p}\to R^n$ and ${\mathit{s}}_l:R^p\times R^{n\times p}\to R^n$ are given by

$$\stackrel{\mathit{ˆ}}{\mathit{l}}(Y,\mathit{v},\mathit{w},\mathit{\tau }{)}_i≔\frac{\sum _j {\tau }_{ij}{Y}_{ij}{v}_j}{\sum _j {\tau }_{ij}{w}_j},$$ (3.19)

$${\mathit{s}}_l(\mathit{w},\mathit{\tau }{)}_i≔{\left(\sum _j {\tau }_{ij}{w}_j\right)}^{-0.5},$$ (3.20)

and $\stackrel{-}{\mathit{f}},\stackrel{-}{{\mathit{f}}^2}\in R^p$ denote the vectors whose elements are the first and second moments of $\mathit{f}$ under $q_f$:

$$\stackrel{-}{\mathit{f}}≔\left(E_{q_{\mathit{f}}}\left(f_j\right)\right)$$ (3.21)

$$\stackrel{-}{{\mathit{f}}^2}≔\left(E_{q_{\mathit{f}}}\left(f_j^2\right)\right).$$ (3.22)

Proof See Appendix A. ■

For intuition into where the EBNM in Proposition 2 comes from, consider estimating $\mathit{l},g_{\mathit{l}}$ in (2.1) with $\mathit{f}$ and $\mathit{\tau }$ known. The model then becomes $n$ independent regressions of the rows of $Y$ on $\mathit{f}$, and the maximum likelihood estimate for $\mathit{l}$ has elements:

$${\hat{l}}_i=\frac{\sum _j {\tau }_{ij}{Y}_{ij}{f}_j}{\sum _j {\tau }_{ij}{f}_j^2},$$ (3.23)

with standard errors

$$s_i={\left(\sum _j {\tau }_{ij}{f}_j^2\right)}^{-0.5}.$$ (3.24)

Further, it is easy to show that

$${\hat{l}}_i~N\left(l_i,s_i^2\right).$$ (3.25)

Combining (3.25) with the prior

$$l_1,\dots ,l_n{~}^{iid}{g}_l,g_l\in 𝒢$$ (3.26)

yields an EBNM problem.

The EBNM in Proposition 2 is the same as the EBNM (3.25)-(3.26), but with the terms $f_j$ and $f_j^2$ replaced with their expectations under $q_{\mathit{f}}$. Thus, the update for $\left(q_{\mathit{l}},g_{\mathit{l}}\right)$ in Algorithm 1, with $\left(q_{\mathit{f}},g_{\mathit{f}},\mathit{\tau }\right)$ fixed, is closely connected to solving the EBMF problem for “known $\mathit{f},\mathit{\tau }$ “.

3.5. Streamlined Implementation Using First and Second Moments

Although Algorithm 1, as written, optimizes over $\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}}\right)$, in practice each step requires only the first and second moments of the distributions $q_{\mathit{l}}$ and $q_{\mathit{f}}$. For example, the EBNM problem in Proposition 1 involves $\stackrel{-}{\mathit{f}}$ and $\stackrel{-}{{\mathit{f}}^2}$ and not $g_{\mathit{f}}$. Consequently, we can simplify implementation by keeping track of only those moments. In particular, when solving the normal means problem, $\mathit{\text{EBNM}}(\mathit{x},\mathit{s})$ in (3.17), we need only return the posterior first and second moments (3.15) and (3.16). This results in a streamlined and intuitive implementation, summarized in Algorithm 2.

Remark 3 Algorithm 2 has a very intuitive form: it has the flavor of an alternating least squares algorithm, which alternates between estimating $\mathit{l}$ given $\mathit{f}$ (Step 6) and $\mathit{f}$ given $\mathit{l}$ (Step 8), but with the addition of the ebnm step (Steps 7 and 9), which can be thought of as regularizing or shrinking the estimates: see Remark 1. This viewpoint highlights connections with related algorithms. For example, the (rank 1 version of the) SSVD algorithm from Yang et al. (2014) has a similar form, but uses a thresholding function in place of the ebnm function to induce shrinkage and/or sparsity.

3.6. The $K$-factor EBMF Model

It is straightforward to extend the variational approach to fit the general $K$ factor model (2.1)-(2.4). In brief, we introduce variational distributions $\left(q_{l_k},q_{f_k}\right)$ for $k=1,\dots ,K$, and then optimize the objective function $F\left(q_{l_1},g_{{\mathit{l}}_1},q_{{\mathit{f}}_1},g_{{\mathit{f}}_1};\dots ;q_{{\mathit{l}}_K},g_{{\mathit{l}}_K},q_{{\mathit{f}}_K},g_{{\mathit{f}}_K};\mathit{\tau }\right)$. Similar to the rank-1 model, this optimization can be done by iteratively updating parameters relating to a single loading or factor, keeping other parameters fixed. And again we simplify implementation by keeping track of only the first and second moments of the distributions $q_{{\mathit{l}}_k}$ and $q_{{\mathit{f}}_k}$, which we denote ${\stackrel{-}{\mathit{l}}}_k,{\stackrel{-}{\mathit{l}}}_k^2,{\stackrel{-}{\mathit{f}}}_k,{\stackrel{-}{\mathit{f}}}_k^2$. The updates to ${\stackrel{-}{\mathit{l}}}_k,{\stackrel{-}{\mathit{l}}}_k{}^{\text{(and }}{\stackrel{-}{\mathit{f}}}_k,{\stackrel{-}{\mathit{f}}}_k^2$) are essentially identical to those for fitting the rank 1 model above, but with $Y_{ij}$ replaced with the residuals obtained by removing the estimated effects of the other $k-1$ factors:

$$R_{ij}^k≔Y_{ij}-\sum _{k^{\prime }\ne k} {\stackrel{-}{\mathit{l}}}_{k^{\prime }i}{\stackrel{-}{\mathit{f}}}_{k^{\prime }j}.$$ (3.27)

Based on this approach we have implemented two algorithms for fitting the $K$-factor model. First, a simple “greedy” algorithm, which starts by fitting the rank 1 model, and then adds factors $k=2,\dots ,K$, one at a time, optimizing over the new factor parameters before moving on to the next factor. Second, a “backfitting” algorithm (Breiman and Friedman, 1985), which iteratively refines the estimates for each factor given the estimates for the other factors. Both algorithms are detailed in Appendix A.

3.7. Selecting $K$

An interesting feature of EB approaches to matrix factorization, noted by Bishop (1999), is that they automatically select the number of factors $K$. This is because the maximum likelihood solution to $g_{{\mathit{l}}_k},g_{{\mathit{f}}_k}$ is sometimes a point mass on 0 (provided $𝒢$ includes this distribution). Furthermore, the same is true of the solution to the variational approximation (see also Bishop, 1999; Stegle et al., 2012). This means that if $K$ is set sufficiently large then some loading/factor combinations will be optimized to be exactly 0. (Or, in the greedy approach, which adds one factor at a time, the algorithm will eventually add a factor that is exactly 0, at which point it terminates.)

Here we note that the variational approximation may be expected to result in conservative estimation (i.e. underestimation) of $K$ compared with the (intractable) use of maximum likelihood to estimate $g_l,g_{\mathit{f}}$. We base our argument on the simplest case: comparing $K=1$ vs $K=0$. Let ${\delta }_0$ denote the degenerate distribution with all its mass at 0. Note that the rank-1 factor model (2.1), with $g_{\mathit{l}}={\delta }_0$ (or $g_{\mathit{f}}={\delta }_0$) is essentially a “rank-0” model. Now note that the variational lower bound, $F$, is exactly equal to the log-likelihood when $g_{\mathit{l}}={\delta }_0$ (or $g_{\mathit{f}}={\delta }_0$). This is because if the prior is a point mass at 0 then the posterior is also a point mass, which trivially factorizes as a product of point masses, and so the variational family $𝒬$ includes the true posterior in this case. Since $F$ is a lower bound to the log-likelihood we have the following simple lemma:

Lemma 4 If $F\left(\hat{q},{\hat{g}}_l,{\hat{g}}_{\mathit{f}},\hat{\tau }\right)>F\left({\delta }_0,{\delta }_0,{\delta }_0,{\hat{\tau }}_0\right)$ then $l\left({\hat{g}}_l,{\hat{g}}_{\mathit{f}},\hat{\tau }\right)>l\left({\delta }_0,{\delta }_0,{\hat{\tau }}_0\right)$.

Proof

$$l\left({\hat{g}}_l,{\hat{g}}_{\mathit{f}},\hat{\tau }\right)\ge F\left(\hat{q},{\hat{g}}_l,{\hat{g}}_{\mathit{f}},\hat{\tau }\right)>F\left({\delta }_0,{\delta }_0,{\delta }_0,{\hat{\tau }}_0\right)=l\left({\delta }_0,{\delta }_0,{\hat{\tau }}_0\right)$$ (3.28)

■

Thus, if the variational approximation $F$ favors ${\hat{g}}_l,{\hat{g}}_{\mathit{f}},\hat{\tau }$ over the rank 0 model, then it is guaranteed that the likelihood would also favor ${\hat{g}}_l,{\hat{g}}_f,\hat{\tau }$ over the rank 0 model. In other words, compared with the likelihood, the variational approximation is conservative in terms of preferring the rank 1 model to the rank 0 model. This conservatism is a double-edged sword. On the one hand it means that if the variational approximation finds structure it should be taken seriously. On the other hand it means that the variational approximation could miss subtle structure.

In practice Algorithm 2 can converge to a local optimum of $F$ that is not as high as the trivial (rank 0) solution, $F\left({\delta }_0,{\delta }_0,{\delta }_0,{\hat{\tau }}_0\right)$. We can add a check for this at the end of Algorithm 2, and set ${\hat{g}}_l={\hat{g}}_f={\delta }_0$ and $\hat{\tau }={\hat{\tau }}_0$ when this occurs.

3.8. Identifiability

In EBMF each loading and factor is identifiable, at best, only up to a multiplicative constant (provided $𝒢$ is a scale family). Specifically, scaling the prior distributions $g_{{\mathit{f}}_k}$ and $g_{l_k}$ by $c_k$ and $1/c_k$ respectively results in the same marginal likelihood, and also results in a corresponding scaling of the posterior distribution on the factors ${\mathit{f}}_k$ and loadings ${\mathit{l}}_k$ (e.g. it scales the posterior first moments by $c_k,1/c_k$ and the second moments by $\left.c_k^2,1/c_k^2\right)$. However, this non-identifiability is not generally a problem, and if necessary it could be dealt with by re-scaling factor estimates to have norm 1.

4. Software Implementation: flash

We have implemented Algorithms 2, 4 and 5 in an R package, flash (“factors and loadings via adaptive shrinkage”). These algorithms can fit the EBMF model for any choice of distributional family ${𝒢}_l,{𝒢}_f$: the user must simply provide a function to solve the EBNM problem for these prior families.

One source of functions for solving the EBNM problem is the “adaptive shrinkage” (ashr) package, which implements methods from Stephens (2017). These methods solve the EBNM problem for several flexible choices of $𝒢$, including:

• $𝒢=𝒮N$, the set of all scale mixtures of zero-centered normals;

• $𝒢=𝒮𝒰$, the set of all symmetric unimodal distributions, with mode at 0;

• $𝒢=𝒰$, the set of all unimodal distributions, with mode at 0;

• $𝒢={𝒰}_{+}$, the set of all non-negative unimodal distributions, with mode at 0.

These methods are computationally stable and efficient, being based on convex optimization methods (Koenker and Mizera, 2014b) and analytic Bayesian posterior computations.

We have also implemented functions to solve the EBNM problem for additional choices of $𝒢$ in the package ebnm (https://github.com/stephenslab/ebnm). These include $𝒢$ being the “point-normal” family:

• $𝒢=𝒫N$, the set of all distributions that are a mixture of a point mass at zero and a normal with mean 0.

This choice is less flexible than those in ashr, and involves non-convex optimizations, but can be faster.

Although in this paper we focus our examples on sparsity-inducing priors with ${𝒢}_l={𝒢}_f=𝒢$ we note that our software makes it easy to experiment with different choices, some of which represent novel methodologies. For example, setting ${𝒢}_l={𝒰}_{+}$ and ${𝒢}_f=𝒮N$ yields an EB version of semi-non-negative matrix factorization (Ding et al., 2008), and we are aware of no existing EB implementations for this problem. Exploring the relative merits of these many possible options in different types of application will be an interesting direction for future work.

4.1. Missing Data

If some elements of $Y$ are missing, then this is easily dealt with. For example, the sums over $j$ in (3.19) and (3.20) are simply computed using only the $j$ for which $Y_{ij}$ is not missing. This corresponds to an assumption that the missing elements of $Y$ are “missing at random” (Rubin, 1976). In practice we implement this by setting ${\tau }_{ij}=0$ whenever $Y_{ij}$ is missing (and filling in the missing entries of $Y$ to an arbitrary number). This allows the implementation to exploit standard fast matrix multiplication routines, which cannot handle missing data. If many data points are missing then it may be helpful to exploit sparse matrix routines.

4.2. Initialization

Both Algorithms 2 and 4 require a rank 1 initialization procedure, init. Here, we use the softImpute function from the package softImpute (Mazumder et al., 2010), with penalty parameter $\lambda =0$, which essentially performs SVD when $Y$ is completely observed, but can also deal with missing values in $Y$.

The backfitting algorithm (Algorithm 5) also requires initialization. One option is to use the greedy algorithm to initialize, which we call “greedy+backfitting”.

5. Numerical Comparisons

We now compare our methods with several competing approaches. To keep these comparisons manageable in scope we focus attention on methods that aim to capture possible sparsity in $L$ and/or $F$. For EBMF we present results for two different shrinkage-oriented prior families, $𝒢$: the scale mixture of normals $(𝒢=𝒮N)$, and the point-normal family $(𝒢=𝒫N)$. We denote these flash and flash_pn respectively when we need to distinguish. In addition we consider Sparse Factor Analysis (SFA) (Engelhardt and Stephens, 2010), SFAmix (Gao et al., 2013), Nonparametric Bayesian Sparse Factor Analysis (NBSFA) (Knowles and Ghahramani, 2011), Penalized Matrix Decomposition (Witten et al., 2009) (PMD, implemented in the R package PMA), and Sparse SVD (Yang et al., 2014) (SSVD, implemented in $R$ package ssvd). Although the methods we compare against involve only a small fraction of the very large number of methods for this problem, the methods were chosen to represent a wide range of different approaches to inducing sparsity: SFA, SFAmix and NBSFA are three Bayesian approaches with quite different approaches to prior specification; PMD is based on a penalized likelihood with $L_1$ penalty on factors and/or loadings; and SSVD is based on iterative thresholding of singular vectors. We also compare with softImpute (Mazumder et al., 2010), which does not explicitly model sparsity in $L$ and $F$, but fits a regularized low-rank matrix using a nuclear-norm penalty. Finally, for reference we also use standard (truncated) SVD.

All of the Bayesian methods (flash, SFA, SFAmix and NBSFA) are “self-tuning”, at least to some extent, and we applied them here with default values. According to Yang et al. (2014) SSVD is robust to choice of tuning parameters, so we also ran SSVD with its default values, using the robust option (method=“method”). The softImpute method has a single tuning parameter ($\lambda$, which controls the nuclear norm penalty), and we chose this penalty by orthogonal cross-validation (OCV; Appendix B). The PMD method can use two tuning parameters (one for $\mathit{l}$ and one for f) to allow different sparsity levels in $\mathit{l}$ vs $\mathit{f}$. However, since tuning two parameters can be inconvenient it also has the option to use a single parameter for both $\mathit{l}$ and $\mathit{f}$. We used OCV to tune parameters in both cases, referring to the methods as PMD.cv2 (2 tuning parameters) and PMD.cv1 (1 tuning parameter).

5.1. Simple Simulations

5.1.1. A Single Factor Example

We simulated data with $n=200,p=300$ under the single-factor model (2.1) with sparse loadings, and a non-sparse factor:

$$l_i~{\pi }_0{\delta }_0+\left(1-{\pi }_0\right)\sum _{m=1}^5\frac{1}{5}N\left(0,{\sigma }_m^2\right)$$ (5.1)

$$f_j~N\left(\mathrm{0,1}\right)$$ (5.2)

where ${\delta }_0$ denotes a point mass on 0, and $\left({\sigma }_1^2,\dots ,{\sigma }_5^2\right)≔(\mathrm{0.25,0.5,1},\mathrm{2,4})$. We simulated using three different levels of sparsity on the loadings, using ${\pi }_0=\mathrm{0.9,0.3,0}$. (We set the noise precision $\tau =\mathrm{1,1}/\mathrm{16,1}/25$ in these three cases to make each problem not too easy and not too hard.)

We applied all methods to this rank-1 problem, specifying the true value $K=1$. (The NBSFA software does not provide the option to fix $K$, so is omitted here.) We compare methods in their accuracy in estimating the true low-rank structure $\left(B≔\mathit{l}{\mathit{f}}^T\right)$ using relative root mean squared error:

$$\text{RRMSE}(\hat{B},B)≔\sqrt{\frac{\sum _{i,j} {\left({\hat{B}}_{ij}-B_{ij}\right)}^2}{\sum _{i,j} B_{ij}^2}}.$$ (5.3)

Despite the simplicity of this simulation, the methods vary greatly in performance (Figure 1). Both versions of flash consistently outperform all the other methods across all scenarios (although softImpute performs similarly in the non-sparse case). The next best performances come from softImpute (SI.cv), PMD.cv2 and SFA, whose relative performances depend on the scenario. All three consistently improve on, or do no worse than, SVD. PMD.cv1 performs similarly to SVD. The SFAmix method performs very variably, sometimes providing very poor estimates, possibly due to poor convergence of the MCMC algorithm (it is the only method here that uses MCMC). The SSVD method consistently performs worse than simple SVD, possibly because it is more adapted to both factors and loadings being sparse (and possibly because, following Yang et al. 2014, we did not use CV to tune its parameters). Inspection of individual results suggests that the poor performance of both SFAmix and SSVD is often due to over-shrinking of non-zero loadings to zero.

Figure 1:

Boxplots comparing accuracy of flash with several other methods in a simple rank-1 simulation. This simulation involves a single dense factor, and a loading that varies from strong sparsity (90% zeros, left) to no sparsity (right). Accuracy is measured by difference in each methods RRMSE from the flash RRMSE, with smaller values indicating highest accuracy. The $y$ axis is plotted on a nonlinear (square-root) scale to avoid the plots being dominated by poorer-performing methods.

5.1.2. A Sparse Bi-cluster Example (Rank 3)

An important feature of our EBMF methods is that they estimate separate distributions $g_{\mathit{l}},g_{\mathit{f}}$ for each factor and each loading, allowing them to adapt to any combination of sparsity in the factors and loadings. This flexibility is not easy to achieve in other ways. For example, methods that use $\text{CV}$ are generally limited to one or two tuning parameters because of the computational difficulties of searching over a larger space.

To illustrate this flexibility we simulated data under the factor model (2.1) with $=$ $150,p=240,K=3,\tau =1/4$, and:

$${\mathit{l}}_{1,i}~N\left(0,2^2\right)i=1,\dots ,10$$ (5.4)

$${\mathit{l}}_{2,i}~N\left(\mathrm{0,1}\right)i=11,\dots ,60$$ (5.5)

$${\mathit{l}}_{3,i}~N\left(\mathrm{0,1}/2^2\right)i=61,\dots ,150$$ (5.6)

$${\mathit{f}}_{1,j}~N\left(\mathrm{0,1}/2^2\right)j=1,\dots ,80$$ (5.7)

$${\mathit{f}}_{2,j}~N\left(\mathrm{0,1}\right)j=81,\dots ,160$$ (5.8)

$${\mathit{f}}_{3,j}~N\left(0,2^2\right)j=161,\dots ,240,$$ (5.9)

with all other elements of ${\mathit{l}}_k$ and ${\mathit{f}}_k$ set to zero for $k=\mathrm{1,2},3$. This example has a sparse bi-cluster structure where distinct groups of samples are each loaded on only one factor (Figure 2a), and both the size of the groups and number of variables in each factor vary.

Figure 2:

Results from simulations with sparse bi-cluster structure $(K=3)$.

We applied flash, softImpute, SSVD and PMD to this example. (We excluded SFA and SFAmix since these methods do not model sparsity in both factors and loadings.) The results (Figure 2) show that again flash consistently outperforms the other methods, and again the next best is softImpute. On this example both SSVD and PMD outperform SVD. Although SSVD and PMD perform similarly on average, their qualitative behavior is different: PMD insufficiently shrink the 0 values, whereas SSVD shrinks the 0 values well but overshrinks some of the signal, essentially removing the smallest of the three loading/factor combinations (Figure 2 b).

5.2. Missing Data Imputation for Real Data Sets

Here we compare methods in their ability to impute missing data using five real data sets. In each case we “hold out” (mask) some of the data points, and then apply the methods to obtain estimates of the missing values. The data sets are as follows:

MovieLens $100K$ data, an (incomplete) $943\times 1682$ matrix of user-movie ratings (integers from 1 to 5) (Harper and Konstan, 2016). Most users do not rate most movies, so the matrix is sparsely observed (94% missing), and contains about 100K observed ratings. We hold out a fraction of the observed entries and assess accuracy of methods in estimating these. We centered and scaled the ratings for each user before analysis.

GTEx eQTL summary data, a $16069\times 44$ matrix of $Z$ scores computed testing association of genetic variants (rows) with gene expression in different human tissues (columns). These data come from the Genotype Tissue Expression (GTEx) project (Consortium et al., 2015), which assessed the effects of thousands of “eQTLs” across 44 human tissues. (An eQTL is a genetic variant that is associated with expression of a gene.) To identify eQTLs, the project tested for association between expression and every near-by genetic variant, each test yielding a $Z$ score. The data used here are the $Z$ scores for the most significant genetic variant for each gene (the “top” eQTL). See Section 5.3 for more detailed analyses of these data.

Brain Tumor data, a $43\times 356$ matrix of gene expression measurements on 4 different types of brain tumor (included in the denoiseR package, Josse et al., 2018). We centered each column before analysis.

Presidential address data, a $13\times 836$ matrix of word counts from the inaugural addresses of 13 US presidents (1940–2009) (also included in the denoiseR package, Josse et al., 2018). Since both row and column means vary greatly we centered and scaled both rows and columns before analysis, using the biScale function from softImpute.

Breast cancer data, a $251\times 226$ matrix of gene expression measurements from Carvalho et al. (2008), which were used as an example in the paper introducing NBSFA (Knowles and Ghahramani, 2011). Following Knowles and Ghahramani (2011) we centered each column (gene) before analysis.

Among the methods considered above, only flash, PMD and softImpute can handle missing data. We add NBSFA (Knowles and Ghahramani, 2011) to these comparisons. To emphasize the importance of parameter tuning we include results for PMD and softImpute with default settings (denoted PMD, SI) as well as using cross-validation (PMD.cv1, SI.cv).

For these real data the appropriate value of $K$ is, of course, unknown. Both flash and NBSFA automatically estimate $K$. For PMD and softImpute we specified $K$ based on the values inferred by flash and NBSFA. (Specifically, we used $K=\mathrm{10,30,20,10,40}$ respectively for the five data sets.)

We applied each method to all 5 data sets, using 10-fold OCV (Appendix B) to mask data points for imputation, repeated 20 times (with different random number seeds) for each data set. We measure imputation accuracy using root mean squared error (RMSE):

$$\text{RMSE}(\hat{Y},Y;\Omega )=\sqrt{\frac{1}{\left|\Omega \right|}\sum _{ij\in \Omega } {\left(Y_{ij}-{\hat{Y}}_{ij}\right)}^2}.$$ (5.10)

where $\Omega$ is the set of indices of the held-out data points.

The results are shown in Figure 3. Although the ranking of methods varies among data sets, flash, PMD.cv1 and SI.cv perform similarly on average, and consistently outperform NBSFA, which in turn typically outperforms (untuned) PMD and unpenalized softImpute. These results highlight the importance of appropriate tuning for the penalized methods, and also the effectiveness of the EB method in flash to provide automatic tuning.

Figure 3:

Comparison of the accuracy of different methods in imputing missing data. Each panel shows a boxplot of error rates (RMSE) for 20 simulations based on masking observed entries in a real data set.

In these comparisons, as in the simulations, the two flash methods typically performed similarly. The exception is the GTEx data, where the scale mixture of normals $(𝒢=𝒮N)$ performed worse. Detailed investigation revealed this to be due to a very small number of very large “outlier” imputed values, well outside the range of the observed data, which grossly inflated RMSE. These outliers were so extreme that it should be possible to implement a filter to avoid them. However, we did not do this here as it seems useful to highlight this unexpected behavior. (Note that this occurs only when data are missing, and even then only in one of the five data sets considered here.)

5.3. Sharing of Genetic Effects on Gene Expression Among Tissues

To illustrate flash in a scientific application, we applied it to the GTEx data described above, a $\mathrm{16,069}\times 44$ matrix of $Z$ scores, with $Z_{ij}$ reflecting the strength (and direction) of effect of eQTL $i$ in tissue $j$. We applied flash with $𝒢=𝒮N$ using the greedy+backfitting algorithm (i.e. the backfitting algorithm, initialized using the greedy algorithm).

The flash results yielded 26 factors (Figure 4–5) which summarize the main patterns of eQTL sharing among tissues (and, conversely, the main patterns of tissue-specificity). For example, the first factor has approximately equal weight for every tissue, and reflects the fact that many eQTLs show similar effects across all 44 tissues. The second factor has strong effects only in the 10 brain tissues, from which we infer that some eQTLs show much stronger effects in brain tissues than other tissues.

Figure 4:

Results from running flash on GTEx data (factors 1 – 8). The pve (“Percentage Variance Explained”) for loading/factor $k$ is defined as ${\text{pve}}_k≔s_k/\left({\sum }_k s_k+\right.\left.{\sum }_{ij} 1/{\tau }_{ij}\right)$ where $s_k≔{\sum }_{ij} {\left({\stackrel{-}{\mathit{l}}}_{ki}{\stackrel{-}{\mathit{f}}}_{kj}\right)}^2$. It is a measure of the amount of signal in the data captured by loading/factor $k$ (but its naming as “percentage variance explained” should be considered loose since the factors are not orthogonal).

Figure 5:

Results from running flash on GTEx data (factors 15 – 26)

Subsequent factors tend to be sparser, and many have a strong effect in only one tissue, capturing “tissue-specific” effects. For example, the 3rd factor shows a strong effect only in whole blood, and captures eQTLs that have much stronger effects in whole blood than other tissues. (Two tissues, “Lung” and “Spleen”, show very small effects in this factor but with the same sign as blood. This is intriguing since the lung has recently been found to make blood cells—see Lefrançais et al. 2017—and a key role of the spleen is storing of blood cells.) Similarly Factors 7, 11 and 14 capture effects specific to “Testis”, “Thyroid” and “Esophagus Mucosa” respectively.

A few other factors show strong effects in a small number of tissues that are known to be biologically related, providing support that the factors identified are scientifically meaningful. For example, factor 10 captures the two tissues related to the cerebellum, “Brain Cerebellar Hemisphere” and “Brain Cerebellum”. Factor 19 captures tissues related to female reproduction, “Ovary”, “Uterus” and “Vagina”. Factor 5 captures “Muscle Skeletal”, with small but concordant effects in the heart tissues (“Heart Atrial Appendage” and “Heart Left Ventricle”). Factor 4, captures the two skin tissues (“Skin Not Sun Exposed Suprapubic”, “Skin Sun Exposed Lower leg”) and also “Esophagus Mucosa”, possibly reflecting the sharing of squamous cells that are found in both the surface of the skin, and the lining of the digestive tract. In factor 24, “Colon Transverse” and “Small Intestine Terminal Ileum” show the strongest effects (and with same sign), reflecting some sharing of effects in these intestinal tissues. Among the 26 factors, only a few are difficult to interpret biologically (e.g. factor 8).

To highlight the benefits of sparsity, we contrast the flash results with those for softImpute, which was the best-performing method in the missing data assessments on these data, but which uses a nuclear norm penalty that does not explicitly reward sparse factors or loadings. The first eight softImpute factors are shown in Figure 6. The softImpute results—except for the first two factors—show little resemblance to the flash results, and in our view are harder to interpret.

Figure 6:

Results from running softImpute on GTEx data (factors 1–8). The factors are both less sparse and less interpretable than the flash results.

5.4. Computational Demands

It is difficult to make general statements about computational demands of our methods, because both the number of factors and number of iterations per factor can vary considerably depending on the data. However, to give a specific example, running our current implementation of the greedy algorithm on the GTEx data (a 16,000 by 44 matrix) takes about 140s (wall time) for $𝒢=𝒫N$ and 650s for $𝒢=𝒮N$ (on a 2015 MacBook Air with a 2.2 GHz Intel Core i7 processor and 8Gb RAM). By comparison, a single run of softImpute without CV takes 2–3s, so a naive implementation of 5-fold CV with 10 different tuning parameters and 10 different values of $K$ would take over 1000s (although one could improve on this by use of warm starts for example).

6. Discussion

Here we discuss some potential extensions or modifications of our work.

6.1. Orthogonality Constraint

Our formulation here does not require the factors or loadings to be orthogonal. In scientific applications we do not see any particular reason to expect underlying factors to be orthogonal. However, imposing such a constraint could have computational or mathematical advantages. Formally adding such a constraint to our objective function seems tricky, but it would be straightforward to modify our algorithms to include an orthogonalization step each update. This would effectively result in an EB version of the SSVD algorithms in Yang et al. (2014), and it seems likely to be computationally faster than our current approach. One disadvantage of this approach is that it is unclear what optimization problem such an algorithm would solve (but the same is true of SSVD, and our algorithms have the advantage that they deal with missing data.)

6.2. Non-negative Matrix Factorization

We focused here on the potential for EBMF to induce sparsity on loadings and factors. However, EBMF can also encode other assumptions. For example, to assume the loadings and factors are non-negative, simply restrict $𝒢$ to be a family of non-negative-valued distributions, yielding “Empirical Bayes non-negative Matrix Factorization” (EBNMF). Indeed, the ashr software can already solve the EBNM problem for some such families $𝒢$, and so flash already implements EBNMF. In preliminary assessments we found that the greedy approach is problematic here: the non-negative constraint makes it harder for later factors to compensate for errors in earlier factors. However, it is straightforward to apply the backfitting algorithm to fit EBNMF, with initialization by any existing NMF method. The performance of this approach is an area for future investigation.

6.3. Tensor Factorization

It is also straightforward to extend EBMF to tensor factorization, specifically a CANDE-COMP/PARAFAC decomposition (Kolda and Bader, 2009):

$$Y_{ijm}=\sum _{k=1}^K l_{ki}{f}_{kj}{h}_{km}+E_{ijm}$$ (6.1)

$$l_{k1},\dots ,l_{kn}{~}^{iid}{g}_{{\mathit{l}}_k},g_{{\mathit{l}}_k}\in 𝒢$$ (6.2)

$$f_{k1},\dots ,f_{kp}{~}^{iid}{g}_{{\mathit{f}}_k},g_{{\mathit{f}}_k}\in 𝒢$$ (6.3)

$$h_{k1},\dots ,h_{kr}~{}^{iid}{g}_{{\mathit{h}}_k},g_{{\mathit{h}}_k}\in 𝒢$$ (6.4)

$$E_{ijm}{~}^{iid}N\left(\mathrm{0,1}/{\tau }_{ijm}\right).$$ (6.5)

The variational approach is easily extended to this case (a generalization of methods in Hore et al., 2016), and updates that increase the objective function can be constructed by solving an EBNM problem, similar to EBMF. It seems likely that issues of convergence to local optima, and the need for good initializations, will need some attention to obtain good practical performance. However, results in Hore et al. (2016) are promising, and the automatic-tuning feature of EB methods seems particularly attractive here. For example, extending PMD to this case—allowing for different sparsity levels in $l,f$ and $h—$would require 3 penalty parameters even in the rank 1 case, making it difficult to tune by CV.

6.4. Non-Gaussian Errors

It is also possible to extend the variational approximations used here to fit non-Gaussian models, such as binomial data; see for example Jaakkola and Jordan (2000); Seeger and Bouchard (2012); Klami (2015). The extension of our EB methods using these ideas is detailed in Wang (2017).

Appendix A. Variational EBMF with $\mathit{K}$ Factors

Here we describe in detail the variational approach to the $K$ factor model, including deriving updates that we use to optimize the variational objective. (These derivations naturally include the $K=1$ model as a special case, and our proof of Proposition 5 below includes Proposition 2 as a special case.)

Let $q_l,q_f$ denote the variational distributions on the $K$ loadings/factors:

$$q_{\mathit{l}}\left({\mathit{l}}_1,\cdots ,{\mathit{l}}_K\right)=\prod _k q_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)$$ (A.1)

$$\begin{array}{rr}& \\ q_{\mathit{f}}\left({\mathit{f}}_1,\dots ,{\mathit{f}}_K\right)& =\prod _k q_{{\mathit{f}}_k}\left({\mathit{f}}_k\right)\end{array}.$$ (A.2)

The objective function $F\left(3.4\right)$ is thus a function of $q_{\mathit{l}}=\left(q_{{\mathit{l}}_1},\dots ,q_{{\mathit{l}}_K}\right),q_{\mathit{f}}=\left(q_{{\mathit{f}}_1},\dots ,q_{{\mathit{f}}_K}\right)$, $g_{\mathit{l}}=\left(g_{{\mathit{l}}_1},\dots ,g_{{\mathit{l}}_K}\right)$ and $g_{\mathit{f}}=\left(g_{{\mathit{f}}_1},\dots ,g_{{\mathit{f}}_K}\right)$, as well as the precision $\mathit{\tau }$:

$$F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)=\int \prod _k{q}_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)q_{{\mathit{f}}_k}\left({\mathit{f}}_k\right)\text{log}\frac{p\left(Y,\mathit{l},\mathit{f};g_{{\mathit{l}}_1},g_{{\mathit{f}}_1},\cdots ,g_{{\mathit{l}}_K},g_{{\mathit{f}}_K},\mathit{\tau }\right)}{{\prod }_k{q}_{{\mathit{l}}_k}({\mathit{l}}_k)q_{{\mathit{f}}_k}\left({\mathit{f}}_k\right)}d{\mathit{l}}_{k}d{\mathit{f}}_k,$$ (A.3)

$$=E_{q_{\mathit{l}},q_{\mathit{f}}}\text{log}p(Y\mid \mathit{l},\mathit{f};\mathit{\tau })+\sum _k E_{q_{{\mathit{l}}_k}}\text{log}\frac{g_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)}{q_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)}+\sum _k E_{q_{{\mathit{f}}_k}}\text{log}\frac{g_{{\mathit{f}}_k}\left({\mathit{f}}_k\right)}{q_{{\mathit{f}}_k}\left({\mathit{f}}_k\right)}.$$ (A.4)

We optimize $F$ by iteratively updating parameters relating to $\mathit{\tau }$, a single loading $k$ $\left(q_{l_k},g_{{\mathit{l}}_k}\right)$ or factor $k\left(q_{{\mathit{f}}_k},g_{{\mathit{f}}_k}\right)$, keeping other parameters fixed. We simplify implementation by keeping track of only the first and second moments of the distributions $q_{l_k}$ and $q_{f_k}$, which we denote ${\stackrel{-}{\mathit{l}}}_k,{\overline{{\mathit{l}}^2}}_k,{\stackrel{-}{\mathit{f}}}_k,{\stackrel{-}{{\mathit{f}}^2}}_k$. We now describe each kind of update in turn.

A.1. Updates for Precision Parameters

Here we derive updates to optimize $F$ over the precision parameters $\mathit{\tau }$. Focusing on the parts of $F$ that depend on $\mathit{\tau }$ gives:

$$F\left(\tau \right)=E_{q_{\mathit{l}}}{E}_{q_{\mathit{f}}}\sum _{ij}0.5\text{log}\left({\tau }_{ij}\right)-0.5{\tau }_{ij}{\left(Y_{ij}-\sum _k{l}_{ki}{f}_{kj}\right)}^2+\text{const}$$ (A.5)

$$=0.5\sum _{ij} \left[\text{log}\left({\tau }_{ij}\right)+{\tau }_{ij}{\overline{R^2}}_{ij}\right]+\text{const}$$ (A.6)

where $\stackrel{-}{R^2}$ is defined by:

$${\overline{R^2}}_{ij}≔E_{q_{\mathit{l}},q_{\mathit{f}}}\left[{\left(Y_{ij}-\sum _{k=1}^K{\mathit{l}}_{ki}{\mathit{f}}_{kj}\right)}^2\right]$$ (A.7)

$$={\left(Y_{ij}-\sum _k{\overline{\mathit{l}}}_{ki}{\overline{\mathit{f}}}_{kj}\right)}^2-\sum _k{\left({\overline{\mathit{l}}}_{ki}\right)}^2{\left({\overline{\mathit{f}}}_{kj}\right)}^2+\sum _k{\overline{{\mathit{l}}^2}}_{ki}{\overline{{\mathit{f}}^2}}_{kj}.$$ (A.8)

If we constrain $\mathit{\tau }\in 𝒯$ then we have

$$\widehat{\tau }=\text{arg}\underset{\tau \in 𝒯}{\text{max}}\sum _{ij}\left[\text{log}\left({\tau }_{ij}\right)-{\tau }_{ij}{\overline{R^2}}_{ij}\right].$$ (A.9)

For example, assuming constant precision ${\mathit{\tau }}_{ij}=\mathit{\tau }$ yields:

$$\widehat{\tau }=\frac{NP}{{\sum }_{ij}{\overline{R^2}}_{ij}}.$$ (A.10)

Assuming column-specific precisions $\left({\mathit{\tau }}_{ij}={\mathit{\tau }}_j\right)$, which is the default in our software, yields:

$${\widehat{\tau }}_j=\frac{N}{{\sum }_i{\overline{R^2}}_{ij}}.$$ (A.11)

Other variance structures are considered in Appendix A.5 of Wang (2017).

A.2. Updating Loadings and Factors

The following Proposition, which generalizes Proposition 2 in the main text, shows how updates for loadings (and factors) for the $K$-factor EBMF model can be achieve by solving an EBNM problem.

Proposition 5 For the $K$-factor model, arg ${\text{max}}_{q_{l_k},g_{l_k}} F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$ is solved by solving an EBNM problem. Specifically

$$\text{arg}\underset{q_{{\mathit{l}}_k},g_{{\mathit{l}}_k}}{\text{max}}F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)=\mathit{\text{EBNM}}\left(\widehat{\mathit{l}}\left(R^k,{\overline{\mathit{f}}}_k,{\overline{{\mathit{f}}^2}}_k,\mathit{\tau }\right),{\mathit{s}}_l\left({\overline{{\mathit{f}}^2}}_k,\mathit{\tau }\right)\right)$$ (A.12)

where the functions $\stackrel{\mathit{ˆ}}{\mathit{l}}$ and ${\mathit{s}}_l$ are given by (3.19) and (3.20), ${\stackrel{-}{\mathit{f}}}_k,{\stackrel{-}{{\mathit{f}}^2}}_k\in R^p$ denote the vectors whose elements are the first and second moments of ${\mathit{f}}_k$ under $q_{{\mathit{f}}_k}$, and $R^k$ denotes the residual matrix (3.27).

Similarly, $\text{arg}{\text{max}}_{q_{{\mathit{f}}_k},g_{f_k}} F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$ is solved by solving an EBNM problem. Specifically,

$$\text{arg}\underset{q_{{\mathit{f}}_k},g_{{\mathit{f}}_k}}{\text{max}}F\left(q_{\mathit{l}},q_{\mathit{f}},g_l,g_f,\mathit{\tau }\right)=\mathit{\text{EBNM}}\left(\widehat{\mathit{f}}\left(R^k,{\overline{\mathit{l}}}_k,{\overline{{\mathit{l}}^2}}_k,\mathit{\tau }\right),{\mathit{s}}_f\left({\overline{{\mathit{l}}^2}}_k,\mathit{\tau }\right)\right)$$ (A.13)

where the functions $\stackrel{\mathit{ˆ}}{\mathit{f}}:R^{n\times p}\times R^n\times R^n\times R^{n\times p}\to R^p$ and ${\mathit{s}}_f:R^n\times R^{n\times p}\to R^p$ are given by

$$\begin{array}{rr}\stackrel{\mathit{ˆ}}{\mathit{f}}(Y,\mathit{v},\mathit{w},\mathit{\tau }{)}_j& ≔\frac{\sum _i {\tau }_{ij}{Y}_{ij}{v}_i}{\sum _i {\tau }_{ij}{w}_i},\\ & \end{array}$$

$${\mathit{s}}_{\mathit{f}}(\mathit{w},\mathit{\tau }{)}_j≔{\left(\sum _i {\tau }_{ij}{w}_i\right)}^{-0.5}.$$ (A.15)

A.2.1. A Lemma on the Normal Means Problem

To prove Proposition 5 we introduce a lemma that characterizes the solution of the normal means problem in terms of an objective that is closely related to the variational objective.

Recall that the EBNM model is:

$$\mathit{x}=\mathit{\theta }+\mathit{e}$$ (A.16)

$$\begin{array}{c}\\ {\theta }_1,\dots ,{\theta }_n{~}^{iid}g,g\in 𝒢\end{array}.$$ (A.17)

where $e_i~N\left(0,s_i^2\right)$.

Solving the EBNM problem involves estimating $g$ by maximum likelihood:

$$\hat{g}=\text{arg}\underset{g\in 𝒢}{\text{max}} l\left(g\right),$$ (A.18)

where

$$l\left(g\right)=\text{log}p(\mathit{x}\mid g).$$ (A.19)

It also involves finding the posterior distributions:

$$p(\mathit{\theta }\mid \mathit{x},\hat{g})=\prod _j p\left({\theta }_j\mid \mathit{x},\hat{g}\right)\propto \prod _j \hat{g}\left({\theta }_j\right)p\left(x_j\mid {\theta }_j,s_j\right).$$ (A.20)

Lemma 6 Solving the EBNM problem also solves:

$$\underset{q_{\theta },g\in 𝒢}{\text{max}} F^{NM}\left(q_{\theta },g\right)$$ (A.21)

where

$$F^{NM}\left(q_{\theta },g\right)=E_{q_{\theta }}\left[-\frac{1}{2}\sum _j \left(A_j{\theta }_j^2-2B_j{\theta }_j\right)\right]+E_{q_{\theta }}\text{log}\frac{g\left(\mathit{\theta }\right)}{q_{\theta }\left(\mathit{\theta }\right)}+\mathit{\text{const}}$$ (A.22)

with $A_j=1/s_j^2$ and $B_j=x_j/s_j^2$, and $g\left(\mathit{\theta }\right)≔{\prod }_j g\left({\theta }_j\right)$.

Equivalently, (A.21)-(A.22) is solved by $g=\hat{g}$ in (A.18) and $q_{\theta }=p(\mathit{\theta }\mid \mathit{x},\hat{g})$ in (A.20), with $x_j=B_j/A_j$ and $s_j^2=1/A_j$.

Proof The log likelihood can be written as

$$l\left(g\right)≔\text{log}\left[p\right(\mathit{x}\mid g\left)\right]$$ (A.23)

$$=\text{log}\left[p\right(\mathit{x},\mathit{\theta }\mid g)/p(\mathit{\theta }\mid \mathit{x},g\left)\right]$$ (A.24)

$$=\int q_{\theta }\left(\mathit{\theta }\right)\text{log}\frac{p(\mathit{x},\mathit{\theta }\mid g)}{p(\mathit{\theta }\mid \mathit{x},g)}d\mathit{\theta }$$ (A.25)

$$=\int q_{\theta }\left(\mathit{\theta }\right)\text{log}\frac{p(\mathit{x},\mathit{\theta }\mid g)}{q_{\theta }\left(\mathit{\theta }\right)}d\mathit{\theta }+\int q_{\theta }\left(\mathit{\theta }\right)\text{log}\frac{q_{\theta }\left(\mathit{\theta }\right)}{p(\mathit{\theta }\mid \mathit{x},g)}d\mathit{\theta }$$ (A.26)

$$=F^{\text{NM}}\left(q_{\theta },g\right)+D_{KL}\left(q_{\theta }\parallel p_{\theta \mid x,g}\right)$$ (A.27)

where

$$F^{\text{NM}}\left(q_{\theta },g\right)=\int q_{\theta }\left(\mathit{\theta }\right)\text{log}\frac{p(\mathit{x},\mathit{\theta }\mid g)}{q_{\theta }\left(\mathit{\theta }\right)}d\mathit{\theta }$$ (A.28)

and

$$D_{KL}\left(q_{\theta }\left|\right|p_{\theta \mid x,g}\right)=-\int q_{\theta }\left(\mathit{\theta }\right)\text{log}\frac{p(\mathit{\theta }\mid \mathit{x},g)}{q_{\theta }\left(\mathit{\theta }\right)}d\mathit{\theta }$$ (A.29)

Here $p_{\theta \mid x,g}$ denotes the posterior distribution $p(\mathit{\theta }\mid \mathit{x},g)$. This identity holds for any distribution $q_{\theta }\left(\mathit{\theta }\right)$.

Rearranging (A.27) gives:

$$F^{\text{NM}}\left(q_{\theta },g\right)=l\left(g\right)-D_{KL}\left(q_{\theta }\parallel p_{\theta \mid x,g}\right).$$ (A.30)

Since $D_{KL}\left(q_{\theta }\left|\right|p_{\theta \mid x,g}\right)\ge 0$, with equality when $q_{\theta }=p_{\theta \mid x,g},F^{\text{NM}}\left(q_{\theta },g\right)$ is maximized over $q_{\theta }$ by setting $q_{\theta }=p_{\theta \mid x,g}$. Further

$$\underset{q_{\theta }}{\text{max}} F^{\text{NM}}\left(q_{\theta },g\right)=l\left(g\right),$$ (A.31)

so

$$\text{arg}\underset{g\in 𝒢}{\text{max}} \underset{q_{\theta }}{\text{max}} F^{\text{NM}}\left(q_{\theta },g\right)=\text{arg}\underset{g\in 𝒢}{\text{max}} l\left(g\right)=\hat{g}.$$ (A.32)

It remains only to show that $F^{\text{NM}}$ has the form (A.22).

By (A.16) and (A.17), we have

$$\text{log}p(\mathit{x},\mathit{\theta }\mid g)=-\frac{1}{2}\sum _j s_j^{-2}{\left(x_j-{\theta }_j\right)}^2+\text{log}g\left(\mathit{\theta }\right)+\text{const}.$$ (A.33)

Thus

$$F^{\text{NM}}\left(q_{\theta },g\right)=E_{q_{\theta }}\left[-\frac{1}{2}\sum _j \left(A_j{\theta }_j^2-2B_j{\theta }_j\right)\right]+{\text{E}}_{q_{\theta }}\text{log}\frac{g\left(\mathit{\theta }\right)}{q_{\theta }\left(\mathit{\theta }\right)}+\text{const}.$$ (A.34)

■

A.2.2. Proof of Proposition 5

We are now ready to prove Proposition 5.

Proof We prove the first part of the proposition since the proof for the second part is essentially the same.

The objective function (A.3) is:

$$F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)=E_{q_{\mathit{l}},q_{\mathit{f}}}\text{log}p\left(Y\mid \mathit{l},\mathit{f};\mathit{\tau }\right)+\sum _k{E}_{q_{{\mathit{l}}_k}}\text{log}\frac{g_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)}{q_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)}+\sum _k{E}_{q_{{\mathit{f}}_k}}\text{log}\frac{g_{{\mathit{f}}_k}\left({\mathit{f}}_k\right)}{q_{{\mathit{f}}_k}\left({\mathit{f}}_k\right)}$$ (A.35)

$$=E_{q_{{\mathit{l}}_k}}\left[-\frac{1}{2}\sum _i \left(A_{ik}{l}_{ki}^2-2B_{ik}{l}_{ki}\right)\right]+E_{q_{l_k}}\text{log}\frac{g_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)}{q_{{\mathit{l}}_k}\left({\mathit{l}}_k\right)}+C_1$$ (A.36)

where $C_1$ is a constant with respect to $q_{l_k},g_{l_k}$ and

$$A_{ik}=\sum _j {\tau }_{ij}{\text{E}}_{q_{\mathit{f}}}\left(f_{kj}^2\right)$$ (A.37)

$$B_{ik}=\sum _j {\tau }_{ij}\left(R_{ij}^k{\text{E}}_{q_{\mathit{f}}}{f}_{kj}\right).$$ (A.38)

Based on Lemma 6, we can solve this optimization problem (A.36) by solving the EBNM problem with:

$$x_i=\frac{\sum _j {\tau }_{ij}\left(R_{ij}^k{\text{E}}_{q_{\mathit{f}}}{f}_{kj}\right)}{\sum _j {\tau }_{ij}{\text{E}}_{q_{\mathit{f}}}\left(f_{kj}^2\right)}$$ (A.39)

$$s_i^2=\frac{1}{\sum _j {\tau }_{ij}{\text{E}}_{q_{\mathit{f}}}\left(f_{kj}^2\right)}.$$ (A.40)

■

A.3. Algorithms

Just as with the rank 1 EBMF model, the updates for the rank $K$ model require only the first and second moments of the variational distributions $q$. Thus we implement the updates in algorithms that keep track of the first moments $\left(\stackrel{-}{\mathit{l}}≔\left({\stackrel{-}{\mathit{l}}}_1,\dots ,{\stackrel{-}{\mathit{l}}}_K\right)\right.$ and $\left.\stackrel{-}{\mathit{f}}≔\left({\stackrel{-}{\mathit{f}}}_1,\dots ,{\stackrel{-}{\mathit{f}}}_K\right)\right)$ and second moments ($\overline{l^2}≔\left({\overline{l^2}}_1,\dots ,{\overline{l^2}}_K\right)$ and $\overline{f^2}≔\left({\overline{f^2}}_1,\dots ,{\overline{f^2}}_K\right)$), and the precision $\mathit{\tau }$.

Algorithm 3 implements a basic update for $\mathit{\tau }$, and for the parameters relating to a single factor $k\left({\overline{l}}_k,{\overline{l^2}}_K,{\overline{f}}_k,{\overline{f^2}}_K\right)$. Note that the latter updates are identical to the updates for fitting the single factor EBMF model, but with $Y_{ij}$ replaced with the residuals obtained by removing the estimated effects of the other $k-1$ factors.

Based on these basic updates we implemented two algorithms for fitting the $K$-factor EBMF model: the greedy algorithm, and the backfitting algorithm, as follows.

A.3.1. Greedy Algorithm

The greedy algorithm is a forward procedure that, at the $k$th step, adds new factors and loadings ${\mathit{l}}_k,{\mathit{f}}_k$ by optimizing over $q_{{\mathit{l}}_k},q_{{\mathit{f}}_k},g_{{\mathit{l}}_k},g_{{\mathit{f}}_k}$ while keeping the distributions related to previous factors fixed. Essentially this involves fitting the single-factor model to the residuals obtained by removing previous factors. The procedure stops adding factors when the estimated new factors (or loadings) are identically zero. The algorithm as follows:

A.3.2. Backfitting Algorithm

The backfitting algorithm iteratively refines a fit of $K$ factors and loadings, by updating them one at a time, at each update keeping the other loadings and factors fixed. The name comes from its connection with the backfitting algorithm in Breiman and Friedman (1985), specifically the fact that it involves iteratively re-fitting to residuals.

A.4. Objective Function Computation

The algorithms above all involve updates that will increase (or, at least, not decrease) the objective function $F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)$. However, these updates do not require computing the objective function itself. In iterative algorithms it can be helpful to compute the objective function to monitor convergence (and as a check on implementation). In this subsection we describe how this can be done. In essence, this involves extending the solver of the EBNM problem to also return the value of the log-likelihood achieved in that problem (which is usually not difficult).

The objective function of the EBMF model is:

$$F\left(q_{\mathit{l}},q_{\mathit{f}},g_{\mathit{l}},g_{\mathit{f}},\mathit{\tau }\right)=E_{q_l,q_{\mathit{f}}}\text{log}p(Y\mid \mathit{l},\mathit{f};\mathit{\tau })+E_{q_l}\text{log}\frac{g_{\mathit{l}}\left(\mathit{l}\right)}{q_{\mathit{l}}\left(\mathit{l}\right)}+E_{q_{\mathit{f}}}\text{log}\frac{g_{\mathit{f}}\left(\mathit{f}\right)}{q_{\mathit{f}}\left(\mathit{f}\right)}$$ (A.41)

The calculation of $E_{q_l,q_f}\text{log}p(Y\mid \mathit{l},\mathit{f};\mathit{\tau })$ is straightforward and $E_{q_l}\text{log}\frac{g_l\left(l\right)}{q_l\left(l\right)}$ and $E_{q_l}\text{log}\frac{g_l\left(l\right)}{q_l\left(l\right)}$ can be calculated using the log-likelihood of the EBNM model using the following Lemma 7.

Lemma 7 Suppose $\hat{g},q$ solves the EBNM problem with data $(\mathit{x},\mathit{s})$:

$$(\hat{g},q)=\mathit{\text{EBNM}}(\mathit{x},\mathit{s}),$$ (A.42)

where $q≔\left(q_1,\dots ,q_n\right)$ are the estimated posterior distributions of the normal means parameters ${\theta }_1,\dots ,{\theta }_n$. Then

$$E_q\left(\text{log}\left(\prod _j \hat{g}\left({\theta }_j\right)/\prod _j q_j\left({\theta }_j\right)\right)\right)=l(\hat{g};\mathit{x},\mathit{s})+\frac{1}{2}\sum _j \text{log}\left(2\pi s_j^2\right)+\left(1/s_j^2\right)\left(x_j^2+E_q\left({\theta }_j^2\right)-2x_j{E}_q\left({\theta }_j\right)\right.$$ (A.43)

where $l(\hat{g};\mathit{x},s)$ is the log of the likelihood for the normal means problem (A.27).

Proof We have from (A.28)

$$F^{\text{NM}}\left(q_{\theta },\widehat{g}\right)=\int q_{\theta }\left(\theta \right)\text{log}\frac{p\left(x,\theta \mid \widehat{g}\right)}{q_{\theta }\left(\theta \right)}d\theta$$ (A.44)

$$=\int q_{\theta }\left(\mathit{\theta }\right)\text{log}\frac{p(\mathit{x}\mid \mathit{\theta })\hat{g}\left(\mathit{\theta }\right)}{q_{\theta }\left(\mathit{\theta }\right)}d\mathit{\theta }$$ (A.45)

$$=E_q\left(\text{log}\left(\prod _j \hat{g}\left({\theta }_j\right)/\prod _j q_j\left({\theta }_j\right)\right)\right)-\frac{1}{2}{E}_q\left[\sum _j \text{log}\left(2\pi s_j^2\right)+\left(1/s_j^2\right){\left(x_j-{\theta }_j\right)}^2\right]$$ (A.46)

And the result follows from noting that $F^{\text{NM}}(\hat{q},\hat{g})=l\left(\hat{g}\right)$. ■

A.5. Inference with Penalty Term

Conceivably, in some settings one might like to encourage solutions to the EBMF problem be sparser than the maximum-likelihood estimates for $g_{\mathit{l}},g_{\mathit{f}}$ would produce. This could be done by extending the EBMF model to introduce a penalty term on the distributions $g_l,g_f$ so that the maximum likelihood estimates are replaced by maximizing a penalized likelihood. We are not advocating for this approach, but it is straightforward given existing machinery, and so we document it here for completeness.

Let $h_{\mathit{l}}\left(g_{\mathit{l}}\right)$ and $h_{\mathit{f}}\left(g_{\mathit{f}}\right)$ denote penalty terms on $g_{\mathit{l}}$ and $g_{\mathit{f}}$, so the penalized log-likelihood would be:

$$\begin{gathered} l\left(g_l,g_f,\tau \right)≔\text{log}\left[p\left(Y\mid g_l,g_f,{\tau }^2\right)\right]+h_l\left(g_l\right)+h_f\left(g_f\right) \\ =F\left(q,g_l,g_f,{\tau }^2\right)+h_l\left(g_l\right)+h_f\left(g_f\right)+D_{KL}\left(q\parallel p\right) \end{gathered}$$ (A.47)

where $F\left(q,g_l,g_{\mathit{f}},{\mathit{\tau }}^2\right)$ and $D_{KL}(q\parallel p)$ are defined in (3.4) and (3.5). And the corresponding penalized variational objective is:

$$\text{max}F\left(q,g_{\mathit{l}},g_{\mathit{f}},{\mathit{\tau }}^2\right)+h_{\mathit{l}}\left(g_{\mathit{l}}\right)+h_{\mathit{f}}\left(g_{\mathit{f}}\right).$$ (A.48)

It is straightforward to modify the algorithms above to maximize this penalized objective: simply modify the EBNM solvers to solve a corresponding penalized normal means problem. That is, instead of estimating the prior $g$ by maximum likelihood, the EBNM solver must now maximize the penalized log-likelihood:

$$\hat{g}=\text{arg}\underset{g\in 𝒢}{\text{max}} l_{\text{EBNM}}\left(g\right)+h\left(g\right),$$ (A.49)

where $l_{\text{EBNM}}$ denote the log-likelihood for the EBNM problem. (The computation of the posterior distributions given $\hat{g}$ is unchanged).

For example, the ashr software (Stephens, 2017) provides the option to include a penalty on $g$ to encourage overestimation of the size of the point mass on zero. This penalty was introduced to ensure conservative behavior in False Discovery Rate applications of the normal means problem. It is unclear that such a penalty is desirable in the matrix factorization application. However, the above discussion shows that using this penalty (e.g. within the ebnm function used by the greedy or backfitting algorithms) can be thought of as solving a penalized version of the EBMF problem.

Appendix B. Orthogonal Cross Validation

Cross-validation assessments involving “holding out” (hiding) data from methods. Here we introduce a novel approach to selecting the data to be held out, which we call Orthogonal Cross Validation (OCV). Although not the main focus of our paper, we believe that OCV is a novel and appealing approach to selecting hold-out data for factor models, e.g. when using CV to select an appropriate dimension $K$ for dimension reduction methods, as in Owen and Wang (2016).

Generic $k$-fold CV involves randomly dividing the data matrix into $k$ parts and then, for each part, training methods on the other $k-1$ parts before assessing error on that part, as in Algorithm 6.

The novel part of OCV is in how to choose the “hold-out” pattern. We randomly divide the columns and rows into $k$ sets. and put these sets into $k$ orthogonal parts, and then take all $Y_{ij}$ with the chosen column and row indices as “hold-out” $Y_{\left(i\right)}$.

To illustrate this scheme, we take 3-fold CV as an example. We randomly divide the columns into 3 sets and the rows into 3 sets as well. The data matrix $Y$ is divided into 9 partition (by row and column permutation):

$$Y=\left(\begin{array}{lll}{Y}_{11}& Y_{12}& Y_{13}\\ Y_{21}& Y_{22}& Y_{23}\\ Y_{31}& Y_{32}& Y_{33}\end{array}\right)$$

Then $Y_{\left(1\right)}=\left\{Y_{11},Y_{22},Y_{33}\right\},Y_{\left(2\right)}=\left\{Y_{12},Y_{23},Y_{31}\right\}$ and $Y_{\left(3\right)}=\left\{Y_{13},Y_{21},Y_{32}\right\}$ are orthogonal to each other. Then the data matrix $Y$ is marked as:

$$Y=\left(\begin{array}{lll}{Y}_{\left(1\right)}& Y_{\left(2\right)}& Y_{\left(3\right)}\\ Y_{\left(3\right)}& Y_{\left(1\right)}& Y_{\left(2\right)}\\ Y_{\left(2\right)}& Y_{\left(3\right)}& Y_{\left(1\right)}\end{array}\right)$$

In OCV, each fold $k,Y_{\left(k\right)}$ contains equally balanced part of data matrix and includes all the row and column indices. This ensures that all i’s and $j$’s are included into each $Y_{-\left(k\right)}$. In 3-fold OCV, we have:

$$Y=\left(\begin{array}{lll}{Y}_{11}\hfill & Y_{12}\hfill & Y_{13}\hfill \\ Y_{21}\hfill & Y_{22}\hfill & Y_{23}\hfill \\ Y_{31}\hfill & Y_{32}\hfill & Y_{33}\hfill \end{array}\right)=\left[\begin{array}{c}{L}^{\left(1\right)}\\ L^{\left(2\right)}\\ L^{\left(3\right)}\end{array}\right]\times \left[\begin{array}{lll}{F}^{\left(1\right)}\hfill & F^{\left(2\right)}\hfill & F^{\left(3\right)}\hfill \end{array}\right]+E$$ (B.1)

$$=\left(\begin{array}{lll}{Y}_{\left(1\right)}& Y_{\left(2\right)}& Y_{\left(3\right)}\\ Y_{\left(3\right)}& Y_{\left(1\right)}& Y_{\left(2\right)}\\ Y_{\left(2\right)}& Y_{\left(3\right)}& Y_{\left(1\right)}\end{array}\right)=\left[\begin{array}{lll}{L}^{\left(1\right)}{F}^{\left(1\right)}& L^{\left(1\right)}{F}^{\left(2\right)}& L^{\left(1\right)}{F}^{\left(3\right)}\\ L^{\left(2\right)}{F}^{\left(1\right)}& L^{\left(2\right)}{F}^{\left(2\right)}& L^{\left(2\right)}{F}^{\left(3\right)}\\ L^{\left(3\right)}{F}^{\left(1\right)}& L^{\left(3\right)}{F}^{\left(2\right)}& L^{\left(3\right)}{F}^{\left(3\right)}\end{array}\right]+E$$ (B.2)

where $Y_{\left(1\right)}=\left\{Y_{11},Y_{22},Y_{33}\right\},Y_{\left(2\right)}=\left\{Y_{12},Y_{23},Y_{31}\right\}$ and $Y_{\left(3\right)}=\left\{Y_{13},Y_{21},Y_{32}\right\}$. We can see for each “hold-out” part, $Y_{\left(k\right)},L^{\left(1\right)},L^{\left(2\right)},L^{\left(3\right)}$ and $F^{\left(1\right)},F^{\left(2\right)},F^{\left(3\right)}$ show up once and only once. In this sense the hold-out pattern is “balanced”.