Learnable Topological Features for Phylogenetic Inference via Graph
  Neural Networks
Structural information of phylogenetic tree topologies plays an important
role in phylogenetic inference. However, finding appropriate topological
structures for specific phylogenetic inference tasks often requires significant
design effort and domain expertise. In this paper, we propose a novel
structural representation method for phylogenetic inference based on learnable
topological features. By combining the raw node features that minimize the
Dirichlet energy with modern graph representation learning techniques, our
learnable topological features can provide efficient structural information of
phylogenetic trees that automatically adapts to different downstream tasks
without requiring domain expertise. We demonstrate the effectiveness and
efficiency of our method on a simulated data tree probability estimation task
and a benchmark of challenging real data variational Bayesian phylogenetic
inference problems.
Cheng Zhang
# Learnable Topological Features for Phylogenetic Inference via Graph Neural Networks

###### Abstract

Structural information of phylogenetic tree topologies plays an important role in phylogenetic inference. However, finding appropriate topological structures for specific phylogenetic inference tasks often requires significant design effort and domain expertise. In this paper, we propose a novel structural representation method for phylogenetic inference based on learnable topological features. By combining the raw node features that minimize the Dirichlet energy with modern graph representation learning techniques, our learnable topological features can provide efficient structural information of phylogenetic trees that automatically adapts to different downstream tasks without requiring domain expertise. We demonstrate the effectiveness and efficiency of our method on a simulated data tree probability estimation task and a benchmark of challenging real data variational Bayesian phylogenetic inference problems.

## 1 Introduction

Phylogenetics is an important discipline of computational biology where the goal is to identify the evolutionary history and relationships among individuals or groups of biological entities. In statistical approaches to phylogenetics, this has been formulated as an inference problem on hypotheses of shared history, i.e., _phylogenetic trees_, based on observed sequence data (e.g., DNA, RNA, or protein sequences) under a model of evolution. The phylogenetic tree defines a probabilistic graphical model, based on which the likelihood of the observed sequences can be efficiently computed (Felsenstein, 2003). Many statistical inference procedures therefore can be applied, including maximum likelihood and Bayesian approaches (Felsenstein, 1981; Yang & Rannala, 1997; Mau et al., 1999; Huelsenbeck et al., 2001).

Phylogenetic inference, however, has been challenging due to the composite parameter space of both continuous and discrete components (i.e., branch lengths and the tree topology) and the combinatorial explosion in the number of tree topologies with the number of sequences. Harnessing the topological information of trees hence becomes crucial in the development of efficient phylogenetic inference algorithms. For example, by assuming conditional independence of separated subtrees, Larget (2013) showed that conditional clade distributions (CCDs) can provide more reliable tree probability estimation that generalizes beyond observed samples. A similar approach was proposed to design more efficient proposals for tree movement when implementing Markov chain Monte Carlo (MCMC) algorithms for Bayesian phylogenetics (Hohna & Drummond, 2012). Utilizing more sophisticated local topological structures, CCDs were later generalized to subsplit Bayesian networks (SBNs) that provide more flexible distributions over tree topologies (Zhang & Matsen IV, 2018). Besides MCMC, variational Bayesian phylogenetics inference (VBPI) was recently proposed that leveraged SBNs and a structured amortization of branch lengths to deliver competitive posterior estimates in a more timely manner (Zhang & Matsen IV, 2019; Zhang, 2020; Zhang & Matsen IV, 2022). Azouri et al. (2021) used a machine learning approach to accelerate maximum likelihood tree-search algorithms by providing more informative topology moves. Topological features have also been found useful for comparison and interpretation of the reconstructed phylogenies (Matsen IV, 2007; Hayati et al., 2022). While these approaches prove effective in practice, they all rely on heuristic features (e.g., clades and subsplits) of phylogenetic trees that often require significant design effort and domain expertise, and may be insufficient for capturing complicated topological information.
Graph Neural Networks (GNNs) are an effective framework for learning representations of graph-structured data. To encode the structural information about graphs, GNNs follow a neighborhood aggregation procedure that computes the representation vector of a node by recursively aggregating and transforming representation vectors of its neighboring nodes. After the final iteration of aggregation, the representation of the entire graph can also be obtained by pooling all the node embeddings together via some permutation invariant operators (Ying et al., 2018). Many GNN variants have been proposed and have achieved superior performance on both node-level and graph-level representation learning tasks (Kipf and Welling, 2017; Hamilton et al., 2017; Li et al., 2016; Zhang et al., 2018; Ying et al., 2018). A natural idea, therefore, is to adapt GNNs to phylogenetic models for automatic topological feature learning. However, the lack of node features for phylogenetic trees makes it challenging as most GNN variants assume fully observed node features at initialization.

In this paper, we propose a novel structural representation method for phylogenetic inference that automatically learns efficient topological features based on GNNs. To obtain the initial node features for phylogenetic trees, we follow previous studies (Zhu and Ghahramani, 2002; Rossi et al., 2021) to minimize the Dirichlet energy, with one not encoding for the tip nodes. Unlike these previous studies, we present a fast linear time algorithm for Dirichlet energy minimization by taking advantage of the hierarchical structure of phylogenetic trees. Moreover, we prove that these features are sufficient for identifying the corresponding tree topology, i.e., there is no information loss in our raw feature representations of phylogenetic trees. These raw node features are then passed to GNNs for more sophisticated structure representation learning required by downstream tasks. Experiments on a synthetic data tree probability estimation problem and a benchmark of challenging real data variational Bayesian phylogenetic inference problems demonstrate the effectiveness and efficiency of our method.

