Computational Approaches for Exponential-Family Factor Analysis

For given data $X\in\mathbb{R}^{n\times p}$, the traditional rank $q$ factor model assumes that $q\ll p$ and that the data is generated by latent factors $\Lambda\in\mathbb{R}^{n\times q}$ with $\Lambda^{\top}=[\Lambda_{1}\cdots\Lambda_{n}]$, and a deterministic projection matrix $V\in\mathbb{R}^{p\times q}$, the loading matrix. We implicitly assume the following data generating process: for each observation $i=1,\ldots,n$:

where $\Phi$ is a $p$-th order symmetric positive definite matrix, a covariance matrix providing potential heterogeneous noise.

Marginally, $X_{i}\stackrel{{\scriptstyle\text{ind}}}{{\sim}}N(0,\Phi+VV^{\top})$, so maximum likelihood estimation of $V$ and $\Phi$ is equivalent to covariance estimation. It is common to assume that $\Phi$ is diagonal so as to not confound the effect of the loadings $V$ (Bartholomew et al., 2011), and so we adopt the same assumption from now on. In this case, the MLE estimator for $V$ can be obtained in closed form using matrix calculus, based on the eigen-decomposition of the data covariance. Alternatively, $V$ and $\Phi$ can be estimated via expectation-maximization, especially if some of the entries in the data matrix $X$ are missing.

The model in general has interesting connections to matrix factorization. For example, probability PCA (Tipping and Bishop, 1998) can be considered as equivalent to the factor model with the only difference that the factor model permits heterogeneous noise structure through the specification of $\Phi$. Under $\Phi=\sigma^{2}I_{p}$, Anderson (1963) established the connection between these two models by demonstrating that the stationary point solution of the factor model likelihood spans the columns of the sample covariance eigenvectors. Drawing further the analogy from the relationship between probability PCA and PCA, the factor model can be considered as the random counterpart of matrix factorization by allowing the factorized components (or latent local factors) $\Lambda$ to be random.

While finding a wide range of applications, the factor model and its deterministic counterpart (matrix factorization) have, however, some limitations. We discuss them below, including a brief summary on some recent improvements, along with our proposed solutions to further generalize the factor model.

The paper is organized as follows: in Section 2 we introduce a more general exponential factor model that addresses the shortcomings listed in 1.1.1 and 1.1.2; next, in Section 3, we discuss our main contributions—a collection of efficient and robust optimization strategies for inference, tackling the points in 1.1.3; Section 4 demonstrates the effectiveness of our factorization with simulated examples and applications on benchmark data from various fields; finally, Section 5 concludes with a summary of the innovations and directions for future work.

We start by addressing the model issues raised in Section 1. Given our concern with computational efficiency, our first issue is non-identifiability; as discussed in 1.1.1, we need to constraint the factors to avoid lack of identifiability due to rotations. From now on we adopt the following standardization of the factors:

1. (i)

   $\Lambda_{i}\stackrel{{\scriptstyle\text{iid}}}{{\sim}}N(0,I_{q})$ for $i\in[n]$ as usual, with the distribution of the rows of $\Lambda$ being invariant to orthogonal transformations;

2. (ii)

   $V$ has scaled pairwise orthogonal columns, that is, $V=UD$ with $U\in\mathcal{S}_{p,q}(\mathbb{R})$, a $p$-frame in the Stiefel manifold of order $q$, and $D=\text{Diag}_{j\in[q]}\{d_{j}\}$ with $d_{1}\geq\cdots\geq d_{q}>0$, so that $V^{\top}V=D^{2}$. We denote this space for $V$ as $\widetilde{\mathcal{S}}_{p,q}(\mathbb{R})$.

This setup makes the factorization model identifiable since for any arbitrary $T\in O(q)$, $V^{*}=VT^{\top}$ can only belong to $\widetilde{\mathcal{S}}_{p,q}(\mathbb{R})$ if $T^{\top}$ commutes with a diagonal matrix, that is, if $T\in O(1)^{q}$, and so $V$ is unique (up to column sign changes, as in the SVD). In practice, given any pair of factors $\widehat{\Lambda}$ and $\widehat{V}$ we just need to find the singular value decomposition of $\widehat{\Lambda}\widehat{V}^{\top}=\Lambda DU^{\top}$ to identify $\Lambda$ and $V=UD$.

