GeoPhy: Differentiable Phylogenetic Inference via Geometric Gradients of Tree
Topologies

Takahiro Mimori 1 2 Michiaki Hamada 1 3

Abstract

As an alternative approach to the conventional MCMCs,
Zhang & Matsen IV (2019) proposed a variational Bayesian
approach termed VBPI, which has subsequently been improved in the expressive powers of topology-dependent
branch length distributions (Zhang, 2020; 2023). Althoughthese methods have presented accurate joint posterior distributions of topology and branch lengths for real datasets, theyrequired reasonable preselection of candidate tree topologies
to avoid a combinatorial explosion in the number of weightparameters beforehand. There have also been proposed variational approaches (Moretti et al., 2021; Koptagel et al.,2022) on top of the combinatorial sequential Monte-Carlo
method (CSMC; Wang et al. (2015)), where topologies andtheir weights were iteratively updated without the need forthe preselection steps. However, the fidelity of the joint
posterior distributions was still largely behind MCMC andVBPI as reported in Koptagel et al. (2022).

Phylogenetic inference, grounded in molecular
evolution models, is essential for understanding evolutionary relationships in biological data.
While variational Bayesian methods offer scalable models for biological analysis, reliable inference for latent tree topology and branch lengths
remains challenging due to the vast possibilities
for topological candidates. In response, we introduce GeoPhy, a novel approach that employs afully differentiable formulation of phylogenetic
inference, representing topological distributions
in continuous geometric spaces without limiting
topological candidates. In experiments using realbenchmark datasets, GeoPhy significantly outperformed other approximate Bayesian methods thatconsidered whole topologies. 1

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 to optimize a variational distribution of the tree topology and branch lengths to approach
the posterior distribution, without preselection of candidate
topologies. In our experiments using real biological datasets,
we demonstrate that GeoPhy significantly outperforms otherapproaches, all without topological restrictions, in terms of
the fidelity of the marginal log-likelihood (MLL) estimatesto gold-standard provided with long-run MCMCs.

1. Introduction
Phylogenetic inference, the reconstruction of tree-structured
evolutionary relationships between biological units, such as
genes, cells, individuals, and species ranging from viruses tomacro-organisms, is a fundamental problem in biology. As
the phylogenetic relationships are often indirectly inferredfrom molecular observations, including DNA, RNA andprotein 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 thephylogenetic inference has typically been an Markov-chain
Monte Carlo (MCMC) method (Ronquist et al., 2012), enhanced with domain-specific techniques such as a mixed
strategy for efficient exploration of tree topologies.

2. Background2.1. Phylogenetic models
Let τ 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 relationshipbetween N species, which is inferred from molecular data,such as DNA, RNA or protein sequences obtained for the
species. Let Y = {Yij ∈ Ω}1≤i≤N,1≤j≤M be a set ofaligned sequences with length M from the species, where
Yij denote a character (base) of the i-th sequence at j-th

1

Waseda University, Tokyo, Japan 2 RIKEN AIP, Tokyo,
Japan 3 AIST, Tokyo, Japan. Correspondence to: Takahiro Mimori <takahiro.mimori@aoni.waseda.jp>, Michiaki Hamada<mhamada@waseda.jp>.
Published at the Differentiable Almost Everything Workshop of the
40 th International Conference on Machine Learning, Honolulu,
Hawaii, USA. July 2023. Copyright 2023 by the author(s).
1This study’s complete version is currently under review.

1

A differentiable approach to Bayesian phylogenetic inference without topological preselections

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 ofthe sequences P (Y |τ, Bτ ) is determined based on evolutionary assumptions. In this study, we follow a common
practice for method evaluations (Zhang & Matsen IV, 2019;
Zhang, 2023) as follows: Y is assumed to be generated
from a Markov process along the branches of τ in a siteindependent manner; The base mutations are assumed tofollow the Jukes-Cantor model (Jukes et al., 1969). The loglikelihood ln P (Y |τ, Bτ ) can be calculated using Felsenstein’s pruning algorithm (Felsenstein, 1973), which is also
known as the sum-product algorithm, and differentiable withrespect to Bτ .

τ(z)

0.4

z1

1.5

0.4

1.6

Bτ ∼ Q(Bτ | τ)
z2

Q(z2)

z2

Figure 1. The proposed scheme for constructing a variational distribution Q(τ, Bτ ) of a tree topology and branch lengths by using
a distribution Q(z) defined on the continuous space. Left: The
marginal distributions of four tip nodes Q(z1 ), . . . , Q(z4 ). Middle: Coordinates of the tip nodes z = {z1 , . . . , z4 } sampled fromthe distribution Q(z), and a tree topology τ (z) determined fromz. Right: A set of branch lengths Bτ (red figures) of the tree τ issampled from a topology dependent distribution Q(Bτ |τ ).

(Fig. 1). Note that we have overloaded τ to represent both avariable and function for notational simplicity. An exampleof the representation space Z is a product of Euclidean
spaces RN ×d or hyperbolic spaces HN ×d (Appendix A),
where d denotes the dimension of each tip’s representation
coordinate. For the link function, we can use τ (z) = TNJ ◦D(z), where D : Z → RN ×N denotes a function thattakes N coordinates and provides a distance matrix between
those based on a geometric measure such as the Euclidean
or hyperbolic distance. TNJ : RN ×N → T denotes a mapthat takes this distance matrix and generates an unrooted
binary tree topology of their phylogeny, determined using
the Neighbor-Joining (NJ) algorithm (Saitou & Nei, 1987).
While the NJ algorithm offers a rooted binary tree topology
accompanied by estimated branch lengths, we only use thetopology information and remove the root node from it toobtain the unrooted tree topology τ ∈ T .

(1)

where DKL , Q(τ ) and Q(Bτ |τ ) denote the KullbackLeibler 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 & Matsen
IV (2019), which has been successively improved for the
expressiveness of Q(Bτ |τ ) (Zhang, 2020; 2023). For the
expression of variational topology mass function Q(τ ), they
all rely on a subsplit Bayesian network (SBN) (Zhang &Matsen IV, 2018), 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 setof tree topologies, which hinders an end-to-end optimization strategy of the distribution over all the topologies andbranch lengths.

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

3. Proposed methods
3.1. Geometric representations of tree topology
ensembles

L[Q] = EQ(z)Q(Bτ |τ (z))

