Title: ARTree: A Deep Autoregressive Model for Phylogenetic Inference
Venue: NeurIPS 2023
Accepted: yes
Keywords: phylogenetic inference

autoregressive model

graph neural network

density estimation

variational inference
Abstract: Designing flexible probabilistic models over tree topologies is important for developing efficient phylogenetic inference methods. To do that, previous works often leverage the similarity of tree topologies via hand-engineered heuristic features which would require domain expertise and may suffer from limited approximation capability. In this paper, we propose a deep autoregressive model for phylogenetic inference based on graph neural networks (GNNs), called ARTree. By decomposing a tree topology into a sequence of leaf node addition operations and modeling the involved conditional distributions based on learnable topological features via GNNs, ARTree can provide a rich family of distributions over tree topologies that have simple sampling algorithms, without using heuristic features. We demonstrate the effectiveness and efficiency of our method on a benchmark of challenging real data tree topology density estimation and variational Bayesian phylogenetic inference problems.

ARTree: A Deep Autoregressive Model for Phylogenetic Inference

Abstract

Designing flexible probabilistic models over tree topologies is important for developing efficient phylogenetic inference methods. To do that, previous works often leverage the similarity of tree topologies via hand-engineered heuristic features which would require pre-sampled tree topologies and may suffer from limited approximation capability. In this paper, we propose a deep autoregressive model for phylogenetic inference based on graph neural networks (GNNs), called ARTree. By decomposing a tree topology into a sequence of leaf node addition operations and modeling the involved conditional distributions based on learnable topological features via GNNs, ARTree can provide a rich family of distributions over the entire tree topology space that have simple sampling algorithms and density estimation procedures, without using heuristic features. We demonstrate the effectiveness and efficiency of our method on a benchmark of challenging real data tree topology density estimation and variational Bayesian phylogenetic inference problems.

1 Introduction

Reconstructing the evolutionary relationships among species has been one of the central problems in computational biology, with a wide range of applications such as genomic epidemiology (Dudas et al., 2017; du Plessis et al., 2021; Attwood et al., 2022) and conservation genetics (DeSalle & Amato, 2004). Based on molecular sequence data (e.g. DNA, RNA, or protein sequences) of the observed species and a model of evolution, this has been formulated as a statistical inference problem on the hypotheses of shared history, i.e., phylogenetic trees , where maximum likelihood and Bayesian approaches are the most popular methods (Felsenstein, 1981; Yang & Rannala, 1997; Mau et al., 1999; Larget & Simon, 1999; Huelsenbeck et al., 2001). However, phylogenetic inference can be challenging due to the composite structure of tree space which contains both continuous and discrete components (e.g., the branch lengths and the tree topologies) and the large search space of tree topologies that explodes combinatorially as the number of species increases (Whidden & Matsen IV, 2015; Dinh et al., 2017).

Recently, several efforts have been made to improve the efficiency of phylogenetic inference algorithms by designing flexible probabilistic models over the tree topology space (Höhna & Drummond, 2012; Larget, 2013; Zhang & Matsen IV, 2018). One typical example is subsplit Bayesian networks (SBNs) (Zhang & Matsen IV, 2018), which is a powerful probabilistic graphical model that provides a flexible family of distributions over tree topologies. Given a sample of tree topologies (e.g., sampled tree topologies from an MCMC run), SBNs have proved effective for accurate tree topology density estimation that generalizes beyond observed samples by leveraging the similarity of hand-engineered subsplit structures among tree topologies. Moreover, SBNs also allow fast ancestral sampling and

hence were later on integrated into a variational Bayesian phylogenetic inference (VBPI) framework to provide variational posteriors over tree topologies (Zhang & Matsen IV, 2019). However, due to the limited parent-child subsplit patterns in the observed samples, SBNs can not provide distributions whose support spans the entire tree topology space (Zhang & Matsen IV, 2022). Furthermore, when used as variational distributions over tree topologies in VBPI, SBNs often rely on subsplit support estimation for variational parameterization, which requires high-quality pre-sampled tree topologies that would become challenging to obtain when the posterior is diffuse.

While SBNs suffer from the aforementioned limitations due to their hand-engineered design, a number of deep learning methods have been proposed for probabilistic modeling of graphs (Jin et al., 2018; You et al., 2018a; Cao & Kipf, 2018; Simonovsky & Komodakis, 2018). Instead of using hand-engineered features, these approaches use neural networks to define probabilistic models for the connections between graph nodes which allow for learnable distributions over graphs. Due to the flexibility of neural networks, the resulting models are capable of learning complex graph patterns automatically. Among these deep graph models, graph autoregressive models (You et al., 2018b; Li et al., 2018; Liao et al., 2019; Dai et al., 2020; Shi et al., 2020) are designed to learn flexible graph distributions that also allow easy sampling procedures by sequentially adding nodes and edges. Therefore, they serve as an ideal substitution of SBNs for phylogenetic inference that can provide more expressive distributions over tree topologies.

In this paper, we propose a novel deep autoregressive model for phylogenetic inference, called ARTree, which allows for more flexible distributions over tree topologies without using heuristic features than SBNs. With a pre-selected order of leaf nodes (i.e., species or taxa), ARTree generates a tree topology by recursively adding new leaf nodes to the edges of the current tree topology, starting from a star-shaped tree topology with the first three leaf nodes (Figure 1). The edge to which a new leaf node connects is determined according to a conditional distribution based on learnable topological features of the current tree topology via GNNs (Zhang, 2023). This way, probability distributions provided by ARTree all have full support that spans the entire tree topology space. Unlike SBNs, ARTree can be readily used in VBPI without requiring subsplit support estimation for parameterization. In experiments, we show that ARTree outperforms SBNs on a benchmark of challenging real data tree topology density estimation and variational Bayesian phylogenetic inference problems.

2 Background

Phylogenetic likelihoods Aphylogenetic tree is commonly described by a bifurcating tree topology τ and the associated non-negative branch lengths q . The tree topology τ represents the evolutionary relationship of the species and the branch lengths q quantify the evolutionary intensity along the edges of τ . The leaf nodes of τ correspond to the observed species and the internal nodes of τ represent the unobserved ancestor species. A continuous time Markov model is often used to describe the transition probabilities of the characters along the edges of the tree (Felsenstein, 2004). Concretely, let Y = { Y 1 , . . . , Y M } ∈ Ω N × M be the observed sequences (with characters in Ω ) of length M over N species. Under the assumption that different sites evolve independently and identically, the likelihood of Y given τ, q takes the form

where a i ranges over all extensions of Y i to the internal nodes with a i u being the character assignment of node u ( r represents the root node), E ( τ ) is the set of edges of τ , q uv is the branch length of the edge ( u, v ) ∈ E ( τ ) , P jk ( q ) is the transition probability from character j to k through a branch of length q , and η is the stationary distribution of the Markov model.

Subsplit Bayesian networks Let X be the set of leaf labels representing the existing species. A non-empty subset of X is called a clade and the set of all clades C ( X ) is equipped with a total order ≻ (e.g., lexicographical order). An ordered clade pair ( W,Z ) satisfying W ∩ Z = ∅ and W ≻ Z is called a subsplit . A subsplit Bayesian network (SBN) is then defined as a Bayesian network whose nodes take subsplit values or singleton clade values that describe the local topological structures of tree topologies. For a rooted tree topology, one can find the corresponding node assignment of SBNs

by following its splitting processes (Figure 4 in Appendix A). The SBN based probability of a rooted tree topology τ then takes the following form

where S i denotes the subsplit- or singleton-clade-valued random varaibles at node i (node 1 is the root node), π i is the index set of the parents of node i and { s i } i ≥ 1 is the corresponding node assignment. For unrooted tree topologies, we can also define their SBN based probabilities by viewing them as rooted tree topologies with unobserved roots and integrating out the positions of the root node as follows: p sbn ( T u = τ ) = ∑ e ∈ E ( τ ) p sbn ( τ e ) , where τ e is the resulting rooted tree topology when the rooting position is on edge e . In practice, the conditional probability tables (CPTs) of SBNs are often parameterized based on a sample of tree topologies (e.g., the observed data for density estimation (Zhang & Matsen IV, 2018) or fast bootstrap/MCMC samples (Minh et al., 2013; Zhang, 2020) for VBPI). As a result, the supports of SBN-induced distributions are often limited by the splitting patterns in the observed samples and could not span the entire tree topology space (Zhang & Matsen IV, 2022). More details on SBNs can be found in Appendix A.

Variational Bayesian phylogenetic inference Given a prior distribution p ( τ, q ) , the phylogenetic posterior distribution takes the form