Next, we relax the convenient but often unrealistic Gaussian assumptions in the likelihood and settle with a more general mean and variance specification in the spirit of quasi-likelihood (Wedderburn, 1974). We assume that $X_{i}\,|\,\Lambda_{i}\sim F(V\Lambda_{i},\Phi_{i})$ where $F$ belongs to the exponential family with link function $g$ and variance function $\mathcal{V}$, that is,

where $\mathbb{V}(\mu_{i})=\text{Diag}\{\mathcal{V}(\mu_{i})\}$ is the diagonal variance and $\eta_{0}\in\mathbb{R}^{p}$ is the latent center of the factor model. This way, we can more naturally represent data $X$ belonging to fields other than real numbers; common cases are binary data with $F$ being Bernoulli or binomial (with weights) and count data with $F$ being Poisson or negative binomial. In particular, the negative binomial distribution offers enhancements over the Poisson distribution by effectively accommodating the over-dispersion characteristic often observed in count data; see, e.g., (Xia, 2020) for a detailed treatment of negative binomial distributions as a compound Poisson type. To accommodate entry-wise weights, as motivated in Section 1.1.2, we set $\Phi_{i}=\Phi W_{i}^{-1}$ where $W_{i}=\text{Diag}_{j=1,\ldots,p}\{w_{ij}\}$ are the known weights, that is, $\Phi_{i}=\text{Diag}_{j=1,\ldots,p}\{\phi_{j}/w_{ij}\}$. This setup implies ${\mathbb{E}}(X_{ij}|\Lambda_{i})=\mu_{ij}$ and $\text{Var}(X_{ij}|\Lambda_{i})=\phi_{j}\mathcal{V}(\mu_{ij})/w_{ij}$. From these two moment conditions, we can adopt the extended quasi-likelihood (Nelder and Pregibon, 1987) to define:

We call this the exponential factor model (EFM).

The MLE estimate of $\theta=(V,\eta_{0},\Phi)$ then requires access to the marginal density for each observation $i\in[n]$,

where $f$ is the density of the standard multivariate normal density of order $q$:

Such a generalized factor framework can be used to efficiently estimate the covariance structure of high dimensional data. Specifically, applying the total variance formula we can derive the covariance of a new observation $X\,|\,\lambda\sim F(V\lambda,\Phi)$ given $\lambda\sim N(0,I_{q})$,

Here, the first term is a diagonal matrix but the second term requires an outer product that induces correlations among the entries of $X$,

where $\mu_{\lambda}={\mathbb{E}}_{\lambda}\big{(}g^{-1}(V\lambda)\big{)}$.

In the special case of the Gaussian distribution, we have $g(\mu)=\mu$ and $\mathcal{V}(\mu)=1$ and so, as expected, $\text{Cov}_{\theta}(X)=\Phi+VV^{\top}$. Using this result, Fan et al. (2008) demonstrated the efficiency of the plug-in covariance estimator via estimation of $V,\phi,\text{Cov}(\Lambda)$. Similarly, we demonstrate that the plug-in efficiency of $(V,\Phi)$ in high dimensional covariance estimation is preserved under our generalized quasi-factor setup using Eq (8).

As in the Gaussian factor model, one common estimation procedure for such a latent space model should be Expectation-Maximization (EM; Dempster et al., 1977a). In our setup, the EM algorithm can be formulated as follows:

E-Step:

Given the parameters $\theta^{(t)}$ at the $t$-th iteration step, we compute

$$Q(\theta;\theta^{(t)})={\mathbb{E}}_{\Lambda|X;\theta^{(t)}}\Bigg{[}\sum_{i=1}^{n}\log f_{\theta}(X_{i},\Lambda_{i})\Bigg{]}\approx\sum_{i=1}^{n}\int\log f_{\theta}(X_{i},\Lambda_{i})\widetilde{f}_{\theta^{(t)}}(\Lambda_{i}|X_{i})d\Lambda_{i},$$

(9)

where $\widetilde{f}_{\theta^{(t)}}(\Lambda_{i}|X_{i})$ is a Laplace approximation to the actual conditional posterior on $\Lambda_{i}|X_{i}$ using a Gaussian density centered at $\widehat{\Lambda}_{i}$ and with precision matrix $\widehat{H}_{i}$. These parameters are found by Fisher scoring on a regularized GLM regression with quasi-likelihood

