Parallelizing MASH v FLASH backfits

Fitting priors to the random dataset

The first attempt to parallelize the backfit was a disaster. The objective for the first five iterations was:

res_random_bad <- readRDS("./data/parallel/MASHvFLASHrandom_bad.rds")

knitr::kable(data.frame("Iteration" = 1:5, "Objective" = res_random_bad$parallel_obj))

The problem, I think, is that since (as mentioned above) the loadings form an overcomplete basis, an update to the equal effects factor could also be achieved (roughly speaking) by updates to each of the unique effects factors. When all are updated simultaneously, the updates overshoot their goal, causing the objective to spiral out of control.

If this reasoning is correct, then better results might be obtained by splitting up the loadings into subsets of loadings that are mutually orthogonal (or nearly so), and then parallelizing the updates to the loadings within each subset. To test this hypothesis, I divided the loadings into an equal effects loading, 44 unique effects loadings (which are of course mutually orthogonal), two data-driven loadings that are distributed across several tissues (loadings 2 and 11 shown here), and the remaining 11 data-driven loadings (each of which is primarily loaded on two or three tissues).

To update the loadings, then, I updated the equal effects loading, then I updated the unique effects loadings in parallel, then I updated data-driven loadings 2 and 11, and finally I updated the remaining data-driven loadings in parallel. Surprisingly, I was able to achieve a better objective than I achieved using a serial backfit.

res_random <- readRDS("./data/parallel/MASHvFLASHrandom.rds")

all_obj <- c(res_random$backfit_obj, res_random$parallel_obj)
plot(res_random$backfit_obj, pch=19, col="blue",
     xlim=c(1, 20), ylim=c(min(all_obj), max(all_obj)),
     xlab="Update", ylab="Objective")
points(res_random$parallel_obj, pch=19, col="red")
legend("bottomright", legend=c("standard", "parallel"), 
       pch=c(19, 19), col=c("blue", "red"))

The difference in objective attained (that is, the serial objective minus the parallel objective) is as follows:

y <- res_random$backfit_obj - res_random$parallel_obj
plot(1:length(y), y, type="l", xlim=c(1, 20), ylim=c(min(y), max(y)),
     xlab="Update", ylab="Difference")

As expected, the parallel updates are much faster (even after dividing the loadings into four subsets).

data <- data.frame(standard = res_random$backfit_t, 
                   lapply = res_random$parallel_t, 
                   mclapply = res_random$multicore_t)
boxplot(data, ylim=c(0, max(data)), ylab="Time per iter (s)")

The total time (in seconds) required for 20 backfitting iterations is:

colSums(data)

standard   lapply mclapply 
 573.111  286.869  222.333 

Backfitting the strong dataset

The same trick is also needed to backfit the “strong” dataset using the same loadings as above and using the priors obtained by backfitting the “random” dataset. If I try to update the full set of loadings in parallel, then the objective again diverges to \(-\infty\). But if I do the parallel updates in four chunks (as described above), I obtain an objective that beats the serial method (after 20 iterations, at least):

res_final <- readRDS("./data/parallel/MASHvFLASHfinal.rds")

all_obj <- c(res_final$backfit_obj, res_final$parallel_obj)
plot(res_final$backfit_obj, pch=19, col="blue",
     xlim=c(1, 20), ylim=c(min(all_obj), max(all_obj)),
     xlab="Update", ylab="Objective")
points(res_final$parallel_obj, pch=19, col="red")
legend("bottomright", legend=c("standard", "parallel"), 
       pch=c(19, 19), col=c("blue", "red"))

The difference in objective attained is as follows:

y <- res_final$backfit_obj - res_final$parallel_obj
plot(1:length(y), y, type="l", xlim=c(1, 20), ylim=c(min(y), max(y)),
     xlab="Update", ylab="Difference")

The difference in elapsed time is especially dramatic here. The parallel updates take less than a third of the time required by the serial updates:

data <- data.frame(standard = res_final$backfit_t, 
                   lapply = res_final$parallel_t, 
                   mclapply = res_final$multicore_t)
boxplot(data, ylim=c(0, max(data)), ylab="Time per iter (s)")

The total time (in seconds) required is:

colSums(data)

standard   lapply mclapply 
 322.891  103.770   93.890