Title: GeoPhy: Differentiable Phylogenetic Inference via Geometric Gradients of Tree Topologies
Venue: NeurIPS 2023
Accepted: yes
Keywords: phylogenetic inference

variational inference

control variates

hyperbolic space
Abstract: Phylogenetic inference, grounded in molecular evolution models, is essential for understanding the evolutionary relationships in biological data. Accounting for the uncertainty of phylogenetic tree variables, which include tree topologies and evolutionary distances on branches, is crucial for accurately inferring species relationships from molecular data and tasks requiring variable marginalization. Variational Bayesian methods are key to developing scalable, practical models; however, it remains challenging to conduct phylogenetic inference without restricting the combinatorially vast number of possible tree topologies. In this work, we introduce a novel, fully differentiable formulation of phylogenetic inference that leverages a unique representation of topological distributions in continuous geometric spaces. Through practical considerations on design spaces and control variates for gradient estimations, our approach, GeoPhy, enables variational inference without limiting the topological candidates. In experiments using real benchmark datasets, GeoPhy significantly outperformed other approximate Bayesian methods that considered whole topologies.

GeoPhy: Differentiable Phylogenetic Inference via Geometric Gradients of Tree Topologies

Abstract

Phylogenetic inference, grounded in molecular evolution models, is essential for understanding the evolutionary relationships in biological data. Accounting for the uncertainty of phylogenetic tree variables, which include tree topologies and evolutionary distances on branches, is crucial for accurately inferring species relationships from molecular data and tasks requiring variable marginalization. Variational Bayesian methods are key to developing scalable, practical models; however, it remains challenging to conduct phylogenetic inference without restricting the combinatorially vast number of possible tree topologies. In this work, we introduce a novel, fully differentiable formulation of phylogenetic inference that leverages a unique representation of topological distributions in continuous geometric spaces. Through practical considerations on design spaces and control variates for gradient estimations, our approach, GeoPhy, enables variational inference without limiting the topological candidates. In experiments using real benchmark datasets, GeoPhy significantly outperformed other approximate Bayesian methods that considered whole topologies.

1 Introduction

Phylogenetic inference, the reconstruction of tree-structured evolutionary relationships between biological units, such as genes, cells, individuals, and species ranging from viruses to macro-organisms, is a fundamental problem in biology. As the phylogenetic relationships are often indirectly inferred from molecular observations, including DNA, RNA and protein sequences, Bayesian inference has been an essential tool to quantify the uncertainty of phylogeny. However, due to the complex nature of the phylogenetic tree object, which involve both a discrete topology and dependent continuous variables for branch lengths, the default approach for the phylogenetic inference has typically been an Markov-chain Monte Carlo (MCMC) method [29], enhanced with domain-specific techniques such as a mixed strategy for efficient exploration of tree topologies.

As an alternative approach to the conventional MCMCs, a variational Bayesian approach for phylogenetic inference was proposed in [42] (VBPI), which has subsequently been improved in the expressive powers of topology-dependent branch length distributions [39, 40]. Although these methods have presented accurate joint posterior distributions of topology and branch lengths for real datasets, they required reasonable preselection of candidate tree topologies to avoid a combinatorial explosion in the number of weight parameters beforehand. There have also been proposed variational approaches [24, 16] on top of the combinatorial sequential Monte Carlo method (CSMC[34]), which iteratively updated weighted topologies without the need for the preselection steps. However, the fidelity of the joint posterior distributions was still largely behind that of MCMC and VBPI as reported in [16].

Michiaki Hamada

In this work, we propose a simple yet effective scheme for parameterizing a binary tree topological distribution with a transformation of continuous distributions. We further formulate a novel differentiable variational Bayesian approach named GeoPhy 1 to approximate the posterior distribution of the tree topology and branch lengths without preselection of candidate topologies. By constructing practical topological models based on Euclidean and hyperbolic spaces, and employing variance reduction techniques for the stochastic gradient estimates of the variational lower bound, we have developed a robust algorithm that enables stable, gradient-based phylogenetic inference. In our experiments using real biological datasets, we demonstrate that GeoPhy significantly outperforms other approaches, all without topological restrictions, in terms of the fidelity of the marginal log-likelihood (MLL) estimates to gold-standard provided with long-run MCMCs.

Our contributions are summarized as follows:

We propose a representation of tree topological distributions based on continuous distributions, thereby we construct a fully differentiable method named GeoPhy for variational Bayesian phylogenetic inference without restricting the support of the topologies.

By considering variance reduction in stochastic gradient estimates and simple construction of topological distributions on Euclidean and hyperbolic spaces, GeoPhy offers unprecedently close marginal log-likelihood estimates to gold standard MCMC runs, outperforming comparable approaches without topological preselections.

2 Background

2.1 Phylogenetic models

Let τ be represent an unrooted binary tree topology with N leaf nodes (tips), and let B τ denote a set of evolutionary distances defined on each of the branches of τ . A phylogenetic tree ( τ, B τ ) represents an evolutionary relationship between N species, which is inferred from molecular data, such as DNA, RNA or protein sequences obtained for the species. Let Y = { Y ij ∈ Ω } 1 ≤ i ≤ N, 1 ≤ j ≤ M be a set of aligned sequences with length M from the species, where Y ij denote a character (base) of the i -th sequence at j -th site, and is contained in a set of possible bases Ω . For DNA sequences, Ω represents a set of 4-bit vectors, where each bit represents 'A', 'T', 'G', or 'C'. A likelihood model of the sequences P ( Y | τ, B τ ) is determined based on evolutionary assumptions. In this study, we follow a common practice for method evaluations [42, 40] as follows: Y is assumed to be generated from a Markov process along the branches of τ in a site-independent manner; The base mutations are assumed to follow the Jukes-Cantor model [12]. The log-likelihood ln P ( Y | τ, B τ ) can be calculated using Felsenstein's pruning algorithm [5], which is also known as the sum-product algorithm, and differentiable with respect to B τ .

2.2 Variational Bayesian phylogenetic inference

The variational inference problem for phylogenetic trees, which seeks to approximate the posterior probability P ( τ, B τ | Y ) , is formulated as follows:

where D KL , Q ( τ ) and Q ( B τ | τ ) denote the Kullback-Leibler divergence, a variational tree topology distribution, and a variational branch length distribution, respectively. The first variational Bayesian phylogenetic inference method (VBPI) was proposed by Zhang and Matsen IV [42], which has been successively improved for the expressiveness of Q ( B τ | τ ) [39, 40]. For the expression of variational topology mass function Q ( τ ) , they all rely on a subsplit Bayesian network (SBN) [41], which represents tree topology mass function Q ( τ ) as a product of conditional probabilities of splits (i.e., bipartition of the tip nodes) given their parent splits. However, SBN necessitates a preselection of the likely set of tree topologies, which hinders an end-to-end optimization strategy of the distribution over all the topologies and branch lengths.

2

τ

(

z

)

Bτ

∼

Q

(

Bτ

|

τ

)

z

∼

Q

(

z

1

, …,

z

4

)

1.6

1.6

1.5

0.4

0.4

z

1

z

2

3

z

4

z

1

z

2

3

z

4

Q

(

z

1

)

Q

z

3

Q

(

z

4

)

Q

(

z

)

2.3 Learnable topological features (LTFs)

Zhang [40] introduced a deterministic embedding method for tree topologies, termed learnable topological features (LTFs), which provide unique signatures for each tree topology. Specifically, given a set of nodes V and edges E of a tree topology τ , a node embedding function f emb : V → R d is determined to minimize the Dirichlet energy defined as follows:

When feature vectors for the tip nodes are determined, Zhang [40] showed that the optimal f emb for the interior nodes can be efficiently computed with two traversals of the tree topology. They also utilized f emb to parameterize the conditional distribution of branch length given tree topology Q ( B τ | τ ) by using graph neural networks with the input node features { f emb ( v ) | v ∈ V } .

2.4 Hyperbolic spaces and probability distributions

Hyperbolic spaces are known to be able to embed hierarchical data with less distortion in considerably fewer dimensions than Euclidean spaces [32]. This property has been exploited to develop various data analysis applications such as link predictions for relational data [26], as well as phylogenetic analyses [21, 19]. For representing the uncertainty of embedded data on hyperbolic spaces, the wrapped normal distribution proposed in [25] has appealing characteristics: it is easy to sample from and evaluate the density of the distribution.

The Lorentz model H d is a representation of the d -dimensional hyperbolic space, which is a submanifold of d +1 dimensional Euclidean space. Denote µ ∈ H d and Σ ∈ R d × d be a location and scale parameters, respectively, a random variable z ∈ H d that follows a wrapped normal distribution WN ( µ, Σ) is defined as follows:

where µ o , exp µ : T µ H d → H d , and PT µ o → µ : T µ o H d → T µ H d denote the origin, the exponential map, and the parallel transport defined on the Lorentz model. Also, the probability density WN ( µ, Σ) has a closed form and is differentiable with respect to the parameters µ, Σ . More details and properties are summarized in Appendix A.

(

)

z

z

3 Proposed methods

3.1 Geometric representations of tree topology ensembles

Considering the typically infeasible task of parameterizing the probability mass function of unrooted tree topologies T , which requires (2 N -5)!! -1 degrees of freedom, we propose an alternative approach. We suggest constructing the mass function Q ( τ ) through a transformation of a certain probability density Q ( z ) over a continuous domain Z , as follows:

where τ : Z → T denotes a deterministic link function that maps N coordinates to the corresponding tree topology (Fig. 1). Note that we have overloaded τ to represent both a variable and function for notational simplicity. An example of the representation space Z is a product of the Euclidean space R N × d or hyperbolic space H N × d , where d denotes the dimension of each tip's representation coordinate. For the link function, we can use τ ( z ) = T NJ ◦ D ( z ) , where D : Z → R N × N denotes a function that takes N coordinates and provides a distance matrix between those based on a geometric measure such as the Euclidean or hyperbolic distance. T NJ : R N × N →T denotes a map that takes this distance matrix and generates an unrooted binary tree topology of their phylogeny, determined using the Neighbor-Joining (NJ) algorithm [31]. While the NJ algorithm offers a rooted binary tree topology accompanied by estimated branch lengths, we only use the topology information and remove the root node from it to obtain the unrooted tree topology τ ∈ T .

3.2 Derivation of variational lower bound

Given a distribution of tip coordinates Q ( z ) and an induced tree topology distribution Q ( τ ) according to equation (4), the variational lower bound (1) is evaluated as follows:

Thanks to the deterministic mapping τ ( z ) , we can obtain an unbiased estimator of L [ Q ] by sampling from Q ( z ) without summing over the combinatorial many topologies T . However, even when the density Q ( z ) is computable, the evaluation of ln Q ( τ ) remains still infeasible according to the definition (4). We resolve this issue by introducing the second lower bound with respect to a conditional variational distribution R ( z | τ ) as follows:

This formulation is similar in structure to the variational auto-encoders [14], where Q ( τ | z ) and R ( z | τ ) are viewed as a probabilistic decoder and encoder of data τ , respectively. However, unlike typical auto-encoders, we design Q ( τ | z ) to be deterministic, and instead adjust Q ( z ) to approximate Q ( τ ) .

Proposition 1. Assuming that supp R ( z | τ ) ⊇ supp Q ( z | τ ) , the inequality ln P ( Y ) ≥ L [ Q ] ≥ L [ Q,R ] holds, where supp denotes the support of distribution. The first and the second equality holds when Q ( τ, B τ ) = P ( τ, B τ | Y ) and R ( z | τ ) = Q ( z | τ ) , respectively.

Proof. The first variational lower bound of the marginal log-likelihood is given as follows:

where the equality condition of the first inequality holds when Q ( τ, B τ ) = P ( τ, B τ | Y ) . Since we have defined Q ( τ ) in equation (4), we can further transform the lower bound as L [ Q ]

from which equation (5) immediately follows. Hence the inequality ln P ( Y ) ≥ L [ Q ] and its equality condition Q ( τ, B τ ) = P ( τ, B τ | Y ) have been proven.

Next, the entropy term -E Q ( z ) [ln Q ( τ ( z ))] = -E Q ( τ ) [ln Q ( τ )] in equation (5) can be transformed to derive further lower bound as follows:

where the last equality is derived by using the relation E Q ( τ ) Q ( z | τ ) [ · ] = E Q ( z ) [ ∑ τ I [ τ = τ ( z )] · ] and Q ( τ ) Q ( z | τ ) = Q ( z ) I [ τ = τ ( z )] . The equality condition of the first inequality holds when R ( z | τ ) = Q ( z | τ ) . Hence, the inequality L [ Q ] ≥ L [ Q,R ] and the equality condition is proven.

Similar to Burda et al. [1], we can also derive a tractable importance-weighted lower-bound of the model evidence (IW-ELBO), which is used for estimating the marginal log-likelihood (MLL), ln P ( Y ) , or an alternative lower-bound objective for maximization. The details and derivations are described in Appendix B.

3.3 Differentiable phylogenetic inference with GeoPhy

Algorithm 1 GeoPhy algorithm

1: θ, ϕ, ψ ← Initialize variational parameters

2: while not converged do

3: ϵ (1: K ) z , ϵ (1: K ) B ← Random samples from distributions p z , p B

4: z (1: K ) ← h θ ( ϵ (1: K ) z )

5: B (1: K ) τ ← h ϕ ( ϵ (1: K ) B , τ ( z (1: K ) ))

̂ ̂ ̂ 7: θ, ϕ, ψ ← Update parameters using an SGA algorithm, given gradients ̂ g θ , ̂ g ϕ , ̂ g ψ

6: g θ , g ϕ , g ψ ← Estimate the gradients ∇ θ L , ∇ ϕ L , ∇ ψ L

8: end while

9: return θ, ϕ, ψ

Here, we introduce parameterized distributions, Q θ ( z ) , Q ϕ ( B τ | τ ) , and R ψ ( z | τ ) , aiming to optimize the variational objective L [ Q θ,ϕ , R ψ ] using stochastic gradient ascent (SGA). The framework, which we term GeoPhy, is summarized in Algorithm 1. While aligning with black-box variational inference [27] and its extension to phylogenetic inference [42], our unique lower bound objective enables us to optimize the distribution over entire phylogenetic trees within continuous geometric spaces.

Gradients of lower bound We derive the gradient terms for stochastic optimization of L , omitting parameters θ, ϕ, ψ for notational simplicity. We assume that Q θ ( z ) and Q ϕ ( B τ | τ ( z )) are both reparameterizable; namely, we can sample from these distributions as follows:

where the terms ϵ z ∼ p z ( ϵ z ) and ϵ B ∼ p B ( ϵ B ) represent sampling from parameter-free distributions p z and p B , respectively, and the functions h θ and h ϕ are differentiable with respect to θ and ϕ , respectively. Although we cannot take the derivative of τ ( z ) with respect to z , the gradient of L with respect to θ is evaluated as follows:

where H denotes the differential entropy. The gradient of L with respect to ϕ and ψ are evaluated as follows:

where we assume a tractable density model for R ψ ( z | τ ) .

Gradient estimators and variance reduction From equation (10), an unbiased estimator of ∇ θ L , using K -sample Monte Carlo samples, can be derived as follows:

where we denote ϵ ( k ) z ∼ p z z ( k ) = h θ ( ϵ ( k ) z ) , ϵ ( k ) B ∼ p B , B ( k ) τ = h ϕ ( ϵ ( k ) B , τ ( z ( k ) )) , and

Note that we explicitly distinguish z and h θ ( ϵ z ) to indicate the target of differentiation with respect to θ .

In practice, it is known to be crucial to reducing the variance of the gradient estimators proportional to the score function ∇ θ ln Q θ to make optimization feasible. This issue is addressed by introducing a control variate c ( z ) , which has a zero expectation E z [ c ( z )] = 0 and works to reduce the variance of the term by subtracting it from the gradient estimator. For the case of K > 1 , it is known to be effective to use simple Leave-one-out (LOO) control variates [15, 28] as follows:

̸

where we defined f k ( z ( \ k ) , B ( \ k ) τ ) := 1 K -1 ∑ K k ′ =1 ,k ′ = k f ( z ( k ) , B ( k ) τ ) , which is a constant with respect to the k -the sample.

We can also adopt a gradient estimator called LAX [7], which uses a learnable surrogate function s χ ( z ) differentiable in z to form control variates as follows:

where we write the case with K = 1 for simplicity. To optimize the surrogate function s χ , we use the estimator ̂ g χ = ∇ χ 〈 ̂ g 2 θ 〉 2 , where ⟨·⟩ denotes the average of the vector elements. We summarize the procedure with LAX estimator in Algorithm 2.

We also consider other options for the gradient estimators, such as the combinations of the LOO and LAX estimators, the gradient of IW-ELBO, and its variance-reduced estimator called VIMCO [22] in our experiments. For the remaining gradients terms ∇ ϕ L and ∇ ψ L , we can simply employ reparameterized estimators as follows:

where we denote ϵ B ∼ p B and z ∼ Q θ . More details of the gradient estimators are summarized in Appendix B.

Variational distributions To investigate the basic effectiveness of GeoPhy algorithm, we employ simple constructions for the variational distributions Q θ ( z ) , Q ϕ ( B τ | τ ) , and R ψ ( z | τ ) . We use an independent distribution for each tip node coordinate, i.e. Q θ ( z ) = ∏ N i =1 Q θ i ( z i ) , where we use a d -dimensional normal or wrapped normal distribution for the coordinates of each tip node z i . For the conditional distribution of branch lengths given tree topology, Q ϕ ( B τ | τ ) , we use the diagonal lognormal distribution whose location and scale parameters are given as a function of the LTFs of the topology τ , which is comprised of GNNs as proposed in [40]. For the model of R ψ ( z | τ ) , we also employ an independent distribution: R ψ ( z | τ ) = ∏ N i =1 R ψ i ( z i | τ ) , where, we use the same type of distribution as Q θ i ( z i ) , independent of τ .

4 Related work

Differentiability for discrete optimization Discrete optimization problems often suffer from the lack of informative gradients of the objective functions. To address this issue, continuous relaxation for discrete optimization has been actively studied, such as a widely-used reparameterization trick with the Gumbel-softmax distribution [10, 20]. Beyond categorical variables, recent approaches have further advanced the continuous relaxation techniques to more complex discrete objects, including spanning trees [33]. However, it is still nontrivial to extend such techniques to the case of binary tree topologies. As outlined in equation (4), we have introduced a distribution over binary tree topologies T derived from continuous distributions Q ( z ) . This method facilitates a gradient-based optimization further aided by variance reduction techniques.

