Minibatch FLASH for large p

Introduction

In my MASH v FLASH application I use a subset of strong tests to identify loadings (i.e., covariance structures) and a random subset of tests to fit priors on the corresponding factors. (My workflow is similar to the workflow for MASH described here.) I suggest that we can instead use all of the tests in each step by borrowing various online optimization techniques.

Idea

Fitting the “strong” and “random” datasets involves fitting a \(44 \times p\) matrix, where \(p\) is on the order of tens of thousands. We subsample the complete data because there are in fact millions of tests, and it is currently not feasible to fit a matrix of this scale.

I suggest the following procedure. Recall that we are primarily interested in obtaining a good loadings matrix and accurate priors on the factors. With loadings and priors on factors fixed, estimating posteriors for individual tests is entirely straightforward.

Here I outline an idea for using all of the tests to fit the (fixed) loadings matrix LL and priors on factors \(g_f\). Let \(Y \in \mathbb{R}^{n \times p}\) be the complete data matrix, so that \(n = 44\) and \(p\) is on the order of millions.

1. Randomly permute the columns of \(Y\) and split the resulting matrix into mini-batches \(Y_1, Y_2, \ldots, Y_m\), so that each \(Y_j\) is of manageable dimension (say, \(44 \times 1000\) or \(44 \times 10000\)).

2. Fit a FLASH object \(f^{(1)}\) to the first mini-batch \(Y_1\).

3. For each successive mini-batch \(Y_j\):

   1. Take the loadings from the previous FLASH object \(f^{(j - 1)}\). Fix them and fit a new FLASH object \(\tilde{f}^{(j)}\) to the new mini-batch \(Y_j\), fixing the grids for the priors on the factors to be the same as the grids used for the previous FLASH object \(f^{(j - 1)}\).

   2. Greedily add as many new loadings as possible to \(\tilde{f}^{(j)}\). (This helps pick up covariance structures that weren’t seen in previous mini-batches, or that were below the threshold of detection.)

   3. Now unfix all of the loadings and backfit \(\tilde{f}^{(j)}\). (Note that this step is only reasonably fast when \(p\) is small!)

   4. Finally, create a new FLASH object by taking weighted averages; in essence, set \(f^{(j)} = \frac{j - 1}{j} f^{(j - 1)} + \frac{1}{j} \tilde{f}^{(j)}\). Only the (expected value of) the loadings of \(f^{(j)}\) and the priors on factors need to be calculated. One can simply take the new loadings to be \(\frac{j - 1}{j} L^{(j - 1)} + \frac{1}{j} \tilde{L}^{(j)}\), where \(L^{(j - 1)}\) is the loadings matrix from \(f^{(j - 1)}\) and \(\tilde{L}^{(j)}\) is the (posterior mean of) the loadings matrix obtained during the \(j\)th iteration. (Newly added loadings can simply be carried over as is.) Since the ASH grid on the priors is fixed, it is similarly straightforward to calculate the priors \(g_{f_k}\); for example, if \(g_f^{(j - 1)} \sim \pi_0^{(j - 1)} \delta_0 + \pi_1^{(j - 1)} N(0, \sigma_1^2) + \ldots + \pi_K^{(j - 1)} N(0, \sigma_K^2)\) and \(\tilde{g}_f^{(j)} \sim \tilde{\pi}_0^{(j)} \delta_0 + \tilde{\pi}_1^{(j)} N(0, \sigma_1^2) + \ldots + \tilde{\pi}_K^{(j)} N(0, \sigma_K^2)\), then set \(g_f^{(j)} \sim \left( \frac{j - 1}{j} \pi_0^{(j - 1)} + \frac{1}{j} \tilde{\pi}_0^{(j)} \right) \delta_0 + \ldots + \left( \frac{j - 1}{j} \pi_K^{(j - 1)} + \frac{1}{j} \tilde{\pi}_K^{(j)} \right) N(0, \sigma_K^2)\).

4. If desired, one can iterate through the complete dataset (or parts of it) multiple times (optionally, re-permuting the dataset after each iteration). If this is done, then the weights in the above averages can be fixed at \((m - 1) / m\) and \(1 / m\) after the first iteration (where \(m\) is the number of mini-batches).

See here for an implementation.

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_0.12.18      digest_0.6.15    
#>  [4] rprojroot_1.3-2   R.methodsS3_1.7.1 backports_1.1.2  
#>  [7] git2r_0.21.0      magrittr_1.5      evaluate_0.10.1  
#> [10] stringi_1.1.6     whisker_0.3-2     R.oo_1.21.0      
#> [13] R.utils_2.6.0     rmarkdown_1.8     tools_3.4.3      
#> [16] stringr_1.3.0     yaml_2.1.17       compiler_3.4.3   
#> [19] htmltools_0.3.6   knitr_1.20