Nonparametric modelling of continuous-time Markov chains with scalable exact and approximate gradients

Abstract

Inferring the infinitesimal rates of continuous-time Markov chains (CTMCs) is a central challenge in many scientific domains. This task is hindered by 3 factors: quadratic growth in the number of rates as the CTMC state space expands, strong dependencies among rates, and incomplete information for many transitions. We introduce a new Bayesian framework that flexibly models CTMC rates by incorporating covariates through Gaussian processes. This approach improves inference by integrating new information and deepens our understanding of CTMC stochastic behaviour by shedding light on potential external drivers. Unlike previous approaches limited to linear covariate effects, our method captures complex nonlinear relationships, enabling fuller use of covariate information and more accurate characterization of their influence. For efficient posterior inference, we employ Hamiltonian Monte Carlo. To avoid the prohibitive cost of repeated matrix-exponential derivatives under existing approaches, we develop 2 scalable strategies for exact and approximate gradient computation. Given N observations and K CTMC states, these strategies reduce the cost from O(NK3), with a substantial constant, to O(K3+NK2), with cheaper constants. We demonstrate the method in Bayesian phylogenetic and phylogeographic inference, showing strong performance on synthetic and real datasets and empirical quadratic scaling in K, even when N<K.