Gradient-based algorithms for tree optimization For the hierarchical clustering (HC), which reconstructs a tree relationship based on the distance measures between samples, gradient-based algorithms [23, 2, 3] have been proposed based on Dasgupta's cost function [4]. In particular, Chami et al. [2] proposed to decode tree topology from hyperbolic coordinates while the optimization is performed for a relaxed cost function, which is differentiable with respect to the coordinates. However, these approaches are not readily applicable to more general problems, including phylogenetic inference, as their formulations depend on the specific form of the cost functions.

Phylogenetic analysis in hyperbolic space The approach of embedding phylogenetic trees into hyperbolic spaces has been explored for visualization and an interpretation of novel samples with existing phylogeny [21, 11]. For the inference task, a maximum-likelihood approach was proposed in [35], which however assumed a simplified likelihood function of pairwise distances. A recent study [19] proposed an MCMC-based algorithm for sampling ( τ, B τ ) , which were linked from coordinates z using the NJ algorithm [31]. However, there remained the issue of an unevaluated Jacobian determinant, which posed a challenge in evaluating inference objectives. Given that we only use topology τ as described in equation (4), the variational lower bound for the inference can be unbiasedly evaluated through sampling, as shown in Proposition 1.

5 Experiments

We applied GeoPhy to perform phylogenetic inference on biological sequence datasets of 27 to 64 species compiled in Lakner et al. [17].

Models and training As a prior distribution of P ( τ ) and P ( B τ | τ ) , we assumed a uniform distribution over all topologies, and an exponential distribution Exp(10) independent for all branches, respectively, as commonly used in the literature [42, 16]. For the neural network used in the parameterization of Q ϕ ( B τ | τ ) , we employed edge convolutional operation (EDGE), which was well-performed architecture in [40]. For the stochastic gradient optimizations, we used the Adam optimizer [13] with a learning rate of 0.001. We trained for GeoPhy until one million Monte Carlo tree samples were consumed for the gradient estimation of our loss function. This number equals the number of likelihood evaluations (NLEs) and is used for a standardized comparison of experimental runs [34, 42]. The marginal log-likelihood (MLL) value is estimated with 1000 MC samples. More details of the experimental setup are found in Appendix C.

Initialization of coordinates We initialized the mean parameters of the tip coordinate distribution Q θ ( z ) with the multi-dimensional scaling (MDS) algorithm when Q θ was given as normal distributions. For Q θ comprised of wrapped normal distributions, we used the hyperbolic MDS algorithm (hMDS) proposed in [32] for the initialization. For a distance matrix used for MDS and hMDS, we used the Hamming distance between each pair of the input sequences Y as similar to [19]. For the scale parameters, we used 0 . 1 for all experiments. For R ψ ( z ) , we used the same mean parameters as Q θ ( z ) and 1.0 for the scale parameters.

5.1 Exploration of stable learning conditions

To achieve a stable and fast convergence in the posterior approximations, we compared several control variates (CVs) for the variance reduction of the gradients by using DS1 dataset [8]. We

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=2, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=2, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=2, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

demonstrate that the choice of control variates is crucial for optimization (Fig. 2). In particular, adaptive control variates (LAX) for individual Monte Carlo samples and the leave-one-out (LOO) CVs for multiple Monte Carlo samples (with K = 3 ) significantly accelerate the convergence of the marginal log-likelihood (MLL) estimates, yielding promising results comparable to VBPI-GNN. Although IW-ELBO was effective when Q ( z ) was comprised of diagonal Normal distributions, no advantages were observed for the case with full covariance matrices. Similar tendencies in the MLL estimates were observed for the other dimension ( d = 3 , 4) and wrapped normal distributions (Appendix D.1).

5.2 Comparison of different topological distributions

We investigated effective choices of tip coordinate distributions Q ( z ) , which yielded tree topologies, by comparing different combinations of distribution types (normal N or wrapped normal WN ), space dimensions ( d = 2 , 3 , 4 ), and covariance matrix types (diagonal or full), across selected CVs (LAX, LOO, and LOO+LAX). While overall performance was stable and comparable for the range of configurations, the most flexible ones: the normal and wrapped normal distributions with d = 4 and a full covariance matrix, indicated relatively high MLL estimates within their respective groups (Table 2 in Appendix D), implying the importance of flexibility in the variational distributions.

5.3 Performance evaluation across eight benchmark datasets

To demonstrate the inference performance of GeoPhy, we compared the marginal log-likelihood (MLL) estimates for the eight real datasets (DS1-8) [8, 6, 37, 9, 17, 43, 38, 30]. The gold-standard values were obtained using the stepping-stone (SS) algorithm [36] in MrBayes [29], wherein each evaluation involved ten independent runs, each with four chains of 10,000,000 iterations, as reported in [42]. In Table 1, we summarize the MLL estimates of GeoPhy and other approximate Bayesian inference approaches for the eight datasets. While VBPI-GNN [40] employs a preselected set of tree topologies as its support set before execution, it is known to provide reasonable MLL estimates near the reference values. The other approaches including CSMC [34], VCSMC [24], ϕ -CSMC [16] and GeoPhy (ours) tackles the more challenging general problem of model optimization by considering all candidate topologies without preselection. We compared the GeoPhy for two configurations of Q ( z ) : a wrapped normal distribution WN with a 4-dimensional full covariance matrix, and a 2-dimensional diagonal matrix, with three choices of CVs for optimization. Results for other settings, including those of the normal distributions in Euclidean spaces, are provided in Table 3 in Appendix D. There, we noted a slight advantage of wrapped normal distributions over their Euclidean counterparts. In

Table 1, we observe that the superior performance of the larger model WN (full,4) over WN (diag,2), observed in Table 2, is consistently reproduced across all datasets. Moreover, although WN (full,2) is relatively less performant in our comparison, almost all the estimates outperformed the other CSMC-based methods, implying the stability and efficiency of our approach.

Table 1: Comparison of the marginal log-likelihood (MLL) estimates with different approaches in eight benchmark datasets. The MLL values for MrBayes SS and VBPI-GNN, which employs the VIMCO estimator of 10 samples and EDGE GNN, are sourced from [40], The values of CSMC, VCSMC, and ϕ -CSMSC are referenced from [16]. We adopt the nomenclature from the literature [42] and refer to the dataset named DS7 in [16] as DS8. The MLL values for our approach (GeoPhy) are shown for two different Q ( z ) configurations: a wrapped normal distribution WN with 4 -dimensional full covariance matrix, and 2 -dimensional diagonal matrix. In each training, we used three different sets of CVs: LAX with K = 1 , LOO with K = 3 which we labeled LOO(3), and a combination of LOO and LAX, denoted as LOO(3)+. The bold figures are the best (highest) values obtained with GeoPhy and the tree CSMC-based methods, all of which perform an approximate Bayesian inference without the preselection of topologies. We underlined GeoPhy's MLL estimates that outperformed the other CSMC-based methods, demonstrating superior performance across various configurations.

0.2

0.4

0.6

0.8

1.0

1.2

1.4

1.6

1.8

BalPhy

BosTar

CalJac

CanFam

CheMed

DauMad

DidVir

EquCab

EulMon

FelCat

GalCra

GorGor

HapGri

HipAmp

HomSap

LemCat

LepEdw

LorTar

MacMul

MicBer

MicGri MicMur

MicMyo

MicRav

MicRuf1

MicRuf2

MicSam

MicTav

MirCoq

PanTro

PhyMac

PonPyg

ProTat

RhiUni

UrsAme

VarVar

0.4

0.8

1.2

1.6

2.0

2.4

2.8

3.2

3.6

BalPhy

BosTar

CalJac

CanFam

CheMed

DauMad

DidVir

EquCab

EulMon

FelCat

GalCra

GorGor

HapGri

HipAmp

HomSap

LemCat

LepEdw

LorTar

MacMul

MicBer

MicGri

MicMur

MicMyo

MicRav

MicRuf1

MicRuf2

MicSam

MicTav

MirCoq

PanTro

PhyMac

PonPyg

ProTat

RhiUni

UrsAme

VarVar

0.0

0.0

7.92

7.04

6.16

5.28

4.40

3.52

2.64

1.76

0.88

0.00

We further analyze the accuracy and expressivity of tree topology distribution Q ( τ ) in comparison with MCMC tree samples. As the training progresses, the majority consensus tree from GeoPhy

7130

7120

7110

MLL estimate

0

5

10

15

20

DS1

dim

2

4

5

6

K

1

3

33800

33750

MLL estimate

0

2

4

6

8

10

12

14

DS3

37600

37400

MLL estimate

0

5

10

15

20

25

30

DS7

0

100

Bipartitions

0.0

0.2

0.4

0.6

0.8

1.0

DS1

MrBayes

VBPI-GNN

Ours

progressively aligns with those obtained via MCMC (Fig. 3). Furthermore, we also observed that the Robinson-Foulds (RF) distance between the consensus trees derived from the topological distribution Q ( τ ) and MCMC (MrBayes) are highly aligned with the MLL estimates (Fig. 4 First to Third), underscoring the validity of MLL-based evaluations. While GeoPhy provides a tree topology mode akin to MCMCs with a close MLL value, it may require more expressivity in Q ( z ) to fully represent tree topology distributions. In Fig. 4 Fourth, we illustrate the room for improvement in the expressivity of Q ( τ ) in representing the bipartition frequency of tree samples, where VBPI-GNN aligns more closely with MrBayes. Further analysis is found in Appendix D.3.