Let Q ϕ ( τ ) and Q ψ ( q | τ ) be variational families over the spaces of tree topologies and branch lengths respectively. The VBPI approach uses Q ϕ , ψ ( τ, q ) = Q ϕ ( τ ) Q ψ ( q | τ ) to approximate the posterior p ( τ, q | Y ) by maximizing the following multi-sample ( K > 1 ) lower bound

The tree topology distribution Q ϕ ( τ ) is often SBNs which in this case rely on subsplit support estimation for parameterization that requires high-quality pre-sampled tree topologies and would become challenging for diffuse posteriors (Zhang & Matsen IV, 2022). The branch lengths distribution Q ψ ( q | τ ) can be diagonal lognormal distribution parametrized via heuristic features or learnable topological features of τ (Zhang & Matsen IV, 2019; Zhang, 2020, 2023). Compared to the singlesample lower bound, the multi-sample lower bound in (4) allows efficient variance-reduced stochastic gradient estimators (e.g. VIMCO (Mnih & Rezende, 2016)) for tree topology variational parameters. Moreover, using multiple samples would encourage exploration over the vast tree topology space, albeit it may also deteriorates the training of the variational approximation (Rainforth et al., 2019). In practice, a moderate K is often suggested (Zhang & Matsen IV, 2022). See more details on VBPI in Appendix B.

Graph autoregressive models By decomposing a graph as a sequence of components (nodes, edges, motifs, etc), graph autoregressive models generate the full graph by adding one component at a time, until some stopping criteria are satisfied (You et al., 2018b; Jin et al., 2018; Liao et al., 2019). In previous works, recurrent neural networks (RNNs) for graphs are usually utilized to predict new graph components conditioned on the sub-graphs generated so far (You et al., 2018b). The key of graph autoregressive models is to find a way to efficiently sequentialize graph structures, which is often domain-specific.

3 Proposed method

In this section, we propose ARTree, a deep autoregressive model for phylogenetic inference that can provide flexible distributions whose support spans the entire tree topology space and can be naturally parameterized without using heuristic approaches such as subsplit support estimation. We first describe a particular autoregressive generating process of phylogenetic tree topologies. We then develop powerful GNNs to parameterize learnable conditional distributions of this generating process. We consider unrooted tree topologies in this section, but the method developed here can be easily adapted to rooted tree topologies.

B

2

Leaf

D

ARTree node

Interior

Edge

4

E

ARTree

D

Probability vector

0 1

Sampling

One

-

hot vector

5

1

3

4

3

7

6

4

2

3

6

5

2

1

3

4

2

1

1

2

GNN

1

2

3

x

1

x

2

x

3

Sampling

x

4

GNN

5

2

3

x

2

x

3

1

4

x

1

x

4

Sampling

x

5

6

2

3

x

2

x

3

1

7

x

1

x

4

x

5

Sampling

x

6

4

5

Leaf

node

Interior

node

index

Edge

Probability

vector

0

1

One

-

hot

vector

Adding

leaf

node

Ordinal

tree

topologies

of

higher

rank

GNN

1

2

3

4

5

7

1

2

5

3

3.1 A sequential generating process of tree topologies

To better illustrate our approach, we begin with some notations. Let τ n = ( V n , E n ) be a tree topology with n leaf nodes and V n , E n are the sets of nodes and edges respectively. Note that | V n | = 2 n -2 and | E n | = 2 n -3 due to the unrooted and bifurcating structure of τ n . The leaf nodes in V n are treated as labeled nodes and the interior nodes in V n are treated as unlabeled nodes. Let us assume a pre-selected order for the leaf nodes X = { x 1 , . . . , x N } , which is called taxa order for short by us. Now, consider a sequential generating process for all possible tree topologies that have leaf nodes X . We start with a definition below.

Definition 1 (Ordinal Tree Topology) . Let X = { x 1 , . . . , x N } be a set of N ( N ≥ 3) leaf nodes. Let τ n = ( V n , E n ) be a tree topology with n ( n ≤ N ) leaf nodes in X . We say τ n is an ordinal tree topology of rank n , if its leaf nodes are the first n elements of X , i.e., V n ∩ X = { x 1 , . . . , x n } .

We now describe a procedure that constructs ordinal tree topologies of rank N recursively by adding one leaf node at a time as follows. We first start from τ 3 , the ordinal tree topology of rank 3, which is the smallest ordinal tree topology and is unique due to its unrooted and bifurcating structure. Suppose now we have an ordinal tree topology τ n = ( V n , E n ) of rank n . To add the leaf node x n +1 to τ n , we i) select an edge e n = ( u, v ) ∈ E n and remove it from E n ; ii) add a new node w and two new edges ( u, w ) , ( w,v ) to the tree topology; iii) add the leaf node x n +1 and an edge ( w,x n +1 ) to the tree topology. This way, we obtain an ordinal tree topology τ n +1 of rank n +1 . Intuitively, the leaf node x n +1 is added to the tree topology τ n by attaching it to an existing edge e n ∈ E n . The position of the selected edge represents the evolutionary relationship between this new species and others. After performing this procedure for n = 3 , . . . , N -1 , we finally obtain an ordinal tree topology τ = τ N of rank N . See Figure 1 for an illustration.

During the generating process described above, the selected edges at each time step form a sequence D = ( e 3 , . . . , e N -1 ) . This sequence D of length N -3 records all the decisions we have made for autoregressively generating a tree topology τ and thus we call D a decision sequence. In fact, there is a one-to-one mapping between decision sequences and ordinal tree topologies of rank N , which is formalized in Theorem 1. Note that a similar process is also used in online phylogenetic sequential Monte Carlo (OPSMC) (Dinh et al., 2016), where the leaf node addition operation is incorporated into the design of the proposal distributions.

Theorem 1. Let D = { D | D = ( e 3 , . . . , e N -1 ) , e n ∈ E n , ∀ 3 ≤ n ≤ N -1 } be the set of all decision sequences of length N -3 and T be the set of all ordinal tree topologies of rank N . Let the map

be the generating process described above. Then g is a bijection between D and T .

According to Theorem 1, for each tree topology τ ∈ T , there is a unique decision sequence given by g -1 ( τ ) . We call this process of finding the decision sequences of tree topologies the decomposition

4

B

index

2

6

5

Algorithm 1: ARTree: An autoregressive model for phylogenetic tree topologies

Input: a set X = { x 1 , . . . , x N } of leaf nodes. Output: an ordinal tree topology τ of rank N ; the ARTree probability Q ( τ ) of τ . τ 3 = ( V 3 , E 3 ) ← the unique ordinal tree topology of rank 3 ; for n = 3 , . . . , N -1 do Calculate the probability vector q n ∈ R | E n | using the current GNN model; Sample an edge decision e n from Discrete ( q n ) and assume e n = ( u, v ) ; Create a new node w ; E n +1 ← ( E n \{ e n } ) ∪ { ( u, w ) , ( w,v ) , ( w,x n +1 ) } ; V n +1 ← V n ∪ { w,x n +1 } ; τ n +1 ← ( V n +1 , E n +1 ) ; end τ ← τ N ; Q ( τ ) ← q 3 ( e 3 ) q 4 ( e 4 ) · · · q N -1 ( e N -1 ) .

process . See more details on the decomposition process in Appendix C. The following lemma shows that one can find g -1 ( τ ) in linear time.

Lemma 1. The time complexity of the decomposition process induced by g -1 ( · ) is O ( N ) .

The proofs of Theorem 1 and Lemma 1 can be found in Appendix D. Based on the bijection g defined in Theorem 1, we can model the distribution Q ( D ) over the space of decision sequences D instead of modeling the distribution Q ( τ ) over T . Due to the sequential nature of D , we can decompose Q ( D ) as the product of conditional distributions over the elements:

In what follows, we simplify Q ( e n | e 3 , . . . , e n -1 ) as Q ( e n | e <n ) and let e < 3 be the empty set.

3.2 Graph neural networks for edge decision distribution

By Theorem 1, the sequence e <n corresponds to a sequence of ordinal tree topologies of increasing ranks ( τ 3 , . . . , τ n ) (the empty set e < 3 corresponds to the unique ordinal tree topology τ 3 of rank 3). Therefore, the discrete distribution Q ( e n | e <n ) in equation (5) defines the probability of adding the leaf node x n +1 to the edge e n of τ n , conditioned on all the ordinal tree topologies ( τ 3 , . . . , τ n ) generated so far. In what follows, we will show step by step how to use graph neural networks (GNNs) to parameterize such a conditional distribution given tree topologies.