$$\ell_{i}^{(t)}(\Lambda_{i})=\log f_{\theta^{(t)}}(X_{i}|\Lambda_{i})+\log f(\Lambda_{i}),$$

(10)

that is, $\Lambda_{i}|X_{i}\approx N(\widehat{\Lambda}_{i},\widehat{H}_{i}^{-1})$ with $\widehat{\Lambda}_{i}=\operatornamewithlimits{argmax}_{\Lambda_{i}}\ell_{i}^{(t)}(\Lambda_{i})$ and $\widehat{H}_{i}={\mathbb{E}}_{X_{i}}[-\partial^{2}\ell_{i}^{(t)}(\widehat{\Lambda}_{i})/\partial\Lambda_{i}\partial\Lambda_{i}^{\top}]$ the negative expected Hessian. When the latent dimension $q$ is small, we can then evaluate the integral in (9) using multivariate Gauss-Hermite cubature (Golub and Welsch, 1969) with $m$ nodes for each $\Lambda_{i}$,

$$Q(\theta;\theta^{(t)})\approx\sum_{i=1}^{n}\sum_{l=1}^{m}\log f_{\theta}(X_{i},\Lambda_{il})w_{il}^{(t)},$$

(11)

where the cubature weights $w_{il}^{(t)}$ are computed based on the Laplace approximation $\Lambda_{i}|X_{i};\theta^{(t)}\approx N(\widehat{\Lambda}_{i},\widehat{H}_{i}^{-1})$.

M-Step:

We can then update the parameters to $\theta^{(t+1)}$ by maximizing the expected complete data likelihood,

$$\theta^{(t+1)}=\operatornamewithlimits{argmax}_{\theta=(V,\Phi,\eta_{0})}Q(\theta;\theta^{(t)})=\operatornamewithlimits{argmax}_{\theta=(V,\Phi,\eta_{0})}\sum_{i=1}^{n}\sum_{l=1}^{m}\log f_{\theta}(X_{i},\Lambda_{il})w_{il}^{(t)}$$

(12)

The M-step is then easily seen to be a quasi-GLM weighted regression of $X$ on the cubature nodes $\Lambda_{il}$. In particular, we estimate the dispersion parameters using weighted Pearson residuals,

$$\widehat{\phi}_{j}^{(t+1)}=\frac{1}{nm}\sum_{i=1}^{n}\sum_{l=1}^{m}\frac{(X_{ij}-\widehat{\mu}_{ijl}^{(t+1)})^{2}}{w_{ij}w_{il}^{*}\mathcal{V}(\widehat{\mu}_{ijl}^{(t+1)})},\quad j\in[p],$$

(13)

where $w_{il}^{*}$ are the cubature weights at the last EM iteration.

The algorithm for this EM optimization is summarized in Algorithm 1. While this computational scheme is robust and efficient, both its complexity and approximation errors are proportional to the latent dimension $q$. It is common in the literature (Fan et al., 2020) to assume that $q\ll p$ so that the factorization is parsimonious. For cases when we need a larger $q$, we explore two alternative optimization methods utilizing Stochastic Gradient Descent (SGD) next.

Under some regularity conditions to allow the exchange of differentiation and integration, the gradient of likelihood can be written as:

If we can evaluate Eq (14) efficiently and accurately, we could then update the EFM parameters through gradient descent with step size $\alpha$:

Ignoring for now the expectation ${\mathbb{E}}_{\Lambda_{i}|X_{i}}(\cdot)$, $\nabla_{\theta}[f_{\theta}(X_{i},\Lambda_{i})]$ are in fact available in closed form given a specification of our model in Eq (3). We provide a summary of these gradients below:

• •

  Gradient for $V_{j}$:

  	$\displaystyle-\nabla_{V_{j}}\log f_{\theta}(X_{i},\Lambda_{i})$	$\displaystyle=\frac{w_{ij}}{\phi_{j}}\nabla_{V_{j}}\bigg{(}\int_{\mu_{ij}}^{X_{ij}}\frac{1}{\mathcal{V}(t)}(X_{ij}-t)dt\bigg{)}=-\frac{w_{ij}}{\phi_{j}}\frac{1}{\mathcal{V}(\mu_{ij})}(X_{ij}-\mu_{ij})\frac{\partial\mu_{ij}}{\partial\eta_{ij}}\frac{\partial\eta_{ij}}{\partial V_{j}}$		(16)
  		$\displaystyle=-\frac{w_{ij}}{\phi_{j}}\frac{1}{\mathcal{V}(\mu_{ij})}(X_{ij}-\mu_{ij})\frac{1}{g^{\prime}(\mu_{ij})}\Lambda_{i}.$