## 2 Background

NotationA phylogenetic tree is denoted as \((\tau,\mathbf{q})\) where \(\tau\) is a bifurcating tree that represents the evolutionary relationship of the species and \(\mathbf{q}\) is a non-negative branch length vector that characterizes the amount of evolution along the edges of \(\tau\). The tip nodes of \(\tau\) correspond to the observed species and the internal nodes of \(\tau\) represent the unobserved characters (e.g., DNA bases) of the ancestral species. The transition probability \(P_{ij}(t)\) from character \(i\) to character \(j\) along an edge of length \(t\) is often defined by a continuous-time substitution model (e.g., Jukes and Cantor (1969)), whose stationary distribution is denoted as \(\eta\). Let \(E(\tau)\) be the set of edges of \(\tau\), \(r\) be the root node (or any internal node if the tree is unrooted and the substitution model is reversible). Let \(\mathbf{Y}=\{Y_{1},Y_{2},\ldots,Y_{M}\}\in\Omega^{N\times M}\) be the observed sequences (with characters in \(\Omega\)) of length \(M\) over \(N\) species.

Phylogenetic posteriorAssuming different sites \(Y_{i},i=1,\ldots,M\) are independent and identically distributed, the likelihood of observing \(\mathbf{Y}\) given the phylogenetic tree \((\tau,\mathbf{q})\) takes the form

\[p(\mathbf{Y}|\tau,\mathbf{q})=\prod_{i=1}^{M}p(Y_{i}|\tau,\mathbf{q})=\prod_{i=1}^{M}\sum _{a^{i}}\eta(a^{i}_{r})\prod_{(u,v)E(\tau)}P_{a^{i}_{u}a^{i}_{v}}(q_{uv}), \tag{1}\]

where \(a^{i}\) ranges over all extensions of \(Y_{i}\) to the internal nodes with \(a^{i}_{u}\) being the assigned character of node \(u\). The above phylogenetic likelihood function can be computed efficiently through the pruning algorithm (Felsenstein, 2003). Given a prior distribution \(p(\tau,\mathbf{q})\) of the tree topology and the branch lengths, Bayesian phylogenetics then amounts to properly estimating the phylogenetic posterior \(p(\tau,\mathbf{q}|\mathbf{Y})\propto p(\mathbf{Y}|\tau,\mathbf{q})p(\tau,\mathbf{q})\).

Variational Bayesian phylogenetic inferenceLet \(Q_{\mathbf{\phi}}(\tau)\) be an SBN-based distribution over the tree topologies and \(Q_{\mathbf{\psi}}(\mathbf{q}|\tau)\) be a non-negative distribution over the branch lengths. VBPI finds the best approximation to \(p(\tau,\mathbf{q}|\mathbf{Y})\) from the family of products of \(Q_{\mathbf{\phi}}(\tau)\) and \(Q_{\mathbf{\psi}}(\mathbf{q}|\tau)\) by maximizing the following multi-sample lower bound

\[L^{K}(\mathbf{\phi},\mathbf{\psi})=\mathbb{E}_{Q_{\mathbf{\phi},\mathbf{\psi}}(\tau^{1:K}, \mathbf{q}^{1:K})}\log\left(\frac{1}{K}\sum_{i=1}^{K}\frac{p(\mathbf{Y}|\tau^{i},\mathbf{ q}^{i})p(\tau^{i},\mathbf{q}^{i})}{Q_{\mathbf{\phi}}(\tau^{i})Q_{\mathbf{\psi}}(\mathbf{q}^{i}| \tau^{i})}\right)\leq\log p(\mathbf{Y}) \tag{2}\]

where \(Q_{\mathbf{\phi},\mathbf{\psi}}(\tau^{1:K},\mathbf{q}^{1:K})=\prod_{i=1}^{K}Q_{\mathbf{\phi}}( \tau^{i})Q_{\mathbf{\psi}}(\mathbf{q}^{i}|\tau^{i})\). To properly parameterize the variational distributions, a support of the conditional probability tables (CPTs) is often acquired from a sample
of tree topologies via fast heuristic bootstrap methods (Minh et al., 2013; Zhang & Matsen IV, 2019). The branch length approximation \(Q_{\mathbf{\psi}}(\mathbf{q}|\tau)\) is taken to be the diagonal Lognormal distribution

\[Q_{\mathbf{\psi}}(\mathbf{q}|\tau)=\prod\nolimits_{e\in E(\tau)}p^{\mathrm{Lognormal}} \left(q_{e}\mid\mu(e,\tau),\sigma(e,\tau)\right)\]

where \(\mu(e,\tau),\sigma(e,\tau)\) are amortized over the tree topology space via shared local structures (i.e., split and primary subsplit pairs (PSPs)), which are available from the support of CPTs. More details about structured amortization, VBPI and SBNs can be found in section 3.2.2 and Appendix A.

Graph neural networksLet \(G=(V,E)\) denote a graph with node feature vectors \(\mathbf{X}_{v}\) for node \(v\in V\), and \(\mathcal{N}(v)\) denote the set of nodes adjacent to \(v\). GNNs iteratively update the representation of a node by running a message passing (MP) scheme for \(T\) time steps. During each MP time step, the representation vectors of each node are updated based on the aggregated messages from its neighbors as follows