Topological node embeddings At the n -th time step of the generating process, we first find the node embeddings of the current tree topology τ n = ( V n , E n ) , which is a set { f n ( u ) ∈ R N : u ∈ V n } that assigns each node with an encoding vector in R N . Following Zhang (2023), we first assign one hot encoding to the leaf nodes, i.e.

where δ is Kronecker delta function; we then get the embeddings for the interior nodes by minimizing the Dirichlet energy ℓ ( f n , τ n ) := ∑ ( u,v ) ∈ E n || f n ( u ) -f n ( v ) || 2 using the efficient two-pass algorithm described in Zhang (2023). One should note that the embeddings for interior nodes may change as new leaf nodes are added to the ordinal tree topologies, which is a main difference between our model and other graph autoregressive models.

Message passing networks Using these topological node embeddings as the initial node features, GNNs apply message passing steps to compute the representation vector of nodes that encode topological information of τ n , where the node features are updated with the information from their neighborhoods in a convolutional manner (Gilmer et al., 2017). More concretely, the l -th round of message passing is implemented by

where F l message and F l updating are the message function and updating function in the l -th round, and N ( v ) is the neighborhood of the node v . In our implementations, the choices of F l message and F l updating follow the edge convolution operator (Wang et al., 2018), while other variants of GNNs can also be applied. The final node features of τ n are given by { f L n ( v ) : v ∈ V n } after L rounds of message passing.

Node hidden states The conditional distribution Q ( ·| e <n ) is highly complicated as it has to capture how x n +1 can be added to τ n based on how previous leaf nodes are added to form the tree topologies. A common approach is to use RNNs to model this complex distribution that strikes a good balance between expressiveness and scalability (You et al., 2018b; Liao et al., 2019). In our model, after obtaining the final node features of τ n , a gated recurrent unit (GRU) (Cho et al., 2014) follows, i.e.

where h n ( v ) is the hidden state of v at the n -th generation step and is initialized to zero for the newly added nodes including those in τ 3 . The node hidden states { h n ( v ); v ∈ V n } , therefore, contain the information of all the tree topologies generated so far which can be used for conditional distribution modeling.

Time guided readout We now construct the distribution Q ( ·| e <n ) over edge decisions based on the node hidden states. As mentioned before, a main difference between our model and other graph autoregressive models is that the node embedding f 0 n ( v ) of a node v may vary with the time step n . We, therefore, incorporate time embeddings into the readout step which first forms the edge features r n ( e ) ∈ R of e = ( u, v ) using

where b n is the sinusoidal positional embedding of time step n that is widely used in Transformers (Vaswani et al., 2017), F pooling is the pooling function implemented as 2-layer MLPs followed by an elementwise maximum operator, and F readout is the readout function implemented as 2-layer MLPs with a scalar output. Then the conditional distribution for edge decision is

where probability vector q n ∈ R | E n | for parametrizing Q ( ·| e <n ) is obtained by applying a softmax function to all the time guided edge features in equation (9b).

As the last step, we sample an edge e n ∈ E n from the discrete distribution in equation (10) and add the leaf node x n +1 to e n as described in Section 3.1. This way, we update the ordinal tree topology from τ n of rank n to τ n +1 of rank n +1 . We can repeat this procedure until an ordinal tree topology τ = τ N of rank N is reached. The probability of τ then takes the form

where D is the decision sequence and ϕ are the learnable parameters in the model. We call this autoregressive model for tree topologies ARTree, and summarize it in Algorithm 1. Note that equation (11) can also be used for tree topology probability evaluation where the decision sequence D = g -1 ( τ ) is obtained from the decomposition process (Appendix C) that enjoys a linear time complexity (Lemma 1). Compared to SBNs, ARTree does not rely on heuristic features for parameterization and can provide distributions whose support spans the entire tree topology space as all possible decisions would have nonzero probabilities due to the softmax parameterization in equation (10). Although different taxa orders may affect the performance of ARTree, we find this effect is negligible in our experiments.

There exist other VI approaches for phylogenetic inference that also use unconfined models over the space of tree topologies. Moretti et al. (2021) proposed to sample tree topologies through subtree merging and resampling following CSMC (Wang et al., 2015), but employed a parametrized proposal distribution. Koptagel et al. (2022) proposed a novel multifurcating tree topology sampler named SLANTIS, which makes decisions on adding edges in a specific order based on a simply parameterized weight matrix and maximum spanning trees, and is integrated with CSMC to sample bifurcating tree

10 4

10 5

M

10

2

10

1

KL divergence

DS1

SBN-SGA

SBN-EM

SBN-EM-

ARTree

10

8

10

6

10

4

10

2

10 0

ground truth

10

8

10

6

10

4

10

2

10 0

estimated probability

SBN-EM

peak 1

peak 2

outliers

10

8

10

6

10

4

10

2

10 0

ground truth

10

8

10

6

10

4

10

2

10 0

estimated probability

ARTree

peak 1

peak 2

outliers

topologies ( ϕ -CSMC). Unlike these methods, ARTree employs GNNs for an autoregressive model that builds up the tree topology sequentially through leaf node addition operations, which not only allows fast sampling of trees, but also provides straightforward density estimation procedures. We demonstrate the advantage of ARTree over these baselines in the experiments.

4 Experiments

In this section, we test the effectiveness and efficiency of ARTree for phylogenetic inference on two benchmark tasks: tree topology density estimation (TDE) and variational Bayesian phylogenetic inference (VBPI). In all experiments, we report the inclusive KL divergence from posterior estimates to the ground truth to measure the approximation error of different methods. We will use 'KL divergence' for inclusive KL divergence throughout this section unless otherwise specified. The code is available at https://github.com/tyuxie/ARTree .

Experimental setup We perform experiments on eight data sets which we will call DS1-8. These data sets, consisting of sequences from 27 to 64 eukaryote species with 378 to 2520 site observations, are commonly used to benchmark phylogenetic MCMC methods (Hedges et al., 1990; Garey et al., 1996; Yang & Yoder, 2003; Henk et al., 2003; Lakner et al., 2008; Zhang & Blackwell, 2001; Yoder & Yang, 2004; Rossman et al., 2001; Höhna & Drummond, 2012; Larget, 2013; Whidden & Matsen IV, 2015). For the Bayesian setting, we focus on the joint posterior distribution of the tree topologies and the branch lengths and assume a uniform prior on the tree topologies, an i.i.d. exponential prior Exp(10) on branch lengths, and the simple JC substitution model (Jukes et al., 1969). For each of these data sets, we run 10 single-chain MrBayes (Ronquist et al., 2012) for one billion iterations, collect samples every 1000 iterations, and discard the first 25% samples as burn-in. These samples form the ground truth of the marginal distribution of tree topologies to which we will compare the posterior estimates obtained by different methods. All GNNs have L = 2 rounds in the message passing step. All the activation functions in MLPs are exponential linear units (ELUs) (Clevert et al., 2015). The taxa order is set to the lexicographical order of the corresponding species names in all experiments except the ablation studies. All the experiments are run on an Intel Xeon Platinum 9242 processor. All models are implemented in PyTorch (Paszke et al., 2019) and trained with the Adam (Kingma & Ba, 2015) optimizer. The learning rate is 0.001 for SBNs, 0.0001 for ARTree, and 0.001 for the branch length model.

4.1 Tree topology density estimation

We first investigate the performance of ARTree for tree topology density estimation given the MCMC posterior samples on DS1-8. Following Zhang & Matsen IV (2018), we run MrBayes on each data set with 10 replicates of 4 chains and 8 runs until the runs have ASDSF (the standard convergence criteria used in MrBayes) less than 0.01 or a maximum of 100 million iterations. The training data

sets are formed by collecting samples every 100 iterations and discarding the first 25%. Now, given a training data set M = { τ m } M m =1 , we train ARTree via maximum likelihood estimation using stochastic gradient ascent. In each iteration, the stochastic gradient is obtained as follows

where a minibatch { τ m b } B b =1 is randomly sampled from M . We compare ARTree to SBN baselines including SBN-EM, SBN-EMα , and SBN-SGA. For SBN-EM and SBN-EMα , we use the same setting as previously done in Zhang & Matsen IV (2018) (see Appendix A for more details). In addition to these EM variants, a gradient based method for SBNs called SBN-SGA is considered, where SBNs are reparametrized with the latent parameters initialized as zero (see equation (18) in Appendix B) and optimized via stochastic gradient ascent, similarly to ARTree. For both ARTree and SBN-SGA, the results are collected after 200000 parameter updates with batch size B = 10 .