• •

  Gradient for $\phi_{j}$:

  We repeat the derivation for parameter $\phi_{j}\in\mathbb{R}_{+}$, for $j\in[p]$:

  	$\displaystyle\nabla_{\phi_{j}}(-\log f_{\theta}(X_{i},\Lambda_{i}))$	$\displaystyle=-\frac{w_{ij}}{\phi_{j}^{2}}\bigg{(}\int_{\mu_{ij}}^{X_{ij}}\frac{1}{\mathcal{V}(t)}(X_{i}-t)dt\bigg{)}+\frac{1}{2\phi_{j}}=-\frac{w_{ij}}{2\phi_{j}^{2}}Q(X_{ij};\mu_{ij})+\frac{1}{2\phi_{j}}$		(17)
  		$\displaystyle=\frac{1}{2\phi_{j}}\bigg{(}-\frac{w_{ij}}{\phi_{j}}Q(X_{ij};\mu_{ij})+1\bigg{)}\approx\frac{1}{2\phi_{j}}\bigg{(}-\frac{w_{ij}(X_{ij}-\mu_{ij})^{2}}{\phi_{j}\mathcal{V}(\mu_{ij})}+1\bigg{)}.$

  where $Q(X_{ij};\mu_{ij})=-2\int_{X_{ij}}^{\mu_{ij}}\frac{1}{\mathcal{V}(t)}(X_{i}-t)dt$ is the quasi-deviance function.

• •

  Gradient for $\eta_{0}$:

  We repeat the derivation for $\eta_{0j}\in\mathbb{R}$, for $j\in[p]$:

  	$\displaystyle\nabla_{\eta_{0j}}(\log f_{\theta}(X_{i},\Lambda_{i}))$	$\displaystyle=\frac{w_{ij}}{\phi_{j}}\nabla_{\eta_{0j}}\bigg{(}\int_{\mu_{ij}}^{X_{ij}}\frac{1}{\mathcal{V}(t)}(X_{i}-t)dt\bigg{)}=-\frac{w_{ij}}{\phi_{j}}\frac{1}{\mathcal{V}(\mu_{ij})}(X_{ij}-\mu_{ij})\frac{\partial\mu_{ij}}{\partial\eta_{ij}}\frac{\partial\eta_{ij}}{\partial\eta_{0j}}$		(18)
  		$\displaystyle=\frac{1}{\Phi_{jj}}\frac{1}{\mathcal{V}(\mu_{ij})g^{\prime}(\mu_{ij})}(\mu_{ij}-X_{ij}).$

