ARTreeFormer: A Faster Attention-based Autoregressive Model for Phylogenetic Inference

TL;DR

ARTreeFormer is a novel attention-based autoregressive model that significantly accelerates phylogenetic inference by employing fixed-point iteration and attention mechanisms, enabling scalable and efficient analysis of large datasets.

Contribution

It introduces ARTreeFormer, combining fixed-point iteration and attention to improve speed and scalability over previous autoregressive models in phylogenetics.

Findings

Achieves faster computation on CUDA devices.

Maintains high approximation accuracy.

Effective on challenging real data phylogenetic problems.

Abstract

Probabilistic modeling over the combinatorially large space of tree topologies remains a central challenge in phylogenetic inference. Previous approaches often necessitate pre-sampled tree topologies, limiting their modeling capability to a subset of the entire tree space. A recent advancement is ARTree, a deep autoregressive model that offers unrestricted distributions for tree topologies. However, its reliance on repetitive tree traversals and inefficient local message passing for computing topological node representations may hamper the scalability to large datasets. This paper proposes ARTreeFormer, a novel approach that harnesses fixed-point iteration and attention mechanisms to accelerate ARTree. By introducing a fixed-point iteration algorithm for computing the topological node embeddings, ARTreeFormer allows fast vectorized computation, especially on CUDA devices. This, together with an attention-based global message passing scheme, significantly improves the computation speed of ARTree while maintaining great approximation performance. We demonstrate the effectiveness and efficiency of our method on a benchmark of challenging real data phylogenetic inference problems.