ARTreeFormer: Fast Phylogenetic Inference

• The paper introduces ARTreeFormer as an autoregressive model that builds unrooted bifurcating tree topologies via sequential edge decisions.
• It replaces traversal-heavy message passing with fixed-point iterations for topological embeddings and a global multi-head attention scheme to enhance GPU batching.
• Empirical results show ARTreeFormer achieves significant speedups and competitive performance in parsimony, tree topology density estimation, and variational Bayesian inference.

ARTreeFormer is an attention-based autoregressive model for phylogenetic inference that performs probabilistic modeling over the space of unrooted bifurcating tree topologies. It was introduced as an acceleration of ARTree, preserving unrestricted support over tree topologies while replacing traversal-heavy components with a fixed-point iteration for topological node embeddings and an attention-based global message passing scheme (Xie et al., 24 Jul 2025). The name is potentially confusable with several unrelated methods called “TreeFormer” in plant skeleton estimation and remote-sensing tree counting, and with the autoregressive “HourglassTree” architecture for geometric tree generation; those are distinct models with different problem settings, data modalities, and objectives (Liu et al., 2024, Amirkolaee et al., 2023, Wang et al., 7 Feb 2025).

1. Nomenclature and scope

ARTreeFormer was proposed for phylogenetic inference, not for plant image analysis, tree counting, or geometric tree synthesis. In the phylogenetic setting, the object of interest is a posterior over evolutionary trees, each consisting of a discrete topology $\tau$ and continuous branch lengths $q \ge 0$, conditioned on aligned sequence data $Y \in \Omega^{N \times M}$ under a continuous-time Markov model of evolution (Xie et al., 24 Jul 2025).

The name requires disambiguation because multiple arXiv papers use closely related terminology. “TreeFormer: Single-view Plant Skeleton Estimation via Tree-constrained Graph Generation” addresses single-view plant skeleton estimation using a non-autoregressive RelationFormer backbone with an MST-based tree constraint, and explicitly states that its method is not autoregressive (Liu et al., 2024). “TreeFormer: a Semi-Supervised Transformer-based Framework for Tree Counting from a Single High Resolution Image” concerns density estimation and counting in aerial or satellite imagery (Amirkolaee et al., 2023). “Autoregressive Generation of Static and Growing Trees” describes an autoregressive transformer called HourglassTree; the paper explicitly notes that it does not introduce the name ARTreeFormer (Wang et al., 7 Feb 2025).

Name	Domain	Core formulation
ARTreeFormer	Phylogenetic inference	Autoregressive topology model over unrooted bifurcating trees
TreeFormer	Plant skeleton estimation	Non-autoregressive image-to-graph generation with MST-constrained training
TreeFormer	Tree counting in remote sensing	Semi-supervised transformer for density estimation
HourglassTree	Geometric tree generation	Autoregressive hourglass transformer for static and growing trees

This suggests that the most precise use of “ARTreeFormer” is the phylogenetic model of Xie and collaborators rather than a generic label for transformer models involving tree-structured outputs.

2. Problem formulation in phylogenetic inference

Phylogenetic inference seeks a posterior distribution over evolutionary trees from aligned sequences. For a tree with topology $\tau$ and branch lengths $q$, the likelihood factorizes over sites as

$$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$

where $a^i$ is an assignment to internal nodes, $\eta$ is the stationary distribution, and $P_{jk}(q)$ are edge transition probabilities (Xie et al., 24 Jul 2025). Bayesian phylogenetic inference targets

$$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$

but exact computation is infeasible because the number of topologies grows as $q \ge 0$0 for $q \ge 0$1 taxa (Xie et al., 24 Jul 2025).

ARTreeFormer addresses the topology component of this problem. It learns a distribution $q \ge 0$2 over tree topologies and is typically paired with a conditional model $q \ge 0$3 for branch lengths (Xie et al., 24 Jul 2025). The model is intended to retain ARTree’s “unconfined support over all tree topologies” while improving scalability. The paper frames this as a response to limitations of conditional clade distributions and subsplit Bayesian networks, whose support is tied to subsplits observed in training and may fail to cover the full topology space when the posterior is diffuse (Xie et al., 24 Jul 2025).