Considering the typically infeasible task of parameterizing
the probability mass function of unrooted tree topologies
T , which requires (2N − 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, asfollows:
Q(τ ) := EQ(z) [I[τ = τ (z)]],

z1

z3
z4

1.6

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

Q(z4)

z ∼ Q(z1, …, z4)

2.2. Variational Bayesian phylogenetic inference

min DKL (Q(τ )Q(Bτ |τ )∥P (τ, Bτ |Y )) ,

z3

Q(z3)
Q(z1)




P (Y, Bτ |τ (z))P (τ (z)).
ln
Q(Bτ |τ (z))Q(τ (z))(3)

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, theevaluation of ln Q(τ ) remains still infeasible with small
samples according to the definition (2). We resolve thisissue by introducing the second lower bound with respect to(2)

where τ : Z → T denotes a deterministic link functionthat maps N coordinates to the corresponding tree topology2z4

A differentiable approach to Bayesian phylogenetic inference without topological preselections

a conditional variational distribution R(z|τ ) as follows:

Variational distributions To investigate the basic effectiveness of GeoPhy algorithm, we employ simple constructions for the variational distributions Qθ (z), Qϕ (Bτ |τ ), andRψ (z|τ ). We use an independent distribution for each tip
QNnode coordinate, i.e. Qθ (z) = i=1 Qθi (zi ), where we use
a d-dimensional normal or wrapped normal distribution (Appendix A) for the coordinates of each tip node zi . For theconditional distribution of branch lengths given tree topology, Qϕ (Bτ |τ ), we use the diagonal lognormal distribution
whose location and scale parameters are parameterized asfunctions of the unique features defined for each topology τ(Appendix B), as proposed in Zhang (2023). For the modelof Rψ (z|τ ), we also employ an independent distribution:
QN
Rψ (z|τ ) = i=1 Rψi (zi |τ ), where, we use the same typeof distribution as Qθi (zi ), independent of τ .

L[Q, R] =
P (Y, Bτ |τ (z))P (τ (z))R(z|τ (z)).
EQ(z)Q(Bτ |τ (z)) ln
Q(Bτ |τ (z))Q(z)(4)

Proposition 3.1. The relation ln P (Y ) ≥ L[Q] ≥ L[Q, R]
holds, where the first and second equality holds whenQ(τ, Bτ ) = P (τ, Bτ |Y ) and R(z|τ ) = Q(z|τ ), respectively.
The proof is provided in Appendix C. Similar to Burdaet al. (2016), we can also derive a tractable importanceweighted 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. For the MLL estimates, we use 1000
Monte Carlo samples, similar to (Zhang & Matsen IV, 2019;
Zhang, 2023). The details and derivations are described inAppendix C.

4. Related work
Differentiability for discrete optimization Discrete optimization problems often suffer from the lack of informative gradients of the objective functions. To address thisissue, continuous relaxation for discrete optimization has
been actively studied, such as a widely-used reparameterization trick with the Gumbel-softmax distribution (Jang et al.,2016; Maddison et al., 2016). Beyond categorical variables,
recent approaches have further advanced the continuous
relaxation techniques to more complex discrete objects, including spanning trees (Struminsky et al., 2021). However,
it is still nontrivial to extend such techniques to the casewith binary tree topologies. As outlined in equation (2), 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 byvariance reduction techniques.

3.3. Differentiable phylogenetic inference with GeoPhy
Here, we introduce the parameterized variational distributions, Qθ (z), Qϕ (Bτ |τ ), Rψ (z|τ ), to facilitate thegradient-based maximization of the variational objective
L[Qθ,ϕ , Rψ ]. We term this framework GeoPhy, since it allows us to optimize the distribution over entire phylogenetic
trees within continuous geometric spaces.
Gradient estimators and variance reduction An unbiased estimator of the gradient ∇θ L is derived as follows:gbθ = ∇θ ln Qθ (z) · f (z, Bτ ) − ∇θ ln Qθ (hθ (ϵz ))
(k)

(k)

(k)

where we denote ϵz ∼ pz z (k) = hθ (ϵz ), ϵB(k)(k)
pB , Bτ = hϕ (ϵB , τ (z (k) )), andf (z, Bτ ) := ln(5)

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 (Monath et al., 2019; Chami
et al., 2020; Chien et al., 2022) have been proposed basedon Dasgupta’s cost function (Dasgupta, 2016). In particular,
Chami et al. (2020) proposed to decode tree topology fromhyperbolic 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 thespecific form of the cost functions.

∼

P (Y, Bτ |τ (z))+ ln P (τ (z))Rψ (z|τ (z)).
Qϕ (Bτ |τ (z))

Note that we explicitly distinguish z and hθ (ϵz ) to indicate
the target of differentiation with respect to θ. Unlike thegradient terms associated with ϕ and ψ, the optimization of
L suffers from the high variance of the gradient gbθ , whicharises from the term being proportional to the score function∇θ Qθ (z). To address this high gradient variance issue, wehave explored several control variates. These are based on aleave-one-out (LOO) baseline estimation that uses multiple
Monte-Carlo samples (Kool et al., 2019; Richter et al., 2020)and surrogate functions (Grathwohl et al., 2018), as detailedin Appendix C.

Phylogenetic analysis in hyperbolic space The approachof embedding phylogenetic trees into hyperbolic spaces
has been explored for visualization and an interpretation ofnovel samples with existing phylogeny (Matsumoto et al.,2021; Jiang et al., 2022). For the inference task, a maximum3

A differentiable approach to Bayesian phylogenetic inference without topological preselections

MicMurMicBer

HomSap
PanTroMicMyoMicGri MicMur

GorGor
PonPyg

MicRav

MirCoq

DauMadDidVirCanFam
FelCatEquCabUrsAmeRhiUni
BosTar

GalCraLorTar

MicRav
MacMulCalJacHapGri LemCat

MicTav
MicRuf2
MicSamCheMedLepEdw

UrsAme CanFam
FelCatDauMadDidVir
ProTat

VarVar
EulMonHapGri
LemCatBosTar

HipAmp
BalPhyMirCoq
GalCraLorTarLepEdw

VarVar
ProTat
EulMon

MicRuf1

MicBer

PonPyg

MicTav

CheMedCalJacPhyMac

HomSap
GorGor
PanTroMicRuf1

MicRuf2

MicSamMacMultwo Q(z) configurations: a wrapped normal distribution
WN with a full 4-dimensional covariance matrix and a2-dimensional diagonal matrix, along with three control
variate choices. The consistently superior performance ofthe larger model, WN (full,4), compared to WN (diag,2) isevident across all datasets. Details of other settings, including those of the normal distributions in Euclidean spaces,
are provided in Table 2. There, we noted a slight advantage of wrapped normal distributions over their Euclidean
counterparts. While VBPI-GNN (Zhang, 2023) 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. Other methods including CSMC (Wang et al., 2015), VCSMC (Moretti et al.,2021), ϕ-CSMC (Koptagel et al., 2022), and GeoPhy (ourapproach), address the more challenging general problem
of model optimization by considering all candidate topologies without preselection. Among these methods, GeoPhy
consistently outperforms other CSMC-based methods, evenwhen a less performant WN (full,2) is used for Q(z). This
demonstrates the stability and efficiency of GeoPhy.

MicMyoMicGri

HipAmp
BalPhyPhyMacRhiUni
EquCabFigure 2. P
Superimposed probability densities of topological distributions N
i=1 Q(zi ) up to 100,000, 200,000, and 500,000 steps(MC samples) from the left. For Q(z), we employed a wrapped
normal distribution with a two-dimensional full covariance matrix. The experiments used the DS3 dataset (N = 36). The
majority-rule consensus phylogenetic tree obtained with MrBayesand each step of GeoPhy are shown in blue and red lines, respectively. The center area is magnified by transforming the radius rof the Poincaré coordinates into tanh 2.1r.

likelihood approach was proposed in (Wilson, 2021), which
however assumed a simplified likelihood function of pairwise distances. A recent study (Macaulay et al., 2023) proposed an MCMC-based algorithm for sampling (τ, Bτ ),which were linked from coordinates z using the NJ algorithm (Saitou & Nei, 1987). However, there remained theissue of an unevaluated Jacobian determinant, which poseda challenge in evaluating inference objectives. Given that
we only use topology τ as described in equation (2), thevariational lower bound for the inference can be unbiasedly
evaluated through sampling, as shown in Proposition 3.1.

6. Conclusion
We developed a differential phylogenetic inference algorithm named GeoPhy that performed an approximate
Bayesian inference without the preselection of candidate
topologies. In experiments conducted with real sequence
datasets, GeoPhy consistently outperformed other methods
that took whole topologies into consideration.

5. Experiments

Acknowledgements

We applied GeoPhy for an approximate posterior inferenceof phylogenetic models given biological sequence data of
species. We trained for GeoPhy until one million MonteCarlo tree samples were consumed for the gradient estimation of the lower bound objectives. This number equals
the number of likelihood evaluations (NLEs) and is used
for a standardized comparison of experimental runs (Wanget al., 2015; Zhang & Matsen IV, 2019). More details of theexperimental setup are found in Appendix D.

We appreciate the valuable comments on our study provided
by Dr. Tsukasa Fukunaga of Waseda University and theanonymous reviewers.

References
Burda, Y., Grosse, R. B., and Salakhutdinov, R. Importanceweighted autoencoders. In Bengio, Y. and LeCun, Y.
(eds.), 4th International Conference on Learning Representations, ICLR 2016, San Juan, Puerto Rico, May 2-4,
2016, Conference Track Proceedings, 2016.

To demonstrate the inference performance of GeoPhy, we
compared the marginal log-likelihood (MLL) estimates forthe eight real datasets (DS1-8) compiled in (Lakner et al.,2008) against gold-standard values obtained in the study(Zhang & Matsen IV, 2019) using MrBayes stepping-stone(SS) algorithm (Ronquist et al., 2012; Xie et al., 2011).
Alongside the training steps, a majority consensus tree procured with GeoPhy progressively aligns with those obtained
via MCMC, as demonstrated in Fig. 2.

Chami, I., Gu, A., Chatziafratis, V., and Ré, C. From trees
to continuous embeddings and back: Hyperbolic hierarchical clustering. Advances in Neural Information Processing Systems, 33:15065–15076, 2020.
Chien, E., Tabaghi, P., and Milenkovic, O. Hyperaid: Denoising in hyperbolic spaces for tree-fitting and hierarchical clustering. arXiv preprint arXiv:2205.09721, 2022.

In Table 1, we compiled the MLL estimates for GeoPhy andother approximate Bayesian inference approaches acrosseight datasets. Specifically, we present GeoPhy results for

Dasgupta, S. A cost function for similarity-based hierarchi4

A differentiable approach to Bayesian phylogenetic inference without topological preselections

Table 1. Comparison of the MLL estimates with different approaches in eight benchmark datasets. The MLL values for MrBayes SS andVBPI-GNN are sourced from (Zhang, 2023), while those for CSMC, VCSMC, and φ-CSMSC are referenced from (Koptagel et al., 2022).
Following the literature (Zhang & Matsen IV, 2019), we refer to the dataset DS7 in (Koptagel et al., 2022) as DS8. For each configurationof GeoPhy (WN (full,4) and WN (diag,2)), 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)+. For GeoPhy, the mean MLL values of five independent runs areshown with the standard deviation indicated in parentheses. The bold figures are the best (highest) values obtained with GeoPhy andthe 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 acrossvarious configurations.
Dataset#Taxa (N )#Sites (M )

