Torchtree: Flexible Phylogenetic Model Development and Inference Using PyTorch

Bayesian inference has predominantly relied on the Markov chain Monte Carlo (MCMC) algorithm for many years. However, MCMC is computationally laborious, especially for complex phylogenetic models of time trees. This bottleneck has led to the search for alternatives, such as variational Bayes, which can scale better to large data sets. In this paper, we introduce torchtree, a framework written in Python that allows developers to easily implement rich phylogenetic models and algorithms using a fixed tree topology. One can either use automatic differentiation or leverage torchtree's plug-in system to compute gradients analytically for model components for which automatic differentiation is slow. We demonstrate that the torchtree variational inference framework performs similarly to BEAST in terms of speed, and delivers promising approximation results, though accuracy varies across scenarios. Furthermore, we explore the use of the forward Kullback-Leibler (KL) divergence as an optimizing criterion for variational inference, which can handle discontinuous and nondifferentiable models. Our experiments show that inference using the forward KL divergence is frequently faster per iteration compared with the evidence lower bound (ELBO) criterion, although the ELBO-based inference may converge faster in some cases. Overall, torchtree provides a flexible and efficient framework for phylogenetic model development and inference using PyTorch.