A central design choice is to model topology construction autoregressively. Given a fixed taxa order $q \ge 0$4, the tree is built by sequentially adding leaf $q \ge 0$5 to an edge $q \ge 0$6 of the current subtree $q \ge 0$7. If $q \ge 0$8 denotes the sequence of edge decisions, then

$q \ge 0$9

with $Y \in \Omega^{N \times M}$0 (Xie et al., 24 Jul 2025). The paper states that performance is empirically insensitive to the fixed taxa order.

3. Autoregressive topology construction

At each autoregressive step, ARTreeFormer operates on the current partial tree $Y \in \Omega^{N \times M}$1 and scores every candidate edge in $Y \in \Omega^{N \times M}$2 as a possible insertion site for the next taxon (Xie et al., 24 Jul 2025). This yields a categorical action distribution over edges rather than over clades or pre-enumerated topologies. The model therefore defines topology probabilities through construction actions, and the final topology probability is the product of stepwise edge-selection probabilities.

The scoring pipeline begins with node representations derived from the current topology. These representations are globally contextualized, pooled into a graph representation $Y \in \Omega^{N \times M}$3, and then combined with endpoint features for each candidate edge. For an edge $Y \in \Omega^{N \times M}$4, the invariant edge pooling is defined by

$Y \in \Omega^{N \times M}$5

using elementwise max on the two endpoint features. The readout then forms

$Y \in \Omega^{N \times M}$6

where $Y \in \Omega^{N \times M}$7 is a sinusoidal positional embedding encoding the autoregressive time step (Xie et al., 24 Jul 2025). The action distribution is

$Y \in \Omega^{N \times M}$8

with $Y \in \Omega^{N \times M}$9 (Xie et al., 24 Jul 2025).

This formulation differs from local message-passing approaches in ARTree. The earlier model used graph neural networks and repeated tree traversals to compute node embeddings and propagate information across the current partial tree. ARTreeFormer replaces those components with a fixed-point solver and global multi-head attention, which the paper presents as a drop-in acceleration (Xie et al., 24 Jul 2025).

A plausible implication is that the model’s action space remains combinatorial, but its internal computation is substantially more amenable to GPU batching than traversal-based solvers and multi-round local propagation.

4. Fixed-point topological embeddings and global attention

A defining technical feature of ARTreeFormer is its fixed-point method for topological node embeddings. For the current ordinal tree $\tau$0 with leaf set $\tau$1 and internal nodes $\tau$2, leaves are assigned one-hot embeddings $\tau$3, and internal embeddings are defined as the minimizer of the Dirichlet energy

$\tau$4

The global minimum satisfies the discrete harmonic condition

$\tau$5

with fixed one-hot values at the leaves (Xie et al., 24 Jul 2025).

Stacking internal-node embeddings into $\tau$6 and defining internal adjacency $\tau$7 and leaf–internal adjacency $\tau$8, the harmonicity conditions become

$\tau$9

ARTreeFormer solves this by fixed-point iteration,

$q$0

initialized with uniform values $q$1 (Xie et al., 24 Jul 2025). The paper states that for any $q$2, the spectral radius satisfies $q$3, implying $q$4, so the iteration converges linearly to a unique solution with a topology- and $q$5-independent rate bound (Xie et al., 24 Jul 2025). A norm-based stopping rule is used, and the required number of iterations $q$6 is stated to be constant independent of $q$7 and $q$8.

The “power trick” further reduces sequential depth. Writing the full node stack $q$9 and corresponding matrix $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$0, the method repeatedly squares $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$1 so that

$$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$2

with $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$3 (Xie et al., 24 Jul 2025). After convergence, embeddings are right-padded to dimension $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$4 and projected by an MLP $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$5.

ARTreeFormer then replaces ARTree’s local message passing with a single-query global multi-head attention block. With node features $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$6 and a learned query $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$7, the attention uses $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$8, $$p(Y \mid \tau, q)=\prod_{i=1}^{M}\sum_{a^i}\eta(a^i_r)\prod_{(u,v)\in E(\tau)}P_{a^i_u a^i_v}(q_{uv}),$$9, and $a^i$0, followed by a 2-layer MLP $a^i$1 to obtain $a^i$2 (Xie et al., 24 Jul 2025). Because the query length is 1, the pooling cost is stated as $a^i$3 rather than $a^i$4.