DS1
27
1949

DS2
29
2520

DS3
36
1812

DS4
41
1137

DS5
50
378

DS6
50
1133

DS7
59
1824

DS8
64
1008

MrBayes SS−7108.42(0.18)−26367.57(0.48)−33735.44(0.5)−13330.06(0.54)−8214.51(0.28)−6724.07(0.86)−37332.76(2.42)−8649.88(1.75)

VBPI-GNN−7108.41(0.14)−26367.73(0.07)−33735.12(0.09)−13329.94(0.19)−8214.64(0.38)−6724.37(0.4)−37332.04(0.26)−8650.65(0.45)

CSMC−8306.76(166.27)−9180.34(170.27)−7290.36(7.23)−27884.37(226.6)−28700.7(4892.67)−30568.49(31.34)−35381.01(218.18)−37211.2(397.97)−33798.06(6.62)−15019.21(100.61)−17106.1(362.74)−13582.24(35.08)−8940.62(46.44)−9449.65(2578.58)−8367.51(8.87)−8029.51(83.67)−9296.66(2046.7)−7013.83(16.99)−11013.57(113.49)−−−9209.18(18.03)−7111.55(0.07)−7119.77(11.80)−7116.09(10.67)−7126.89(10.06)−7130.67(10.67)−7128.40(9.78)−26379.48(11.60)−26368.44(0.13)−26368.54(0.12)−26444.84(27.91)−26380.41(14.40)−26375.28(11.78)−33757.79(8.07)−33736.01(0.03)−33735.85(0.12)−33823.74(15.62)−33737.75(2.48)−33736.91(1.91)−13342.71(1.61)−13339.26(3.19)−13337.42(1.32)−13358.16(9.79)−13346.94(4.25)−13347.32(4.42)−8240.87(9.80)−8234.06(7.53)−8233.89(6.63)−8251.45(9.72)−8239.36(4.62)−8235.41(5.70)−6735.14(2.64)−6733.91(0.57)−6735.90(1.13)−6745.60(8.36)−6741.63(3.23)−6742.40(1.94)

−
−
−
−
−
−

VCSMC
ϕ-CSMCWN (full,4)LOO(3)
LOO(3)+WN (diag,2)
LOO(3)
LOO(3)+

cal clustering. In Proceedings of the forty-eighth annual
ACM symposium on Theory of Computing, pp. 118–127,
2016.

−37377.86(29.48)−37350.77(11.74)−37358.96(13.06)−37516.88(69.88)−37382.28(31.96)−37411.28(56.74)−8663.51(6.85)−8671.32(5.99)−8660.48(0.78)−8719.44(60.54)−8690.41(15.92)−8683.22(13.13)

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

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

Jang, E., Gu, S., and Poole, B. Categorical reparameterization with gumbel-softmax. arXiv preprint
arXiv:1611.01144, 2016.

Garey, J. R., Near, T. J., Nonnemacher, M. R., and Nadler,
S. A. Molecular evidence for acanthocephala as a subtaxon of rotifera. Journal of Molecular Evolution, 43:
287–292, 1996.

Jiang, Y., Tabaghi, P., and Mirarab, S. 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, pp. 68–85. Springer, 2022.

Grathwohl, W., Choi, D., Wu, Y., Roeder, G., and Duvenaud, D. Backpropagation through the void: Optimizing
control variates for black-box gradient estimation. In
International Conference on Learning Representations,2018.

Jukes, T. H., Cantor, C. R., et al. Evolution of protein
molecules. Mammalian protein metabolism, 3:21–132,
1969.

Hedges, S. B., Moberg, K. D., and Maxson, L. R. 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.

Kingma, D. P. and Ba, J. Adam: A method for stochastic
optimization. arXiv preprint arXiv:1412.6980, 2014.
Kool, W., van Hoof, H., and Welling, M. Buy 4 reinforce
samples, get a baseline for free! 2019.
5

A differentiable approach to Bayesian phylogenetic inference without topological preselections

Koptagel, H., Kviman, O., Melin, H., Safinianaini, N., andLagergren, J. Vaiphy: a variational inference based algorithm for phylogeny. In Advances in Neural Information
Processing Systems, 2022.

Saitou, N. and Nei, M. The neighbor-joining method: a new
method for reconstructing phylogenetic trees. Molecularbiology and evolution, 4(4):406–425, 1987.
Sala, F., De Sa, C., Gu, A., and Ré, C. Representation
tradeoffs for hyperbolic embeddings. In International
conference on machine learning, pp. 4460–4469. PMLR,
2018.

Lakner, C., Van Der Mark, P., Huelsenbeck, J. P., Larget,
B., and Ronquist, F. Efficiency of markov chain montecarlo tree proposals in bayesian phylogenetics. Systematic
biology, 57(1):86–103, 2008.

Struminsky, K., Gadetsky, A., Rakitin, D., Karpushkin, D.,
and Vetrov, D. P. Leveraging recursive gumbel-max trick
for approximate inference in combinatorial spaces. Advances in Neural Information Processing Systems, 34:
10999–11011, 2021.

Macaulay, M., Darling, A. E., and Fourment, M. Fidelityof hyperbolic space for bayesian phylogenetic inference.
PLoS computational biology, 2023.
Maddison, C. J., Mnih, A., and Teh, Y. W. The concrete
distribution: A continuous relaxation of discrete randomvariables. arXiv preprint arXiv:1611.00712, 2016.

Wang, L., Bouchard-Côté, A., and Doucet, A. Bayesian
phylogenetic inference using a combinatorial sequential
monte carlo method. Journal of the American Statistical
Association, 110(512):1362–1374, 2015.

Matsumoto, H., Mimori, T., and Fukunaga, T. Novel metricfor hyperbolic phylogenetic tree embeddings. Biology
Methods and Protocols, 6(1):bpab006, 2021.

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

