Parallel backfits

Introduction

I’ve rewritten flashier’s parallel backfitting algorithm. As before, factors can be backfit in parallel by setting parameter backfit.order = "parallel". The number of cores and type of cluster (socket or fork) can be set using global options (here, I use options(cl.type = "FORK", cl.cores = parallel::detectCores())).

Each worker is responsible for \(\frac{K}{\text{n.cores}}\) calls to ebnm (where \(K\) is the total number of factors), so we can only really expect performance benefits from parallelization when each call to ebnm is fairly computationally intensive and when \(K\) is somewhat large. Further, since parallel updates are not guaranteed to produce a monotonic increase in the objective function, serial updates should generally be preferred when the dataset is not large.

Test Cases

For large problems, parallelization can provide a noticeable speedup. As a first test case, I greedily fit 50 factors to the droplet-based 3’ scRNA-seq dataset from Montoro et al., which is, I think, on the smaller end of the scale of dataset that could benefit from parallel backfitting updates. The dataset, which I also used to benchmark flashier, is described here. I backfit using both the default serial method and the new implementation of the parallel approach. As shown below, parallel updates are able to attain the same ELBO as the serial method about three times more quickly.

library(ggplot2)

timing <- readRDS("./output/parallel_v2/sp_trachea_timing.rds")

serial.res <- data.table::fread("./output/parallel_v2/sp_trachea_serial.txt")
serial.res$method <- "serial"
serial.res$time <- timing$serial * 1:nrow(serial.res) / nrow(serial.res)

parallel.res <- data.table::fread("./output/parallel_v2/sp_trachea_parallel.txt")
parallel.res$method <- "parallel"
parallel.res$time <- timing$parallel * 1:nrow(parallel.res) / nrow(parallel.res)

all.res <- rbind(serial.res, parallel.res)
all.res$time <- as.numeric(all.res$time)

ggplot(all.res, aes(x = time, y = Obj, color = method)) + geom_line() +
  labs(x = "Elapsed time (min)", y = "ELBO", title = "Montoro et al. droplet-based dataset")

Next, I repeat the experiment on the larger PulseSeq dataset from Montoro et al. (described here). I get bus errors when I try to use all available cores (using parallel::detectCores()), so I lower cl.cores to 16. (For both of these datasets, however, there is not much additional speedup after the first 8 cores.) Here, parallel updates are only about twice as fast as serial ones.

Do note, however, that the default backfitting method already uses some tricks to speed up backfits. In particular, it stops updating any given factor when that factor appears to have converged. For example, factor 42 only gets updated 10 times, while factor 6 must be updated 249 times before it is considered to have converged. As a result, any changes that occur in the final 239 backfitting iterations will not be reflected in the final form of factor 42. In contrast, parallel backfits update every factor every time, so with the tolerance parameter set appropriately, parallel backfits will get one close to a local maximum much more reliably. It’s possible to get the same reliability with serial backfits by setting backfit.order = "sequential", but this will be much slower than the default method.

library(ggplot2)

timing <- readRDS("./output/parallel_v2/pulseseq_timing.rds")

serial.res <- data.table::fread("./output/parallel_v2/pulseseq_serial.txt")
serial.res$method <- "serial"
serial.res$time <- timing$serial * 1:nrow(serial.res) / nrow(serial.res)

parallel.res <- data.table::fread("./output/parallel_v2/pulseseq_parallel.txt")
parallel.res$method <- "parallel"
parallel.res$time <- timing$parallel * 1:nrow(parallel.res) / nrow(parallel.res)

all.res <- rbind(serial.res, parallel.res)
all.res$time <- as.numeric(all.res$time)

ggplot(all.res, aes(x = time, y = Obj, color = method)) + geom_line() +
  labs(x = "Elapsed time (hr)", y = "ELBO", title = "Montoro et al. PulseSeq dataset")

Limitations

Code

Session information

sessionInfo()

#> R version 3.5.3 (2019-03-11)
#> Platform: x86_64-apple-darwin15.6.0 (64-bit)
#> Running under: macOS Mojave 10.14.5
#> 
#> Matrix products: default
#> BLAS: /Library/Frameworks/R.framework/Versions/3.5/Resources/lib/libRblas.0.dylib
#> LAPACK: /Library/Frameworks/R.framework/Versions/3.5/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     
#> 
#> other attached packages:
#> [1] ggplot2_3.2.0  mixsqp_0.1-119 ashr_2.2-38    ebnm_0.1-24   
#> [5] flashier_0.1.7
#> 
#> loaded via a namespace (and not attached):
#>  [1] Rcpp_1.0.1        pillar_1.3.1      compiler_3.5.3   
#>  [4] git2r_0.25.2      workflowr_1.2.0   iterators_1.0.10 
#>  [7] tools_3.5.3       digest_0.6.18     tibble_2.1.1     
#> [10] evaluate_0.13     gtable_0.3.0      lattice_0.20-38  
#> [13] pkgconfig_2.0.2   rlang_0.3.1       Matrix_1.2-15    
#> [16] foreach_1.4.4     yaml_2.2.0        parallel_3.5.3   
#> [19] xfun_0.6          withr_2.1.2       dplyr_0.8.0.1    
#> [22] stringr_1.4.0     knitr_1.22        fs_1.2.7         
#> [25] tidyselect_0.2.5  rprojroot_1.3-2   grid_3.5.3       
#> [28] data.table_1.12.2 glue_1.3.1        R6_2.4.0         
#> [31] rmarkdown_1.12    purrr_0.3.2       magrittr_1.5     
#> [34] whisker_0.3-2     backports_1.1.3   scales_1.0.0     
#> [37] codetools_0.2-16  htmltools_0.3.6   MASS_7.3-51.1    
#> [40] assertthat_0.2.1  colorspace_1.4-1  labeling_0.3     
#> [43] stringi_1.4.3     lazyeval_0.2.2    doParallel_1.0.14
#> [46] pscl_1.5.2        munsell_0.5.0     truncnorm_1.0-8  
#> [49] SQUAREM_2017.10-1 crayon_1.3.4