The left and middle plots of Figure 2 show a comparison between ARTree and SBN-EM on DS1, which has a peaky posterior distribution. Compared to SBN-EM, ARTree provides more accurate probability estimates for tree topologies on the peaks and significantly reduces the large biases in the low probability region (the crimson dots). The right plot of Figure 2 shows the KL divergence of different methods as a function of the sample size of the training data. We see that ARTree consistently outperforms SBN based methods for all M s. Moreover, as the sample size M increases, ARTree keeps providing better approximation while SBNs start to level off when M is large. This indicates the superior flexibility of ARTree over SBNs for tree topology density estimation.

Table 1 shows the KL divergences of different methods on DS1-8. We see that ARTree outperforms SBN based methods on all data sets. The gradient based method SBN-SGA is better than SBN-EM on most of the data sets because SBN-EM is well initialized (Zhang & Matsen IV, 2018) and more likely to get trapped in local modes. From this point of view, the comparison between ARTree and SBN-SGA is fair because they both use a uniform initialization that facilitates exploration.

4.2 Variational Bayesian phylogenetic inference

Our second experiment is on VBPI, where we compare ARTree to SBNs for tree topology variational approximations. Both methods are evaluated on the aforementioned benchmark data sets DS1-8. Following Zhang & Matsen IV (2019), we use the simplest SBN and gather the subsplit support from 10 replicates of 10000 ultrafast maximum likelihood bootstrap trees (Minh et al., 2013). For both ARTree and SBNs, the collaborative branch lengths are parametrized using the learnable topological features with the edge convolution operator (EDGE) for GNNs (Zhang, 2023). We set K = 10 for the multi-sample lower bound (4) and use the following annealed unnormalized posterior at the i -th iteration

0

10

20

30

40

50

order index

10

2

10

1

KL divergence

SBN

ARTree

10 4

10 5

iterations

8200

8000

7800

7600

7400

7200

ELBO

10 5

-7111.0

-7110.0

ARTree(10)

ARTree(5)

ARTree*(10)

SBN(10)

SBN(5)

10

4

10

3

10

2

10

1

10 0

ground truth

10

4

10

3

10

2

10

1

10 0

variational approximation

SBN

ARTree

where β i = min { 1 . 0 , 0 . 001 + i/H } is the inverse temperature that goes from 0.001 to 1 after H iterations. For ARTree, a long annealing period H = 200000 is used for DS6 and DS7 due to the highly multimodal posterior distributions on these two data sets (Whidden & Matsen IV, 2015) and H = 100000 is used for the other data sets. For SBNs, we set H = 100000 for all data sets. The Monte Carlo gradient estimates for the tree topology parameters and the branch lengths parameters are obtained via VIMCO (Mnih & Rezende, 2016) and the reparametrization trick (Zhang & Matsen IV, 2019) respectively. The results are collected after 400000 parameter updates.

The left plot in Figure 3 shows the evidence lower bound (ELBO) as a function of the number of iterations on DS1. Although the larger support of ARTree adds to the complexity of training for tree topology variational approximation, we see that by the time SBN based methods converge, ARTree based methods achieve comparable (if not better) lower bounds and finally surpass the SBN baselines in the end. We also find that using fewer particles ( K = 5) in the training objective tends

to provide larger ELBO. Moreover, time guidance turns out to be crucial for ARTree, as evidenced by the significant performance drop when it is turned off. As shown in the middle plot, compared to SBNs, ARTree can provide a more accurate variational approximation of the tree topology posterior. To investigate the effect of taxa orders on ARTree, we randomly sample 50 taxa orders and report the KL divergence for each order in the right plot of Figure 3. We find that ARTree exhibits weak randomness as the taxa order varies and consistently outperforms SBNs by a large margin.

Table 2 shows the KL divergences to the ground truth, evidence lower bound (ELBO), 10-sample lower bound (LB-10), and marginal likelihood (ML) estimates obtained by different methods on DS1-8. We find that ARTree achieves smaller KL divergences than SBNs across all data sets and performs on par or better than SBNs for lower bound and marginal likelihood estimation. Compared to SBNs, the ELBOs provided by ARTree tend to have larger variances, especially on DS2, DS3, and DS7, which is partly due to the larger support of ARTree that spans the entire tree topology space (see more discussions in Appendix E).

5 Conclusion

In this paper, we introduced ARTree, a deep autoregressive model over tree topologies for phylogenetic inference. Unlike SBNs that rely on hand-engineered features for parameterization and require presampled tree topologies, ARTree is built solely on top of learnable topological features (Zhang, 2023) via GNNs which allows for a rich family of distributions over the entire phylogenetic tree topology space. Moreover, as an autoregressive model, ARTree also allows simple forward sampling procedures and straightforward density computation, which make it readily usable for tree topology density estimation and variational Bayesian phylogenetic inference. In experiments, we showed that ARTree outperforms SBNs on a benchmark of challenging real data tree topology density estimation and variational Bayesian phylogenetic inference problems, especially in terms of tree topology posterior approximation accuracy.

References

Stephen W. Attwood, Sarah C. Hill, David M. Aanensen, Thomas R. Connor, and Oliver G. Pybus. Phylogenetic and phylodynamic approaches to understanding and combating the early SARS-CoV2 pandemic. Nature Reviews. Genetics , 23:547 - 562, 2022.

Nicola De Cao and Thomas Kipf. MolGAN: An implicit generative model for small molecular graphs, 2018. URL https://arxiv.org/abs/1805.11973 .

Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using RNN encoder-decoder for statistical machine translation. arXiv preprint arXiv:1406.1078 , 2014.

Djork-Arné Clevert, Thomas Unterthiner, and Sepp Hochreiter. Fast and accurate deep network learning by exponential linear units (ELUs). arXiv: Learning , 2015.

Hanjun Dai, Azade Nazi, Yujia Li, Bo Dai, and Dale Schuurmans. Scalable deep generative modeling for sparse graphs. In International Conference on Machine Learning , pp. 2302-2312, 2020.

Rob DeSalle and George Amato. The expansion of conservation genetics. Nat. Rev. Genet. , 5(9): 702-712, September 2004. ISSN 1471-0056. doi: 10.1038/nrg1425. URL http://dx.doi.org/ 10.1038/nrg1425 .

Vu Dinh, Arman Bilge, Cheng Zhang, and Frederick A Matsen IV. Probabilistic path Hamiltonian Monte Carlo. In Proceedings of the 34th International Conference on Machine Learning , pp. 1009-1018, July 2017. URL http://proceedings.mlr.press/v70/dinh17a.html .

Vu C. Dinh, Aaron E. Darling, and Frederick A. Matsen IV. Online Bayesian phylogenetic inference: Theoretical foundations via sequential Monte Carlo. Systematic Biology , 67:503 - 517, 2016.

Louis du Plessis, John T McCrone, Alexander E Zarebski, Verity Hill, Christopher Ruis, Bernardo Gutierrez, Jayna Raghwani, Jordan Ashworth, Rachel Colquhoun, Thomas R Connor, Nuno R Faria, Ben Jackson, Nicholas J Loman, Áine O'Toole, Samuel M Nicholls, Kris V Parag, Emily Scher, Tetyana I Vasylyeva, Erik M Volz, Alexander Watts, Isaac I Bogoch, Kamran Khan, COVID-19 Genomics UK (COG-UK) Consortium†, David M Aanensen, Moritz U G Kraemer, Andrew Rambaut, and Oliver G Pybus. Establishment and lineage dynamics of the SARS-CoV2 epidemic in the UK. Science , January 2021. ISSN 0036-8075, 1095-9203. doi: 10.1126/ science.abf2946. URL https://science.sciencemag.org/content/early/2021/01/07/ science.abf2946 .