Monath, N., Zaheer, M., Silva, D., McCallum, A., andAhmed, A. Gradient-based hierarchical clustering using
continuous representations of trees in hyperbolic space.
In Proceedings of the 25th ACM SIGKDD International
Conference on Knowledge Discovery & Data Mining, pp.
714–722, 2019.

Xie, W., Lewis, P. O., Fan, Y., Kuo, L., and Chen, M.-H.
Improving marginal likelihood estimation for bayesian
phylogenetic model selection. Systematic biology, 60(2):150–160, 2011.
Yang, Z. and Yoder, A. D. Comparison of likelihood andbayesian 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.

Moretti, A. K., Zhang, L., Naesseth, C. A., Venner, H., Blei,D., and Pe’er, I. Variational combinatorial sequential
monte carlo methods for bayesian phylogenetic inference.
In Uncertainty in Artificial Intelligence, pp. 971–981.
PMLR, 2021.

Yoder, A. D. and Yang, Z. Divergence dates for malagasy
lemurs estimated from multiple gene loci: geological andevolutionary context. Molecular Ecology, 13(4):757–773,
2004.

Nagano, Y., Yamaguchi, S., Fujita, Y., and Koyama, M.
A wrapped normal distribution on hyperbolic space forgradient-based learning. In International Conference on
Machine Learning, pp. 4693–4702. PMLR, 2019.

Zhang, C. Improved variational bayesian phylogenetic inference with normalizing flows. In Larochelle, H., Ranzato, M., Hadsell, R., Balcan, M. F., and Lin, H. (eds.),
Advances in Neural Information Processing Systems, volume 33, pp. 18760–18771. Curran Associates, Inc., 2020.

Richter, L., Boustati, A., Nüsken, N., Ruiz, F., and Akyildiz,
O. D. Vargrad: a low-variance gradient estimator forvariational inference. Advances in Neural Information
Processing Systems, 33:13481–13492, 2020.
Ronquist, F., Teslenko, M., Van Der Mark, P., Ayres, D. L.,
Darling, A., Höhna, S., Larget, B., Liu, L., Suchard,
M. A., and Huelsenbeck, J. P. Mrbayes 3.2: efficientbayesian phylogenetic inference and model choice acrossa large model space. Systematic biology, 61(3):539–542,
2012.

Zhang, C. Learnable topological features for phylogenetic
inference via graph neural networks. In The EleventhInternational Conference on Learning Representations,2023.
Zhang, C. and Matsen IV, F. A. Generalizing tree probability
estimation via bayesian networks. In Bengio, S., Wallach,
H., Larochelle, H., Grauman, K., Cesa-Bianchi, N., andGarnett, R. (eds.), Advances in Neural Information Processing Systems 31, pp. 1449–1458. Curran Associates,
Inc., 2018.

Rossman, A. Y., McKemy, J. M., Pardo-Schultheiss, R. A.,
and Schroers, H.-J. Molecular studies of the bionectriaceae using large subunit rdna sequences. Mycologia, 93(1):100–110, 2001.
6

A differentiable approach to Bayesian phylogenetic inference without topological preselections

Zhang, C. and Matsen IV, F. A. Variational bayesian phylogenetic inference. In International Conference on Learning Representations, 2019.
Zhang, N. and Blackwell, M. Molecular phylogeny of
dogwood anthracnose fungus (discula destructiva) andthe diaporthales. Mycologia, 93(2):355–365, 2001.

7

A differentiable approach to Bayesian phylogenetic inference without topological preselections

A. Hyperbolic spaces and wrapped normal distributions
Hyperbolic spaces are known to be able to embed hierarchical data with less distortion in considerably fewer dimensions
than Euclidean spaces (Sala et al., 2018). The wrapped normal distribution, as proposed in (Nagano et al., 2019), 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 Hd , represents the d-dimenstional hyperbolic space as a submanifold of a d + 1 dimensional
Euclidean space. Given u, v ∈ Rd+1 , we can define the pseudo-inner product and pseudo-norm as follows:⟨u, v⟩L := −u0 v0 +dX

uj v j ,j=1∥u∥L :=

q⟨u, u⟩L .

(6)

The distance between hyperbolic coordinates ν, µ ∈ Hd is defined as follows:d(ν, µ) := cosh−1 (− ⟨ν, µ⟩L ),(7)

where cosh−1 denotes the inverse function of the hyperbolic cosine function. Consider hyperbolic coordinates ν, µ ∈ Hd
and tangent vectors u ∈ Tµ Hd and v ∈ Tν Hd . An exponential map expµ (u) ∈ Hd , a logarithm map logµ (ν) ∈ Tµ Hd , anda parallel transport map PTν→µ (v) ∈ Tµ Hd , can be calculated as follows:u,(8)
expµ (u) = cosh(∥u∥L )µ + sinh(∥u∥L )∥u∥Lcosh−1 (α)(ν − αµ) ,logµ (ν) = √α2 − 1⟨µ − αν, v⟩LPTν→µ (v) = v +(ν + µ),
α+1

(9)
(10)

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

Given location and scale parameters denoted as µ ∈ Hd and Σ ∈ Rd×d , respectively, the procedure for sampling from awrapped normal distribution z ∼ WN (µ, Σ) defined over Hd is given as follows:z = expµ ◦PTµo →µ (u),u1:d ∼ N (0, Σ),(11)where µo = (1, 0, . . . , 0)⊤ denotes the origin of the Hd . Note that we set u0 = 0, and u := u0:d ∈ Tµo Hd represents atangent vector at the origin Tµo Hd .
From the sampling definition in equation (11), the probability density function WN (z; µ, Σ) can be derived as follows:sinh ∥u∥L
log WN (z; µ, Σ) = log N (u; 0, Σ) − (d − 1) ln
,
(12)
∥u∥L
where u is defined as u = PTµ→µo ◦ logµ (z). For detailed derivation, we refer to Appendix A of (Nagano et al., 2019).

B. 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 (Zhang, 2023). 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 topologicalfeatures (LTFs) of the topology τ , which are computed with a method described in (Zhang, 2023). Below, we summarize anarchitecture 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(0)
unrooted binary tree topology τ . The input to the GNN consists of node features represented by LTFs denoted as {hv }v∈Vτ .
These features undergo transformation L times as follows:(L)◦ · · · ◦ g (1) ({h(0) }v∈Vτ ),{hv(L) }v∈Vτ = GNN({h(0)v }v∈Vτ ) = g

8

(13)

A differentiable approach to Bayesian phylogenetic inference without topological preselections

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:ehv = MLPV (hv(L) ),me (v,u) = MAX(ehv , ehu ),(∀v ∈ Vτ )(14)(∀(v, u) ∈ Eτ )(15)(∀(v, u) ∈ Eτ )µ(v,u) , log σ(v,u) = MLPE (me (v,u) )(16)

where MLPN , MAX, and MLPE denotes a multi-layer perceptron for node features with two hidden layers, the elementwise 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 MLPN and MLPE , 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 (Zhang, 2023), 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 {hv }v∈Vτ according to the following scheme:{hv(ℓ+1) }v∈Vτ = g (ℓ) ({h(ℓ)
v }v∈Vτ ),(17)

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:


(ℓ)
(ℓ)
(ℓ)
e(ℓ)
, ∀u ∈ Nτ (v)
(18)
h(ℓ)
u→v = MLPv ∥hu − hvh′(ℓ+1)= AGG(ℓ) e(ℓ)
u→v ,vu∈Nτ (v)
hv(ℓ+1) = ELU hv′(ℓ+1) ,

(19)

(20)

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 eu→v , ∀u ∈ N (v)for each element.

C. Variational Lower Bounds and Gradient Estimators
C.1. Variational lower boundIn Proposition C.1 we will demonstrate that the following functional is a lower bound of the marginal log-likelihoodln P (Y ):L[Q, R] := EQ(z,Bτ ) [ln F ′ (z, Bτ )] = EQ(z) EQ(Bτ |z) [ln F (z, Bτ )] − ln Q(z)(21)≤ ln P (Y ),(22)
where F and F ′ are respectively defined as follows:
F (z, Bτ ) :=P (Y, Bτ |τ (z))
P (τ (z))R(z|τ (z)),
Q(Bτ |τ (z))