The paper specifies layer normalization before attention and MLP, ELU activations, a residual connection after attention, node feature dimension $a^i$5, and $a^i$6 attention heads (Xie et al., 24 Jul 2025). Topological structure is encoded through the fixed-point embeddings and their projection, while temporal conditioning enters through the sinusoidal positional embedding $a^i$7 (Xie et al., 24 Jul 2025).

5. Objectives, optimization, and empirical results

ARTreeFormer is evaluated under three training regimes: maximum parsimony, tree topology density estimation (TDE), and variational Bayesian phylogenetic inference (VBPI) (Xie et al., 24 Jul 2025).

For maximum parsimony, the target distribution is

$a^i$8

where $a^i$9 is the parsimony score. Optimization uses an annealed multi-sample objective with $\eta$0 particles and VIMCO for low-variance gradients, with $\eta$1 increasing linearly until 1 and results reported after 400k updates (Xie et al., 24 Jul 2025).

For TDE, posterior samples from MrBayes are used as training data, and the objective is maximum likelihood,

$\eta$2

equivalently cross-entropy over autoregressive actions (Xie et al., 24 Jul 2025).

For VBPI, the variational family factorizes as $\eta$3, and the optimized $\eta$4-sample ELBO is

$\eta$5

with annealed likelihood in early training, VIMCO for tree gradients, and reparameterization for branch lengths (Xie et al., 24 Jul 2025). The experiments use the JC substitution model and an $\eta$6 prior on branch lengths in MrBayes comparisons. The implementation is in PyTorch with Adam at learning rate $\eta$7 (Xie et al., 24 Jul 2025).

The evaluation covers eight benchmarks, DS1–DS8, with 27–64 taxa and 378–2520 sites (Xie et al., 24 Jul 2025). The main empirical findings are as follows.

Setting	Result
Speed	Approximately $\eta$8 faster generation and approximately $\eta$9 faster training on GPU versus ARTree
Parsimony	On DS1, reaches the same lower bound/score approximately $P_{jk}(q)$0 faster
TDE	Matches or outperforms ARTree on several datasets and consistently improves over SBN variants
VBPI	Marginal likelihood estimates nearly identical to VBPI-ARTree and VBPI-SBN, and on par with MrBayes SS

The TDE KL divergences reported in the paper are particularly informative. On DS1, SBN-EM, SBN-EM-$P_{jk}(q)$1, SBN-SGA, ARTree, and ARTreeFormer achieve 0.0136, 0.0130, 0.0504, 0.0045, and 0.0065, respectively; on DS3, 0.1243, 0.0882, 0.0922, 0.0548, and 0.0474; on DS4, 0.0763, 0.0637, 0.0739, 0.0299, and 0.0267; on DS7, 0.0483, 0.0399, 0.0301, 0.0191, and 0.0152; and on DS8, 0.1415, 0.1236, 0.1177, 0.0741, and 0.0563 (Xie et al., 24 Jul 2025). The paper therefore characterizes ARTreeFormer as matching or outperforming ARTree while remaining consistently stronger than SBN variants.

For VBPI marginal likelihood, the paper reports near-identity with ARTree-based VBPI and with stepping-stone estimates from MrBayes. Examples include DS1, where ARTreeFormer gives $P_{jk}(q)$2 versus $P_{jk}(q)$3 for ARTree, and DS6, where it gives $P_{jk}(q)$4 versus $P_{jk}(q)$5 (Xie et al., 24 Jul 2025). On influenza HA subsets with $P_{jk}(q)$6 up to 100, ARTreeFormer is reported to achieve estimates closer to SBN than ARTree for large $P_{jk}(q)$7 while maintaining low variance (Xie et al., 24 Jul 2025).

6. Complexity, memory behavior, and implementation characteristics

The computational comparison to ARTree is one of the paper’s central contributions. ARTree’s node embeddings require repeated tree traversals and its local message passing requires multiple rounds, leading to total complexity

$P_{jk}(q)$8

where $P_{jk}(q)$9 is batch size, $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$0 is the number of taxa, $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$1 is the number of message-passing rounds, and $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$2 is feature dimension (Xie et al., 24 Jul 2025). ARTreeFormer instead uses batched fixed-point iteration with the power trick and single-query global attention, yielding