Gytis Dudas, Luiz Max Carvalho, Trevor Bedford, Andrew J Tatem, Guy Baele, Nuno R Faria, Daniel J Park, Jason T Ladner, Armando Arias, Danny Asogun, Filip Bielejec, Sarah L Caddy, Matthew Cotten, Jonathan D'Ambrozio, Simon Dellicour, Antonino Di Caro, Joseph W Diclaro, Sophie Duraffour, Michael J Elmore, Lawrence S Fakoli, Ousmane Faye, Merle L Gilbert, Sahr M Gevao, Stephen Gire, Adrianne Gladden-Young, Andreas Gnirke, Augustine Goba, Donald S Grant, Bart L Haagmans, Julian A Hiscox, Umaru Jah, Jeffrey R Kugelman, Di Liu, Jia Lu, Christine M Malboeuf, Suzanne Mate, David A Matthews, Christian B Matranga, Luke W Meredith, James Qu, Joshua Quick, Suzan D Pas, My V T Phan, Georgios Pollakis, Chantal B Reusken, Mariano Sanchez-Lockhart, Stephen F Schaffner, John S Schieffelin, Rachel S Sealfon, Etienne SimonLoriere, Saskia L Smits, Kilian Stoecker, Lucy Thorne, Ekaete Alice Tobin, Mohamed A Vandi, Simon J Watson, Kendra West, Shannon Whitmer, Michael R Wiley, Sarah M Winnicki, Shirlee Wohl, Roman Wölfel, Nathan L Yozwiak, Kristian G Andersen, Sylvia O Blyden, Fatorma Bolay, Miles W Carroll, Bernice Dahn, Boubacar Diallo, Pierre Formenty, Christophe Fraser, George F Gao, Robert F Garry, Ian Goodfellow, Stephan Günther, Christian T Happi, Edward C Holmes, Brima Kargbo, Sakoba Keïta, Paul Kellam, Marion P G Koopmans, Jens H Kuhn, Nicholas J Loman, N'faly Magassouba, Dhamari Naidoo, Stuart T Nichol, Tolbert Nyenswah, Gustavo Palacios, Oliver G Pybus, Pardis C Sabeti, Amadou Sall, Ute Ströher, Isatta Wurie, Marc A Suchard, Philippe Lemey, and Andrew Rambaut. Virus genomes reveal factors that spread and sustained the Ebola epidemic. Nature , April 2017. ISSN 0028-0836, 1476-4687. doi: 10.1038/nature22040. URL http://dx.doi.org/10.1038/nature22040 .

Felsenstein. Evolutionary trees from DNA sequences: A maximum likelihood approach. Journal of Molecular Evolution , 17:268-276, 1981.

Joseph Felsenstein. Inferring Phylogenies . Sinauer associates, 2 edition, 2004.

R. Garey, T. J. Near, M. R. Nonnemacher, and S. A. Nadler. Molecular evidence for Acanthocephala as a subtaxon of Rotifera. Mol. Evol. , 43:287-292, 1996.

Justin Gilmer, Samuel S. Schoenholz, Patrick F. Riley, Oriol Vinyals, and George E. Dahl. Neural message passing for quantum chemistry. ArXiv , abs/1704.01212, 2017.

B. Hedges, K. D. Moberg, and L. R. Maxson. Tetrapod phylogeny inferred from 18S and 28S ribosomal RNA sequences and review of the evidence for amniote relationships. Mol. Biol. Evol. , 7:607-633, 1990.

A. Henk, A. Weir, and M. Blackwell. Laboulbeniopsis termitarius, an ectoparasite of termites newly recognized as a member of the Laboulbeniomycetes. Mycologia , 95:561-564, 2003.

Sebastian Höhna and Alexei J. Drummond. Guided tree topology proposals for Bayesian phylogenetic inference. Syst. Biol. , 61(1):1-11, January 2012. ISSN 1063-5157. doi: 10.1093/sysbio/syr074. URL http://dx.doi.org/10.1093/sysbio/syr074 .

P. Huelsenbeck, F. Ronquist, R. Nielsen, and J. P. Bollback. Bayesian inference of phylogeny and its impact on evolutionary biology. Science , 294:2310-2314, 2001.

Wengong Jin, Regina Barzilay, and T. Jaakkola. Junction tree variational autoencoder for molecular graph generation. ArXiv , abs/1802.04364, 2018.

Thomas H Jukes, Charles R Cantor, et al. Evolution of protein molecules. Mammalian protein metabolism , 3:21-132, 1969.

P. Kingma and J. Ba. Adam: A method for stochastic optimization. In ICLR , 2015.

Hazal Koptagel, Oskar Kviman, Harald Melin, Negar Safinianaini, and Jens Lagergren. VaiPhy: a variational inference based algorithm for phylogeny. In Advances in Neural Information Processing Systems , 2022.

Lakner, P. van der Mark, J. P. Huelsenbeck, B. Larget, and F. Ronquist. Efficiency of Markov chain Monte Carlo tree proposals in Bayesian phylogenetics. Syst. Biol. , 57:86-103, 2008.

Bret Larget. The estimation of tree posterior probabilities using conditional clade probability distributions. Syst. Biol. , 62(4):501-511, July 2013. ISSN 1063-5157. doi: 10.1093/sysbio/syt014. URL http://dx.doi.org/10.1093/sysbio/syt014 .

Bret R. Larget and D. L. Simon. Markov chain Monte Carlo algorithms for the Bayesian analysis of phylogenetic trees. Molecular Biology and Evolution , 16:750-750, 1999.

Li, O. Vinyals, C. Dyer, R. Pascanu, and P. Battaglia. Learning deep generative models of graphs, 2018. URL https://arxiv.org/abs/1803.03324 .

Renjie Liao, Yujia Li, Yang Song, Shenlong Wang, Charlie Nash, William L. Hamilton, David Kristjanson Duvenaud, Raquel Urtasun, and Richard S. Zemel. Efficient graph generation with graph recurrent attention networks. ArXiv , abs/1910.00760, 2019.

Mau, M. Newton, and B. Larget. Bayesian phylogenetic inference via Markov chain Monte Carlo methods. Biometrics , 55:1-12, 1999.

Bui Quang Minh, Minh Anh Nguyen, and Arndt von Haeseler. Ultrafast approximation for phylogenetic bootstrap. Molecular Biology and Evolution , 30:1188 - 1195, 2013.

Andriy Mnih and Danilo Jimenez Rezende. Variational inference for monte carlo objectives. In International Conference on Machine Learning , 2016.

Antonio Khalil Moretti, Liyi Zhang, Christian Andersson Naesseth, Hadiah Venner, David M. Blei, and Itsik Pe'er. Variational combinatorial sequential Monte Carlo methods for Bayesian phylogenetic inference. In Conference on Uncertainty in Artificial Intelligence , 2021.

Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. PyTorch: An imperative style, high-performance deep learning library. In Neural Information Processing Systems , 2019.

Tom Rainforth, Adam R. Kosioreck, Tuan Anh Le, Chris J. Maddison, Maximilian Igl, Frank Wood, and Yee Whye Teh. Tighter variational bounds are not necessarily better. In Proceedings of the 36th International Conference on Machine Learning , 2019.

Fredrik Ronquist, Maxim Teslenko, Paul Van Der Mark, Daniel L Ayres, Aaron Darling, Sebastian Höhna, Bret Larget, Liang Liu, Marc A Suchard, and John P Huelsenbeck. MrBayes 3.2: Efficient Bayesian phylogenetic inference and model choice across a large model space. Systematic Biology , 61(3):539-542, 2012.

Y. Rossman, J. M. Mckemy, R. A. Pardo-Schultheiss, and H. J. Schroers. Molecular studies of the Bionectriaceae using large subunit rDNA sequences. Mycologia , 93:100-110, 2001.

Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. GraphAF: a flow-based autoregressive model for molecular graph generation. In International Conference on Learning Representations , 2020. URL https://openreview.net/forum?id=S1esMkHYPr .

Martin Simonovsky and Nikos Komodakis. GraphVAE: Towards generation of small graphs using variational autoencoders. In International Conference on Artificial Neural Networks , pp. 412-422, 2018.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Systems , volume 30, 2017.

Liangliang Wang, Alexandre Bouchard-Côté, and A. Doucet. Bayesian phylogenetic inference using a combinatorial sequential Monte Carlo method. Journal of the American Statistical Association , 110:1362 - 1374, 2015. URL https://api.semanticscholar.org/CorpusID:4495539 .

Yue Wang, Yongbin Sun, Ziwei Liu, Sanjay E. Sarma, Michael M. Bronstein, and Justin M. Solomon. Dynamic graph CNN for learning on point clouds. ACM Transactions on Graphics (TOG) , 38:1 12, 2018.

Chris Whidden and Frederick A Matsen IV. Quantifying MCMC exploration of phylogenetic tree space. Syst. Biol. , 64(3):472-491, May 2015. ISSN 1063-5157, 1076-836X. doi: 10.1093/sysbio/ syv006. URL http://dx.doi.org/10.1093/sysbio/syv006 .

Yang and A. D. Yoder. Comparison of likelihood and Bayesian methods for estimating divergence times using multiple gene loci and calibration points, with application to a radiation of cute-looking mouse lemur species. Syst. Biol. , 52:705-716, 2003.

Ziheng Yang and Bruce Rannala. Bayesian phylogenetic inference using DNA sequences: a Markov chain Monte Carlo method. Molecular Biology and Evolution , 14(7):717-724, 1997.

D. Yoder and Z. Yang. Divergence datas for Malagasy lemurs estimated from multiple gene loci: geological and evolutionary context. Mol. Ecol. , 13:757-773, 2004.

Jiaxuan You, Bowen Liu, Zhitao Ying, Vijay Pande, and Jure Leskovec. Graph convolutional policy network for goal-directed molecular graph generation. In Advances in Neural Information Processing Systems , pp. 6410-6421, 2018a.