F ′ (z, Bτ ) :=F (z, Bτ )
.
Q(z)

(23)

Proposition C.1 (Restatement of Proposition 3.1). The relation ln P (Y ) ≥ L[Q] ≥ L[Q, R] holds, where the first andsecond 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:
P (Y, τ, Bτ )ln P (Y ) ≥ ln P (Y ) − DKL [Q(τ, Bτ )∥P (τ, Bτ |Y )] = EQ(τ,Bτ ) ln:= L[Q],
Q(τ, Bτ )9(24)

A differentiable approach to Bayesian phylogenetic inference without topological preselections

where the equality condition of the first inequality holds when Q(τ, Bτ ) = P (τ, Bτ |Y ). Since we have defined Q(τ ) inequation (2), we can further transform the lower bound as L[Q]"#X
P (Y, Bτ |τ )P (τ )+ ln
I[τ = τ (z)] EQ(Bτ |τ ) ln,
(25)
L[Q] = EQ(z)
Q(Bτ |τ )
Q(τ )
τ ∈T

from which equation (3) 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 − EQ(z) [ln Q(τ (z))] = − EQ(τ ) [ln Q(τ )] in equation (3) can be transformed to derive further lower
bound as follows:− EQ(τ ) [ln Q(τ )] ≥ − EQ(τ ) [ln Q(τ ) + DKL (Q(z|τ )∥R(z|τ ))]




Q(z)
Q(τ )Q(z|τ )= − EQ(z) ln,= − EQ(τ )Q(z|τ ) ln
R(τ |z)R(z|τ (z))
P
where the last equality is derived by using the relation EQ(τ )Q(z|τ ) [·] = EQ(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.
C.2. Gradient estimators for variational lower boundThe gradient of L[Qθ,ϕ , Rψ ] with respect to θ is given by∇θ L = ∇θ EQθ (z) EQφ (Bτ |z) [ln F (z, Bτ )] − ln Qθ (z)





P (Y, Bτ |τ (z))= EQθ (z) (∇θ ln Qθ (z)) EQφ (Bτ |τ (z)) ln+ ln P (τ (z))Rψ (z|τ (z))
Qϕ (Bτ |τ (z))+ ∇θ H[Qθ (z)],

(26)

(27)

where H denotes the differential entropy. We assume that Qϕ (Bτ |τ ) is reparameterizable as in (Zhang, 2023): namely,Bτ can be sampled through Bτ = hϕ (ϵB , τ ), where ϵB ∼ pB (ϵB ), where pB (ϵ) and hϕ denote a parameter-free basedistribution and a differentiable function with ϕ, respectively. Consequently, the gradient of L with respect to ϕ is evaluatedas follows:∇ϕ L = ∇ϕ EQθ (z) EQφ (Bτ |z) [ln F (z, Bτ )]
P (Y, Bτ = hϕ (ϵB , τ )|τ (z))= EQθ (z) EpB (ϵB ) ∇ϕ ln.
Qϕ (Bτ = hϕ (ϵB , τ )|τ (z))

(28)
(29)

Lastly, the gradient of L with respect to ψ can be evaluated with a tractable density model Rψ (z|τ ) as follows:∇ψ L = ∇ψ EQθ (z) EQφ (Bτ |z) [ln F (z, Bτ )] = EQθ (z) [∇ψ ln Rψ (z|τ (z))] .

(30)

Given samples ϵz ∼ pz and ϵB ∼ pB , 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:gbθ = ∇θ ln Qθ (z) · ln F (z, Bτ ) − ∇θ ln Qθ (hθ (ϵz )),(31)

gbψ = ∇ψ ln F (z, hϕ (ϵB , τ (z))) = ∇ψ ln Rψ (z|τ (z)).

(33)

gbϕ = ∇ϕ ln F (z, hϕ (ϵB , τ (z))) = ∇ϕ ln

P (Y, hϕ (ϵB , τ (z))|τ (z)),
Qϕ (hϕ (ϵB , τ (z))|τ (z))(32)

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

Lb′ = ln Qθ (z) · detach[f (z, Bτ )] + f (z, hϕ (ϵB , τ (z))) − ln Qθ (hθ (ϵz )),(34)

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))).
10

A differentiable approach to Bayesian phylogenetic inference without topological preselections

C.3. Multi-sample gradient estimators
(k)

(k)

(k)

(k)

Given a K set of Monte-Carlo (MC) samples from Qθ,ϕ (z, Bτ ), i.e. {ϵZ , z (k) = hθ (ϵZ )}K
k=1 and {ϵB , Bτ(k)hϕ (ϵB , τ (z (k) ))}K,
wecansimplyestimate
∇L
[Q
,
R
]
as
follows:
θ
θ,ϕ
ψ
k=1

1 X
∇θ ln Qθ (z (k) ) · f (z (k) , Bτ(k) ) − ∇θ ln Qθ (hθ (ϵ(k)
z )) .
K

=

K

(K)
gbθ =

(35)

k=1

As a simple extension of equation (34), the gradients are obtained through an auto-gradient computation of the following
target:
(K)
1
Lb′
=
K

K 
X

k=1


(k)
ln Qθ (z (k) ) · detach[f (z (k) , Bτ(k) )] + f (z (k) , hϕ (ϵB , τ (z (k) ))) − ln Qθ (hθ (ϵ(k)
z )) ,(36)

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

For the term of K-sample gradient estimator gbθ proportional to the score function ∇θ ln Qθ , a leave-one-out (LOO)
variance reduction is known to be effective (Kool et al., 2019; Richter et al., 2020), which is denoted as follows:(K)

gbLOO,θ =

K


i
1 Xh
∇θ ln Qθ (z (k) ) · f (z (k) , Bτ(k) ) − fk (z (\k) , Bτ(\k) ) − ∇θ ln Qθ (hθ (ϵ(k)
z )) ,
K

(37)

k=1

where fk denotes:

fk (z (\k) , Bτ(\k) ) :=

1
K −1

K
X

′

′

f (z (k ) , Bτ(k ) ).

(38)

k′ =1,k′ ̸=k

To employ the LOO gradient estimator for θ, the target of auto-gradient computation in equation (36) needs to be adjusted
as follows:(K)
1
Lb′ LOO =
K

K 
X

k=1

ln Qθ (z (k) ) · detach[f (z (k) , Bτ(k) ) − fk (z (\k) , Bτ(\k) )](k)+f (z (k) , hϕ (ϵB , τ (z (k) ))) − ln Qθ (hθ (ϵ(k)
))
,
z

(39)

C.5. LAX estimators for adaptive variance reduction
The LAX estimator (Grathwohl et al., 2018) is a stochastic gradient estimator based on a surrogate function, which canbe adaptively learned to reduce the variance regarding the term ∇θ ln Qθ (z). In our case, the LAX estimator is given as
follows:
gbLAX,θ := ∇θ ln Qθ (z) · (f (z, Bτ ) − sχ (z)) + ∇θ sχ (hθ (ϵz )).

(40)

gbLAX,θ := ∇θ ln Qθ (z) · (f (z, Bτ ) − sχ (z)) + ∇θ sχ (hθ (ϵz )) − ∇θ ln Qθ (hθ (ϵz )).

(41)

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 gbLAX,θ , we optimize χ to minimize the following objective as proposed in(Grathwohl et al., 2018):
VQθ (z) [bgθ ] :=

nθ
nθ

1 X
1 X
VQθ (z) [bg θi ] =EQθ (z) [bgθ2i ] − EQθ (z) [bg θi ] 2 ,nθ i=1nθ i=111(42)

A differentiable approach to Bayesian phylogenetic inference without topological preselections

