Fast cross-validation for Bayesian inference using proposals on a linear subspace
Mohamed Tarek [1, 2]; Many Francis [3]; Anastasios Panagiotelis [2]
[1] Pumas-AI Inc., USA; [2] Business School, University of Sydney, Australia; [3] Formerly at Pumas-AI Inc., USA
Objectives: Hierarchical models typically found in the field of pharmacometrics have a number of parameters that scales with the number of subjects in the population. In the Bayesian inference workflow, model comparison and evaluation is often done using the expected log predictive density (ELPD) estimated using some form of cross-validation (CV), e.g. leave-1-out CV [1,2,3,4].
Methods: The Pareto smoothed importance sampling (PSIS) [5,6], leave-one-out CV method has gained significant popularity in the recent years for its ability to efficiently estimate the ELPD without running the Markov Chain Monte Carlo (MCMC) algorithm multiple times. However, the PSIS CV method can fail due to highly influential data points causing a significant bias in the ELPD estimate. In this paper, another approximate CV algorithm is proposed that offers better accuracy than PSIS CV while being significantly faster than re-running the inference. The proposed method relies on running approximate MCMC. The approximate MCMC works by re-parameterizing the log probability function such that MCMC proposals are only made on a pre-defined subspace of the full parameter space. The subspace is obtained using principal component analysis of the MCMC samples using the initial full data run. The idea of using parameter subspaces in MCMC was explored before in [7] in the context of Bayesian neural networks and in [8] in the context of pharmacometrics. This work is different in that it applies the subspace MCMC approach in [8] to CV while providing additional mathematical justification not previously presented in any work to our knowledge.
Results: A mathematical justification and analysis of the proposed method is presented in this work together with experimental results. The test model used was a pharmacokinetic (PK) model with 2 depot compartments, 1 central compartment and linear absorption and clearance. A synthetic population of 12 subjects and 11 observations per subject was simulated and used for the inference. Leave-future-1-observation CV was run using the re-running CV, PSIS CV, and the subspace CV methods. No less than 4 observations per subject was allowed in each CV run. The following table shows the maximum mean discrepancy (MMD) comparing the pointwise ELPD using the subspace CV and PSIS CV methods to the re-running inference CV method. A subspace size of 5 was used in the subspace CV method. The running time is also shown. As shown in the table, the proposed subspace CV method can achieve more than 500 times lower MMD using approximately 1% of the running time.
|
|
Re-running Inference CV |
PSIS CV |
Subspace CV |
|
MMD |
0 |
1.32061 |
0.00178 |
|
Time (sec) |
1158.40 |
1.01 |
12.47 |
Conclusions: Performing MCMC by generating proposals on a subspace is a cost effective alternative for estimating the ELPD compared to re-running the full MCMC and is more accurate than pure PSIS CV. The proposed method can also be used to enhance the accuracy of PSIS CV instead of replacing it, running the proposed method only when the pareto shape parameter in PSIS CV is higher than a certain threshold.
References:
[1] Gelman, A., Vehtari, A., Simpson, D., Margossian, C.C., Carpenter, B., Yao, Y., Kennedy, L., Gabry, J., Bürkner, P.-C., Modrák, M.: Bayesian Workflow. arXiv:2011.01808 [stat] (2020). arXiv: 2011.01808.
[2] Vehtari, A., Ojanen, J.: A survey of Bayesian predictive methods for model assessment, selection and comparison (2012)
[3] Piironen, J., Vehtari, A.: Comparison of Bayesian predictive methods for model selection. Statistics and Computing 27, 711–735 (2017)
[4] Gneiting, T., Raftery, A.E.: Strictly proper scoring rules, prediction, and estimation. Journal of the American Statistical Association 102(477), 359–378 (2007) https://arxiv.org/abs/https://doi.org/10.1198/016214506000001437. https://doi.org/10.1198/016214506000001437
[5] Vehtari, A., Gelman, A., Gabry, J.: Practical Bayesian model evaluation using leave-one-out cross-validation and WAIC. Statistics and computing 27, 1413–1432 (2017)
[6] Vehtari, A., Simpson, D., Gelman, A., Yao, Y., Gabry, J.: Pareto Smoothed Importance Sampling. arXiv:1507.02646 [stat] (2021). arXiv: 1507.02646.
[7] Izmailov, P., Maddox, W.J., Kirichenko, P., Garipov, T., Vetrov, D., Wilson, A.G.: Subspace inference for bayesian deep learning. In: Uncertainty in Artificial Intelligence, pp. 1169–1179 (2020). PMLR
[8] Manu Francis, Vijay Ivaturi, Mohamed Tarek: Subspace MCMC algorithm for Bayesian parameter estimation of hierarchical PK/PD models in Pumas. In: Population Approach Group in Europe (2021).