Jiaxuan You, Rex Ying, Xiang Ren, William L. Hamilton, and Jure Leskovec. GraphRNN: Generating realistic graphs with deep auto-regressive models. In International Conference on Machine Learning , 2018b.

Cheng Zhang. Improved variational Bayesian phylogenetic inference with normalizing flows. In Neural Information Processing Systems , 2020.

Cheng Zhang. Learnable topological features for phylogenetic inference via graph neural networks. In International Conference on Learning Representations , 2023.

Cheng Zhang and Frederick A Matsen IV. Generalizing tree probability estimation via Bayesian networks. In Neural Information Processing Systems , 2018.

Cheng Zhang and Frederick A Matsen IV. Variational Bayesian phylogenetic inference. In International Conference on Learning Representations , 2019.

Cheng Zhang and Frederick A Matsen IV. A variational approach to Bayesian phylogenetic inference, 2022. URL https://arxiv.org/abs/2204.07747 .

Zhang and M. Blackwell. Molecular phylogeny of dogwood anthracnose fungus (Discula destructiva) and the Diaporthales. Mycologia , 93:355-365, 2001.

S

1

S

2

S

3

S

4

S

5

S

6

S

7

· · ·

· · ·

· · ·

· · ·

· · ·

· · ·

· · ·

· · ·

D

A

B

C

A

B

C

D

ABC

D

A

BC

D

A

B

C

D

D

AB

CD

A

B

C

D

A

B

C

D

assign

assign

S

4

S

5

S

6

S

7

S

2

S

3

S

1

A Details of subsplit Bayesian networks

One recent and expressive graphical model that provides a flexible family of distributions over tree topologies is the subsplit Bayesian network, as proposed by Zhang & Matsen IV (2018). Let X be the set of N labeled leaf nodes. A non-empty set C of X is referred to as a clade and the set of all clades of X , denoted by C ( X ) , is a totally ordered set with a partial order ≻ (e.g., lexicographical order) defined on it. An ordered pair of clades ( W,Z ) is called a subsplit of a clade C if it is a bipartition of C , i.e., W ≻ Z , W ∩ Z = ∅ , and W ∪ Z = C .

Definition 2 (Subsplit Bayesian Network) . A subsplit Bayesian network (SBN) B X on a leaf node set X of size N is defined as a Bayesian network whose nodes take on subsplit or singleton clade values of X and has the following properties: (a) The root node of B X takes on subsplits of the entire labeled leaf node set X ; (b) B X contains a full and complete binary tree network B ∗ X as a subnetwork; (c) The depth of B X is N -1 , with the root counted as depth 1 .

Due to the binary structure of B ∗ X , the nodes in SBNs can be indexed by denoting the root node with S 1 and two children of S i with S 2 i and S 2 i +1 recursively where S i is an internal node (see the left plot in Figure 4). For any rooted tree topology, by assigning the corresponding subsplits or singleton clades values { S i = s i } i ≥ 1 to its nodes, one can uniquely map it into an SBN node assignment (see the middle and right plots in Figure 4).

As Bayesian networks, the SBN based probability of a rooted tree topology τ takes the following form

where π i is the index set of the parents of node i . For unrooted tree topologies, we can also define their SBN based probabilities by viewing them as rooted tree topologies with unobserved roots and integrating out the positions of the root node as follows:

where τ e is the resulting rooted tree topology when the rooting position is on edge e .

In practice, SBNs are parameterized according to the conditional probability sharing principle where the conditional probability for parent-child subsplit pairs are shared across the SBN network, regardless of their locations. The set of all conditional probabilities are called conditional probability tables (CPTs). Parameterizing SBNs, therefore, often requires finding an appropriate support of CPTs. For tree topology density estimation, this can be done using the sample of tree topologies that is given as the data set. For variational Bayesian phylogenetic inference, as no sample of tree topologies is available, one often resorts to fast bootstrap or MCMC methods (Minh et al., 2013; Zhang, 2020). Let S r denotes the root subsplits and S ch | pa denotes the child-parent subsplit pairs in the support. The parameters of SBNs are then p = { p s 1 ; s 1 ∈ S r } ∪ { p s | t ; s | t ∈ S ch | pa } where

As a result, the supports of SBN-induced distributions are often limited by the splitting patterns in the observed samples and could not span the entire tree topology space (Zhang & Matsen IV, 2022).

The SBN-EM Algorithm For unrooted tree topologies, the SBN based probability (15) can be viewed as a hidden variable model where the root subsplit is the hidden variable. In this case, SBNs can be trained using the expectation-maximization (EM) algorithm, as proposed by Zhang & Matsen IV (2018). Given a training set { τ k } M k =1 , we first initialize the parameter estimates as ˆ p EM, (0) (i.e., the simple average estimates as in Zhang & Matsen IV (2018)). In the i -th step, we run the E-step and M-step as follows

E-step : ∀ 1 ≤ k ≤ M , compute q ( i ) k ( s 1 ) = p ( τ k ,s 1 | ˆ p EM , ( i ) ) ∑ s 1 ∼ τ k p ( τ k ,s 1 | ˆ p EM , ( i ) ) where s 1 ∼ τ k means

the subsplit s 1 can be achieved by placing a 'virtual root' on an edge of τ .

M-step : update the parameter estimates by

where I is the indicator function, s e j,k is the node value of S j in τ e k , ˜ m s and ˜ m s,t are equivalent counts and α is the regularization coefficient that encourages generalization.

When α > 0 , this algorithm is called SBN-EMα .

B Details of variational Bayesian phylogenetic inference

With two variational families Q ϕ ( τ ) and Q ψ ( q | τ ) over the space of tree topologies and branch lengths, the variational Bayesian phylogenetic inference (VBPI) approach forces Q ϕ , ψ ( τ, q ) = Q ϕ ( τ ) Q ψ ( q | τ ) to approximate the posterior p ( τ, q | Y ) by maximizing the following multi-sample lower bound

Gradients of the objective (17) w.r.t. ϕ and ψ can be estimated by the VIMCO estimator (Mnih & Rezende, 2016) and the reparameterization trick respectively. In the following, we introduce some common choices of Q ϕ ( τ ) and Q ψ ( q | τ ) .

Choice of Q ϕ ( τ ) Before the proposed ARTree framework in this article, SBNs is the common choice of Q ϕ ( τ ) . As introduced in Appendix A, SBNs provide a probability distribution over unrooted tree topologies in equation (15). Given a subsplit support of CPTs, SBNs can be parameterized as follows

The parameters ϕ = { ϕ s 1 ; s 1 ∈ S r } ∪ { ϕ s | t ; s | t ∈ S ch | pa } are called latent parameters of SBNs.

Choice of Q ψ ( q | τ ) The distribution Q ψ ( q | τ ) is often taken to be a diagonal lognormal distribution, which can be parametrized using some heuristic features (Zhang & Matsen IV, 2019) or the recently proposed learnable topological features (Zhang, 2023) of τ as follows. For each edge e = ( u, v ) in τ , one can first obtain the edge features using h e = f ( h u , h v ) where h u is the GNN output at node u and f is a permutation invariant function. Then the mean and standard deviation parameters are given by

where MLP µ and MLP σ are two multi-layer perceptrons (MLPs).

C Details of tree topology decomposition process

The tree topology decomposition process, which maps a tree topology to a corresponding decision sequence, is indeed the inverse operation of Algorithm 1. Also, the decomposition process is implemented in a recursive way starting from the tree topology τ N of rank N . Intuitively, given an ordinal tree topology τ n +1 of rank n +1 , one can detach the leaf node x n +1 as well as its unique neighbor w and reconnect the two neighbors of w . The remaining graph, denoted by τ n , is an ordinal tree topology of rank n ; the edge decision e n is exactly the reconnected edge. This process continues until the unique ordinal tree topology τ 3 of rank 3 is reached. We summarize the sketch of tree topology decomposition process in Algorithm 2.

Given a tree topology τ with all of the N leaf nodes, we can evaluate its ARTree based probability by first mapping it to a decision sequence D = ( e 3 , . . . , e N -1 ) following Algorithm 2 and then calculate the probability as the product of conditionals

D The proofs of Theorem 1 and Lemma 1

Theorem 1. Let D = { D | D = ( e 3 , . . . , e N -1 ) , e n ∈ E n , ∀ 3 ≤ n ≤ N -1 } be the set of all decision sequences of length N -3 and T be the set of all ordinal tree topologies of rank N . Let the map

be the generating process used in ARTree. Then g is a bijection between D and T .

Proof of Theorem 1 It is obvious that g is a well-defined map from D to T . To prove it is a bijection, it suffices to show g is injective and surjective.

