[Paper Review] TreeFlow: probabilistic programming and automatic differentiation for phylogenetics

Christiaan Swanepoel, Mathieu Fourment+5 more|arXiv (Cornell University)|Nov 9, 2022

Genomics and Phylogenetic Studies4 citations

TL;DR

TreeFlow is a probabilistic programming and automatic differentiation library for phylogenetics that enables scalable gradient-based inference on phylogenetic trees by efficiently computing gradients of the phylogenetic likelihood using TensorFlow. It achieves performance within a factor of 2–3 of specialized libraries like bito/BEAGLE for complex models, demonstrating linear computational scaling and enabling flexible, rapid implementation of new evolutionary models.

ABSTRACT

Probabilistic programming frameworks are powerful tools for statistical modelling and inference. They are not immediately generalisable to phylogenetic problems due to the particular computational properties of the phylogenetic tree object. TreeFlow is a software library for probabilistic programming and automatic differentiation with phylogenetic trees. It implements inference algorithms for phylogenetic tree times and model parameters given a tree topology. We demonstrate how TreeFlow can be used to quickly implement and assess new models. We also show that it provides reasonable performance for gradient-based inference algorithms compared to specialized computational libraries for phylogenetics.

Motivation & Objective

• To bridge the gap between general-purpose probabilistic programming and phylogenetic inference by enabling gradient-based inference on phylogenetic trees.
• To address the computational challenges of combining discrete tree topologies with continuous branch lengths in probabilistic models.
• To provide a flexible yet efficient framework for implementing and benchmarking new phylogenetic models using automatic differentiation.
• To enable scalable inference for complex models involving evolutionary rates, population dynamics, and time-varying parameters.
• To lay the foundation for future support of tree topology inference within a differentiable probabilistic programming framework.

Proposed method

• TreeFlow implements a differentiable phylogenetic likelihood using dynamic programming with TensorArray-based control flow to support recursion in computational graphs.
• It leverages TensorFlow’s automatic differentiation to compute exact gradients of the log-likelihood with respect to branch lengths and model parameters.
• The library uses a hierarchical parameterization of tree topology and branch lengths, allowing joint inference over continuous and discrete components.
• It supports gradient-based inference methods such as Hamiltonian Monte Carlo and variational inference via automatic differentiation.
• TreeFlow’s implementation is optimized for computational efficiency, with performance comparable to specialized libraries like bito/BEAGLE for GTR/Weibull models.
• The framework supports extensible model specification through Python-based probabilistic programming, enabling rapid prototyping of new evolutionary models.

Experimental results

Research questions

• RQ1 Can automatic differentiation be effectively applied to the phylogenetic likelihood computation, which involves recursive and control-flow-heavy operations?
• RQ2 How does the performance of a general-purpose differentiable phylogenetic library compare to specialized high-performance libraries like bito/BEAGLE?
• RQ3 Can TreeFlow scale efficiently to large datasets, and what is its empirical computational complexity?
• RQ4 To what extent can TreeFlow support complex models involving evolutionary rates and time-varying population dynamics?
• RQ5 What are the limitations of mean-field variational inference in capturing posterior correlations in phylogenetic parameters, and how can structured approximations improve this?

Key findings

• TreeFlow achieves computational performance within a factor of 2–3 of bito/BEAGLE for the GTR/Weibull model, indicating strong competitiveness with specialized libraries.
• The computational scaling of TreeFlow is approximately linear with respect to the number of sequences, as indicated by slope estimates below 2 in log-log runtime vs. data size fits.
• The library’s use of TensorFlow’s control flow and execution model enables efficient computation of gradients for recursive likelihood calculations, outperforming JAX’s eager execution in this context.
• Mean-field variational inference in TreeFlow fails to capture posterior correlations in divergence times, particularly in complex datasets like influenza, highlighting the need for structured approximations.
• TreeFlow supports rapid prototyping of new models, including those with time-varying evolutionary rates and population dynamics, through its flexible probabilistic programming interface.
• The framework is extensible and well-suited for future integration of topology inference via differentiable tree search, though this remains a significant computational challenge.