\[\mathbf{h}_{v}^{(t+1)}=\mathrm{UPDATE}^{(t)}\left(\mathbf{h}_{v}^{(t)},\mathbf{m}_{v}^{(t +1)}\right),\quad\mathbf{m}_{v}^{(t+1)}=\mathrm{AGG}^{(t)}\left(\left\{\mathbf{h}_{u} ^{(t)}:u\in\mathcal{N}(v)\right\}\right)\]

where \(\mathbf{h}_{v}^{(t)}\) is the feature vector of node \(v\) at time step \(t\), with initialization \(\mathbf{h}_{v}^{(0)}=\mathbf{X}_{v}\), \(\mathrm{UPDATE}^{(t)}\) is the update function, and \(\mathrm{AGG}^{(t)}\) is the aggregation function. A number of powerful GNNs with different implementations of the update and aggregation functions have been proposed (Kipf & Welling, 2017; Hamilton et al., 2017; Li et al., 2016; Velickovic et al., 2018; Xu et al., 2019; Wang et al., 2019). In additional to the local node-level features, GNNs can also provide features for the entire graph. To learn these global features, an additional \(\mathrm{READOUT}\) function is often introduced to aggregate node features from the final iteration

\[\mathbf{h}_{G}=\mathrm{READOUT}\left(\left\{\mathbf{h}_{v}^{(T)}:v\in V\right\} \right).\]

\(\mathrm{READOUT}\) can be any function that is permutation invariant to the node features.

## 3 Proposed Method

In this section, we propose a general approach that automatically learns topological features directly from phylogenetic trees. We first introduce a simple embedding method that provides raw features for the nodes of phylogenetic trees, together with an efficient linear time algorithm for obtaining these raw features and a discussion on some of their theoretical properties regarding tree topology representation. We then describe how these raw features can be adapted to learn efficient representations of certain structures of trees (e.g., edges) for downstream tasks.

### Interior Node Embedding

Learning tree structure features directly from tree topologies often requires raw node/edge features, as typically assumed in most GNN models. Unfortunately, this is not the case for phylogenetic models.

Figure 1: An overview of the proposed topological feature learning framework for phylogenetic inference. **Left**: A phylogenetic tree topology with one hot encoding for the tip nodes and missing features for the interior nodes. **Middle**: Interior node embedding via Dirichlet energy minimization. **Right**: Subsequently, the tree topology with embedded node features are fed into a GNN model for more sophisticated tree structure representation learning required by downstream tasks.
Although we can use one hot encoding for the tip nodes according to their corresponding species (taxa names only, not the sequences), the interior nodes still lack original features. The first step of tree structure representation learning for phylogenetic models, therefore, is to properly input those missing features for the interior nodes. Following previous studies (Zhu & Ghahramani, 2002; Rossi et al., 2021), we make a common assumption that the node features change smoothly across the tree topologies (i.e., the features of every node are similar to those of the neighbors). A widely used criterion of smoothness for functions defined on nodes of a graph is the _Dirichlet energy_. Given a tree topology \(\tau=(V,E)\) and a function \(f:V\mapsto\mathbb{R}^{d}\), the Dirichlet energy is defined as

\[\ell(f,\tau)=\sum_{(u,v)\in E}\|f(u)-f(v)\|^{2}.\]

Let \(V=V^{b}\cup V^{o}\), where \(V^{b}\) denotes the set of leaf nodes and \(V^{o}\) denotes the set of interior nodes. Let \(\mathbf{X}^{b}=\{\mathbf{x}_{v}|v\in V^{b}\}\) be the set of one hot embeddings for the leaf nodes. The interior node features \(\mathbf{X}^{o}=\{\mathbf{x}_{v}|v\in V^{o}\}\) then can be obtained by minimizing the Dirichlet energy

\[\widehat{\mathbf{X}^{o}}=\mathop{\arg\min}_{\mathbf{X}^{o}}\ell(\mathbf{X}^{o},\mathbf{X}^{b},\tau)=\mathop{\arg\min}_{\mathbf{X}^{o}}\sum_{(u,v)\in E}\|\mathbf{x}_{u}-\mathbf{x}_{v} \|^{2}.\]

#### 3.1.1 A Linear Time Two-pass Algorithm

Note that the above Dirichlet energy function is convex, its minimizer therefore can be obtained by solving the following optimality condition

\[\frac{\partial\ell(\mathbf{X}^{o},\mathbf{X}^{b},\tau)}{\partial\mathbf{X}^{o}}(\widehat{ \mathbf{X}^{o}})=\mathbf{0}. \tag{3}\]

It turns out that equation 3 has a close-form solution based on matrix inversion. However, as matrix inversion scales cubically in general, it is infeasible for graphs with many nodes. Fortunately, by leveraging the hierarchical structure of phylogenetic trees, we can design a more efficient linear time algorithm for the solution of equation 3 as follows. We first rewrite equation 3 as a system of linear equations

\[\sum\nolimits_{v\in\mathcal{N}(u)}(\widehat{\mathbf{x}}_{u}-\widehat{\mathbf{x}}_{v} )=\mathbf{0},\quad\forall u\in V^{o},\qquad\widehat{\mathbf{x}}_{v}=\mathbf{x}_{v},\quad \forall v\in V^{b}, \tag{4}\]