In terms of execution time, we compared GeoPhy with VBPI-GNN and MrBayes, obtaining results that are promising and comparable to VBPI-GNN (Fig. 10 Left). Though we have not included CSMC-based methods in the comparison due to deviations in MLL values, they tend to capitalize on parallelism for faster execution as demonstrated in [16].

6 Conclusion

We developed a novel differential phylogenetic inference framework named GeoPhy, which optimized a variational distribution of tree topology and branch lengths without the preselection of candidate topologies. We also proposed a practical implementation for stable model optimization through choices of distribution models and control variates. In experiments conducted with real sequence datasets, GeoPhy consistently outperformed other approximate Bayesian methods that considered whole topologies.

7 Limitations and Future work

Although GeoPhy exhibits remarkable performance on standard benchmarks without preselecting topologies, it may be necessary to use more expressive distributions for Q ( z ) than independent parametric distributions to reach a level of performance comparable to state-of-the-art VBPI or gold-standard MCMC evaluations. Another limitation of our study lies in its exclusive focus on efficiency, measured in terms of the number of likelihood evaluations (NLEs), without paying significant attention to optimizing the computational cost per iteration. The general design of our variational distribution allows us to replace the tree link function such as UPGMA as presented in Fig. 9 Right. Future work should explore alternative design choices in continuous tree topology representations and the link functions to address these limitations. Extending our framework to include complex models remains crucial. Given the potential of scalable phylogenetic methods to enhance our understanding of viral and bacterial evolution, further exploration could substantially impact public health and disease control.

Acknowledgments and Disclosure of Funding

We appreciate the valuable comments on our study provided by Dr. Tsukasa Fukunaga of Waseda University and the anonymous reviewers. TM was partially supported by JSPS KAKENHI JP20H04239 and JST ACT-X JPMJAX22AI. MH was partially supported by JST CREST JPMJCR21F1 and AMED JP22ama121055, JP21ae0121049 and JP21gm0010008.

References

Yuri Burda, Roger B. Grosse, and Ruslan Salakhutdinov. Importance weighted autoencoders. In Yoshua Bengio and Yann LeCun, editors, 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4, 2016, Conference Track Proceedings , 2016. URL http://arxiv.org/abs/1509.00519 .

Ines Chami, Albert Gu, Vaggos Chatziafratis, and Christopher Ré. From trees to continuous embeddings and back: Hyperbolic hierarchical clustering. Advances in Neural Information Processing Systems , 33:15065-15076, 2020.

Eli Chien, Puoya Tabaghi, and Olgica Milenkovic. Hyperaid: Denoising in hyperbolic spaces for tree-fitting and hierarchical clustering. arXiv preprint arXiv:2205.09721 , 2022.

Sanjoy Dasgupta. A cost function for similarity-based hierarchical clustering. In Proceedings of the forty-eighth annual ACM symposium on Theory of Computing , pages 118-127, 2016.

Joseph Felsenstein. Maximum likelihood and minimum-steps methods for estimating evolutionary trees from data on discrete characters. Systematic Biology , 22(3):240-249, 1973.

James R Garey, Thomas J Near, Michael R Nonnemacher, and Steven A Nadler. Molecular evidence for acanthocephala as a subtaxon of rotifera. Journal of Molecular Evolution , 43: 287-292, 1996.

Will Grathwohl, Dami Choi, Yuhuai Wu, Geoffrey Roeder, and David Duvenaud. Backpropagation through the void: Optimizing control variates for black-box gradient estimation. In International Conference on Learning Representations , 2018.

S Blair Hedges, Kirk D Moberg, and Linda R Maxson. Tetrapod phylogeny inferred from 18s and 28s ribosomal rna sequences and a review of the evidence for amniote relationships. Molecular Biology and Evolution , 7(6):607-633, 1990.

Daniel A Henk, Alex Weir, and Meredith Blackwell. Laboulbeniopsis termitarius, an ectoparasite of termites newly recognized as a member of the laboulbeniomycetes. Mycologia , 95(4): 561-564, 2003.

Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterization with gumbel-softmax. arXiv preprint arXiv:1611.01144 , 2016.

Yueyu Jiang, Puoya Tabaghi, and Siavash Mirarab. Phylogenetic placement problem: A hyperbolic embedding approach. In Comparative Genomics: 19th International Conference, RECOMB-CG 2022, La Jolla, CA, USA, May 20-21, 2022, Proceedings , pages 68-85. Springer, 2022.

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

Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980 , 2014.

Diederik P Kingma and Max Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

Wouter Kool, Herke van Hoof, and Max Welling. Buy 4 reinforce samples, get a baseline for free! 2019.

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

Clemens Lakner, Paul Van Der Mark, John P Huelsenbeck, Bret Larget, and Fredrik Ronquist. Efficiency of markov chain monte carlo tree proposals in bayesian phylogenetics. Systematic biology , 57(1):86-103, 2008.

Marc Law, Renjie Liao, Jake Snell, and Richard Zemel. Lorentzian distance learning for hyperbolic representations. In International Conference on Machine Learning , pages 36723681. PMLR, 2019.

Matthew Macaulay, Aaron E Darling, and Mathieu Fourment. Fidelity of hyperbolic space for bayesian phylogenetic inference. PLoS computational biology , 2023.

Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint arXiv:1611.00712 , 2016.

Hirotaka Matsumoto, Takahiro Mimori, and Tsukasa Fukunaga. Novel metric for hyperbolic phylogenetic tree embeddings. Biology Methods and Protocols , 6(1):bpab006, 2021.

Andriy Mnih and Danilo Rezende. Variational inference for monte carlo objectives. In International Conference on Machine Learning , pages 2188-2196. PMLR, 2016.

Nicholas Monath, Manzil Zaheer, Daniel Silva, Andrew McCallum, and Amr Ahmed. Gradientbased hierarchical clustering using continuous representations of trees in hyperbolic space. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 714-722, 2019.

Antonio Khalil Moretti, Liyi Zhang, Christian A Naesseth, Hadiah Venner, David Blei, and Itsik Pe'er. Variational combinatorial sequential monte carlo methods for bayesian phylogenetic inference. In Uncertainty in Artificial Intelligence , pages 971-981. PMLR, 2021.

Yoshihiro Nagano, Shoichiro Yamaguchi, Yasuhiro Fujita, and Masanori Koyama. A wrapped normal distribution on hyperbolic space for gradient-based learning. In International Conference on Machine Learning , pages 4693-4702. PMLR, 2019.

Maximillian Nickel and Douwe Kiela. Poincaré embeddings for learning hierarchical representations. Advances in neural information processing systems , 30, 2017.

Rajesh Ranganath, Sean Gerrish, and David Blei. Black box variational inference. In Artificial intelligence and statistics , pages 814-822. PMLR, 2014.

Lorenz Richter, Ayman Boustati, Nikolas Nüsken, Francisco Ruiz, and Omer Deniz Akyildiz. Vargrad: a low-variance gradient estimator for variational inference. Advances in Neural Information Processing Systems , 33:13481-13492, 2020.

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

Amy Y Rossman, John M McKemy, Rebecca A Pardo-Schultheiss, and Hans-Josef Schroers. Molecular studies of the bionectriaceae using large subunit rdna sequences. Mycologia , 93(1): 100-110, 2001.

Naruya Saitou and Masatoshi Nei. The neighbor-joining method: a new method for reconstructing phylogenetic trees. Molecular biology and evolution , 4(4):406-425, 1987.

Frederic Sala, Chris De Sa, Albert Gu, and Christopher Ré. Representation tradeoffs for hyperbolic embeddings. In International conference on machine learning , pages 4460-4469. PMLR, 2018.

Kirill Struminsky, Artyom Gadetsky, Denis Rakitin, Danil Karpushkin, and Dmitry P Vetrov. Leveraging recursive gumbel-max trick for approximate inference in combinatorial spaces. Advances in Neural Information Processing Systems , 34:10999-11011, 2021.

Liangliang Wang, Alexandre Bouchard-Côté, and Arnaud Doucet. Bayesian phylogenetic inference using a combinatorial sequential monte carlo method. Journal of the American Statistical Association , 110(512):1362-1374, 2015.

Benjamin Wilson. Learning phylogenetic trees as hyperbolic point configurations. arXiv preprint arXiv:2104.11430 , 2021.

Wangang Xie, Paul O Lewis, Yu Fan, Lynn Kuo, and Ming-Hui Chen. Improving marginal likelihood estimation for bayesian phylogenetic model selection. Systematic biology , 60(2): 150-160, 2011.

Ziheng Yang and Anne D Yoder. Comparison of likelihood and bayesian methods for estimating divergence times using multiple gene loci and calibration points, with application to a radiation of cute-looking mouse lemur species. Systematic biology , 52(5):705-716, 2003.

Anne D Yoder and Ziheng Yang. Divergence dates for malagasy lemurs estimated from multiple gene loci: geological and evolutionary context. Molecular Ecology , 13(4):757-773, 2004.

Cheng Zhang. Improved variational bayesian phylogenetic inference with normalizing flows. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 18760-18771. Curran Associates, Inc., 2020. URL https://proceedings.neurips.cc/paper/2020/file/ d96409bf894217686ba124d7356686c9-Paper.pdf .

