Skip to content

Review: torchtree: Flexible Phylogenetic Model Development and Inference Using PyTorch

Citation

  • Fourment, M., Macaulay, M., Swanepoel, C. J., Ji, X., Suchard, M. A., & Matsen IV, F. A. (2025). torchtree: flexible phylogenetic model development and inference using PyTorch. Systematic Biology, 75(1), 39–51.
  • DOI

Abstract

torchtree is a Python framework for phylogenetic model development and inference built on PyTorch. By expressing phylogenetic models as differentiable computational graphs, it enables gradient-based optimization, variational Bayesian inference, and automatic differentiation of complex model components. The framework is highly modular, allowing researchers to combine substitution models, tree priors, and inference algorithms from a library of components.

TorchTree represents an alternative architectural approach to GPU-accelerated phylogenetics compared to Hifuku. Where Hifuku writes custom CUDA kernels for the pruning recursion and fills an elite-archive (MAP-Elites) survey with them, TorchTree expresses the entire likelihood computation as a PyTorch computational graph and uses automatic differentiation for gradient-based inference (variational Bayes, Hamiltonian Monte Carlo, or MAP optimization).

The two approaches make different tradeoffs. TorchTree's automatic differentiation enables methods that require gradients with respect to branch lengths (HMC, ADVI), which are difficult to implement in hand-written CUDA kernels. Hifuku's hand-written kernels achieve tighter control over memory layout and floating-point precision (enforcing float32 partials with float64 accumulators exactly as specified in CONSTRAINTS.md) and integrate the tree metric and the log-likelihood in the same CUDA kernel without framework overhead.

TorchTree is referenced in Hifuku's documentation as a related approach in the space of GPU-accelerated Bayesian phylogenetics. Both systems address the same computational bottleneck (phylogenetic likelihood on GPU) but from different starting points (framework-driven autodiff vs. hand-tuned CUDA). The modular substitution model library in TorchTree includes the same GTR, LG, WAG, and JTT models as Hifuku, providing a useful cross-reference for model implementation correctness.