L EARNABLE T OPOLOGICAL F EATURES FOR P HYLOGENETIC I NFERENCE VIA G RAPH N EURAL N ETWORKS
A BSTRACT
Structural information of phylogenetic tree topologies plays an important role in phylogenetic inference.
However, ﬁnding appropriate topological structures for speciﬁc phylogenetic inference tasks often requires signiﬁcant 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 efﬁcient structural information of phylogenetic trees that automatically adapts to different downstream tasks without requiring domain expertise.
We demonstrate the effectiveness and efﬁciency of our method on a simulated data tree probability estimation task and a benchmark of challenging real data variational
Bayesian phylogenetic inference problems
I NTRODUCTION
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 deﬁnes a probabilistic graphical model, based on which the likelihood of the observed sequences can be efﬁciently 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 efﬁcient 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 efﬁcient proposals for tree movement when implementing Markov chain Monte Carlo (MCMC) algorithms for Bayesian phylogenetics (Höhna & Drummond, 2012).
Utilizing more sophisticated local topological structures, CCDs were later generalized to subsplit Bayesian networks (SBNs) that provide more ﬂexible 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 signiﬁcant design effort and domain expertise, and may be insufﬁcient for capturing complicated topological information.
Graph Neural Networks (GNNs) are an effective framework for learning representations of graphstructured 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 ﬁnal 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 & 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 efﬁcient topological features based on GNNs.
To obtain the initial node features for phylogenetic trees, we follow previous studies (Zhu & Ghahramani, 2002; Rossi et al., 2021) to minimize the Dirichlet energy, with one hot 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 sufﬁcient 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 efﬁciency of our method.
B ACKGROUND
Notation A phylogenetic tree is denoted as ( τ, q ) where τ is a bifurcating tree that represents the evolutionary relationship of the species and q is a non-negative branch length vector that characterizes the amount of evolution along the edges of τ .
The tip nodes of τ correspond to the observed species and the internal nodes of τ 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 deﬁned by a continuous-time substitution model (e.g., Jukes & Cantor (1969)), whose stationary distribution is denoted as η .
Let E ( τ ) be the set of edges of τ , r be the root node (or any internal node if the tree is unrooted and the substitution model is reversible).
Let Y =
{ Y 1 , Y 2 , . . .
, Y M } ∈ Ω N × M be the observed sequences (with characters in Ω ) of length M over N species.
Phylogenetic posterior
Assuming different sites
Y i , i = 1 , . . .
, M are independent and identically distributed, the likelihood of observing Y given the phylogenetic tree ( τ, q ) takes the form
[Equation]
where a i ranges over all extensions of Y i to the internal nodes with a iu being the assigned character of node u .
The above phylogenetic likelihood function can be computed efﬁciently through the pruning algorithm (Felsenstein, 2003).
Given a prior distribution p ( τ, q ) of the tree topology and the branch lengths, Bayesian phylogenetics then amounts to properly estimating the phylogenetic posterior p ( τ, q | Y )
∝ p ( Y | τ, q ) p ( τ, q ) .
Variational Bayesian phylogenetic inference Let Q φ ( τ ) be an SBN-based distribution over the tree topologies and Q ψ ( q | τ ) be a non-negative distribution over the branch lengths.
VBPI ﬁnds the best approximation to p ( τ, q | Y ) from the family of products of Q φ ( τ ) and Q ψ ( q | τ ) by maximizing the following multi-sample lower bound
[Equation]
where Q φ , ψ ( τ 1: K , q 1: K ) = (cid:81)
Ki
=
1 Q φ ( τ i ) Q ψ ( q i | τ i ) .
To properly parameterize the variational distributions, a support of the conditional probability tables (CPTs) is often acquired from a sample
[Figure]
tree The branch length approximation Q ψ ( q | τ ) is taken to be the diagonal Lognormal distribution
[Equation]
where µ ( e, τ ) , σ ( e, τ ) 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 networks Let G =
( V, E ) denote a graph with node feature vectors X v for node v ∈ V , and 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
[Equation]
where h ( t ) v is the feature vector of node v at time step t , with initialization h (0) v = X v , UPDATE ( t ) is the update function, and 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; Veliˇckovi´c 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 READOUT function is often introduced to aggregate node features from the ﬁnal iteration
[Equation]
READOUT can be any function that is permutation invariant to the node features.
P ROPOSED M ETHOD
In this section, we propose a general approach that automatically learns topological features directly from phylogenetic trees.
We ﬁrst introduce a simple embedding method that provides raw features for the nodes of phylogenetic trees, together with an efﬁcient 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 efﬁcient representations of certain structures of trees (e.g., edges) for downstream tasks.
3. I NTERIOR N ODE E MBEDDING
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.
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 ﬁrst 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 deﬁned on nodes of a graph is the Dirichlet energy .
Given a tree topology τ = ( V, E ) and a function f : V (cid:55)→ R d , the Dirichlet energy is deﬁned as
[Equation]
Let V =
V b ∪ V o , where V b denotes the set of leaf nodes and V o denotes the set of interior nodes.
Let X b = { x v | v ∈ V b } be the set of one hot embeddings for the leaf nodes.
The interior node features X o =
{ x v | v ∈ V o } then can be obtained by minimizing the Dirichlet energy
[Equation]
3.1. A L INEAR T IME T WO - PASS A LGORITHM
Note that the above Dirichlet energy function is convex, its minimizer therefore can be obtained by solving the following optimality condition
[Equation]
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 efﬁcient linear time algorithm for the solution of equation 3 as follows.
We ﬁrst rewrite equation 3 as a system of linear equations
[Equation]
where N ( u ) is the set of neighbors of node u .
Given a topological ordering induced by the tree 1 , 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 ﬁrst pass, we traverse the tree in a postorder fashion and express the node features as a linear function of those of their parents,
[Equation]
for all the nodes expect the root node, where π u denotes the parent node of u .
More speciﬁcally, we ﬁrst initialize c u = 0 , d u = x u for all leaf nodes u ∈ V b .
For all the interior nodes except the root node, we compute c u , d u recursively as follows (see a detailed derivation in Appendix B)
[Equation]
where 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 equationfor all the child nodes from the ﬁrst pass, we can compute the node feature directly from equationas below
[Equation]
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 λ = min u ∈ V o \{ r } |N ( u ) | .
For all interior node u ∈ V o \{ r } , 0 ≤ c u ≤ 1 λ − 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. T REE T OPOLOGY R EPRESENTATION P OWER
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 { (cid:98) x u ∈ R d | u ∈ V } be a set of d -dimensional node features that satisﬁes equations 4.
∀ 1 ≤ n ≤ d , let (cid:99) X
[ n ] = { (cid:98) x u [ n ] | u ∈ V } be the set of the n -th components of node features.
Then, ∀ 1 ≤ n ≤ d , we have: (i) the extremum values (i.e., maximum and minimum) of (cid:99) X [ n ] can be achieved at some tip nodes; (ii) if the extremum values are achieved at some interior nodes, then (cid:99) X
[ n ] has only one member, or in other words, (cid:98) x u [ n ] is the same ∀ u ∈ V .
Theorem 1.
Let N be the number of tip nodes.
Let { (cid:98) x u ∈ R N | u ∈ V } be the solution to the linear system 4 with one hot encoding for the tip nodes.
Then, ∀ u ∈ V o , we have
( i )
< (cid:98) x u [ n ] < 1 , ∀ 1 ≤ n ≤ N, and ( ii ) (cid:88) N n =1 (cid:98) 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 coefﬁcients 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 { (cid:98) x u | u ∈ V } be the solution to the linear system 4, with linearly independent tip node features.
Let { (cid:98) x u
= (cid:80) v ∈ V b a vu x v | u ∈ V o } be the convex combination representations of the interior node features.
For any tip node v ∈ V b , we have
[Equation]
Theorem 2 (Identiﬁability) .
Let X o = { (cid:98) x u | u ∈ V o } and Z o = { (cid:98)
z u | u ∈ V o } be the sets of interior node features that minimizes the Dirichlet energy for phylogenetic tree topologies τ x and τ z respectively, given the same linearly independent tip node features.
If X o =
Z o , then τ x = τ 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.
3. S TRUCTURAL R EPRESENTATION L EARNING VIA G RAPH N EURAL
N ETWORKS
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 τ , let { h (0) v : v ∈ V } be the raw features and { h ( T ) v : v ∈ V } be the output features after the ﬁnal 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
[Equation]
before adapting to different downstream tasks, as demonstrated in the following examples.
3.2. E NERGY B ASED M ODELS FOR T REE P ROBABILITY E STIMATION
Our ﬁrst example is on graph-level representation learning of phylogenetic tree topologies.
Let 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 φ : T (cid:55)→ R that assigns each tree topology a scalar value as follows
[Equation]
where g ◦ MLP (0) can be viewed as a READOUT function in section 2.
This allows us to construct energy based models (EBMs) for tree probability estimation
[Equation]
As Z ( φ ) is usually intractable, we can employ noise contrastive estimation (NCE) (Gutmann &
Hyvärinen, 2010) to train these energy based models.
Let p n be some noise distribution that has tractable density and allows efﬁcient sampling procedures.
Let D φ ( τ ) = log q
φ ( τ ) − log p n ( τ ) .
We can train D φ 2 to minimize the following objective function (NCE loss)
[Equation]
where S ( x ) = 1 1+exp( − x ) is the sigmoid function.
It is easy to verify that the minimum is achieved at
D φ ∗ ( τ ) = log p data ( τ ) − log p n ( τ ) .
Therefore, q φ ∗ ( τ ) =
p data ( τ ) =
p n ( τ ) exp ( D φ ∗ ( τ )) .
3.2. B RANCH L ENGTH P ARAMETERIZATION 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 S r denote the set of splits and S psp denote the set of PSPs.
The simple split-based parameterization assigns parameters ψ µ , ψ σ for splits in S r .
The mean and standard deviation for each edge e on τ are then given by the associated parameters of the corresponding split e/τ as follows
[Equation]
The more ﬂexible PSP parameterization assigns additional parameters for PSPs in S psp and adds the associated parameters of the corresponding PSPs e//τ to equation 8 to reﬁne the mean and standard deviation parameterization
[Equation]
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 difﬁcult for amortized inference over different tree topologies.
Based on the learnable node features, we can design a more ﬂexible 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 τ , 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
[Equation]
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.
E XPERIMENTS
In this section, we test the effectiveness and efﬁciency 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 & Matsen IV (2019), in
VBPI we used the simplest SBN for the tree topology variational distribution, and the CPT supports
[Figure]
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 .
Experimental setup.
We evaluate ﬁve 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.
4. S IMULATED D ATA T REE P ROBABILITY E STIMATION
We ﬁrst 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 ( τ ) by drawing a sample from the symmetric Dirichlet distribution Dir( β 1 ) of order 10395 with a pre-selected arbitrary order of trees.
The concentration parameter β 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 ( τ ) to be the uniform distribution.
Results were collected after 200,000 parameter updates.
Note that the minimum NCE loss in this case is
[Equation]
where JSD( ·(cid:107)· ) 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 ﬂexibility provided by GNN iterations is crucial for tree probability estimation that would beneﬁt from more informative graph-level features.
Although the
[Table]
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 ( τ ) .
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.
4. R EAL D ATA V ARIATIONAL B AYESIAN P HYLOGENETIC I NFERENCE
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; Höhna & 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 ( Exp (10)) for the branch lengths and the simple Jukes & Cantor (1969) substitution model.
We gathered the support of CPTs from 10 replicates of 10000 ultrafast maximum likelihood bootstrap trees (Minh et al., 2013).
We set K = 10 for the multi-sample lower bound, with a schedule λ 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
[Figure]
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 ﬂexible 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 EDGE 3 achieve comparable (if not better) lower bounds and quickly surpass these baselines as the number of iteration increases.
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 gaps 4 of different methods on individual trees in the 95% credible set of DS1 as in Zhang (2020).
The middle and right plots in Figureshow the amortization gaps of different parameterization methods on each tree topology τ .
We see the amortization gaps of MLP and EDGE are considerably smaller than those of Split and PSP respectively, showing the efﬁciency of learnable topological features for amortized branch length distributions.
Again, incorporating more local topological information is beneﬁcial, as evidenced by the signiﬁcant improvement of EDGE over MLP.
More results about the amortization gaps can be found in Table 2 in the appendix.
C ONCLUSION
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 efﬁcient 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 efﬁciently leveraging structural information of tree topologies is of great importance.
B T HE T WO - PASS A LGORITHM FOR I NTERIOR N ODE E MBEDDING
Due to the hierarchical structure of phylogenetic trees
[Equation]
we derive a linear time two-pass algorithm for solving the linear equations 4.
In the ﬁrst pass
[Equation]
we traverse the tree in a postorder fashion and express the node features as a linear function of those of their parents as follows
[Equation]
for all the nodes except the root node, where π u is the parent node of u .
For all leaf nodes u ∈ V b , this is straightforward: c u = 0 , d u = x u .
For any interior node u that is not the root node, as a result of postorder traversal, c v , d v is available ∀ v ∈ ch( u ) when u is visited.
Therefore, we can rewrite the equation about u in 4 as
[Equation]
This implies
[Equation]
which gives the recursive updating formula in equation 6.
In the second pass
[Equation]
we traverse the tree in a preorder fashion.
We ﬁrst visit the root node r .
From the ﬁrst pass, c v , d v is available ∀ v ∈ ch( r ) .
Similarly, from equation 4, we have
[Equation]
x r + d v )
Therefore,
[Equation]
For any other interior node u , as a result of preorder traversal, (cid:98) x π u is available when u is visited.
We can compute its node feature (cid:98) x u via equation 5.
C P ROOFS FOR L EMMAS , C OROLLARIES , AND T HEOREMS
Proof for Lemma
Lemma 1.
Let λ = min u ∈ V o \{ r } |N ( u ) | .
For all interior node u ∈ V o \{ r } , 0 ≤ c u ≤ 1 λ − 1 .
Proof.
We prove by induction.
Since |N ( u ) | ≥ 2 , ∀ u ∈ V o , we have λ ≥ 2 .
Note that c u = 0 , ∀ u ∈ V b .
Suppose 0 ≤ c v ≤ 1 λ − 1 , ∀ v ∈ ch( u ) , it sufﬁces to show that 0 ≤ c u ≤ 1 λ − 1 as long as u is not the root node.
From the recursive updating formula, we have
[Equation]
As 2 ≤ λ ≤ |N ( u ) | ,
[Equation]
Therefore, 0 ≤ c u ≤ 1 λ − 1 .
Remark.
For bifurcating phylogenetic trees, we have λ = 3 ⇒ 0 ≤ c u ≤ 1 2 .
Proof for Lemma
Lemma 2 (Extremum Principle) .
Let { (cid:98) x u ∈ R d | u ∈ V } be a set of d -dimensional node features that satisﬁes equations 4.
∀ 1 ≤ n ≤ d , let (cid:99) X
[ n ] = { (cid:98) x u [ n ] | u ∈ V } be the set of the n -th components of node features.
Then, ∀ 1 ≤ n ≤ d , we have: (i) the extremum values (i.e., maximum and minimum) of (cid:99) X [ n ] can be achieved at some tip nodes; (ii) if the extremum values are achieved at some interior nodes, then (cid:99) X
[ n ] has only one member, or in other words, (cid:98) x u [ n ] is the same ∀ u ∈ V .
Proof.
From equations 4, we have
[Equation]
In other words, for any interior node, its node feature is the mean of those of its neighbors.
Therefore, ∀ 1 ≤ n ≤ d , the extremum values of (cid:99) X
[ n ] can be achieved at the boundary of the graph, i.e., the tip nodes.
On the other hand, if the extremum value of (cid:99) X [ n ] is achieved at some interior node u , then the extremum is also achieved at all the neighbors of u .
Since the tree topology is connected, this implies that (cid:98) x u [ n ] is a constant ∀ u ∈ V .
Proof for Theorem
Theorem 1.
Let N be the number of tip nodes.
Let { (cid:98) x u ∈ R N | u ∈ V } be the solution to the linear system 4 with one hot encoding for the tip nodes.
Then, ∀ u ∈ V o , we have ( i ) 0 < (cid:98) x u [ n ] < 1 , ∀ 1 ≤ n ≤ N, and ( ii ) (cid:88)
[Equation]
Proof.
As the tip node features are one hot vectors, ( i ) follows immediately from Lemma 2.
Let
S = { x ∈ R N | (cid:80) Nn
=1 x [ n ] = 1 } and P S be the projection onto S .
Since { (cid:98) x u ∈ R N | u ∈ V } solves the linear system 4, it minimizes the Dirichlet energy.
Now consider its projection { P S ( (cid:98) x u )
∈
R N | u ∈ V } .
Since one hot vectors are in S , P S ( (cid:98) x u ) = (cid:98) x u , ∀ u ∈ V b .
Note that P S is a projection operator, we have
[Equation]
Note that { (cid:98) x u ∈ R N | u ∈ V } minimizes the Dirichlet energy, the equality has to hold which implies that (cid:98) x u − (cid:98) x v is parallel to the hyperplane S , ∀ ( u, v ) ∈ E .
Therefore, (cid:98) x u ∈ S, ∀ u ∈ V
o as the tree topology is connected and the tip node features are already in S .
Proof for Corollary
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.
Proof.
Let A be a matrix whose columns corresponds to the tip node features.
From the linear system 4, we see that all interior node features lie in the column space of A and hence can be uniquely represented as
[Equation]
as the columns are linearly independent.
Moreover, { (cid:98) y u ∈
R
N | u ∈ V } satisﬁes the same linear system(cid:88)
[Equation]
where { y v ∈ R N | v ∈ V b } are one hot encoding vectors.
By Theorem 1, we have ∀ u ∈ V
o
( i )
< (cid:98) y u [ n ] < 1 , ∀ 1 ≤ n ≤ N, and ( ii )
N (cid:88) n =1 (cid:98) y u [ n ] = 1 . which concludes the proof.
Proof for Lemma
Lemma 3.
Let { (cid:98) x u | u ∈ V } be the solution to the linear system 4, with linearly independent tip node features.
Let { (cid:98) x u
= (cid:80) v ∈ V b a vu x v | u ∈ V o } be the convex combination representations of the interior node features.
For any tip node v ∈ V b , we have
[Equation]
Proof.
For any tip node v ∈ V b , let ¯ v ∈ N ( v ) be the adjacent interior node of v .
It sufﬁces to show that a v ¯ v > a vu , ∀ u ∈ V o \{ ¯ v } .
As { x v | v ∈ V b } are linearly independent, from equations 4 we have
[Equation]
Now consider a new tree topology τ (cid:48) that has v removed from τ .
Note that equation 11 still holds for all interior nodes except ¯ v , the maximum value of { a vu | u ∈ V \{ v }} therefore can be achieved at either the boundary of τ (cid:48) or ¯ v .
As a vu = 0 , ∀ u ∈ V b \{ v } and a v ¯ v > 0 (by Corollary 1), the maximum value hence is achieved at a v ¯ v , i.e., a v ¯ v ≥ a vu , ∀ u ∈ V
o .
Now suppose there exists an interior node u (cid:54) =
¯ v such that a vu = a v ¯ v .
Similarly to Lemma 2, this implies a vu = a v ¯ v , ∀ u ∈ V \{ v } , which is an obvious contradiction.
Proof for Theorem
Theorem 2 (Identiﬁability) .
Let X o = { (cid:98) x u | u ∈ V o } and Z o = { (cid:98)
z u | u ∈ V o } be the sets of interior node features that minimizes the Dirichlet energy for phylogenetic tree topologies τ x and τ z respectively, given the same linearly independent tip node features.
If X o =
Z o , then τ x = τ z .
Proof.
Consider the case for unrooted trees ﬁrst.
We prove by induction on the number of tip nodes
N of the tree topology.
For N = 3 , the tree topology is trivial.
If N > 3 , then for each tip node,
Lemma 3 identiﬁes its adjacent node by convex combination coefﬁcient maximization.
As the number of interior nodes is less than the number of tip nodes, there must be two tip nodes that connect to the same interior node.
Merging the two tip nodes into their shared neighbor reduces the problem to size
N − 1 , with the shared neighbor being the new tip node.
It is easy to check that for this new set of node features, the tip node features are also linearly independent.
Therefore, the remaining part of the tree topology can be recovered by induction hypothesis.
The proof for rooted trees is similar.
D M ORE R ESULTS ON A MORTIZATION G APS
[Table]
E R ELATED W ORKS
Many deep learning approaches have also been proposed for phylogenetic models recently.
SolisLemus et al.
(2022) use symmetry-preserving neural networks to explicitly encode permutation invariance of phylogenetic trees.
However, their method requires enumeration of all permutations of tip nodes that would leave tree topologies unchanged and hence is challenging to extend to datasets with many sequences.
Jiang et al.
(2022) learn hyperbolic embeddings of gene sequences for phylogenetic tree placement and updates.
While these embeddings maybe useful for sequentially updating species trees with only a handful of genes, they are not suitable for learning representations of tree topologies and their local structures with a given taxa set.
Fioravanti et al.
(2018) use
CNNs to incorporate phylogenetic information for the classiﬁcation of metagenomics data where the phylogenetic tree is assumed to be known.
Their method, therefore, is also not suitable for phylogenetic inference where the tree topology is unknown and needs to be inferred from the data.
F D ETAILS ON G RAPH C ONVOLUTIONAL O PERATORS
The followings are the update and aggregate functions for the graph convolutional operators used in our experiments.
• Graph convolutional networks (GCN)
[Equation]
where d u is the degree of node u . • Graph isomorphism networks (GIN)
[Equation]
where (cid:15) ( t ) can be either a learnable parameter or a ﬁxed scalar.
• GraphSAGE operator (SAGE)
[Equation]
• Gated graph convolutional operator (GGNN)
[Equation]
where GRU ( t ) is a Gated recurrent unit, a gating mechanism in recurrent neural networks.
• Edge convolutional operator (EDGE)
[Equation]
where || means concatenation and e ( t +1)
u
→ v is the edge feature from node u to node v .