where \(\mathcal{N}(u)\) is the set of neighbors of node \(u\). Given a topological ordering induced by the tree1, we can obtain the solution within a two-pass sweep through the tree topology, similar to the Thomas algorithm for solving tridiagonal systems of linear equations (Thomas, 1949). In the first pass, we traverse the tree in a postorder fashion and express the node features as a linear function of those of their parents,

Footnote 1: This is trivial for rooted trees since they are directed. For unrooted trees, we can choose an interior node as the root node and use the topological ordering of the corresponding rooted trees.

\[\widehat{\mathbf{x}}_{u}=c_{u}\widehat{\mathbf{x}}_{\pi_{u}}+\mathbf{d}_{u}, \tag{5}\]

for all the nodes expect the root node, where \(\pi_{u}\) denotes the parent node of \(u\). More specifically, we first initialize \(c_{u}=0,\mathbf{d}_{u}=\mathbf{x}_{u}\) for all leaf nodes \(u\in V^{b}\). For all the interior nodes except the root node, we compute \(c_{u},\mathbf{d}_{u}\) recursively as follows (see a detailed derivation in Appendix B)

\[c_{u}=\frac{1}{|\mathcal{N}(u)|-\sum_{v\in\mathrm{ch}(u)}c_{v}},\quad\mathbf{d}_{u }=\frac{\sum_{v\in\mathrm{ch}(u)}\mathbf{d}_{v}}{|\mathcal{N}(u)|-\sum_{v\in \mathrm{ch}(u)}c_{v}}, \tag{6}\]

where \(\mathrm{ch}(u)\) denotes the set of child nodes of \(u\). In the second pass, we traverse the tree in a preorder fashion and compute the solution by back substitution. Concretely, at the root node \(r\), given equation 5 for all the child nodes from the first pass, we can compute the node feature directly from equation 4 as below

\[\widehat{\mathbf{x}}_{r}=\frac{\sum_{v\in\mathrm{ch}(r)}\mathbf{d}_{v}}{|\mathcal{N}(r )|-\sum_{v\in\mathrm{ch}(r)}c_{v}}. \tag{7}\]

For all the other interior nodes, the node features can be obtained via equation 5 by substituting the learned features for the parent nodes. We summarize our two-pass algorithm in Algorithm 1. Moreover, the algorithm is numerically stable due to the following lemma (proof in Appendix C).

**Lemma 1**.: _Let \(\lambda=\min_{u\in V^{o}\setminus\{r\}}|\mathcal{N}(u)|\). For all interior node \(u\in V^{o}\setminus\{r\}\), \(0\leq c_{u}\leq\frac{1}{\lambda-1}\)._

Besides bifurcating phylogenetic trees, the above two-pass algorithm can be easily adapted to interior node embedding for general tree-shaped graphs with given tip node features.
#### 3.1.2 Tree Topology Representation Power

In this section, we discuss some theoretical properties regarding the tree topology presentation power of the node features introduced above. We start with a useful lemma that elucidates an important behavior of the solution to the linear system 4, which is similar to the solutions to elliptic equations.

**Lemma 2** (Extremum Principle).: _Let \(\{\widehat{\mathbf{x}}_{u}\in\mathbb{R}^{d}|u\in V\}\) be a set of \(d\)-dimensional node features that satisfies equations 4. \(\forall 1\leq n\leq d\), let \(\widehat{\mathbf{X}}[n]=\{\widehat{\mathbf{x}}_{u}[n]|u\in V\}\) be the set of the \(n\)-th components of node features. Then, \(\forall 1\leq n\leq d\), we have: (i) the extremum values (i.e., maximum and minimum) of \(\widehat{\mathbf{X}}[n]\) can be achieved at some tip nodes; (ii) if the extremum values are achieved at some interior nodes, then \(\widehat{\mathbf{X}}[n]\) has only one member, or in other words, \(\widehat{\mathbf{x}}_{u}[n]\) is the same \(\forall u\in V\)._

**Theorem 1**.: _Let \(N\) be the number of tip nodes. Let \(\{\widehat{\mathbf{x}}_{u}\in\mathbb{R}^{N}|u\in V\}\) be the solution to the linear system 4 with one hot encoding for the tip nodes. Then, \(\forall u\in V^{o}\), we have_

\[(i)\ 0<\widehat{\mathbf{x}}_{u}[n]<1,\quad\forall 1\leq n\leq N,\quad\text{and} \quad(ii)\ \sum\nolimits_{n=1}^{N}\widehat{\mathbf{x}}_{u}[n]=1.\]

The complete proofs of Lemma 2 and Theorem 1 are provided in Appendix C. When the tip node features are linearly independent, a similar proposition holds when we consider the coefficients of the linear combination of the tip node features for the interior node features instead.

**Corollary 1**.: _Suppose that the tip node features are linearly independent, the interior node features obtained from the solution to the linear system 4 all lie in the interior of the convex hull of all tip node features._

The proof is provided in Appendix C. The following lemma reveals a key property of the nodes that are adjacent to the boundary of the tree topology in the embedded feature space.

**Lemma 3**.: _Let \(\{\widehat{\mathbf{x}}_{u}|u\in V\}\) be the solution to the linear system 4, with linearly independent tip node features. Let \(\{\widehat{\mathbf{x}}_{u}=\sum_{v\in V^{h}}a^{u}_{u}\mathbf{x}_{v}|u\in V^{o}\}\) be the convex combination representations of the interior node features. For any tip node \(v\in V^{b}\), we have_