$$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$3

with ideal vectorization reducing the sequential dependence to $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$4 (Xie et al., 24 Jul 2025).

The acceleration mechanisms stated in the paper are dense batched matrix multiplication for fixed-point updates, repeated squaring to reduce sequential depth, $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$5 updates to $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$6 and $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$7 after each attachment, and vectorization across trees and nodes for efficient CUDA GEMMs (Xie et al., 24 Jul 2025). In practical guidance, the paper recommends stacking $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$8 and $$p(\tau, q \mid Y) \propto p(Y \mid \tau, q)\, p(\tau, q),$$9 across the batch dimension, keeping tensors contiguous, and preallocating buffers for powers of $q \ge 0$00 (Xie et al., 24 Jul 2025).

Runtime and memory results reinforce this systems-oriented redesign. ARTreeFormer delivers approximately $q \ge 0$01 faster generation and approximately $q \ge 0$02 faster training on GPU overall, with node embedding time less than 10% of ARTree’s and message-passing time approximately 50% of ARTree’s on CPU/GPU (Xie et al., 24 Jul 2025). In VBPI experiments, the runtime memory is lower across DS1–DS8; for example, on DS8 the paper reports 1044MB for ARTreeFormer versus 2148MB for ARTree, with comparable parameter counts of approximately 219k versus approximately 209k (Xie et al., 24 Jul 2025).

The implementation uses a single NVIDIA A100 GPU in the reported experiments, and large batch sizes are emphasized as important for exploiting vectorization benefits (Xie et al., 24 Jul 2025). The model is publicly available at the repository specified by the authors (Xie et al., 24 Jul 2025).

7. Limitations, interpretation, and relation to similarly named models

The paper identifies several limitations. First, although the convergence rate of the fixed-point iteration is uniform, batched matrix multiplication over internal-node blocks still scales as $q \ge 0$03 per step, so very large trees may remain limited by GPU memory and time (Xie et al., 24 Jul 2025). Second, the power trick reduces sequential depth to $q \ge 0$04, but does not eliminate sequential dependence entirely; very small stopping thresholds $q \ge 0$05 increase compute (Xie et al., 24 Jul 2025). Third, the single global attention layer with a single query vector may underfit subtle structural dependencies; the paper suggests that deeper attention stacks or structural biases such as relation-aware attention could help (Xie et al., 24 Jul 2025). Fourth, the model assumes a fixed taxa order, even though empirical performance is reported to be insensitive to that order (Xie et al., 24 Jul 2025). Finally, ARTreeFormer models only topologies, so full posterior fidelity also depends on the separate variational family used for branch lengths (Xie et al., 24 Jul 2025).

A likely misconception is to treat ARTreeFormer as a generic “tree transformer.” The cited arXiv literature shows that the label “TreeFormer” spans several unrelated formulations. In plant skeleton estimation, “TreeFormer” denotes a non-autoregressive image-to-graph method built on RelationFormer, Deformable DETR, an MST projection, and a selective feature suppression layer; the paper explicitly distinguishes this from autoregressive graph generators such as GGT and states that TreeFormer itself is not AR (Liu et al., 2024). In remote sensing, “TreeFormer” denotes a semi-supervised framework for density estimation and counting with a Pyramid Vision Transformer encoder, contextual attention-based feature fusion, a tree density regressor, and a tree counter token (Amirkolaee et al., 2023). In geometric generation, the relevant autoregressive transformer is named HourglassTree rather than ARTreeFormer, with a multi-resolution hourglass architecture for static trees, point-cloud conditioning, and 4D growth sequences (Wang et al., 7 Feb 2025).

ARTreeFormer therefore occupies a specific position within the broader family of transformer methods on tree-structured domains: it is an autoregressive topology model for phylogenetic inference whose technical novelty lies in replacing ARTree’s traversal-based Dirichlet solver and local GNN propagation with CUDA-friendly fixed-point embedding computation and global attention pooling (Xie et al., 24 Jul 2025). This suggests that its significance is as much algorithmic and systems-oriented as it is probabilistic: the method preserves the expressive autoregressive topology distribution while restructuring the computation so that large-batch vectorization becomes practical.