where nθ denotes the dimension of θ. As the gradient in equation (40) is given as an unbiased estimator of ∇θ L, which is
not dependent on χ, we can use the relation ∇χ EQθ (z) [bgLAX,θi ] = 0. Therefore, the unbiased estimator of the gradient∇χ VQθ (z) [bgLAX,θ ] is given as follows:
n

gbχ =

θ
1 X
2
∇χ gbLAX,θ
.
inθ i=1(43)

As we require the gradient of ∇θ L with respect to χ for the optimization, we use different objectives for auto-gradientcomputation with respect to θ and the other parameters ϕ and ψ as follows:
Lb′ LAX,θ = ln Qθ (z) · (detach[f (z, Bτ )] − sχ (z)) + sχ (hθ (ϵz )) − ln Qθ (hθ (ϵz )),Lb′ ϕ,ψ = f (z, hϕ (ϵB , τ (z))).

(44)
(45)

C.6. LAX estimators with multiple MC-samples

For the cases with K MC-samples, we use LAX estimators by differentiating the following objectives:(K)
1
Lb′ LAX,θ =K(K)
1
Lb′ ϕ,ψ =
K

K 
X

k=1
K
X




(k)
ln Qθ (z (k) ) · detach[f (z (k) , Bτ(k) )] − sχ (z (k) ) + sχ (hθ (ϵ(k)
))
−
ln
Q
(h
(ϵ
))
,
θ θ z
z
(k)

f (z (k) , hϕ (ϵB , τ (z (k) ))).

(46)

(47)

k=1

When we combine LAX estimators with LOO control variates. the target for auto-gradient computation changes to thefollowing:(K)
1
Lb′ LOO+LAX,θ =
K

K 
X

k=1



ln Qθ (z (k) ) · detach[f (z (k) , Bτ(k) ) − fk (z (\k) , Bτ(\k) )] − sχ (z (k) )(k)+sχ (hθ (ϵ(k)
))
−
ln
Q
(h
(ϵ
))
.
θ θ z
z

(48)

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

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

An importance-weighted evidence lower bound (IW-ELBO) (Burda et al., 2016), is a tighter lower bound of the log-likelihoodln P (Y ) than ELBO. For our model, a conventional K-sample IW-ELBO is given as follows:
#
"
K
(k)
1 X P (Y, Bτ , τ (z (k) ))(K)
LIW [Q] := EQ(z(1) ,B (1) )···Q(z(K) ,B (K) ) ln
.
(49)
(k)
τ
τ
K
Q(Bτ , τ (z (k) ))
k=1

(K)

The fact that LIW [Q] is the lower bound of ln P (Y ) is directly followed from Theorem 1 in (Burda et al., 2016). 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 3.1. We define the K-sample tractable IW-ELBO as follows:"#K
1 X ′ (k) (k)(K)
LIW [Q, R] := EQ(z(1) ,B (1) )···Q(z(K) ,B (K) ) ln
F (z , Bτ ) ,
(50)
τ
τ
K
k=1

(K)