• •

  Hessian for $V_{j}$:

  In our later optimization, we additionally need the Hessian of the $\log f_{\theta}(X_{i},\Lambda_{i})$. The Hessian for $V_{j}$ and $\Lambda_{i}$ are symmetric with respect to each other with a simple notation change; below we use the former as an example. We need the following definitions:

  • –

    $S_{ij}=\frac{w_{ij}}{\phi_{j}}\frac{{{g^{-1}}^{\prime}(\eta_{ij})}^{2}}{V(\mu_{ij})}$ and $G_{ij}=\frac{{g^{-1}}^{\prime}(\eta_{ij})}{V(\mu_{ij})}\frac{w_{ij}}{\phi_{j}}(X_{ij}-\mu_{ij})$

  • –

    $D_{i\cdot}=\text{Diag}\{S_{i\cdot}^{(t)}\}$ with $S_{i\cdot}$ denotes the $i$-th row of $S$. Similarly, $\text{Diag}\{S_{\cdot j}^{(t)}\}$ with $S_{\cdot j}$ denotes the $j$-th column of $S$.

  The Hessian for $V_{j}$ is then

  	$\displaystyle{\mathbb{E}}_{X_{i}}$	$\displaystyle\bigg{[}\nabla_{V_{j}}^{2}(-\log f_{\theta}(X_{i},\Lambda_{i}))\bigg{]}=$		(19)
  		$\displaystyle{\mathbb{E}}_{X_{i}}\bigg{[}\bigg{(}\nabla_{V_{j}}\big{(}G(\mu_{i})^{\top}[S^{\frac{1}{2}}(\mu_{i})\Phi S^{\frac{1}{2}}(\mu_{i})]^{-1}\big{)}(\mu_{i}-X_{i})+G(\mu_{i})^{\top}[S^{\frac{1}{2}}(\mu_{i})\Phi S^{\frac{1}{2}}(\mu_{i})]^{-1}\nabla_{V_{j}}\big{(}\mu_{i}-X_{i}\big{)}\bigg{)}\Lambda_{i}\bigg{]}$
  		$\displaystyle=\sum_{j=1}^{p}G(\mu_{ij})^{2}[S^{\frac{1}{2}}(\mu_{ij})\Phi_{jj}S^{\frac{1}{2}}(\mu_{ij})]^{-1}\frac{\partial\mu_{ij}}{\partial\eta_{ij}}\frac{\partial\eta_{ij}}{\partial V_{j}}\Lambda_{i}$
  		$\displaystyle=\Lambda_{i}\bigg{(}\text{Diag}(G(\mu_{i})^{2})[S^{\frac{1}{2}}(\mu_{i})\Phi S^{\frac{1}{2}}(\mu_{i})]^{-1}\bigg{)}\Lambda_{i}^{\top}$

  where the first components is 0 because ${\mathbb{E}}(\mu_{i}-X_{i})=\mathbf{0}_{p}$.

The random sampling for the gradient of every observation $i\in[n]$ in Eq (14) is however expensive. We here additionally leverage modern stochastic optimization to randomly sample partial data to approximate the optimization gradient. The optimization is termed Stochastic Gradient Descent (SGD; Robbins and Monro, 1951) due to its interpretation as a stochastic approximation of the actual gradient function. Since the method effectively reduces the sample size by applying a sub-sampling on the original dataset, the SGD can be used to accelerate all optimization algorithms with explicit gradient formulation. We here use the SML optimization as an example to illustrate the implementation.

Taking the last equality from Eq (14) and applying the law of large numbers, we can compute the gradient using stochastic sample of size $B$ and $S$:

The batch size $B$ is chosen to be smaller than sample size $n$, which scales the original complexity with a factor of $B/n$ per iteration. To maintain distributional assumptions, we sample $X^{(b)}$ with replacement from the original data $X$. In addition, since this stochastic sampled gradient is proposed to maximize the true likelihood instead of the simulated likelihood, the sample size $S$ required to compute the optimization gradients does not need to grow in an order of the actual sample size (e.g. $S=n^{1/2}$ as required for SML (Lee, 1995)).

As it is compared to second-order optimization such as Newton’s method, step size selection is of crucial importance for stochastic gradient descent. A large step size will make the algorithm oscillate while a small step size hardly improves our likelihood function. The theoretical analysis states that we should choose the step size according to the conditional number of the parameter Hessian matrix (Bertsekas, 1999). When the Hessian matrix is not available, recent researchers have refined the step size selection by utilizing the momentum and the scale of the parameters. Specifically, AdaGrad (Duchi et al., 2011) scales the gradient update with the gradient’s second moment while RMSProp (Ieleman and Hinton, 2012) employed exponential decay to smooth out the gradient direction. More recently, combining both RMSprop and AdaGrad, the Adam (Kingma and Ba, 2015) method has gained its well-deserved attention in modern stochastic optimization research. Here we adopt the Adam method with the parameter recommended in the original reference ($\beta_{1}=0.9,\beta_{2}=0.999,\epsilon=10^{-8}$). To avoid the oscillation around the minimum, we also employed a decay learning rate with $\gamma_{t}=\frac{\alpha}{1+0.5t}$. When it is necessary, the tuning on the hyperparameter $\alpha$ can be conducted by randomly sampling $\alpha$ on a log grid.

To tackle specifically the expectation ${\mathbb{E}}_{\widetilde{\Lambda}_{i}}(\cdot)$ with $\widetilde{\Lambda}_{i}\stackrel{{\scriptstyle d}}{{=}}\Lambda_{i}|X_{i}$, we propose two optimization algorithms in the following subsections.