More efficient FLASH algorithms

Introduction

Currently, FLASH does not perform well on large datasets. The reason is that it maintains and manipulates at least three matrices (Rk, R2, and R2k) which are each of the same size as the data matrix Y. When var_type is by_row or by_column, a fourth large matrix tau is maintained, and when there is missing data, still another large matrix is stored (a copy of Y, with missing data imputed to be zero).

Let \(Y \in \mathbb{R}^{n \times p}\), and let \(k\) be the rank of the flash fit. I show here that when \(k^2 < \min(n, p)\), FLASH can be implemented without maintaining any matrices of the same dimensions as Y. Even when \(k\) is large, such a matrix only needs to be created once during factor updates and once during loading updates, and it is only temporarily needed.

I outline two improved implementations. In the first, the only \(n \times p\) matrices that are required are the data matrix \(Y\) with NAs set to zero and, if there is missing data, a matrix \(Z\) with zeroes where data is missing and ones elsewhere. The original data matrix is not needed. I never modify or transpose \(Y\) or \(Z\), so the total memory requirements should not be much larger than the memory needed to store \(Y\) and \(Z\).

The second assumes that one is willing to maintain a matrix of residuals \(R\), initialized at \(Y\) (\(Z\) is kept to ensure that \(R\) has zeroes where data is missing). The primary disadvantage is that one cannot take advantage of sparsity in \(Y\), and it requires some care to update \(R\) correctly and at the right time, but it is potentially much faster than the first implementation.

Initializing factors