where F ′ is defined in equation (23). We will prove in Theorem C.3 that LIW [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 3.1, the relationship isless obvious when K > 1. Before delving into that, we prepare the following proposition.
Proposition C.2. Given Q(z, τ ) as defined in equation 2 and an arbitrary conditional distribution R(z|τ ), it follows that
ER(z|τ ) [I[τ = τ (z)]] ≤ 1,
where setting R(z|τ ) = Q(z|τ ) is a sufficient condition for the equality to hold.
12

(51)

A differentiable approach to Bayesian phylogenetic inference without topological preselections

Proof. The inequality immediately follows from the definition as follows:
ER(z|τ ) [I[τ = τ (z)]] ≤ ER(z|τ ) [1] = 1.

(52)

Next, when we set R(z|τ ) = Q(z|τ ), the condition for equality is satisfied as follows:
EQ(z|τ ) [I[τ = τ (z)]] =EQ(z) [I[τ = τ (z)]2 ]
Q(τ )
=
= 1,
Q(τ )
Q(τ )

(53)

where we have used the definition of Q(τ ) := EQ(z) [I[τ = τ (z)]] from equation (2) and the resulting relation Q(z|τ )Q(τ ) =I[τ = τ (z)]Q(z).
Theorem C.3. Given Q(z, τ ) as defined in equation 2 and an arbitrary conditional distribution R(z|τ ), for any natural
number K > 1, the following relation holds:
(K)

(K−1)

ln P (Y ) ≥ LIW [Q, R] ≥ LIW[Q, R].

(54)(K)

Additionally, if F ′ (z, Bτ ) is bounded and ∀τ, ER(z|τ ) [I[τ = τ (z)]] = 1, then LIW [Q, R] approaches ln P (Y ) as K → ∞.
Proof. We first show that for any natural number K > M ,
(K)

(M )

LIW [Q, R] ≥ LIW [Q, R].
(k)

(55)

(k)

For simplicity, we denote Q(z (k) , Bτ ) and F ′ (z (k) , Bτ ) as Qk and Fk′ , respectively, in the following discussion. LetK
UMrepresent a uniform distribution over a subset with M distinct indices chosen from the K indices {1, . . . , K}. Similar
to the approach used in (Burda et al., 2016), we will utilize the following relationship:
K

1 X ′
K
Fk = E{i1 ,...,iM }∼UM
K
k=1

"

#
M
1 X ′
.
F
M m=1 im

(56)

Now, the inequality (55) is derived as follows:"

EQ1 ···QK ln

K

1 X ′
Fk
K
k=1

!#

"

K
= EQ1 ···QK ln E{i1 ,...,im }∼UM

"

"

"

M

1 X ′
F
M m=1 imM

1 X ′F
M m=1 imKln
≥ EQ1 ···QK E{i1 ,...,im }∼UM"

= EQ1 ···QM ln

M

1 X ′
F
M m=1 m

!#

,

!##
!##

(57)

(58)

(59)

where we have also used Jensen’s inequality.
(K)

Next, we show that ln P (Y ) ≥ LIW [Q, R]. We again use Jensen’s inequality as follows:
"

#
K
1 X ′
ln
Fk
K
k=1
#
"
K
1 X ′
Fk = ln EQ(z,Bτ ) [F ′ (z, Bτ )] .
≤ ln EQ1 ···QK
K(K)
LIW [Q, R] = EQ1 ···QK

k=1

13

(60)

(61)

A differentiable approach to Bayesian phylogenetic inference without topological preselections

The last term is further transformed as follows:P (Y, Bτ |τ (z))R(z|τ (z))
ln EQ(z,Bτ ) [F (z, Bτ )] = ln EQ(z,Bτ )Q(Bτ |τ (z))Q(z)P (Y, Bτ |τ (z))R(z|τ (z))= ln EQ(z,Bτ )Q(z, Bτ )
P (Y, τ (z))R(z|τ (z))= ln EQ(z)
Q(z)


X
P (Y, τ ′ )R(z|τ ′ ) ′
EQ(z)
= ln
I[τ = τ (z)]Q(z)τ ′ ∈TX
P (Y, τ ′ ) ER(z|τ ′ ) [I[τ ′ = τ (z)]]
= ln


′

(62)
(63)
(64)
(65)
(66)

τ ′ ∈T≤ lnX

P (Y, τ ′ ) = ln P (Y ),(67)

τ ′ ∈T

where, in the transition from the first to the second row, we employed the following relation:
Q(Bτ |τ (z)) =

X

τ ∈T

Q(Bτ |τ )I[τ = τ (z)] =

X

τ ∈T

Q(Bτ |τ )Q(τ |z) = Q(Bτ |z),(68)

and we have used Proposition C.2 for the last inequality.
Finally, we will show that the following convergence property assuming that F (z, Bτ ) is bounded:(K)
LIW [Q, R] → ln

X

P (Y, τ ′ ) ER(z|τ ′ ) I[τ ′ = τ (z)]τ ′ ∈T

1
From the strong law of large numbers, it follows that K

PK

EQ(zk ,Bk ) [F ′ (zk , Bk )] =!

(K → ∞).

(69)′
k=1 Fk converges to the following term almost surely as K → ∞:X

P (Y, τ ′ ) ER(z|τ ′ ) I[τ ′ = τ (z)],(70)τ ′ ∈T

where we have employed the same transformations as used from equation (62) to (66). Observe that the r.h.s term of theequation (69) equals to ln P (Y ) when ∀τ ′ ∈ T , ER(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(Zhang, 2023). From Theorem C.3, 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 ∀τ, ER(z|τ ) [I[τ = τ (z)]] = 1. According to Proposition C.2, this equality condition issatisfied when we set R(z|τ ) = Q(z|τ ), which is approached by maximizing L[Q, R] with respect to R as indicated in
Proposition 3.1.
C.8. Gradient estimators for IW-ELBO(K)

The gradient of IW-ELBO LIW [Qθ,ϕ , Rψ ] with respect to θ is given by(K)∇θ LIW = EQ (z(1) ,B (1) )···Q (z(K) ,B (K) )
θ,φ
θ,φ
τ
τ

+ EQ

θ,φ (z(K) ,B (K) )(1) ,B (1) )···Q
θ,φ (z
τ
τ

"K
X

wk (z

(1:K)

,B

(1:K)

′

(k)

,B

(1:K)

)∇θ ln F (z

k=1

"K
X

k=1

14

∇θ ln Qθ (z(k)) · ℓ(z(1:K), Bτ(k) )
#

) ,

#
(71)

A differentiable approach to Bayesian phylogenetic inference without topological preselections

where we have defined(k)wk (z (1:K) , Bτ(1:K) ) := PKF ′ (z (k) , Bτ ),′′ (k′ ) , B (k ) )τ
k′ =1 F (z!
K
1 X ′ (k′ ) (k′ )
(1:K)
(1:K)
ℓ(z
, Bτ
) := ln
F (z , Bτ ) .
K ′
k =1

(72)

(73)

Similarly, as F ′ is differentiable with respect to Bτ , and Bτ = hϕ (ϵB , τ ) is differentiable with respect to ϕ, the gradient of(K)
LIW with respect to ϕ can be evaluated as follows:
"K
#
X
(K)
(1:K)
(1:K)
′ (k)
(k)
∇ϕ LIW = EQθ (z(1) )···Qθ (z(K) ) EpB (ϵ(1) )···pB (ϵ(K) )
wk (z
, Bτ
)∇ϕ ln F (z , hϕ (ϵ , τ )) .
(74)
B

B

k=1

(k)

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

gbIW,θ :=

K
X

k=1

h
i
∇θ ln Qθ (z (k) ) · −wk (z (1:K) , Bτ(1:K) ) + ℓ(z (1:K) , Bτ(1:K) ) .

(75)

The remaining gradient estimators are given as follows:(K)

gbIW,ϕ =(K)

gbIW,ϕ =KX

wk (z (1:K) , Bτ(1:K) )∇ϕ ln F (z (k) , hϕ (ϵ(k) , τ )),(76)

wk (z (1:K) , Bτ(1:K) )∇ψ ln F (z (k) , hϕ (ϵ(k) , τ )).

(77)

k=1
K
X

k=1

In total, the target for computing auto-gradient for these gradients is given as follows:′(K)
LbIW =

+

K
X

k=1
K
X

i
h
ln Qθ (z (k) ) · detach −wk (z (1:K) , Bτ(1:K) ) + ℓ(z (1:K) , Bτ(1:K) )(k)detach[wk (z (1:K) , Bτ(1:K) )] ln F (z (k) , hϕ (ϵB , τ (z (k) ))).

(78)

k=1

D. Experimental Details
D.1. 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 (Zhang & Matsen
IV, 2019; Koptagel et al., 2022). For the neural network used in the parameterization of Qϕ (Bτ |τ ), we employed edge
convolutional operation (EDGE), which was well-performed architecture in (Zhang, 2023). For the stochastic gradient
optimizations, we used the Adam optimizer (Kingma & Ba, 2014) with a learning rate of 0.0001.
D.2. 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 thehyperbolic MDS algorithm (hMDS) proposed in (Sala et al., 2018) 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 (Macaulay et al., 2023).
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.
15

A differentiable approach to Bayesian phylogenetic inference without topological preselections

D.3. 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.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 & Matsen IV (2019), 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.
D.4. Variational branch length distribuions
For the variational branch length distribution Qϕ (Bτ |τ ), we followed an architecture of (Zhang, 2023); namely, each branch
length independently followed a lognormal distribution which was parameterized with a graph neural network (GNN).
Details are described in Appendix B.
D.5. LAX estimators
As input features of a surrogate function sχ (z) used in the LAX estimators, we employed a flattened vector of coordinatesz ∈ RN ×d when z resides in Euclidean space. In cases where the coordinates were z ∈ Hd , we first transformed z with alogarithm map logµo z ∈ Tµo Hd , then omitted their constant value 0-th elements and subsequently flattened the result. We
implemented a simple multi-layer perceptron (MLP) network with a single hidden layer of width 10N d and a subsequent
sigmoid linear unit (SiLU) activation function as the neural network to output sχ (z).
D.6. 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 (Zhang & Matsen IV, 2019) and (Koptagel et al., 2022), we replicated the MrBayes SS
runs using MrBayes version 3.2.7a. The script we used is provided below.
BEGIN MRBAYES;
set autoclose=yes nowarn=yes Seed=123 Swapseed=123;
lset nst=1;
prset statefreqpr=fixed(equal);prset brlenspr=Unconstrained:exp(10.0);
ss ngen=10000000 nruns=10 nchains=4 printfreq=1000 samplefreq=100 \savebrlens=yes filename=mrbayes_ss_out;
END;
We incorporated the results in the row named Replication in Table 2, where the values aligned more closely with thosefound in (Zhang & Matsen IV, 2019). We deduced that the prior distribution used in (Koptagel et al., 2022) might havebeen set differently as the current default values of brlenspr are Unconstrained : GammaDir(1.0, 0.100, 1.0, 1.0) 2 , which
deviates from the model assumption used for the benchmarks. We observed that the line brlenspr was not included in thecode provided in Appendix F of (Koptagel et al., 2022). Having been able to replicate the results found in (Zhang & Matsen
IV, 2019), we opted to use their values as a reference in Table 1.
D.7. Visualization of tree topologies
PN
We visualized the sum of the probability densities for tip node distribution i=1 Q(zi ) in Fig. 2 by projecting a hyperbolic
coordinate zi ∈ Hd onto the Poincaré coordinates z ik = zik / (1 + zi0 ) (k = 1, . . . , d), then applying a transformationz i 7→ z ′i = tanh (a ∥z i ∥2 ) · z i / ∥z i ∥2 with a = 2.1 to close up the central region. To display the density Q in the newcoordinates, the Jacobian term was also considered to evaluate the density Q(z ′i ).
2

https://github.com/NBISweden/MrBayes/blob/develop/doc/manual/Manual_MrBayes_v3.2.pdf

16

A differentiable approach to Bayesian phylogenetic inference without topological preselections

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 Hd of the i-th tip node was determined as the location parameterµi ∈ Hd of the wrapped normal distribution Q(zi ) = WN (zi ; µi , Σi ). Let ξu ∈ Hd denotes the coordinate of an interior
node u, we defined ξu by using the Lorentzian centroid operation C (?) as follows:ξeuξu := C({cs }s∈Sτ (u) , {νs }s∈Sτ (u) ) = r DE ,ee− ξu , ξu

(79)

L

P
where ξeu := s∈Sτ (u) νs cs 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 τ , cs := 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 swhere |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 (79).

E. Additional Results
E.1. Marginal log-likelihood estimates for eight datasets
We present more comprehensive results in Table 2, extending upon the data from Table 1. This table presents the marginal
log-likelihood (MLL) estimates obtained with various GeoPhy configurations and other conventional methods for thedatasets DS1-DS8 (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). Once again, GeoPhy demonstrates its superior
performance, consistently outperforming CSMC-based approaches that do not require preselection of tree topologies acrossthe 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 themost effective among the tested configurations.

17A differentiable approach to Bayesian phylogenetic inference without topological preselections

Table 2. Extended results of Table 1 comparing the marginal log-likelihood (MLL) estimates with different approaches in eight benchmark
datasets. The MLL values for MrBayes SS and VBPI-GNN were obtained from (Zhang, 2023), while CSMC, VCSMC, and φ-CSMSC
are referenced from (Koptagel et al., 2022). We also included replicated results for MrBayes SS. The MLL values for our approach(GeoPhy) are presented for a variety of Q(z) configurations, encompassing distribution types (normal N or wrapped normal WN ),embedding dimensions (2 or 4), and the covariance matrix (full or diagonal). Each result features various CVs: LAX with K = 1, LOO
with K = 3 denoted as LOO(3), and a combination of LOO and LAX, denoted as LOO(3)+. The figures highlighted in bold represent
the highest values obtained with GeoPhy and the three CSMC-based methods, all of which perform an approximate Bayesian inference
without the preselection of topologies. We have underlined MLL estimates where GeoPhy outperformed the other CSMC-based methods.
Dataset#Taxa (N )#Sites (M )

DS1
27
1949

DS2
29
2520

DS3
36
1812

DS4
41
1137

DS5
50
378

DS6
50
1133

DS7
59
1824

DS8
64
1008

MrBayes SS−7108.42(0.18)−7107.81(0.25)−26367.57(0.48)−26366.45(0.40)−33735.44(0.5)−33732.79(0.63)−13330.06(0.54)−13328.40(0.48)−8214.51(0.28)−8209.17(0.46)−6724.07(0.86)−6721.54(0.77)−37332.76(2.42)−37331.85(3.08)−8649.88(1.75)−8646.18(1.19)

VBPI-GNN−7108.41(0.14)−26367.73(0.07)−33735.12(0.09)−13329.94(0.19)−8214.64(0.38)−6724.37(0.4)−37332.04(0.26)−8650.65(0.45)

CSMC−8306.76(166.27)−9180.34(170.27)−7290.36(7.23)−27884.37(226.6)−28700.7(4892.67)−30568.49(31.34)−35381.01(218.18)−37211.2(397.97)−33798.06(6.62)−15019.21(100.61)−17106.1(362.74)−13582.24(35.08)−8940.62(46.44)−9449.65(2578.58)−8367.51(8.87)−8029.51(83.67)−9296.66(2046.7)−7013.83(16.99)−11013.57(113.49)−−−9209.18(18.03)−7126.79(10.06)−7130.60(10.66)−7128.36(9.77)−7123.95(12.03)−7120.88(13.15)−7119.81(11.71)−26440.54(28.78)−26375.10(11.75)−26369.93(0.25)−26382.91(16.89)−26368.53(0.05)−26368.49(0.10)−33814.98(20.31)−33737.71(2.32)−33735.91(0.13)−33762.45(10.68)−33736.04(0.09)−33735.92(0.15)−13356.21(9.48)−13345.55(3.26)−13346.03(4.54)−13341.62(3.64)−13338.99(6.08)−13339.79(4.55)−8251.99(9.43)−8236.99(6.29)−8236.13(5.71)−8241.07(6.56)−8238.16(0.52)−8236.69(4.70)−6747.15(15.21)−6747.46(6.90)−6751.97(11.24)−6735.78(5.25)−6735.59(4.51)−6736.74(3.55)

−
−
−
−
−
−

−37526.41(66.28)−37375.93(28.86)−37430.82(67.19)−37396.04(26.39)−37357.86(10.76)−37353.08(16.97)−8727.93(43.33)−8716.61(26.32)−8691.38(10.70)−8679.48(27.55)−8665.54(5.66)−8665.99(7.53)−7129.70(6.14)−7132.35(6.89)−7122.76(10.81)−7120.03(11.92)−7124.62(12.33)−7123.37(11.28)−26487.71(54.79)−26391.00(11.88)−26380.59(14.39)−26378.55(11.05)−26368.49(0.13)−26368.51(0.09)−33807.05(22.97)−33736.98(1.89)−33736.93(2.10)−33753.20(3.03)−33736.03(0.16)−33735.99(0.05)−13353.30(5.92)−13347.17(7.77)−13343.21(2.14)−13342.27(2.71)−13337.74(1.71)−13337.06(1.45)−8251.01(10.34)−8237.75(6.08)−8239.96(4.84)−8237.33(5.41)−8234.18(6.11)−8241.25(8.15)−6750.00(11.91)−6752.46(8.64)−6753.84(14.30)−6734.51(1.95)−6734.49(3.14)−6734.63(2.18)−37487.49(50.43)−37462.07(54.40)−37419.02(35.94)−37373.32(10.08)−37347.46(11.93)−37352.30(12.32)−8736.81(52.38)−8684.38(7.98)−8691.96(13.51)−8662.53(4.58)−8666.63(7.86)−8666.39(7.54)−7126.89(10.06)−7130.67(10.67)−7128.40(9.78)−7122.10(12.29)−7120.94(13.11)−7125.78(13.10)−26444.84(27.91)−26380.41(14.40)−26375.28(11.78)−26381.84(17.18)−26368.52(0.03)−26368.51(0.10)−33823.74(15.62)−33737.75(2.48)−33736.91(1.91)−33759.19(9.95)−33735.98(0.08)−33736.00(0.18)−13358.16(9.79)−13346.94(4.25)−13347.32(4.42)−13342.81(3.45)−13339.77(3.84)−13342.38(6.35)−8251.45(9.72)−8239.36(4.62)−8235.41(5.70)−8243.92(6.74)−8236.42(3.63)−8235.03(5.36)−6745.60(8.36)−6741.63(3.23)−6742.40(1.94)−6733.38(0.79)−6735.12(2.52)−6736.20(1.91)−37516.88(69.88)−37382.28(31.96)−37411.28(56.74)−37369.36(13.45)−37341.92(9.15)−37345.80(11.13)−8719.44(60.54)−8690.41(15.92)−8683.22(13.13)−8666.85(10.63)−8673.15(0.97)−8666.68(5.78)−7124.63(8.09)−7125.94(13.07)−7115.19(8.16)−7111.55(0.07)−7119.77(11.80)−7116.09(10.67)−26458.50(32.42)−26391.02(11.89)−26385.21(13.77)−26379.48(11.60)−26368.44(0.13)−26368.54(0.12)−33804.63(20.76)−33736.98(1.96)−33736.97(1.93)−33757.79(8.07)−33736.01(0.03)−33735.85(0.12)−13358.16(1.20)−13344.13(0.22)−13343.95(1.76)−13342.71(1.61)−13339.26(3.19)−13337.42(1.32)−8251.09(7.46)−8236.90(5.13)−8239.55(4.72)−8240.87(9.80)−8234.06(7.53)−8233.89(6.63)−6748.78(7.38)−6753.86(10.68)−6747.61(6.87)−6735.14(2.64)−6733.91(0.57)−6735.90(1.13)−37484.98(34.91)−37416.00(3.12)−37431.76(43.65)−37377.86(29.48)−37350.77(11.74)−37358.96(13.06)−8717.27(28.49)−8684.90(12.81)−8683.54(3.57)−8663.51(6.85)−8671.32(5.99)−8660.48(0.78)(Replication)

VCSMC
ϕ-CSMCN (diag,2)
LOO(3)
LOO(3)+N (diag,4)
LOO(3)
LOO(3)+
N (full,2)
LOO(3)
LOO(3)+
N (full,4)LOO(3)
LOO(3)+WN (diag,2)
LOO(3)
LOO(3)+WN (diag,4)
LOO(3)
LOO(3)+WN (full,2)
LOO(3)
LOO(3)+WN (full,4)LOO(3)
LOO(3)+

18