\[u^{*}=\arg\max_{u\in V^{o}}a^{u}_{u}\quad\Leftrightarrow\quad u^{*}\in \mathcal{N}(v).\]

**Theorem 2** (Identifiability).: _Let \(\mathbf{X}^{o}=\{\widehat{\mathbf{x}}_{u}|u\in V^{o}\}\) and \(\mathbf{Z}^{o}=\{\widehat{\mathbf{x}}_{u}|u\in V^{o}\}\) be the sets of interior node features that minimizes the Dirichlet energy for phylogenetic tree topologies \(\tau_{x}\) and \(\tau_{z}\) respectively, given the same linearly independent tip node features. If \(\mathbf{X}^{o}=\mathbf{Z}^{o}\), then \(\tau_{x}=\tau_{z}\)._

The proofs of Lemma 3 and Theorem 2 are provided in Appendix C. By Theorem 2, we see that the proposed node embeddings are complete representations of phylogenetic tree topologies with no information loss.

### Structural Representation Learning via Graph Neural Networks

Using node embeddings introduced in section 3.1 as raw features, we now show how to learn more sophisticated representations of tree structures for different phylogenetic inference tasks via GNNs. Given a tree topology \(\tau\), let \(\{\mathbf{h}^{(0)}_{v}:v\in V\}\) be the raw features and \(\{\mathbf{h}^{(T)}_{v}:v\in V\}\) be the output features after the final iteration of GNNs. We feed these output features of GNNs into a multi-layer perceptron (MLP) to get a set of learnable features for each node

\[\mathbf{h}_{v}=\operatorname{MLP}^{(0)}\left(\mathbf{h}^{(T)}_{v}\right),\quad\forall \ v\in V,\]

before adapting to different downstream tasks, as demonstrated in the following examples.
#### 3.2.1 Energy Based Models for Tree Probability Estimation

Our first example is on graph-level representation learning of phylogenetic tree topologies. Let \(\mathcal{T}\) denote the entire tree topology space. Given learnable node features of tree topologies, one can use a permutation invariant function \(g\) to obtain graph-level features and hence create an energy function \(F_{\mathbf{\phi}}:\mathcal{T}\mapsto\mathbb{R}\) that assigns each tree topology a scalar value as follows

\[F_{\mathbf{\phi}}(\tau)=\mathrm{MLP}^{(1)}(\mathbf{h}_{G}),\quad\mathbf{h}_{G}=g\left(\{ \mathbf{h}_{v}:v\in V\}\right).\]

where \(g\circ\mathrm{MLP}^{(0)}\) can be viewed as a \(\mathrm{READOUT}\) function in section 2. This allows us to construct energy based models (EBMs) for tree probability estimation

\[q_{\mathbf{\phi}}(\tau)=\frac{\exp\left(-F_{\mathbf{\phi}}(\tau)\right)}{Z(\mathbf{\phi})},\quad Z(\mathbf{\phi})=\sum\nolimits_{\tau\in\mathcal{T}}\exp\left(-F_{\mathbf{\phi} }(\tau)\right).\]

As \(Z(\mathbf{\phi})\) is usually intractable, we can employ noise contrastive estimation (NCE) (Gutmann and Hyvarinen, 2010) to train these energy based models. Let \(p_{n}\) be some noise distribution that has tractable density and allows efficient sampling procedures. Let \(D_{\mathbf{\phi}}(\tau)=\log q_{\mathbf{\phi}}(\tau)-\log p_{n}(\tau).\) We can train \(D_{\mathbf{\phi}}\)2 to minimize the following objective function (NCE loss)

Footnote 2: Here \(Z(\mathbf{\phi})\) is taken as a free parameter and is included into \(\mathbf{\phi}\).

\[J(\mathbf{\phi})=-\left(\mathbb{E}_{\tau\sim p_{\mathrm{data}}(\tau)}\log\left(S \left(D_{\mathbf{\phi}}(\tau)\right)\right)+\mathbb{E}_{\tau\sim p_{n}(\tau)}\log \left(1-S\left(D_{\mathbf{\phi}}(\tau)\right)\right)\right),\]

where \(S(x)=\frac{1}{1+\exp(-x)}\) is the sigmoid function. It is easy to verify that the minimum is achieved at \(D_{\mathbf{\phi}^{*}}(\tau)=\log p_{\mathrm{data}}(\tau)-\log p_{n}(\tau).\) Therefore, \(q_{\mathbf{\phi}^{*}}(\tau)=p_{\mathrm{data}}(\tau)=p_{n}(\tau)\exp\left(D_{\mathbf{ \phi}^{*}}(\tau)\right)\).

#### 3.2.2 Branch Length Parameterization for Vbpi

The branch length parameterization in VBPI so far has relied on hand-engineered features (i.e., splits and PSPs) for the edges on tree topologies. Let \(\mathbb{S}_{\mathrm{r}}\) denote the set of splits and \(\mathbb{S}_{\mathrm{psp}}\) denote the set of PSPs. The simple split-based parameterization assigns parameters \(\mathbf{\psi}^{\mu},\mathbf{\psi}^{\sigma}\) for splits in \(\mathbb{S}_{\mathrm{r}}\). The mean and standard deviation for each edge \(e\) on \(\tau\) are then given by the associated parameters of the corresponding split \(e/\tau\) as follows

\[\mu(e,\tau)=\psi_{e/\tau}^{\mu},\quad\sigma(e,\tau)=\psi_{e/\tau}^{\sigma}. \tag{8}\]