To initialize a new factor, we currently form a matrix of residuals, run an SVD-type algorithm on it, and then initialize using the leading singular vector. However, we don’t need to recourse to an external package here. Since we only require one singular vector, the minimization problem is simple. When there is no missing data, the problem is to solve: \[ (l_k, f_k) = {\arg \min}_{l_k, f_k} \| R_{-k} - l_k f_k' \|_F^2 = {\arg \min}_{l_k, f_k} \| Y - L_{-k}F_{-k}' - l_k f_k' \|_F^2 \] It is easy to derive an alternating algorithm by differentiating with respect to \(l_k\) and setting the result equal to zero: \[ l_k = \frac{R_{-k}f_k}{f_k^Tf_k} = \frac{Yf_k - L_{-k}(F_{-k}'f_k)}{f_k^Tf_k} \] And similarly: \[ f_k^T = \frac{l_k^TR_{-k}}{l_k^Tl_k} = \frac{l_k^TY - (l_k^TL_{-k})F_{-k}'}{l_k^Tl_k} \]

Each of these updates can be performed very quickly; further, since we don’t need to solve the problem exactly, we can stop early, and the same procedure is likely sufficient to (approximately) initialize factors when there is not a lot of missing data.

But we can also derive efficient updates that are exact for missing data. In this case, the problem is to solve: \[ (l_k, f_k) = {\arg \min}_{l_k, f_k} \| Z \odot (R_{-k} - l_k f_k') \|_F^2 = {\arg \min}_{l_k, f_k} \| Y - Z \odot L_{-k}F_{-k}' - Z \odot l_k f_k' \|_F^2 \]

The terms on the right-hand side that depend on \(l_k^{(i)}\) are \[\begin{aligned} -2 &(Z \odot R_{-k})_{i \bullet}(Z_{i \bullet}^T \odot l_k^{(i)} f_k) + (Z_{i \bullet}^T \odot l_k^{(i)} f_k)^T (Z_{i \bullet}^T \odot l_k^{(i)} f_k) \\ &= -2l_k^{(i)} (Z_{i \bullet} \odot (R_{-k})_{i \bullet})f_k + l_k^{(i)2} (Z_{i \bullet}^T \odot f_k)^T f_k \end{aligned}\] Differentiating and setting the result equal to zero gives (in vector form) \[ l_k = (Z \odot R_{-k})f_k / Zf_k^2, \] where division and squaring are elementwise. Write \[ (Z \odot R_{-k})f_k = Yf_k - (Z \odot L_{-k}F_{-k}')f_k\] and note that the \(i\)th entry of \((Z \odot L_{-k}F_{-k}')f_k\) is \[\sum_j z_{ij} \sum_{\ell: \ell \ne k} l_{i \ell} f_{j \ell} f_k = \sum_{\ell: \ell \ne k} l_{i \ell} \sum_j z_{ij} f_{j \ell} f_k = \sum_{\ell: \ell \ne k} l_{i \ell} \sum_j z_{ij} \tilde{f}_{j \ell},\] where \(\tilde{F} = f_k \odot_b F_{-k}\), with \(\odot_b\) denoting elementwise multiplication using broadcasting (that is, the \(\ell\)th column of \(\tilde{F}\) is formed by taking the elementwise product of \(f_k\) and the \(\ell\)th column of \(F_{-k}\)). Thus, \[ (Z \odot R_{-k})f_k = Yf_k - \text{rowSums}(L_{-k} \odot Z(f_k \odot_b F_{-k}))\]

In sum, exact updates can be obtained for missing data by alternately setting \[ l_k = (Yf_k - \text{rowSums}(L_{-k} \odot Z(f_k \odot_b F_{-k})) / Zf_k^2\] and \[ f_k^T = (l_k^T Y - \text{colSums}(F_{-k}^T \odot (l_k \odot_b L_{-k})^TZ)) / l_k^{2T}Z \]

This algorithm requires three matrix multiplies involving a \(n \times p\) matrix, with a flop count of \(2(k + 1)np\) per update. The algorithm for no missing data only requires one, with a flop count of \(2np\) per update, so we can expect this initialization to be slower by a factor of \(k + 1\). Still, note that the matrix of residuals never needs to be formed. Further, if there is a lot of missing data, then \(Y\) and \(Z\) can be stored as sparse matrices. In that case, the largest dense matrix that needs to be manipulated will be of dimension \(\max(n, p) \times k - 1\).

Updating loadings

Updating factors

Factor updates can easily be obtained by reversing the roles of, on the one hand, \(EL\) and \(EF\) and, on the other, var_type = by_row and var_type = by_column.

Updating \(\tau\)

\(\tau\) can be updated easily and efficiently if it is done every time a single factor or loading is updated. Recall that most updates to \(\tau\) involve some operation on the matrix of squared residuals: var_type = constant takes the overall mean, by_row takes row means, and by_column takes column means. Write \[\text{R2} = R^2 + Z \odot (EL^2) (EF^2)' - Z \odot (EL)^2 ((EF)^2)' \]

For var_type = constant, \[\begin{aligned} \frac{n_{\text{nonnmissing}}}{\tau} = \text{sum}(\text{R2}) &= \text{sum}(R^2) + \text{sum} ((EL^2) \odot Z(EF^2)) - \text{sum} ((EL)^2 \odot Z((EF)^2)) \\ &= \text{sum}(R^2) + \text{sum} ((EF^2)^T \odot (EL^2)^T Z) - \text{sum} (((EF)^2)^T \odot ((EL)^2)^T Z) \end{aligned}\]

Similar expressions can be derived for by_row (viz. by_column) by replacing the sums with rowSums in the first equality (viz. with colSums in the second equality) and by replacing the scalar number of non-missing values \(n_{\text{nonnmissing}}\) with a vector of the number of non-missing values per row (viz. per column).

Since at most one factor or loading has changed between updates of \(\tau\), one can further simplify. For example, take the case var_type = constant. After a single loading update, \[\begin{aligned} n_{\text{nonmissing}} \left( \frac{1}{\tau} - \frac{1}{\tilde{\tau}} \right) = \Delta (\text{sum}(R^2)) + (\Delta EL_k^2)^T Z(EF_k^2) - (\Delta (EL_k)^2)^T Z((EF_k)^2), \end{aligned}\] where \(\Delta\) denotes the new value minus the old value.

If \(R\) is stored, then \(\Delta (\text{sum}(R^2))\) is probably most easily obtained by calculating and storing sum\((R^2)\) every time \(R\) is updated (at a cost of \(2np\) flops). Alternatively, writing \[ \begin{aligned} \Delta R^2 &= (R_{-k} - Z \odot L_k^{\text{new}}F_k')^2 - (R_{-k} - Z \odot L_k^{\text{old}}F_k')^2 \\ &= -2R_{-k} \odot (L_k^{\text{new}} - L_k^{\text{old}}) F_k' + (Z \odot L_k^{\text{new}}F_k')^2 - (Z \odot L_k^{\text{old}}F_k')^2 \end{aligned}\] and noting that \[ \text{sum}((Z \odot L_k^{\text{new}}F_k')^2) = ((L_k^{\text{new}})^2)^T ZF_k^2, \] one has that \[ \Delta (\text{sum}(R^2)) = ((L_k^{\text{new}})^2 - (L_k^{\text{old}})^2)^T ZF_k^2 -2(L_k^{\text{new}} - L_k^{\text{old}})^T R_{-k} F_k \] Finally, replacing \(R_{-k}\) with \(Y - \sum_{i: i \ne k} Z \odot L_i F_i'\): \[\begin{aligned} \Delta(\text{sum}(R^2)) &= ((L_k^{\text{new}})^2 - (L_k^{\text{old}})^2)^T ZF_k^2 - 2(L_k^{\text{new}} - L_k^{\text{old}})^T (Y F_k - \text{rowSums} (L_{-k} \odot Z(F_k \odot_b F_{-k}))) \end{aligned} \] (See “Calculating the EBNM xs” above for a more detailed derivation of the last step.) This calculation requires approximately \(npk\) flops, so I would expect it to be slower than simply calculating sum\((R^2)\) after each update, especially as \(k\) gets large.

I omit the details for by_row and by_column and the simplifications obtained when there is no missing data; these are easy to derive from the above.

Updating objective

The part of the objective that is not stored in the KL fields is calculated as e_loglik = -0.5 * sum(log(2 * pi / tau) + tau * R2) (here, R2 has NAs where data is missing so that the entries do not contribute to the sum).

One should observe, however, that if the objective is calculated immediately after updating \(\tau\), then this expression has a very simple form. With by_constant, for example, \[\frac{1}{\tau} = \frac{\text{sum}(\text{R2})}{n_{\text{nonmissing}}},\] so \[\text{sum}(\tau \cdot \text{R2}) = \tau \cdot \text{sum}(\text{R2}) = n_{\text{nonmissing}}\]

In effect, one can verify that for constant, by_row, and by_column, e_loglik is simply \[ -\frac{1}{2} \text{sum} \left(n_{\text{nonmissing}} \odot \left(\log \left(\frac{2 \pi}{\tau} \right) + 1 \right) \right), \] where \(n_{\text{nonmissing}}\) is a scalar, \(n\)-vector of nonmissing values per row (for by_row), or \(p\)-vector of nonmissing values per column (for by_column).

If var_type is zero and is a constant (say, S = 1), then one can maintain a scalar value of tau as if the precision were constant (using the above updates). Since it will still be true that, for this pseudo-\(\tau\), \(\text{sum(R2)} = n_{\text{nonmissing}} / \tau\), e_loglik can then be calculated as \[ -\frac{1}{2} \text{sum} \left(n_{\text{nonmissing}} \odot \left(\log \left(2 \pi S^2 \right) + \frac{1}{\tau S^2} \right) \right) \] This also works if S is a row or column vector.

If S is a full matrix and \(R\) is being stored, then still another approach is possible; one can store \(\text{sum}(\log(2 \pi S^2))\) and then calculate \(\text{sum(R2} / \text{S}^2)\) as \[ \text{sum} \left( (1 / S^2) \odot R^2 \right) + \text{sum}(EL^2 \cdot (1 / S^2) EF^2) - \text{sum}((EL)^2 \cdot (1 / S^2) (EF)^2),\] where the matrix \((1 / S^2)\) has zeroes where data is missing. This is of course similar to the formula for \(1 / \tau\) when var_type = constant, and similar simple updates that depend only on updated factors and loadings are possible. For example, one can update the objective after a single loading update by calculating \[ \Delta(\text{sum}((1 / S^2) \odot R^2)) + (\Delta EL_k^2)^T (1 / S^2)(EF_k^2) - (\Delta (EL_k)^2)^T (1 / S^2)((EF_k)^2)\] (Due to the form of these updates, I recommend storing 1 / S^2 rather than S, with zeroes where data is missing.)

Session information

sessionInfo()

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

Matrix products: default
BLAS: /Library/Frameworks/R.framework/Versions/3.4/Resources/lib/libRblas.0.dylib
LAPACK: /Library/Frameworks/R.framework/Versions/3.4/Resources/lib/libRlapack.dylib

locale:
[1] en_US.UTF-8/en_US.UTF-8/en_US.UTF-8/C/en_US.UTF-8/en_US.UTF-8

attached base packages:
[1] stats     graphics  grDevices utils     datasets  methods   base     

loaded via a namespace (and not attached):
 [1] workflowr_1.0.1   Rcpp_1.0.0        digest_0.6.18    
 [4] rprojroot_1.3-2   R.methodsS3_1.7.1 backports_1.1.2  
 [7] magrittr_1.5      git2r_0.21.0      evaluate_0.12    
[10] stringi_1.2.4     whisker_0.3-2     R.oo_1.21.0      
[13] R.utils_2.6.0     rmarkdown_1.8     tools_3.4.3      
[16] stringr_1.3.1     xfun_0.4          yaml_2.2.0       
[19] compiler_3.4.3    htmltools_0.3.6   knitr_1.20.22