Yi Zhang (1), William R. Gillespie (1), Ben Bales (2), Aki Vehtari (3)
(1) Metrum Research Group, USA, (2) Columbia University, USA, (3) Aalto university, Finland
Introduction: Stan is a widely used, open-source, probabilistic programming language and Bayesian inference engine [1, 2]. It provides a general model specification language and uses dynamic HMC simulation for fully Bayesian data analysis. Torsten is a library of Stan functions that simplifies implementation of pharmacometric (PMX) models and extends the range of models that may be implemented [3].
Objectives: Demonstrate a multilevel parallelism that combines a cross-chain warmup algorithm we developed recently for Stan and the within-chain parallel functionality of Torsten.
We show that this mechanisim significantly improves large PKPD model simulation efficiency.
Methods: Torsten’s functions are written in C++ and they comply with Stan development requirements. They may be accessed by all Stan interfaces such as RStan and CmdStan. Torsten provides ODE integrator functions and PMX solution functions with MPI-based parallel support [4]. For population modeling, these functions solve each ODE system corresponding to each individual in parallel. In addition, we develop a parallel cross-chain warmup scheme that performs dynamic warmup adaptation by aggregating samples from different chains. During the HMC warmup stage, the joint log posterior samples from multiple chains are collected and used to calculate potential scale reduction coefficients (Rhat) and effective sample sizes (ESS) [5]. Warmup is terminated when they surpass predetermined target values. Also during the warmup we use the aggregated samples to calculate and update the HMC stepsize and mass matrix. This new warmup adaptation terminates warmup sampling automatically to avoid unncessary computation. The MPI implementation of the algorithm includes both inter-chain communication for cross-chain warmup, as well as within-chain communication for population calculation load distribution. With this setup we are able to perform large-scale simulations using a cluster and distribute the computation load evenly across chains.
Results: Models that would benefit most from the proposed scheme are those that involve heavy computation in each HMC leap-frog step, e.g. models with a large number of parameters or that require numerical solution of ODEs. Here we show initial results on two such Stan models. Model 1 is a hierarchical model that has nearly 400 normally distributed parameters. Model 2 involves solving an ODE over a long period of evolution time. We apply the parallel scheme to the models and simulate using default Stan settings of 4 chains, with 1000 post-warmup samples/chain. We also perform regular Stan simulations that require 1000 warmup iterations. The cross-chain warmup scheme gives similar performance in terms of ESS/iteration (Table 1), while reducing the number of warmup iterations and thus wall time significantly (Tables 2 & 3).
======================================================
ESS_bulk/iteration ESS_tail/iteration
Model Regular run Cross-chain run Regular run Cross-chain run
1 0.29 0.29 0.40 0.41
2 0.23 0.28 0.30 0.30
======================================================
Table 1: Sampling performance: Effective sample size (ESS) per iteration
======================================================
ESS_bulk/sec ESS_tail/sec
Model Regular run Cross-chain run Regular run Cross-chain run
1 83 103 118 144
2 10 18 13 20
======================================================
Table 2: Sampling performance: ESS per wall time unit (sec)
======================================================
Regular run Cross-chain run
Model warmup iterations walltime(s) warmup iterations walltime(s)
1 1000 14.1 350 11.1
2 1000 91.5 350 61.9
======================================================
Table 3: Total simulation time (Warmup + Sampling)
Conclusions: Multi-level parallelism based on recent developments in Stan and Torsten significantly improves computational efficiency and extends the range of models that may be practically implemented.
References:
[1] B. Carpenter, A. Gelman, M. Hoffman, D. Lee, B. Goodrich, M. Betancourt, M. Brubaker, J. Guo, P. Li, , and A. Riddell. Stan: A probabilistic programming language. Journal of Statistical Software, 76(1):1–32, 2017.
[2] B. Carpenter, M. D. Hoffman, M. Brubaker, D. Lee, P. Li, and M. Betancourt. The Stan Math Library: Reverse-mode automatic differentiation in C++. arXiv, 1509.07164, 2015.
[3] https://github.com/metrumresearchgroup/Torsten
[4] W. R. Gillespie, Y. Zhang, Torsten: Stan functions for pharmacometric applications. New functionality including within chain parallel computation. PAGE 2019 Abstr IV-29.
[5] A. Vehtari, A. Gelman, D. Simpson, B. Carpenter, and P. Bürkner. Rank-normalization, folding, and localization: An improved Rhat for assessing convergence of MCMC, http://arxiv.org/abs/1903.08008
Reference: PAGE () Abstr 9395 [www.page-meeting.org/?abstract=9395]
Poster: Methodology - Estimation Methods