The more flexible PSP parameterization assigns additional parameters for PSPs in \(\mathbb{S}_{\mathrm{psp}}\) and adds the associated parameters of the corresponding PSPs \(e/\tau\) to equation 8 to refine the mean and standard deviation parameterization

\[\mu(e,\tau)=\psi_{e/\tau}^{\mu}+\sum\nolimits_{s\in e/\tau}\psi_{s}^{\mu},\ \ \sigma(e,\tau)=\psi_{e/\tau}^{\sigma}+\sum\nolimits_{s\in e/\tau}\psi_{s}^{\sigma}. \tag{9}\]

Although these heuristic features prove effective, they often require substantial design effort, a sample of tree topologies for feature collection, and can not adapt themselves during training which makes it difficult for amortized inference over different tree topologies. Based on the learnable node features, we can design a more flexible branch length parameterization that is capable of distilling more effective structural information of tree topologies for variational approximations. For each edge \(e=(u,v)\) on \(\tau\), similarly as in section 3.2.1, one can use a permutation invariant function \(f\) to obtain edge-level features and transform them into the mean and standard deviation parameters as follows

\[\mu(e,\tau)=\mathrm{MLP}^{\mu}\left(\mathbf{h}_{e}\right),\quad\sigma(e,\tau)= \mathrm{MLP}^{\sigma}\left(\mathbf{h}_{e}\right),\quad\mathbf{h}_{e}=f\left(\{\mathbf{h}_ {u},\mathbf{h}_{v}\}\right). \tag{10}\]

Compared to heuristic feature based parameterizations in 8 and 9, learnable topological feature based parameterizations in 10 allow much richer design for the branch length distributions across different tree topologies and do not require pre-sampled tree topologies for feature collection.

## 4 Experiments

In this section, we test the effectiveness and efficiency of learnable topological features for phylogenetic inference on the two aforementioned benchmark tasks: tree probability estimation via energy based models and branch length parameterization for VBPI. Following Zhang and Matsen (2019), in VBPI we used the simplest SBN for the tree topology variational distribution, and the CPT supports
were estimated from ultrafast maximum likelihood phylogenetic bootstrap trees using UFBoot (Minh et al., 2013). The code is available at [https://github.com/zcrabbit/vbpi-gnn](https://github.com/zcrabbit/vbpi-gnn).

Experimental setup.We evaluate five commonly used GNN variants with the following convolution operators: graph convolution networks (GCN), graph isomorphism operator (GIN), GraphSAGE operator (SAGE), gated graph convolution operator (GGNN) and edge convolution operator (EDGE). See more details about these convolution operators in Appendix F. In addition to the above GNN variants, we also considered a simpler model that skips all GNN iterations (i.e., \(T=0\)) and referred to it as MLP in the sequel. All GNN variants have 2 GNN layers (including the input layer), and all involved MLPs have 2 layers. We used summation as our permutation invariant aggregation function for graph-level features and maximization for edge-level features. All models were implemented in Pytorch (Paszke et al., 2019) with the Adam optimizer (Kingma & Ba, 2015).We designed our experiments with the goals of (i) verifying the effectiveness of GNN-based EBMs for tree topology estimation and (ii) verifying the improvement of GNN-based branch length parameterization for VBPI over the baseline approaches (i.e., split and PSP based parameterizations) and investigating how helpful the learnable topological features are for reducing the amortization gaps.

### Simulated Data Tree Probability Estimation

We first investigated the representative power of learnable topological features for approximating distributions on phylogenetic trees using energy based models (EBMs), and conducted experiments on a simulated data set. We used the space of unrooted phylogenetic trees with 8 leaves, which contains 10395 unique trees in total. Similarly as in Zhang & Matsen IV (2019), we generated a target distribution \(p_{0}(\tau)\) by drawing a sample from the symmetric Dirichlet distribution \(\mathrm{Dir}(\beta 1)\) of order 10395 with a pre-selected arbitrary order of trees. The concentration parameter \(\beta\) is used to control the diffuseness of the target distribution and was set to 0.008 to provide enough information for inference while allowing for adequate diffusion in the target. As mentioned earlier in section 3.2.1, we used noise contrastive estimation (NCE) to train our EBMs where we set the noise distribution \(p_{n}(\tau)\) to be the uniform distribution. Results were collected after 200,000 parameter updates. Note that the minimum NCE loss in this case is

\[J^{*}=-2\mathrm{JSD}\left(p_{0}(\tau)\|p_{n}(\tau)\right)+2\log 2,\]

where \(\mathrm{JSD}(\cdot\|\cdot)\) is the Jensen-Shannon divergence.

Figure 2 shows the empirical performance of different methods. From the left plot, we see that the NCE losses converge rapidly and the gaps between NCE losses for the GNN variants and the best NCE loss \(J^{*}\) (dashed red line) are close to zero, demonstrating the representative power of learnable topological features on phylogenetic tree probability estimations. The evolution of KL divergences (middle plot) is consistent with the NCE losses. Compared to MLP, all GNN variants perform better, indicating that the extra flexibility provided by GNN iterations is crucial for tree probability estimation that would benefit from more informative graph-level features. Although the

Figure 2: Comparison of learnable topological feature based EBMs for probability mass estimation of unrooted phylogenetic trees with 8 leaves using NCE. **Left:** NCE loss. **Middle:** KL divergence. **Right:** EBM approximations vs ground truth probabilities. The NCE loss and KL divergence results were obtained from 10 independent runs and the error bars represent one standard deviation.
raw features from interior node embedding contain all information of phylogenetic tree topologies, we see that distilling effective structural information from them is still challenging. This makes GNN models that are by design more capable of learning geometric representations a favorable choice. The right plot compares the probability mass approximations provided by EBMs using MLP and GGNN (which performs the best among all GNN variants), to the ground truth \(p_{0}(\tau)\). We see that EBMs using GGNN consistently provide accurate approximations for trees across a wide range of probabilities. On the other hand, estimates provided by those using MLP are often of large bias, except for a few trees with high probabilities.

### Real Data Variational Bayesian Phylogenetic Inference

The second task we considered is VBPI, where we compared learnable topological feature based branch length parameterizations to heuristic feature based parameterizations (denoted as Split and PSP resepectively) proposed in the original VBPI approach (Zhang & Matsen IV, 2019). All methods were evaluated on 8 real datasets that are commonly used to benchmark Bayesian phylogenetic inference 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). These datasets, which we call DS1-8, consist of sequences from 27 to 64 eukaryote species with 378 to 2520 site observations. We concentrate on the most challenging part of the Bayesian phylogenetics: joint learning of the tree topologies and the branch lengths, and assume a uniform prior on the tree topology, an i.i.d. exponential prior (\(\mathrm{Exp}(10)\)) for the branch lengths and the simple Jukes & Cantor (1969) substitution model. We gathered the support of CPTs from 10 replicates of 1000 ultrafast maximum likelihood bootstrap trees (Minh et al., 2013). We set \(K=10\) for the multi-sample lower bound, with a schedule \(\lambda_{n}=\min(1,0.001+n/100000)\), going from 0.001 to 1 after 100000 iterations. The Monte Carlo gradient estimates for the tree topology parameters and branch length parameters were obtained via VIMCO (Mnih & Rezende, 2016) and the reparameterization trick (Kingma & Welling, 2014) respectively. Results were collected after 400,000 parameter updates.