Cheng Zhang. Learnable topological features for phylogenetic inference via graph neural networks. In The Eleventh International Conference on Learning Representations , 2023. URL https://openreview.net/forum?id=hVVUY7p64WL .

Cheng Zhang and Frederick A Matsen IV. Generalizing tree probability estimation via bayesian networks. In S. Bengio, H. Wallach, H. Larochelle, K. Grauman, N. Cesa-Bianchi, and R. Garnett, editors, Advances in Neural Information Processing Systems 31 , pages 1449-1458. Curran Associates, Inc., 2018. URL http://papers.nips.cc/paper/ 7418-generalizing-tree-probability-estimation-via-bayesian-networks. pdf .

Cheng Zhang and Frederick A. Matsen IV. Variational bayesian phylogenetic inference. In International Conference on Learning Representations , 2019. URL https://openreview. net/forum?id=SJVmjjR9FX .

Ning Zhang and Meredith Blackwell. Molecular phylogeny of dogwood anthracnose fungus (discula destructiva) and the diaporthales. Mycologia , 93(2):355-365, 2001.

A Details of Background

A.1 Summary of calculations for wrapped normal distributions

The wrapped normal distribution, as proposed in [25], is a probability distribution defined on hyperbolic spaces, which is easy to sample from and evaluate its probability density at arbitrary coordinates in the hyperbolic spaces. In the following, we provide a detailed summary of the calculation involved in applying the wrapped normal distributions within the context of the Lorentz model of hyperbolic spaces.

The Lorentz model, denoted as H d , represents the d -dimenstional hyperbolic space as a submanifold of a d +1 dimensional Euclidean space. Given u, v ∈ R d +1 , we can define the pseudo-inner product and pseudo-norm as follows:

The distance between hyperbolic coordinates ν, µ ∈ H d is defined as follows:

where cosh -1 denotes the inverse function of the hyperbolic cosine function. Consider hyperbolic coordinates ν, µ ∈ H d and tangent vectors u ∈ T µ H d and v ∈ T ν H d . An exponential map exp µ ( u ) ∈ H d , a logarithm map log µ ( ν ) ∈ T µ H d , and a parallel transport map PT ν → µ ( v ) ∈ T µ H d , can be calculated as follows:

where we denote α = -⟨ ν, µ ⟩ L , and cosh -1 represents the inverse function of cosh .

Given location and scale parameters denoted as µ ∈ H d and Σ ∈ R d × d , respectively, the procedure for sampling from a wrapped normal distribution z ∼ WN ( µ, Σ) defined over H d is given as follows:

where µ o = (1 , 0 , . . . , 0) ⊤ denotes the origin of the H d . Note that we set u 0 = 0 , and u := u 0: d ∈ T µ o H d represents a tangent vector at the origin T µ o H d .

From the sampling definition in equation (22), the probability density function WN ( z ; µ, Σ) can be derived as follows:

where u is defined as u = PT µ → µ o ◦ log µ ( z ) . For detailed derivation, we refer to Appendix A of [25].

A.2 GNN-based parameterization for variational branch length distributions

In this work, we employ a variational branch length distribution Q ϕ ( B τ | τ ) parameterized with a graph neural network (GNN) as described in [40]. In concrete, each of the branch lengths follows an independent lognormal distribution, where its location and scale parameters are predicted with a GNN that takes the tree topology τ and the learnable topological features (LTFs) of the topology τ , which are computed with a method described in [40]. Below, we summarize an architecture that we use in this study.

Branch length parameterizations Let V τ and E τ respectively represent the sets of nodes and branch edges for a given unrooted binary tree topology τ . The input to the GNN consists of node features represented by LTFs denoted as { h (0) v } v ∈ V τ . These features undergo transformation L times as follows:

where we set L = 2 . The function g ( ℓ ) represents a GNN layer. For this function, we utilize edge convolutional layers, which will be described in more detail in the following paragraph.

Next, the last node features h ( L ) are transformed to output parameters of edge length as follows:

where MLP N , MAX , and MLP E denotes a multi-layer perceptron for node features with two hidden layers, the element-wise max operation, and a multi-layer perceptron with a hidden layer that outputs the location and scale parameter ( µ e , σ e ) of the lognormal distributions for each edge e ∈ E τ . For each of the hidden layers employed in MLP N and MLP E , we set its width to 100 and apply the ELU activation function after the linear transformation of input values.

Edge convolutional layers In a previous study [40], a GNN with edge convolutional layers, referred to as EDGE, demonstrated strong performance when predicting the posterior tree distributions. In EDGE, the function g ( ℓ ) transforms node features { h ( ℓ ) v } v ∈ V τ according to the following scheme:

where g ( ℓ ) is comprised of the edge convolutional operation with the exponential linear unit (ELU) activation function. Specifically, the transformation with the layer g ( ℓ ) is computed as follows:

where N τ ( v ) represents a set of neighboring nodes connected to node v in the tree topology τ , ∥ refers to the concatenation operation of elements, MLP ( ℓ ) denotes a full connection layer and the exponential linear unit (ELU) activation unit, and AGG ( ℓ ) represents an aggregation operation that takes the maximum value of neighboring edge features e ( ℓ ) u → v , ∀ u ∈ N ( v ) for each element.

B Variational Lower Bounds and Gradient Estimators

B.1 Variational lower bound

In Proposition 1, we present that the following functional is a lower bound of the marginal loglikelihood ln P ( Y ) .

where F and F ′ are respectively defined as follows:

B.2 Gradient estimators for variational lower bound

The gradient of L [ Q θ,ϕ , R ψ ] with respect to θ is given by

where H denotes the differential entropy. We assume that Q ϕ ( B τ | τ ) is reparameterizable as in [40]: namely, B τ can be sampled through B τ = h ϕ ( ϵ B , τ ) , where ϵ B ∼ p B ( ϵ B ) , where p B ( ϵ ) and h ϕ denote a parameter-free base distribution and a differentiable function with ϕ , respectively. Consequently, the gradient of L with respect to ϕ is evaluated as follows:

Lastly, the gradient of L with respect to ψ can be evaluated with a tractable density model R ψ ( z | τ ) as follows:

Given samples ϵ z ∼ p z and ϵ B ∼ p B , we can compute z = h θ ( ϵ z ) , τ ( z ) , and B τ = h ϕ ( ϵ B , τ ( z )) . Then, the below equations are estimators of gradients ∇ θ L , ∇ ϕ L , and ∇ ψ L , respectively:

The gradients can be computed through the auto-gradient of the following target:

where we denote f ( z, B τ ) = ln F ( z, B τ ) , and detach[ · ] refers to an operation that blocks backpropagation through its argument. For clarity in terms of differentiability with respect to the parameters, we distinguish between expressions ( z , B τ ) and ( h θ ( ϵ z ) , h ϕ ( ϵ B , τ ( z )) ).

B.3 Multi-sample gradient estimators

Given a K set of Monte Carlo (MC) samples from Q θ,ϕ ( z, B τ ) , i.e. { ϵ ( k ) Z , z ( k ) = h θ ( ϵ ( k ) Z ) } K k =1 and { ϵ ( k ) B , B ( k ) τ = h ϕ ( ϵ ( k ) B , τ ( z ( k ) )) } K k =1 , we can simply estimate ∇ L θ [ Q θ,ϕ , R ψ ] as follows:

As a simple extension of equation (43), the gradients are obtained through an auto-gradient computation of the following target:

B.4 Leave-one-out (LOO) control variates for variance reduction

For the term of K -sample gradient estimator ̂ g ( K ) θ proportional to the score function ∇ θ ln Q θ , a leave-one-out (LOO) variance reduction is known to be effective [15, 28], which is denoted as follows:

where f k denotes:

̸

To employ the LOO gradient estimator for θ , the target of auto-gradient computation in equation (45) needs to be adjusted as follows:

B.5 LAX estimators for adaptive variance reduction

Algorithm 2 GeoPhy algorithm with LAX

1: θ, ϕ, ψ, χ ← Initialize variational parameters

2: while not converged do

3: ϵ (1: K ) z , ϵ (1: K ) B ← Random samples from distributions p z , p B

4: z (1: K ) ← h θ ( ϵ (1: K ) z )

5: B (1: K ) τ ← h ϕ ( ϵ (1: K ) B , τ ( z (1: K ) ))

7: ̂ g χ ←∇ χ 〈 ̂ g 2 θ 〉

6: ̂ g θ , ̂ g ϕ , ̂ g ψ ← Estimate the gradients ∇ θ L , ∇ ϕ L , ∇ ψ L

8: θ, ϕ, ψ, χ ← Update parameters using an SGA algorithm, given gradients ̂ g θ , ̂ g ϕ , ̂ g ψ , ̂ g χ

9: end while

10: return θ, ϕ, ψ

The LAX estimator [7] is a stochastic gradient estimator based on a surrogate function, which can be adaptively learned to reduce the variance regarding the term ∇ θ ln Q θ ( z ) . In our case, the LAX estimator is given as follows:

We summarize the modified procedure of GeoPhy with LAX in Algorithm 2. As we assume Q θ ( z ) is differentiable with respect to z , we can also use a modified estimator as follows:

Since it is favorable to reduce the variance of ̂ g LAX ,θ , we optimize χ to minimize the following objective as proposed in [7]:

where n θ denotes the dimension of θ . As the gradient in equation (49) is given as an unbiased estimator of ∇ θ L , which is not dependent on χ , we can use the relation ∇ χ E Q θ ( z ) [ ̂ g LAX ,θ i ] = 0 . Therefore, the unbiased estimator of the gradient ∇ χ 〈 V Q θ ( z ) [ ̂ g LAX ,θ ] 〉 is given as follows:

As we require the gradient of ∇ θ L with respect to χ for the optimization, we use different objectives for auto-gradient computation with respect to θ and the other parameters ϕ and ψ as follows:

B.6 LAX estimators with multiple MC-samples

For the cases with K MC-samples, we use LAX estimators by differentiating the following objectives:

When we combine LAX estimators with LOO control variates. the target for auto-gradient computation changes to the following:

We note that ̂ L ′ ( K ) ϕ,ψ is not affected by the introduction of LOO control variates.

B.7 Derivation of importance-weighted evidence lower bound (IW-ELBO)

An importance-weighted evidence lower bound (IW-ELBO) [1], is a tighter lower bound of the log-likelihood ln P ( Y ) than ELBO. For our model, a conventional K -sample IW-ELBO is given as follows:

The fact that L ( K ) IW [ Q ] is the lower bound of ln P ( Y ) is directly followed from Theorem 1 in [1]. However, as our model cannot directly evaluate the mass function Q ( τ ) , we must resort to considering the second lower bound, similar to the case of K = 1 as depicted in Proposition 1. We define the K -sample tractable IW-ELBO as follows:

where F ′ is defined in equation (34). We will prove in Theorem 1 that L ( K ) IW [ Q,R ] serves as a lower bound of the ln P ( Y ) . Although this inequality holds when K = 1 , as shown by ln P ( Y ) ≥ L [ Q ] ≥ L [ Q,R ] in Proposition 1, the relationship is less obvious when K > 1 . Before delving into that, we prepare the following proposition.

Proposition 2. Given Q ( z, τ ) as defined in equation 4 and an arbitrary conditional distribution R ( z | τ ) , it follows that

where setting R ( z | τ ) = Q ( z | τ ) is a sufficient condition for the equality to hold.

Proof. The inequality immediately follows from the definition as follows:

Next, when we set R ( z | τ ) = Q ( z | τ ) , the condition for equality is satisfied as follows:

where we have used the definition of Q ( τ ) := E Q ( z ) [ I [ τ = τ ( z )]] from equation (4) and the resulting relation Q ( z | τ ) Q ( τ ) = I [ τ = τ ( z )] Q ( z ) .

Theorem 1. Given Q ( z, τ ) as defined in equation 4 and an arbitrary conditional distribution R ( z | τ ) that satisfies supp R ( z | τ ) ⊇ supp Q ( z | τ ) , for any natural number K > 1 , the following relation holds:

Additionally, if F ′ ( z, B τ ) is bounded and ∀ τ, E R ( z | τ ) [ I [ τ = τ ( z )]] = 1 , then L ( K ) IW [ Q,R ] approaches ln P ( Y ) as K →∞ .

Proof. We first show that for any natural number K > M ,

For simplicity, we denote Q ( z ( k ) , B ( k ) τ ) and F ′ ( z ( k ) , B ( k ) τ ) as Q k and F ′ k , respectively, in the following discussion. Let U K M represent a uniform distribution over a subset with M distinct indices chosen from the K indices { 1 , . . . , K } . Similar to the approach used in [1], we will utilize the following relationship:

Now, the inequality (64) is derived as follows:

where we have also used Jensen's inequality.

Next, we show that ln P ( Y ) ≥ L ( K ) IW [ Q,R ] . We again use Jensen's inequality as follows:

The last term is further transformed as follows:

where, in the transition from the first to the second row, we employed the following relation:

and we have used Proposition 2 for the last inequality.

Finally, we will show that the following convergence property assuming that F ( z, B τ ) is bounded:

From the strong law of large numbers, it follows that 1 K ∑ K k =1 F ′ k converges to the following term almost surely as K →∞ :

where we have employed the same transformations as used from equation (71) to (75). Observe that the r.h.s term of the equation (78) equals to ln P ( Y ) when ∀ τ ′ ∈ T , E R ( z | τ ′ ) [ I [ τ ′ = τ ( z )]] = 1 , which completes the proof.

Estimation of marginal log-likelihood For the estimation of ln P ( Y ) , we employ L ( K ) [ Q,R ] with K = 1 , 000 similar to [40]. From Theorem 1, IW-ELBO L ( K ) [ Q,R ] is at least a better lower bound of ln P ( Y ) than ELBO L [ Q,R ] , and converges to ln P ( Y ) when ∀ τ, E R ( z | τ ) [ I [ τ = τ ( z )]] = 1 . According to Proposition 2, this equality condition is satisfied when we set R ( z | τ ) = Q ( z | τ ) , which is approached by maximizing L [ Q,R ] with respect to R as indicated in Proposition 1.

B.8 Gradient estimators for IW-ELBO

The gradient of IW-ELBO L ( K ) IW [ Q θ,ϕ , R ψ ] with respect to θ is given by

where we have defined

Similarly, as F ′ is differentiable with respect to B τ , and B τ = h ϕ ( ϵ B , τ ) is differentiable with respect to ϕ , the gradient of L ( K ) IW with respect to ϕ can be evaluated as follows:

Since ∇ θ ln F ′ ( z ( k ) , B ( k ) τ ) = -∇ θ ln Q θ ( z ( k ) ) from equation (34), an unbiased estimator of the gradient ∇ θ L ( K ) is given as follows:

The remaining gradient estimators are given as follows:

In total, the target for computing auto-gradient for these gradients is given as follows:

B.9 VIMCO estimators

The VIMCO estimator [22] aims at reducing the variance in each of the terms proportional to ℓ ( z (1: K ) , B (1: K ) ) in equation (84). This is accomprished by forming for a control variate ℓ k ( z ( \ k ) , B ( \ B τ ) ) that consists of only variables other than z ( k ) and B ( k ) for every k as follows:

where we denote that

̸

̸

Note that the unbiasedness of ̂ g θ, VIMCO can be confirmed through the following observation:

To employ the VIMCO estimator for θ , the target for auto-gradient computation as stated in equation (87) needs to be adjusted as follows:

C Experimental Details

C.1 Training process

For the training of GeoPhy, we continued the stochastic gradient descent process until a total of 1,000,000 Monte Carlo (MC) tree samples were consumed. Specifically, if K MC-samples were used

per step, we performed up to 1,000,000 / K steps. It is noteworthy that the number of MC samples equaled the number of likelihood evaluations (NLEs), which provided us with a basis for comparing convergence speed between different runs, as shown in Fig. 2. In all of our experiments, we used Adam optimizer with an initial learning rate of 0 . 0001 . The learning rate was then multiplied by 0 . 75 after every 200,000 steps. Similar to approaches taken by Zhang and Matsen IV [42], we incorporated an annealing procedure during the initial consumption of 100,000 MC samples. Specifically, we replaced the likelihood function in the lower bound with P ( Y | B τ , τ ) β and linearly increased the inverse temperature β from 0 . 001 to 1 throughout the iterations. Note that all the estimations of marginal log-likelihood (MLL) were performed with β set to 1 .

For each estimation of MLL values, GeoPhy and VBPI-GNN used IW-ELBO with 1000 samples, whereas all CSMC-based methods used 2048 particle weights [16].

C.2 Variational branch length distribuions

For the variational branch length distribution Q ϕ ( B τ | τ ) , we followed an architecture of [40]; namely, each branch length independently followed a lognormal distribution which was parameterized with a graph neural network (GNN). Details are described in Appendix A.2.

C.3 LAX estimators

As input features of a surrogate function s χ ( z ) used in the LAX estimators, we employed a flattened vector of coordinates z ∈ R N × d when z resides in Euclidean space. In cases where the coordinates were z ∈ H d , we first transformed z with a logarithm map log µ o z ∈ T µ o H d , then omitted their constant value 0 -th elements and subsequently flattened the result. We implemented a simple multilayer perceptron (MLP) network with a single hidden layer of width 10 Nd and a subsequent sigmoid linear unit (SiLU) activation function as the neural network to output s χ ( z ) .

C.4 Replication of MLL estimates with MrBayes SS

Given the observed discrepancies in marginal log-likelihood (MLL) estimates obtained with the MrBayes stepping-stone (SS) method between references [42] and [16], we replicated the MrBayes SS runs using MrBayes version 3.2.7a. The script we used is provided below.

We incorporated the results in the row named Replication in Table 3, where the values aligned more closely with those found in [42]. We deduced that the prior distribution used in [16] might have been set differently as the current default values of brlenspr are Unconstrained : GammaDir(1 . 0 , 0 . 100 , 1 . 0 , 1 . 0) 3 , which deviates from the model assumption used for the benchmarks. We observed that the line brlenspr was not included in the code provided in Appendix F of [16]. Having been able to replicate the results found in [42], we opted to use their values as a reference in Table 1.

C.5 Visualization of tree topologies

In Fig. 3, we visualized the sum of the probability densities for tip node distribution ∑ N i =1 Q ( z i ) by projecting each hyperbolic coordinate z i ∈ H d onto the Poincaré coordinates z ik = z ik / (1 + z i 0 ) ( k = 1 , . . . , d ) , then applying a transformation z i ↦→ z ′ i = tanh( a ∥ z i ∥ 2 ) · z i / ∥ z i ∥ 2

