SlowGuess/NeurIPS_2023_OCR huggingface dataset:
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  $\mathbb{S}_{\mathrm{r}}$  denotes the root subsplits and  $\mathbb{S}_{\mathrm{ch|pa}}$  denotes the child-parent subsplit pairs in the support. The parameters of SBNs are then  $p = \{p_{s_1}; s_1 \in \mathbb{S}_{\mathrm{r}}\} \cup \{p_{s|t}; s|t \in \mathbb{S}_{\mathrm{ch|pa}}\}$  where
$$
p _ {s _ {1}} = p \left(S _ {1} = s _ {1}\right), \quad p _ {s | t} = p \left(S _ {i} = s \mid S _ {\pi_ {i}} = t\right), \forall i > 1. \tag {16}
$$
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  $\{\tau_k\}_{k=1}^M$ , we first initialize the parameter estimates as  $\hat{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:  $\forall 1 \leq k \leq M$ , compute  $q_k^{(i)}(s_1) = \frac{p(\tau_k, s_1 | \hat{p}^{\mathrm{EM}, (i)})}{\sum_{s_1 \sim \tau_k} p(\tau_k, s_1 | \hat{p}^{\mathrm{EM}, (i)})}$  where  $s_1 \sim \tau_k$  means the subsplit  $s_1$  can be achieved by placing a "virtual root" on an edge of  $\tau$ .
- M-step: update the parameter estimates by
$$
\begin{array}{l} \hat {p} _ {s _ {1}} ^ {\mathrm {E M}, (i + 1)} = \frac {\bar {m} _ {s _ {1}} ^ {(i)} + \alpha \tilde {m} _ {s _ {1}}}{K + \alpha \sum_ {s _ {1} \in \mathbb {S} _ {\mathrm {r}}} \tilde {m} _ {s _ {1}}}, \quad \bar {m} _ {s _ {1}} ^ {(i)} = \sum_ {k = 1} ^ {M} \sum_ {e \in E (\tau_ {k})} q _ {k} ^ {(i)} (s _ {1}) \mathbb {I} (s _ {1, k} ^ {e} = s _ {1}) \\ \tilde {p} _ {s | t} ^ {\mathrm {E M}, (i + 1)} = \frac {\bar {m} _ {s , t} ^ {(i)} + \alpha \tilde {m} _ {s , t}}{\sum_ {s} \left(\bar {m} _ {s , t} ^ {(i)} + \alpha \tilde {m} _ {s , t}\right)}, \quad \bar {m} _ {s, t} ^ {(i)} = \sum_ {k = 1} ^ {M} \sum_ {e \in E (\tau_ {k})} q _ {k} ^ {(i)} \left(s _ {1, k} ^ {e}\right) \sum_ {j > 1} \mathbb {I} \left(s _ {j, k} ^ {e} = s, s _ {\pi_ {j}, k} ^ {e} = t\right) \\ \end{array}
$$
where  $\mathbb{I}$  is the indicator function,  $s_{j,k}^{e}$  is the node value of  $S_{j}$  in  $\tau_k^e$ ,  $\tilde{m}_s$  and  $\tilde{m}_{s,t}$  are equivalent counts and  $\alpha$  is the regularization coefficient that encourages generalization.
When  $\alpha > 0$ , this algorithm is called SBN-EM- $\alpha$ .
B Details of variational Bayesian phylogenetic inference
With two variational families  $Q_{\phi}(\tau)$  and  $Q_{\psi}(q|\tau)$  over the space of tree topologies and branch lengths, the variational Bayesian phylogenetic inference (VBPI) approach forces  $Q_{\phi,\psi}(\tau,q) = Q_{\phi}(\tau)Q_{\psi}(q|\tau)$  to approximate the posterior  $p(\tau,q|\mathbf{Y})$  by maximizing the following multi-sample lower bound
$$
L ^ {K} (\phi , \psi) = \mathbb {E} _ {\left\{(\tau^ {i}, \boldsymbol {q} ^ {i}) \right\} _ {i = 1} ^ {K} \stackrel {\text {i . i . d .}} {\sim} Q _ {\phi , \psi}} \log \left(\frac {1}{K} \sum_ {i = 1} ^ {K} \frac {p (\boldsymbol {Y} | \tau^ {i} , \boldsymbol {q} ^ {i}) p (\tau^ {i} , \boldsymbol {q} ^ {i})}{Q _ {\phi} (\tau^ {i}) Q _ {\psi} (\boldsymbol {q} ^ {i} | \tau^ {i})}\right). \tag {17}
$$
Gradients of the objective (17) w.r.t.  $\phi$  and  $\psi$  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_{\phi}(\tau)$  and  $Q_{\psi}(\pmb{q}|\tau)$ .
Choice of  $Q_{\phi}(\tau)$  Before the proposed ARTree framework in this article, SBNs is the common choice of  $Q_{\phi}(\tau)$ . 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
$$
p _ {s _ {1}} = \frac {\exp \left(\phi_ {s _ {1}}\right)}{\sum_ {s ^ {\prime} \in \mathbb {S} _ {\mathrm {r}}} \exp \left(\phi_ {s ^ {\prime}}\right)}, s _ {1} \in \mathbb {S} _ {\mathrm {r}}; p _ {s | t} = \frac {\exp \left(\phi_ {s | t}\right)}{\sum_ {s ^ {\prime} : s ^ {\prime} | t \in \mathbb {S} _ {\mathrm {c h} | \mathrm {p a}}} \exp \left(\phi_ {s ^ {\prime} | t}\right)}, s | t \in \mathbb {S} _ {\mathrm {c h} | \mathrm {p a}}. \tag {18}
$$
The parameters  $\phi = \{\phi_{s_1}; s_1 \in \mathbb{S}_{\mathrm{r}}\} \cup \{\phi_{s|t}; s|t \in \mathbb{S}_{\mathrm{ch|pa}}\}$  are called latent parameters of SBNs.
Choice of  $Q_{\psi}(q|\tau)$  The distribution  $Q_{\psi}(q|\tau)$  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  $\tau$  as follows. For each edge  $e = (u,v)$  in  $\tau$ , 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
$$
\mu (e, \tau) = \mathrm {M L P} ^ {\mu} \left(h _ {e}\right), \quad \sigma (e, \tau) = \mathrm {M L P} ^ {\sigma} \left(h _ {e}\right)
$$
where  $\mathrm{MLP}^{\mu}$  and  $\mathrm{MLP}^{\sigma}$  are two multi-layer perceptrons (MLPs).
Algorithm 2: Tree topology decomposition process
Input: a tree topology  $\tau$  with all of the  $N$  leaf nodes.  
Output: a decision sequence  $D$ .  
 $\tau_{N} = (V_{N},E_{N})\gets$  the tree topology  $\tau$   
for  $n = N - 1,\dots ,3$  do Determine the unique neighbor  $w$  of the leaf node  $x_{n + 1}$  Determine the two neighbors  $u$  and  $v$  (except  $x_{n + 1}$ ) of  $w$ $V_{n}\leftarrow V_{n + 1}\backslash \{w,x_{n + 1}\}$ $E_{n}\gets (E_{n + 1}\cup \{(u,v)\})\backslash \{(w,x_{n + 1}),(w,u),(w,v)\}$ $\tau_{n}\gets (V_{n},E_{n})$ $e_n\gets (u,v)$    
end  
 $D\gets (e_3,\ldots ,e_{N - 1})$
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  $\tau_{N}$  of rank  $N$ . Intuitively, given an ordinal tree topology  $\tau_{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  $\tau_{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  $\tau_{3}$  of rank 3 is reached. We summarize the sketch of tree topology decomposition process in Algorithm 2.
Given a tree topology  $\tau$  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,\dots ,e_{N - 1})$  following Algorithm 2 and then calculate the probability as the product of conditionals
$$
Q (\tau) = Q (D) = \prod_ {n = 3} ^ {N - 1} Q \left(e _ {n} \mid e _ {<   n}\right).
$$
D The proofs of Theorem 1 and Lemma 1
be the generating process used in ARTree. Then  $g$  is a bijection between  $\mathcal{D}$  and  $\mathcal{T}$ .
Proof of Theorem 1 It is obvious that  $g$  is a well-defined map from  $\mathcal{D}$  to  $\mathcal{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)} \in \mathcal{D}$  and  $D^{(1)} \neq D^{(2)}$ . Let  $k$  be the first position where  $D^{(1)}$  and  $D^{(2)}$  begin to differ, i.e.  $e_i^{(1)} = e_i^{(2)}$  for all  $i < k$  and  $e_k^{(1)} \neq e_k^{(2)}$ . If  $g(D^{(1)}) = g(D^{(2)}) =: \tau$ , one can take the subtree topology of rank  $k$  and  $k + 1$  of  $\tau, \tau_k$  and  $\tau_{k + 1}$ . Noting that  $e_k$  refers to the edge in  $\tau$  where to add the new node  $x_{n + 1}$ , the equation  $e_k^{(1)} \neq e_k^{(2)}$  implies they will induce different  $\tau_{k + 1}$ 's. This contradicts the uniqueness of  $\tau_{k + 1}$ . Therefore, we conclude that  $g$  is injective.
Next, we prove that  $g$  is surjective. For a tree topology  $\tau$  with rank  $N$ , we denote its subtree topology of rank  $k$  by  $\tau_k$ , where  $k = 3, \ldots, n$ . In fact, for each  $k$ , the tree topology  $\tau_{k+1}$  corresponds to adding the leaf node  $x_{k+1}$  to an edge in  $\tau_k$ , and we denote this edge by  $e_k$ . It is easy to verify that the constructed  $D = (e_3, \ldots, e_{N-1})$  is a preimage of  $\tau$ .
Table 3: The ELBO estimates on DS1 obtained by different combinations of tree topology model  $Q(\tau)$  and branch length model  $Q(\mathbf{q}|\tau)$ . The results are averaged over 100 independent runs with standard deviation in the brackets.
Model combination	ELBO
SBN + branch length model trained along with SBN	-7110.24(0.03)
SBN + branch length model trained along with ARTree	-7110.26(0.03)
ARTree + branch length model trained along with ARTree	-7110.09(0.04)
Table 4: Runtime comparison in the variational inference setting on  $\mathrm{DS1}^2$ . SBN* and ARTree* refer to the early stopping of SBN and ARTree that surpass the CSMC baseline in terms of marginal likelihood estimation (-7290.36), respectively. The experiments are run on a single core of MacBook Pro 2019.
VCSMC	VaiPhy	φ-CSMC	SBN	ARTree	SBN*	ARTree*
Total training time (minutes)	248.3	45.1	N/A	659.3	3740.8	10.2	79.5
Evaluation time (one estimate of ML, minutes)	2.4	1.6	102.2	0.15	0.41	0.15	0.41
Proof of Lemma 1 Assume the tree topology  $\tau$  is stored as a binary tree data structure, where each node other than the root node also has a parent node pointer. Before decomposing  $\tau$ , we first build a dictionary of the  $(n,x_{n})$  mappings by a traversal across  $\tau$  that maps  $n$  to the leaf node  $x_{n}$  in  $\tau$ ,  $\forall n\leq 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,\dots ,3$ , resulting in a time complexity of  $O(N)$ . Therefore, the time complexity of the decomposition process induced by  $g^{-1}(\cdot)$  is  $O(N)$ .
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_{\psi}(q|\tau)$ , compared to SBNs. When combined with ARTree, the approximation accuracy of  $Q_{\psi}(q|\tau)$  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_{\psi}(q|\tau)$  trained along with ARTree for those  $\tau$  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.
$^{2}$ Training: We trained all models following the setting in their original papers: VCSMC was trained for 100 iterations with 2048 particles per iteration; VaiPhy was trained for 200 iterations with 128 particles per iteration;  $\phi$ -CSMC directly estimates ML based on VaiPhy, and therefore does not need extra training; both ARTree and SBN were trained for 400000 iterations with 10 particles per iteration.
Evaluation: We find that the evaluation strategies in their original papers are quite different, e.g. VaiPhy, SBN, and ARTree used importance sampling to estimate ML with different repetition times; VCSMC and  $\phi$ -CSMC instead estimated ML with sequential Monte Carlo (SMC), also with different repetition times. To be fair, we report the time for producing one estimate of ML from each of these models (VaiPhy, SBN, and ARTree used importance sampling with 1000 particles; VCSMC and  $\phi$ -CSMC used SMC with 2048 particles).
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(\pmb{q}|\tau)$  instead of the tree topology model  $Q(\tau)$ . 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(\pmb{q}|\tau)$  which is conditioned on tree topology  $\tau$ , 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.