SlowGuess/NeurIPS_2023_OCR huggingface dataset:
where probability vector  $q_{n} \in \mathbb{R}^{|E_{n}|}$  for parametrizing  $Q(\cdot |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 \in 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  $\tau_{n}$  of rank  $n$  to  $\tau_{n + 1}$  of rank  $n + 1$ . We can repeat this procedure until an ordinal tree topology  $\tau = \tau_N$  of rank  $N$  is reached. The probability of  $\tau$  then takes the form
$$
Q _ {\phi} (\tau) = Q _ {\phi} (D) = \prod_ {n = 3} ^ {N - 1} Q _ {\phi} \left(e _ {n} \mid e _ {<   n}\right), \tag {11}
$$
where  $D$  is the decision sequence and  $\phi$  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}(\tau)$  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
Figure 2: Performances of different methods for TDE on DS1. Left/Middle: Comparison of the ground truth and the estimated probabilities using SBN-EM and ARTree. A tree topology is marked as an outlier if it satisfies  $|\log(\text{estimated probability}) - \log(\text{ground truth})| > 2$ . Right: The KL divergence as a function of the sample size. The results are averaged over 10 replicates with one standard deviation as the error bar.
topologies ( $\phi$ -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.
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; Hohna & 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  $\mathrm{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
Table 1: KL divergences to the ground truth of different methods across 8 benchmark data sets. Sampled trees column shows the numbers of unique tree topologies in the training sets formed by MrBayes runs. The results are averaged over 10 replicates. The results of SBN-EM, SBN-EM- $\alpha$  are from Zhang & Matsen IV (2018).
Data set	#Taxa	#Sites	Sampled trees	KL divergence to ground truth
SBN-EM	SBN-EM-α	SBN-SGA	ARTree
DS1	27	1949	1228	0.0136	0.0130	0.0504	0.0045
DS2	29	2520	7	0.0199	0.0128	0.0118	0.0097
DS3	36	1812	43	0.1243	0.0882	0.0922	0.0548
DS4	41	1137	828	0.0763	0.0637	0.0739	0.0299
DS5	50	378	33752	0.8599	0.8218	0.8044	0.6266
DS6	50	1133	35407	0.3016	0.2786	0.2674	0.2360
DS7	59	1824	1125	0.0483	0.0399	0.0301	0.0191
DS8	64	1008	3067	0.1415	0.1236	0.1177	0.0741
sets are formed by collecting samples every 100 iterations and discarding the first  $25\%$ . Now, given a training data set  $\mathcal{M} = \{\tau_m\}_{m=1}^M$ , we train ARTree via maximum likelihood estimation using stochastic gradient ascent. In each iteration, the stochastic gradient is obtained as follows
$$
\nabla_ {\phi} L (\phi ; \mathcal {M}) = \frac {1}{B} \sum_ {b = 1} ^ {B} \nabla_ {\phi} \log Q _ {\phi} \left(\tau_ {m _ {b}}\right), \tag {12}
$$
where a minibatch  $\{\tau_{m_b}\}_{b = 1}^B$  is randomly sampled from  $\mathcal{M}$ . We compare ARTree to SBN baselines including SBN-EM, SBN-EM- $\alpha$ , and SBN-SGA. For SBN-EM and SBN-EM- $\alpha$ , 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
$$
p (\boldsymbol {Y}, \tau , \boldsymbol {q}; \beta_ {i}) = p (\boldsymbol {Y} | \tau , \boldsymbol {q}) ^ {\beta_ {i}} p (\tau , \boldsymbol {q}) \tag {13}
$$
Figure 3: Performances of ARTree and SBN as tree topology variational approximations for VBPI on DS1. Left: the evidence lower bound (ELBO) as a function of iterations. The numbers of particles used in the training objective are in the brackets. The ARTree* method refers to ARTree without time guidance, i.e.  $b_{n} = 0$  for all  $n$  in the readout step. Middle: variational approximations vs ground truth posterior probabilities of the tree topologies. Right: KL divergences across 50 random taxa orders. The KL divergence of SBNs is averaged over 10 independent trainings.
Table 2: KL divergences to the ground truth, evidence lower bound (ELBO), 10-sample lower bound (LB-10), and marginal likelihood (ML) estimates of different methods across 8 benchmark data sets. GT trees row shows the number of unique tree topologies in the ground truth. The marginal likelihood estimates are obtained via importance sampling using 1000 samples. The KL results are averaged over 10 independent trainings. For ELBO, LB-10, and ML, the results are averaged over 100, 100, and 1000 independent runs respectively with standard deviation in the brackets. The results of  $\phi$ -CSMC are from Koptagel et al. (2022). For ELBO and LB-10, a larger mean is better; for ML, a smaller standard deviation is better<sup>1</sup>.
Data set	DS1	DS2	DS3	DS4	DS5	DS6	DS7	DS8
# Taxa	27	29	36	41	50	50	59	64
# Sites	1949	2520	1812	1137	378	1133	1824	1008
GT trees	2784	42	351	11505	1516877	809765	11525	82162
ML	SBN	0.0707	0.0144	0.0554	0.0739	1.2472	0.3795	0.1531	0.3173
ARTree	0.0097	0.0004	0.0064	0.0219	0.8979	0.2216	0.0123	0.1231
SBN	-7110.24(0.03)	-26368.88(0.03)	-33736.22(0.02)	-13331.83(0.03)	-8217.80(0.04)	-6728.65(0.06)	-37334.85(0.04)	-8655.05(0.05)
ARTree	-7110.09(0.04)	-26368.78(0.07)	-33736.17(0.08)	-13331.82(0.05)	-8217.68(0.04)	-6728.65(0.06)	-37334.84(0.13)	-8655.03(0.05)
SBN	-7108.69(0.02)	-26367.87(0.02)	-33735.26(0.02)	-13330.29(0.02)	-8215.42(0.04)	-6725.33(0.04)	-37332.58(0.03)	-8651.78(0.04)
ARTree	-7108.68(0.02)	-26367.86(0.02)	-33735.25(0.02)	-13330.27(0.03)	-8215.34(0.03)	-6725.33(0.04)	-37332.54(0.03)	-8651.73(0.04)
φ-CSMC	-7290.36(7.23)	-30568.49(31.34)	-33798.06(6.62)	-13582.24(35.08)	-8367.51(8.87)	-7013.83(16.99)	N/A	-9209.18(18.03)
SBN	-7108.41(0.15)	-26367.71(0.08)	-33735.09(0.09)	-13329.94(0.20)	-8214.62(0.40)	-6724.37(0.43)	-37331.97(0.28)	-8650.64(0.50)
ARTree	-7108.41(0.19)	-26367.71(0.07)	-33735.09(0.09)	-13329.94(0.17)	-8214.59(0.34)	-6724.37(0.46)	-37331.95(0.27)	-8650.61(0.48)
where  $\beta_{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
<sup>1</sup> We use standard deviation as a criterion because in the context of importance sampling, the variance of an estimator reflects the approximation accuracy of the importance distribution to the target. Note that this is only reasonable if the importance distribution is proper (so that the marginal likelihood is unbiased), and hence does not apply to  $\phi$ -CSMC even though it also has much larger variance.
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).
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.
This work was supported by National Natural Science Foundation of China (grant no. 12201014 and grant no. 12292983), as well as National Institutes of Health grant AI162611. The research of Cheng Zhang was supported in part by National Engineering Laboratory for Big Data Analysis and Applications, the Key Laboratory of Mathematics and Its Applications (LMAM) and the Key Laboratory of Mathematical Economics and Quantitative Finance (LMEQF) of Peking University. The authors are grateful for the computational resources provided by the High-performance Computing Platform of Peking University. The authors appreciate the anonymous NeurIPS reviewers for their constructive feedback.
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-CoV-2 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 Merrienboer, 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.
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, Lukasz 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.
Z. 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.
A. 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.
N. Zhang and M. Blackwell. Molecular phylogeny of dogwood anthracnose fungus (Discula destructiva) and the Diaporthales. Mycologia, 93:355-365, 2001.
Figure 4: Subsplit Bayesian networks and a simple example for a leaf set of 4 taxa (denoted by  $A, B, C, D$  respectively). Left: General subsplit Bayesian networks. The solid full and complete binary tree network is  $B_{\mathcal{X}}^{*}$ . The dashed arrows represent the additional dependence for more expressiveness. Middle Left: Examples of (rooted) phylogenetic trees that are hypothesized to model the evolutionary history of the taxa. Middle Right: The corresponding subsplit assignments for the trees. For ease of illustration, subsplit  $(Y, Z)$  is represented as  $\frac{Y}{Z}$  in the graph. Right: The SBN for this example, which is  $B_{\mathcal{X}}^{*}$  in this case.
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  $\mathcal{X}$  be the set of  $N$  labeled leaf nodes. A non-empty set  $C$  of  $\mathcal{X}$  is referred to as a clade and the set of all clades of  $\mathcal{X}$ , denoted by  $\mathcal{C}(\mathcal{X})$ , is a totally ordered set with a partial order  $\succ$  (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\succ Z$ ,  $W\cap Z = \emptyset$ , and  $W\cup Z = C$ .
Definition 2 (Subsplit Bayesian Network). A subsplit Bayesian network (SBN)  $\mathcal{B}_{\mathcal{X}}$  on a leaf node set  $\mathcal{X}$  of size  $N$  is defined as a Bayesian network whose nodes take on subsplit or singleton clade values of  $\mathcal{X}$  and has the following properties: (a) The root node of  $\mathcal{B}_{\mathcal{X}}$  takes on subsplits of the entire labeled leaf node set  $\mathcal{X}$ ; (b)  $\mathcal{B}_{\mathcal{X}}$  contains a full and complete binary tree network  $B_{\mathcal{X}}^{*}$  as a subnetwork; (c) The depth of  $B_{\mathcal{X}}$  is  $N - 1$ , with the root counted as depth 1.
Due to the binary structure of  $B_{\mathcal{X}}^*$ , the nodes in SBNs can be indexed by denoting the root node with  $S_1$  and two children of  $S_i$  with  $S_{2i}$  and  $S_{2i+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 \geq 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  $\tau$  takes the following form
$$
p _ {\mathrm {s b n}} (T = \tau) = p \left(S _ {1} = s _ {1}\right) \prod_ {i > 1} p \left(S _ {i} = s _ {i} \mid S _ {\pi_ {i}} = s _ {\pi_ {i}}\right), \tag {14}
$$
where  $\pi_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:
$$
p _ {\mathrm {s b n}} \left(T ^ {\mathrm {u}} = \tau\right) = \sum_ {e \in E (\tau)} p _ {\mathrm {s b n}} \left(\tau^ {e}\right) \tag {15}
$$
where  $\tau^e$  is the resulting rooted tree topology when the rooting position is on edge  $e$ .