with a = 2 . 1 to emphasize the central region. To display the density Q in the new coordinates, the Jacobian term was also considered to evaluate the density Q ( z ′ i ) .

For the comparison of consensus tree topologies, we plotted the edges of the tree by connecting each of their end node coordinate pairs with a geodesic line. The coordinate in H d of the i -th tip node was determined as the location parameter µ i ∈ H d of the wrapped normal distribution Q ( z i ) = WN ( z i ; µ i , Σ i ) . Let ξ u ∈ H d denotes the coordinate of an interior node u , we defined ξ u by using the Lorentzian centroid operation C [18] as follows:

where ˜ ξ u := ∑ s ∈S τ ( u ) ν s c s denote an unnormalized sum of weighted coordinates, s ∈ S τ ( u ) denote a subset of tip node indices partitioned by the interior node u in the tree topology τ , c s := C ( { µ i } i ∈ s , { 1 } i ∈ s ) denote the Lorentzian centroid of the tip nodes contained in the subset s , and ν s = N -| s | denote the number of the tip nodes in the complement set of s where | s | represents the number of tip nodes in the subset s . As an unrooted tree topology τ can be identified by the set of tip node partitions introduced by the interior nodes of τ , the same unrooted tree topologies give the same set of interior coordinates { ξ u } u ∈ V according to equation (93).

D Additional Results

D.1 Marginal log-likelihood estimates throughout iterations for DS1 dataset

In Fig. 2, we highlight several key findings of training processes. Specifically, (1) control variates were crucial for optimizations; (2) The LAX estimator for individual Monte Carlo (MC) samples ( K = 1 ) and leave-one-out (LOO) CVs for multiple MC samples were effective; (3) IW-ELBO and VIMCO were not particularly effective when a full diagonal matrix is used for Q ( z ) . In Fig. 5 and 6, we show that these findings hold consistently across various control variates and configurations of Q ( z ) , including, different types of distributions (normal N and wrapped normal WN ), embedding dimensions ( d = 2 , 3 , and 4 ), and varying numbers of MC samples ( K = 1 , 2 , and 3 ).

D.2 Marginal log-likelihood estimates for eight datasets

We present more comprehensive results in Table 3, extending upon the data from Table 1. This table showcases the marginal log-likelihood (MLL) estimates obtained with various GeoPhy configurations and other conventional methods for the datasets DS1-DS8. Once again, GeoPhy demonstrates its superior performance, consistently outperforming CSMC-based approaches that do not require preselection of tree topologies across the majority of configurations and datasets. This reaffirms the stability and excellence of our approach. Additionally, we found that a Q ( z ) configuration using a 4 -dimensional wrapped normal distribution with a full covariance matrix was the most effective among the tested configurations.

D.3 Analysis of tree topology distributions

To assess the accuracy and expressivity of GeoPhy's tree topology distribution, we contrasted Q ( τ ) with MCMC tree samples. For our analysis, we considered experimental runs using LAX estimators with either K = 1 or 3 . For a configuration of Q ( z ) , We employed a diagonal or multivariate wrapped normal distribution with 2 , 4 , 5 , and 6 dimensions. We sampled 1000 MC trees from Q ( τ ) for each experiment and then deduced the majority consensus tree. As depicted in Fig. 8 across the eight datasets, the Robinson-Foulds (RF) distances between the majority consensus tree of Q ( τ ) and the one obtained with MrBayes were well aligned with the MLL values.

To highlight the expressivity limitation of Q ( τ ) obtained from a simple independent distribution of Q ( z ) , we compiled three statistics related to the tree topology samples for DS1, DS4, and DS7 datasets (Table 4), wherein we employed the model with the least RF distance for each dataset. For DS1 and DS4, GeoPhy's tree topologies were less diverse than those of MrBayes and VBPI-GNN and had more concentration on the most frequent topology. For more diffused tree DS7, GeoPhy also showed diverse tree samples with a close diversity index to the other two.

We also summarize the frequency of bipartitions observed for tree topology samples in Fig. 8. The intermediate values between 0 and 1 in this frequency plot reveal topology diversities, where VBPI-GNN aligns more closely with MrBayes compared to GeoPhy. For DS1 and DS4, GeoPhy tends to take values near zero and one in the slope region. For DS7, while GeoPhy traced the curve more accurately, fluctuations around this curve highlight potential room for improvement.

D.4 Runtime comparison