Table 1 shows the estimates of the evidence lower bound (ELBO) and the marginal likelihood using different branch length parameterizations on the 8 benchmark datasets, including the results for the

\begin{table}
\begin{tabular}{l l l l l l l l l l} \hline \hline \multicolumn{1}{l}{Data set} & DS1 & DS2 & DS3 & DS4 & DS5 & DS6 & DS7 & DS8 \\ \# Taka & 27 & 29 & 36 & 41 & 50 & 50 & 59 & 64 \\ \# SITKs & 1949 & 2520 & 1812 & 1137 & 378 & 1133 & 1824 & 1008 \\ \hline \multirow{7}{*}{\begin{tabular}{} \end{tabular} } & Split & -7112.162.00 & -2639.950(0.89) & -33736.87(0.36) & -13332.82(0.75) & -4218.860.20 & -4729.40(0.47) & -37335.440(1.12) & -4861.562(0.00) \\  & PSP & -7111.231.241 & -26340.43 & -0.3473.714(0.45) & -13332.02(0.65) & -3123.840.19 & -4729.20(0.44) & -37335.180(1.03) & -4855.400(0.32) \\  & MLP & -7110.040.10 & -26348.583 & -0.3107.326(0.05) & -13331.900(0.00) & **-3213.012.13** & -0.7288.901(0.17) & -37334.90(1.12) & -4655.330(1.15) \\  & GCN & -7110.320(10) & -2638.585(0.09) & -33736.20(0.05) & -13331.940(0.00) & -4218.000(10.10) & -6728.70(10.16) & -37334.920(12.1) & -4655.170(10.15) \\  & GIN & -7110.270(0.09) & -2638.586(0.09) & -33736.20(0.05) & -13331.920(0.02) & -6717.530(10.17) & -4728.590(10.17) & -4565.00(0.15) \\  & SAGE & -7110.200(10) & **-2638.584.500(0.09)** & -33736.200(0.06) & -13331.600(0.10) & -2417.210(0.11) & -4728.510(10.15) & -37334.910(11.1) & -4655.050(10.14) \\  & GGNN & **-7110.260.100** & **-2638.584.000** & **-33736.200(0.06)** & **-13331.790(0.00)** & **-3217.830.100** & **-4728.560(10.16) & -37334.870(12.1) & -4655.051(01.15) \\  & EDGE & **-7110.260.100** & **-2638.584.000** & **-33736.200(0.08)** & -13331.800(0.10) & **-3217.500.100** & **-3273.570(10.16)** & **-37334.840(10.14)** & -4655.051(01.14) \\ \hline \multirow{7}{*}{
\begin{tabular}{} \end{tabular} } & Split & -710.470(0.27) & -2637.087(0.08) & -33735.125(1.23) & -13330.100(1.34) & -2418.501(0.48) & -6734.080(0.473) & -37332.180(0.48) & -4631.509(0.94) \\  & PSP & -710.108(0.10) & -26367.740(0.08) & -33735.12(10.12) & -13339.920(0.22) & -24314.600(0.44) & -4741.040(0.47) & -37332.000(0.53) & -6366.050(0.51) \\  & MLP & -7108.40(10.24) & -26367.740(0.08) & -33735.120(10.10) & -13332.990(0.22) & -2314.600(0.45) & -6724.470(0.47) & -37332.000(0.34) & -4560.720(0.53) \\  & GCN & -7108.400(10.24) & -26367.730(0.08) & -33735.120(110.10) & -13332.990(0.18) & -2314.600(0.49) & -6724.1400(0.49) & -37332.000(0.34) & -4650.680(0.54) \\  & GIN & -7108.400(10.17) & **-26387.700(0.08)** & **-33735.120(0.100)** & **-13332.990(0.18)** & -2314.600(0.45) & -6724.300(0.42) & -37332.000(0.30) & -3650.608(0.49) \\  & GCN & -7108.400(10.19) & -26367.730(0.100) & -33753.120(0.09) & -13329.90(0.19) & -2414.600(0.36) & -6724.300(0.42) & -37332.000(0.30) & -4650.608(0.45) \\  & EDGE & **-7108.410.140** & **-2630.750(0.07)** & **-3375.120(0.09)** & -13329.900(0.19) & -2314.640(0.38) & **-6724.300(0.49)** & **-37332.600(0.24)** & **-4680.508(0.45)** \\  & SS & -7108.420(10.18) & -26367.570(0.48) & -33775.440(0.50) & -13330.600(0.54) & **-5214.510(0.28)** & -6724.070(0.86) & -37332.762(0.42) & -4649.581(11.75) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Evidence Lower bound (ELBO) and marginal likelihood (ML) estimates of different methods across 8 benchmark datasets for Bayesian phylogenetic inference. The marginal likelihood estimates of all variational methods are obtained via importance sampling using 1000 samples, and the results (in units of nats) are averaged over 100 independent runs with standard deviation in brackets. Results for stepping-stone (SS) are from Zhang & Matsen IV (2019)(using 10 independent MrBayes (Ronquist et al., 2012) runs, each with 4 chains for 10,000,000 iterations
stepping-stone (SS) method (Xie et al., 2011), which is one of the state-of-the-art sampling based methods for marginal likelihood estimation. For each data set, a better approximation would lead to a smaller variance of the marginal likelihood estimates. We see that solely using the raw features, MLP-based parameterization already outperformed the Split and PSP baselines by providing tighter lower bounds. With more expressive representations of local structures enabled by GNN iterations, GNN-based parameterization further improved upon MLP-based methods, indicating the importance of harnessing local topological information for flexible branch length distributions. Moreover, when used as importance distributions for marginal likelihood estimation via importance sampling, MLP and GNN variants provide more steady estimates (less variance) than Split and PSP respectively. All variational approaches compare favorably to SS and require much fewer samples. The left plot in Figure 3 shows the evidence lower bounds as a function of the number of parameter updates on DS1. Although neural networks based parameterization adds to the complexity of training in VI, we see that by the time Split and PSP converge, MLP and EDGE3 achieve comparable (if not better) lower bounds and quickly surpass these baselines as the number of iteration increases.

Footnote 3: We use EDGE as an example here for branch length parameterization since it can learn edge features (see Appendix F). All GNN variants (except the simple GCN) performed similarly in this example (see Table 1).

As diagonal Lognormal branch length distributions were used for all parameterization methods, how these variational distributions were amortized over tree topologies under different parameterizations therefore is crucial for the overall approximation performance. To better understand this effect of amortized inference, we further investigated the amortization gaps4 of different methods on individual trees in the 95% credible set of DS1 as in Zhang (2020). The middle and right plots in Figure 3 show the amortization gaps of different parameterization methods on each tree topology \(\tau\). We see the amortization gaps of MLP and EDGE are considerably smaller than those of Split and PSP respectively, showing the efficiency of learnable topological features for amortized branch length distributions. Again, incorporating more local topological information is beneficial, as evidenced by the significant improvement of EDGE over MLP. More results about the amortization gaps can be found in Table 2 in the appendix.

Footnote 4: The amortization gap on a tree topology \(\tau\) is defined as \(L(Q^{*}|\tau)-L(Q_{\phi}|\tau)\), where \(L(Q_{\psi}|\tau)\) is the ELBO of the approximating distribution \(Q_{\psi}(q|\tau)\) and \(L(Q^{*}|\tau)\) is the maximum lower bound that can be achieved with the same variational family. See more details in Zhang (2020); Cremer et al. (2018).

## 5 Conclusion

We presented a novel approach for phylogenetic inference based on learnable topological features. By combining the raw node features that minimize the Dirichlet energy with modern GNN variants, our learnable topological features can provide efficient structural information without requiring domain expertise. In experiments, we demonstrated the effectiveness of our approach for tree probability estimation on simulated data and showed that our method consistently outperforms the baseline approaches for VBPI on a benchmark of real data sets. Future work would investigate more sophisticated GNNs for phylogenetic trees, and applications to other phylogenetic inference tasks where efficiently leveraging structural information of tree topologies is of great importance.

Figure 3: Performance on DS1. **Left:** Lower bounds. **Middle & Right:** Amortization gaps on trees in the \(95\%\) credible sets.
#### Acknowledgments

This work was supported by National Natural Science Foundation of China (grant no. 12201014), as well as National Institutes of Health grant AI162611. The research of the author was support in part by the Key Laboratory of Mathematics and Its Applications (LMAM) and the Key Laboratory of Mathematical Economics and Quantitative Finance (LMEQF) of Peking University. The author is grateful for the computational resources provided by the High-performance Computing Platform of Peking University. The author appreciates the anonymous ICLR reviewers for their constructive feedback.