̸

We first show that g is injective. Assume there are two decision sequences D (1) , D (2) ∈ D and D (1) = D (2) . Let k be the first position where D (1) and D (2) begin to differ, i.e. e (1) i = e (2) i for all i < k and e (1) k = e (2) k . If g ( D (1) ) = g ( D (2) ) =: τ , one can take the subtree topology of rank k and k +1 of τ , τ k and τ k +1 . Noting that e k refers to the edge in τ where to add the new node x n +1 , the equation e (1) k = e (2) k implies they will induce different τ k +1 s. This contradicts the uniqueness of τ k +1 . Therefore, we conclude that g is injective.

̸

Next, we prove that g is surjective. For a tree topology τ with rank N , we denote its subtree topology of rank k by τ k , where k = 3 , . . . , n . In fact, for each k , the tree topology τ k +1 corresponds to adding the leaf node x k +1 to an edge in τ k , and we denote this edge by e k . It is easy to verify that the constructed D = ( e 3 , . . . , e N -1 ) is a preimage of τ .

̸

;

Lemma 1. The time complexity of the decomposition process induced by g -1 ( · ) is O ( N ) .

Proof of Lemma 1 Assume the tree topology τ is stored as a binary tree data structure, where each node other than the root node also has a parent node pointer. Before decomposing τ , we first build a dictionary of the ( n, x n ) mappings by a traversal across τ that maps n to the leaf node x n in τ , ∀ n ≤ N . The time complexity of building this dictionary is O ( N ) . In the decomposition process, given an ordinal tree topology of rank n +1 , it costs O (1) time to locate x n +1 and determine the unique parent node w of x n +1 and the neighbors of w . It is obvious that the time complexity of detaching and reconnecting operations is O (1) . One can repeat this procedure for n = N -1 , . . . , 3 , resulting in a time complexity of O ( N ) . Therefore, the time complexity of the decomposition process induced by g -1 ( · ) is O ( N ) .

E Limitations

Larger variances of ELBO estimates As an autoregressive model for phylogenetic tree topology, ARTree provides reliable approximations for target distributions over tree topologies, as evidenced by our real data experiments. However, when incorporated with the branch length model for VBPI, the ELBOs provided by ARTree tend to have larger variances, which we find is caused by the occasional occurrence of 'outliers' among samples. In fact, as the support of ARTree spans over the entire tree topology space, it adds to the difficulty of fitting the conditional distribution Q ψ ( q | τ ) , compared to SBNs. When combined with ARTree, the approximation accuracy of Q ψ ( q | τ ) might be related to the cluster structure of peaks in the tree topology posterior. See Figure 3 in Whidden & Matsen IV (2015) for cluster subtree-prune-and-regraft (SPR) graphs of DS1-8. We also examine the approximation accuracies of Q ψ ( q | τ ) trained along with ARTree for those τ in the support of SBNs and find a significant enhancement in ELBO and reduction of variances. This phenomenon raises an important topic: proper design and optimization of branch length model when the support of tree topology model spans the entire space. We leave this for future work.

Minor increasement of lower bounds From Table 2, the estimates of lower bounds (as well as marginal likelihoods) of ARTree do not improve much over those of SBNs. We provide two explanations for this: (i) The lower bounds in VBPI are more sensitive to the quality of branch length model Q ( q | τ ) instead of the tree topology model Q ( τ ) . As SBN and ARTree use the same parametrization of the branch length model, we do not expect a large improvement in lower bounds. (ii) The support of ARTree spans the entire tree topology space. This adds to the difficulty of training Q ( q | τ ) which is conditioned on tree topology τ , as discussed in the above paragraph. Finally, from Table 3, we see that the enhancement on ELBOs of ARTree indeed comes from a better tree topology model as evidenced by the result of the 'SBN + branch length model trained along with ARTree' combination.