We measured CPU runtimes using a single thread on a 20-core Xeon processor. In Fig. 10, we compared single-core runtime of MrBayes, VPBI-GNN, and GeoPhy (ours). As MrBayes is written in C and highly optimized its per-iteration computation time is much faster than current variational methods. In general, the per-step efficiency of GeoPhy is comparable to that of VBPI-GNN. Although the convergence of MLL estimates tends to be slow in terms of the number of consumed samples for larger K (Fig. 9 Left), the total CPU-time slightly decreases as K increases (Fig. 10 Left). We also included the estimated runtime across the eight datasets and the different model configurations, where the number of species ranges from 27 to 64 (Fig. 10 Right).

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=2, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=3, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=4, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=2, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=3, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=4, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=2, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=3, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (diag d=4, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=2, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=3, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=4, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=2, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=3, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=4, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=2, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=3, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Normal (full d=4, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=2, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=3, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=4, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=2, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=3, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=4, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=2, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=3, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (diag d=4, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=2, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=3, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=4, K=1)

NoCV

LAX

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=2, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=3, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=4, K=2)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=2, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=3, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

0.0

0.2

0.4

0.6

0.8

1.0

Iterations

1e6

-9000

-8500

-8000

-7500

MLL estimate

Wrapped Normal (full d=4, K=3)

NoCV

LAX

LOO

LOO+LAX

IW

VIMCO

7130

7120

7110

MLL estimate

0

5

10

15

20

DS1

dim

2

4

5

6

K

1

3

26400 26450 26500

MLL estimate

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

DS2

33800

33750

MLL estimate

0

2

4

6

8

10

12

14

DS3

13360

13340

MLL estimate

0

5

10

15

20

25

DS4

8260

8240

8220

MLL estimate

0

5

10

15

20

25

30

DS5

6760

6740

MLL estimate

0.0

2.5

5.0

7.5

10.0

12.5

15.0

17.5

DS6

37600

37400

MLL estimate

0

5

10

15

20

25

30

DS7

8800

8700

MLL estimate

0

5

10

15

20

25

30

35

40

DS8

0

100

Bipartitions

0.0

0.2

0.4

0.6

0.8

1.0

DS1

MrBayes

VBPI-GNN

Ours

0

50

100

150

200

Bipartitions

0.0

0.2

0.4

0.6

0.8

1.0

DS4

MrBayes

VBPI-GNN

Ours

0

50

100

150

200

Bipartitions

0.0

0.2

0.4

0.6

0.8

1.0

DS7

MrBayes

VBPI-GNN

Ours

0.00

0.25

0.50

0.75

1.00

Iterations

1e6

-8400

-8200

-8000

-7800

-7600

-7400

-7200

MLL estimate

Normal (full d=2)

K=2 (LOO)

K=2 (LOO+LAX)

K=3 (LOO+LAX)

K=5 (LOO)

K=5 (LOO+LAX)

K=10 (LOO)

K=10 (LOO+LAX)

0.00

0.25

0.50

0.75

1.00

Iterations

1e6

-8400

-8200

-8000

-7800

-7600

-7400

-7200

MLL estimate

Normal (full d=2, K=3)

NJ (LOO)

NJ (LOO+LAX)

UPGMA (LOO)

UPGMA (LOO+LAX)

0

25000

50000

75000

Runtime (sec)

Ours(K=10)

Ours(K=5)

Ours(K=3)

Ours(K=1)

VBPI-GNN*

VBPI-GNN

MrBayes SS*

MrBayes SS

MrBayes

DS1

DS1

DS2

DS3

DS4

DS5

DS6

DS7

DS8

0

10

20

30

40

50

Sec / 1k MC samples

N(2,diag)

N(2,diag)+LAX

WN(4,full)

WN(4,full)+LAX

Peer reviews:
comments: SUMMARY:
This work presents a new robust and scalable method for inferring phylogenetic trees based on (variational) Bayesian inference.

STRENGTHS:
Originality
- the originality of this work is in providing a robust and rigorous solution to an important application problem where the application of Bayesian methodology has so far been relatively limited

Quality
- the work is well motivated and appropriately implemented including the relevant derivations and experimental benchmarking 

Clarity
- language is fluent, and intuitive illustrations support reading; the technical details are described in sufficient detail and complemented by intuitive explanation. Conclucions and limitations are clearly stated. 

Significance
- significant, general and timely application problem addressed
- improved robustness, while maintaining scalability that allows practical application

WEAKNESSES:
Treatment of the benchmark data sets is limited compared to the potential of the method for real applications. This could be expanded to highlight the relevance of the work.

QUESTIONS:
Good performance compared to alternatives is shown but is it possible to demonstrate more on practical relevance for this?

LIMITATIONS:
No source code.

No assessment of societal implications.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-06-25

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors propose a method for learning phylogenetic trees from sequence data. \n\nThe key idea is borrowed from [16], which is to represent the tree topology in terms of an embedding of the leafs of the trees in a continuous space, $z \\in \\mathcal{Z}$ from which the topology is extracted via a mapping $\\tau \\;:\\; \\mathcal{Z} \\to \\mathcal{T}$, where $\\mathcal{T}$ is the space of binary trees, using a distance-based technique (in particular the popular Neighbor-Joining (NJ) method). The fact that the representation in terms of $z$ is continuous has benefits in a variational representation of the posterior distribution over trees.\n\nThe theoretical contribution of the present paper is resolving some issues with \"an unevaluated Jacobian determinant between $B_\\tau$ and $z$\" (see Sec. 4 \"Related work\", p. 6).\n\nThe resulting method is demonstrated through simulations, showing that the method yields high-scoring models (high log-likelihoods), but not as high as existing variational and MCMC techniques (Table 2).\n\nSTRENGTHS:\n- I like the idea (which really borrowed from [16], so not really a new idea) of the embedding in terms of continuous coordinates\n- solid theoretical derivation of the variational estimators\n- promising results\n- quite well written (although a bit dense)\n\nWEAKNESSES:\n- not clear what the benefits are compared to existing variational and MCMC methods (which are presented as the gold-standard)\n\nQUESTIONS:\nMy main question is what are the benefits compared to, e.g., MrBayes, VBPI-GNN, and the method proposed in [33], the first two of which are referred to as gold standard. I'm guessing the proposed method is computationally more efficient and scalable, but this should be clearly demonstrated (or am I missing something?).\n\nI should point out that I'm no expert of variational Bayes, so I didn't check the derivations from that point of view. I'm counting on the other reviewers to comment on that side.\n\ndetailed comments:\np. 7: \"marginal log-likelihood (MLL) estimates\": Does this mean MLL values? I mean, \"maximum (log-) likelihood estimate\" usually refers to an estimate of a parameter that is obtained by likelihood maximization, so it'd be good to be unambiguous.\np. 7: typo \"Monte-Calro\"\n\nLIMITATIONS:\nI think so (but see my question about the benefits wrt. existing methods above).", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors present a novel variational distribution for tree topologies, $Q(\\tau)$, in the context of variational inference (VI) in phylogenetics. They construct their $Q(\\tau)$ by introducing a continuous distribution $Q(z)$ in hyperbolic space and define $Q(\\tau)$ by an expectation over the support of $Q(z)$ of an indicator function acting on a mapping from the hyperbolic space to the topology space; this allows sampling from $Q(z)$ and reconstructing \\tau through the neighbour joining algorithm and, by the differentiable $Q(z)$ distribution, enables Monte-Carlo gradient estimation to maximize the Evidence Lower Bound (ELBO). \n\nAn ablation study is conducted to evaluate different $Q(z)$ distributions (Normal and Wrapped-Normal) and control variates. The model is then compared to other VI methods and MrBayes across 8 datasets.\n\nSTRENGTHS:\nThe problem of phylogenetic tree reconstruction is paramount to understanding evolution, e.g., the evolution of species and cancer. The grand size of the tree topology space induced by a set of taxa constitutes a major obstacle for inference; any method successful in efficient exploration of this space is a significant contribution to the field. Here, the authors construct a differentiable $Q(\\tau)$, enabling the use of gradients to indirectly maneuver the tree topology space. Furthermore, the proposed framework does not require restricting the tree topology space based on pre-processing steps as in VBPI. \n\nThe novel, differentiable approach to $Q(\\tau)$, the avoidance of restrictions of the tree support and reporting strong results on datasets DS1-DS8 in terms of the ELBO(Q,R), makes the paper a significant and original contribution to the field of Bayesian phylogenetics.\n\nWEAKNESSES:\nThe standard deviations of the proposed method reported in Table 2 are low for multiple datasets; this may be due to the consistency of the optimization across seeds, but may also indicate that the support of $Q(z)$ collapses to regions containing one/few unique $\\tau(z)$. Since there is no analysis/discussion regarding this, the paper presents limited information regarding the representative power of $Q(\\tau)$, which is a major set back of the paper, especially as the MrBayes \"golden-run\" posteriors over the tree topologies for DS1-DS8 support multiple topologies (Whidden and Matsen 2015, https://www.ncbi.nlm.nih.gov/pmc/articles/PMC4395846/ ).\n\nThe use of $Q(z)$ makes the method extra susceptible to limited topology support by the inherent mode-seeking behaviour in VI; while a mode in the topology space could represent multiple tree topologies, a mode in Z-space might only support reconstruction of, at worst, one tree topology. \n\nThe paper would benefit from reporting runtime compared to MrBayes and other VI methods, especially as one of the strengths in VI is speed. The current report on runtime is limited and is relegated to Appendix C.\n\nFurthermore, an obvious weakness is the need to introduce a lower bound of the ELBO.\n\nQUESTIONS:\nThe paper is a significant, novel and promising contribution to the field and is in its current format acceptable for publication.\n\nHowever, based on my concerns in the weakness section, addressing the following points in the rebuttal could make me increase my score:\n1. A discussion regarding the low variance in Table 2 and limited knowledge of the representative power $Q(\\tau)$.\n2. Clearer reporting of the runtime of the algorithm in the main paper.\n\nAddressing the following points in the rebuttal could increase my score even further:\n1. A experiment demonstrating the representative power/limitations of $Q(\\tau)$\n\n\nSome misprints: \"Monte-Calro\" figure text of Figure 2 line 4.\n\nLIMITATIONS:\nLimitations has been somewhat discussed (see weaknesses for omitted discussion of possible limitation).", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-05"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nAuthors proposed GeoPhy as a fully differentiable approach for phylogenetic inference, addressing a fundamental problem in phylogenetic inference. In experiments with real benchmark datasets, GeoPhy demonstrated its superior performance compared to other methods when considering all topological candidates. This approach is of interest to the general ML community.\n\nSTRENGTHS:\nOverall, the manuscript is well-written.\n\nWEAKNESSES:\nIt is highly recommended to include an algorithm block summarizing the GeoPhy method from input to output for clarity. Additionally, there are some details in the experimental section that need clarification.\n\nQUESTIONS:\nMajor questions:\n\n1. In Section 3.1, the neighbor-joining algorithm is defined as the projection from the distance matrix to the tree topology. However, other standard algorithms, such as UPGMA, could be alternative approaches. Have you investigated if the choice of projection leads to different performance in the benchmark tests?\n\n2. Appendix C.5 provides a rough estimation of the running time, but it is unclear how the other baseline models included in the benchmark tests perform in terms of time. Furthermore, how does the algorithm scale as the number of tips and sites increase? Have you observed any trends in the time curve with respect to the number of taxa and sites?\n\n3. In Table 2, GeoPhy achieves superior performance in the benchmark tests based on likelihoods. Did the optimized topologies look significantly different across the different methods?\n\n4. How were the per-step Monte Carlo samples (K) determined in the benchmark tests? There is no clear winner based on Table 1.\n\n5. Could you elaborate on the approaches used to calculate the marginal log-likelihood (MLL) for all the listed methods?\n\nLIMITATIONS:\nAuthors briefly mentioned two limitations of the proposed model: 1) Q(z) can be improved to enhance its expressive power, and 2) the computation cost per update step can be optimized to speed up the overall inference process. These practical concerns can be further addressed in future studies.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposed a family of implicit distribution over tree topologies which allows support free variational Bayesian phylogenetic inference. The distribution is constructed based on the neighbor joining algorithm which maps a distribution over the tip node vectors to the tree topology space. Both Euclidean and hyperbolic spaces are considered for the tip node distributions. As the resulting distribution over tree topologies is intractable, an auxiliary reverse distribution was introduced which leads to a looser lower bound for optimization. Various gradient estimators were investigated for the tree topology variational parameters optimization. Experiments on a benchmark of real data problems demonstrated the advantage over other variational approaches (mainly sequential monte carlo based) that does not require preselected tree topololgies.\n\nSTRENGTHS:\n1. The paper is written clearly and well organized. The problem of constructing flexible families of distributions over tree topologies is important for the phylogenetic communities and would to of interest to many practitioners.\n2. In addition to a careful derivation to the training objective, the authors also investigated different gradient estimators for tree topology variational parameters.\n\nWEAKNESSES:\n1. As admitted by the authors, the proposed variational approximation $Q(z)$ and the reverse distribution $R(z|\\tau)$ is simple, which may damage the overall approximation quality of the method. \n2. In the derivation of the lower bound, the reverse distribution $R(z|\\tau)$ needs to satisfy the constraint: the support of $R(z|\\tau)$ should be inside the region $\\{z: \\tau(z)=\\tau\\}$.\n2. In the experiments, it seems that the performance would be really sensitive to the choice of gradient estimators. I wonder if that is related to the small sample size $K$? The marginal likelihood estimates also tends to have large variance in most cases.\n3. It may be hard to assess the performance on the tree topology posterior estimation as the distribution is implicit and defined in a rather sutble way.\n\nQUESTIONS:\n1. In equation 10 and equaiton 12, the expection of the last term is actually zero. Any idea why this term is kept there given that other control variates are considered?\n2. How many samples are used to compute the marginal likelihood estimates?\n3. The current construction of the reverse distribution $R(z|\\tau)$ seems not take the constraints $\\tau(z)=\\tau$ into consideration directly. Would this have a negative impact on performance?\n4. I note that only low dimension tip node vectors ($d\\leq 4$) are considered in this paper. Would higher dimension embedding be helpful?\n5. How about the tree topology approximation given by the implicit distribution? Although it may be hard to estimate the density, one can also report summary statistics from the sampled trees, for example.\n\nLIMITATIONS:\nyes, they do.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-18"}