Torchtree: flexible phylogenetic model development and inference using PyTorch

TL;DR

Torchtree is a flexible Python framework built on PyTorch that enables efficient development and inference of complex phylogenetic models, offering alternatives to traditional MCMC methods with comparable speed and accuracy.

Contribution

Introduces torchtree, a novel framework for phylogenetic inference that utilizes variational Bayes and automatic differentiation within PyTorch, improving scalability and flexibility.

Findings

Torchtree performs similarly to BEAST in speed and accuracy.

Using forward KL divergence can speed up inference per iteration.

ELBO-based inference may converge faster in some cases.

Abstract

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 datasets. 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 approximation accuracy. Furthermore, we explore the use of the forward KL divergence as an optimizing criterion for variational inference, which can handle discontinuous and non-differentiable models. Our experiments show that inference using the forward KL divergence tends to be faster per iteration compared to 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.