Time complexity Another limitation of ARTree is its time complexity. We compare the training time and evaluation time in Table 4. Both ARTree and SBNs take more time in training than other methods with unconfined support since they both build machine-learning models with enormous parameters and rely heavily on optimization. However, they take a comparable amount of time to provide good enough approximations for marginal likelihood estimation of similar accuracy, as evidenced by ARTree ∗ and SBN ∗ . ARTree takes more time than SBNs because it relies on several submodules which, although complicated, are designed to promote the expressive power to accommodate the complex tree space and are widely-used strategies in the literature. The inefficiency of autoregressive generative models is also an inherent issue.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis work proposes a new approach, ARTree, for obtaining more complex tree-topology approximations by combining deep autoregressive models and GNNs. The existing works in black-box VI for phylogenetic inference predominantly rely on SBNs as approximations of the tree-topology distribution, but here ARTree is shown to be superior in terms of KL divergence to the true distribution (obtained via MCMC). ARTree, in contrast to SBNs, does not rely on presampled tree topologies and thus explores the full tree-topology space, not a subset.\n\nSTRENGTHS:\nDue to the combinatorial nature of the tree-topology space, designing efficient and powerful density approximation algorithms of the tree-topology posterior is arguably the most complicated aspect of Bayesian phylogenetic inference. In the VI setting, SBNs have been the SOTA algorithms, but, as is mentioned in the paper, SBNs rely on presampled tree topologies. This means SBNs require candidate trees from other tree-sampling algorithms, making them not stand alone. \n\nThe algorithm proposed in the paper is an important contribution to field that is receiving increased interest in the ML community (most VI for phylogenetics papers have been published in NeurIPS, ICLR and UAI). Additionally, there have not been many attempts to improve over SBNs. As such, the proposed work appears to be relevant and could be of significant interest to the NeurIPS community.\n\nThe writing and clarity of the paper is good, although I have some questions below.\n\nThe experiments are performed using appropriate baselines (SBNs and VBPI w. SBNs). However I am missing some important references like VaiPhy (Koptagel et al., selected as oral at NeurIPS 22) and VCSMC (Moretti et al., from UAI 21). More on this below.\n\nWEAKNESSES:\n**Related work**: As this is a paper proposes a new algorithm for doing VI for phylogenetics, VaiPhy by Koptagel et al. (2022) should be referenced, preferably also VCSMC by Moretti et al. (2021). These work are highly related. I do not consider it necessary to experimentally compare ARTree with these algorithms though, as they do not compare favorably with VBPI empirically. \n\nFurthermore, in the VaiPhy paper a sequential algorithm (SLANTIS) is designed for sampling tree topologies. I believe it can be seen as a conditional sampler in the sense that it decided whether to replace (labeled) edges in a topology based on a Bernoulli probability, given the existing edges in the topology. It is does not use GNNs, and the algorithms are clearly distinct as SLANTIS uses precomputed maximum spanning trees. However, I think a conceptual distinction should be included in the paper, nonetheless.\n\n**Evaluating tree topologies**: It is clear that the likelihood of a topology, $\\tau$, simulated by ARTree can be efficiently evaluated. At each decision, the probability of making that decision can be computed and then the probability of proposing $\\tau$ is the product of these intermediate probabilities (as in Eq. 5; is this correct?). \n\nNow, suppose I give you another tree topology, $\\tau'$, that was not simulated by ARTree. Can ARTree compute the likelihood of $\\tau'$, i.e. $Q(\\tau')$? To me it is not clear how this is achieved by reading the paper. This seems to be like an important downstream task which SBNs can handle. If it is possible, I recommend emphasizing this feature, and how a practitioner would achieve it. If it is not possible, I think that this missing feature should be discussed. \n\nNote that I do not regard this to be a crucial feature, ARTree is still an impressive algorithm. However, it would add transparency and promote future work on ARTree.\n\n**\"Domain expertise\"**: It is repeated multiple times as a key flaw of SBNs is the required domain expertise. I do not see, and it is not explained, how ARTree diminishes this requirement? In fact, how do SBNs require more domain expertise? Running MrBAYES or UFBoot to get the presampled tree topologies can be done without understanding these algorithms in depth, as the softwares are very neatly provided. Especially, does not implementing and understanding ARTree require the same domain expertise by the practitioner as implementing and understanding SBNs?\n\nThis should be carefully clarified in the text. Alternatively, I think, the \"domain expertise\" argument should be removed as it does not add information as the paper is written at the moment. There are plenty of compelling arguments for ARTree over SBNs as is.\n\n**Experiments**: In the caption of Table 2: \"The KL results are averaged over 10 independent trainings.\" I was expecting to see standard deviations of these 10 KL numbers. How come they are not included? \n\nAlso in the same caption: \"For ELBO, LB-10, and ML,the results are averaged over 100, 100, and 1000 independent runs\". Do independent runs imply \"independent trainings\" here too? If not, why not use the 10 trained models used for the KL values to get uncertanties wrt the learned model parameters? Finally note that it says \"100, 100 and 1000\", which I figure is a typo.\n\n**Stds of ML results**: I am aware that previous works reward low-variance estimators of the marginal log-likelihood, as is done here in Table 2. My guess is that this is an appropriate way to compare models that provide estimates of lower-bounds of the ML with models that can harshly over-estimate the ML, like the stepping-stone algorithm used with MrBAYES. However, here the comparison is between two VBPI models, using either ARTree or SBNs, which both use estimates of lower bounds. Could the authors please expand on why then the standard deviation is an appropriate measure of the success of the models? \n\nFor instance, I can come up with models with estimators that have zero standard deviation by sacrificing bias. Does this make these models \"better\"?\n\nIf my concerns above are discussed and clarified, I may be willing to raise my score.\n\nQUESTIONS:\n**Critique of SBNs**: Could the authors please point to where in Zhang and Matsen, (2022) I may find the discussion regarding \"the limited parent-child subsplit patterns in the observed samples\" (line 37 in the submission)? This is an important argument for ARTree which I have not seen investigated before. To me it could make sense to include this discussion in the Appendix of this submission, as to make the paper more stand alone.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces ARTree, a deep generative autoregressive model for phylogenetic tree reconstruction. The authors define an autoregressive sequential process for generating tree topologies, $\\tau$, and prove that there is a one-to-one mapping between the resulting topologies and the decision sequence, $D$, instantiated by the process. Utilizing this fact, the authors define a distribution over $D$, letting each decision at time n be drawn from a Categorical distribution given previous decisions 1,..., n-1, and parameterize these Categoricals by calculating and passing learnable topological features to Graph Neural networks (GNN) with a recurrent unit and unit to incorporate time embeddings. \nARTree is then used as variational distribution for tree topologies $Q(\\tau)$, along with a GNN based parameterization of the branch length distribution, within Variational Bayesian phylogenetic inference (VBPI). The representative power of $Q(\\tau)$ is evaluated by comparing it with Subsplit Bayesian networks (SBNs) to a \"ground truth\" posterior distribution (based on the posteriors of long-running MrBayes experiments) and within the context of VBPI on benchmark datasets.\n\nSTRENGTHS:\nThe vastness of the tree topology space in phylogenetic inference is a well-known obstacle in both classical and tumor phylogenetics. The VI approach to Bayesian phylogenetics is an ongoing area of research and is need of more sophisticated tree topology variational distributions and experiment design to evaluate these varaiational distributions. The tree topology density experiment in 4.1 not only shows strong representative power of ARTree, but the experiment design itself is a contribution to the research field. Furthermore, ARTree is a large improvement when compared to previous VI-methods in phylogeny that use unconfined $Q(\\tau)$ (however this is not highlighted in the paper, see Weaknesses section).\n\nWEAKNESSES:\nThe paper fails to mention other works within VI in phylogenetics, e.g., VaiPhy (https://arxiv.org/abs/2203.01121) and VCSMC (https://arxiv.org/abs/2106.00075); these methods do not confine the $Q(\\tau)$ support either and are relevant related work. \n\nThe paper fails to highlight the strong performance of ARTree in VBPI w.r.t. other VI methods with unconfined $Q(\\tau)$ - adding a row to Table 2 with results from, e.g., VCSMC and VaiPhy would greatly accentuate the contribution of the paper.\n\nARTree relies on several subroutines to be able to construct and parameterize the generative decision sequences. This makes the method complicated to grasp and implement, which, given the complex problem at hand, can be regarded as a strength in the authors perseverance rather than a weakness of the paper. However, the amount of steps involved, e.g., L steps of messing passing and calculating the topological node embeddings, naturally invokes questions regarding inference runtime and memory usage. The lack of runtime comparisons between VBPI with ARTree, other methods in VI and MrBayes is a weakness of the paper.\n\nQUESTIONS:\nOverall, a strong contribution to field of Bayesian phylogenetics and a well-written paper. \n\nThe addressing the following points could make me increase my score further:\n1. Experiment on runtime of ARTree in the context of VBPI\n2. Addressing the issue of minimal increase of raised in Weaknesses\n\nFailing to incorporate the related works mentioned in Weaknesses-section could make me decrease my score.\n\nMisprints:\n56 \"edges of current…\" -> \"edges of the current…\"\n\nLIMITATIONS:\nLimitations, except for potential runtime concerns (see Weaknesses-section), are properly addressed.", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nIn this paper, the authors propose a tractable distribution over tree spaces that can be fit for use in density estimation or variational inference.  The key idea is to build a tree by sequentially adding leaves one at a time by adding an additional branch to the tree.  Each step of this process has a reasonable state space, and it is easy to see that such a process generates distributions that span all of tree space.  The authors then parameterize the \"action space\" of this process using (recurrent) graph neural networks, which can be trained using either maximum likelihood in the density estimation case or VIMCO to optimize the ELBO in the VI case.  The authors apply their method to 8 standard phylogenetic benchmarking datasets finding comparable or superior performance to existing density estimation .or VI methods.\n\nSTRENGTHS:\n* The ideas presented in this paper are extremely simple and elegant.\n* The big picture of the approach is easy to describe and conceptually straightforward.\n* The performance of the method seems to be an advance over existing methods, even by metrics that favor existing methods (e.g., inclusive KL is kind to SBNs, as exclusive KL would be infinite for SBNs that do not have support on all of tree space).\n\nWEAKNESSES:\n* A more thorough description of the technical details of the parameterization of the model would be helpful. In particular, it would be useful to have a schematic representing all of the components and how they fit together. Equations 7-11 had a lot of subcomponents (e.g., $P$ and $R$ and $b_n$ and emb, etc... etc...) which were hard to keep track of and see how they all fit together.\n\n\nMinor:\n* I know that it is common in the field, but it is not obvious to me why one would want to take $K$ greater than $1$ in equation (4).  If $K$ is $1$, then (4) is exactly the usual ELBO.  Maximizing the $K=1$ ELBO corresponds to minimizing the KL between the variational and true posteriors, which seems desirable.  Taking $K$ larger than one certainly tightens the lower bound on the evidence, but that doesn't necessarily mean that it will result in a better variational approximation to the posterior.  In fact, as $K \\to \\infty$ equation (4) should become independent of $Q$, which seems undesirable.  See for example https://proceedings.mlr.press/v80/rainforth18b.html\n* Many of the references at the end have minor formatting issues (e.g., lacking capitalization: \"bayesian,  \"markov\", \"monte carlo\", \"Graphrnn\", etc...)\n\nQUESTIONS:\nHave the authors explored how necessary it is to condition the decisions on all previous decisions?  Does a Markov decision process perform substantially worse?  That is, does one need the recurrent GNN, or would a non-recurrent GNN be sufficient?\n\nLIMITATIONS:\nThe authors have adequately addressed the limitations of their study, and I do not foresee any potential negative societal impacts.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-19"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper presents a new way to construct the variational distribution of tree topologies based on autoregressive generation with GNNs, which is used in the problem of variational phylogenetic inference. The paper mainly compares the new construction of the variational distribution with subsplit Bayesian network approaches. The experiments show that the proposed approach outperforms SBN in terms of learning the ground-truth tree topologies.\n\nSTRENGTHS:\n1. The paper seems to address an interesting sub-problem of the task, which is how to flexibly generate tree topologies for the variational distribution. I'm not a domain expert in variational phylogenetic inference, but previous approaches usually follow the way of SBNs and this paper proposes a novel and better alternative to SBNs.\n\n2. The proposed approach of autoregressive generation with GNNs looks intuitive. The paper provides comprehensive experiments in the comparison with SBN, which support the claims of the proposed method.\n\n3. The paper is well-written and easy to follow.\n\nWEAKNESSES:\n1. As stated in the paper, the main drawback of SBN is that it could not span the entire tree topology space. It seems that there is no analysis on how/why the proposed method is better at doing this in addition to the empirical comparison in the experiments.\n\n2. The main contribution of the paper is in Section 3.1, which is a new parameterisation of $Q(r)$. Most of the techniques in Section 3.2 follows Zhang (2023). It might be better to shorten 3.2, as I think it might not the focus of the paper.\n\n3. Although Table 1 shows that ARTree has better number in terms of revealing the ground-truth trees, it seems that ARTree does not improve much on ELBO and ML in Table 2.\n\nQUESTIONS:\nCan the authors talk about the potential use of the proposed method in computational biology in addition to getting better EBLO or other metrics of modelling the data? such as interpretability.\n\nLIMITATIONS:\nN/a", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-21"}