Title: VARIATIONAL BAYESIAN PHYLOGENETIC INFERENCE

Abstract: Bayesian phylogenetic inference is currently done via Markov chain Monte Carlo with simple mechanisms for proposing new states, which hinders exploration efficiency and often requires long runs to deliver accurate posterior estimates. In this paper we present an alternative approach: a variational framework for Bayesian phylogenetic analysis. We approximate the true posterior using an expressive graphical model for tree distributions, called a subsplit Bayesian network, together with appropriate branch length distributions. We train the variational approximation via stochastic gradient ascent and adopt multi-sample based gradient estimators for different latent variables separately to handle the composite latent space of phylogenetic models. We show that our structured variational approximations are flexible enough to provide comparable posterior estimation to MCMC, while requiring less computation due to a more efficient tree exploration mechanism enabled by variational inference. Moreover, the variational approximations can be readily used for further statistical analysis such as marginal likelihood estimation for model comparison via importance sampling. Experiments on both synthetic data and real data Bayesian phylogenetic inference problems demonstrate the effectiveness and efficiency of our methods.

1. INTRODUCTION:

Bayesian phylogenetic inference is an essential tool in modern evolutionary biology. Given an alignment of nucleotide or amino acid sequences and appropriate prior distributions, Bayesian methods provide principled ways to assess the phylogenetic uncertainty by positing and approximating a posterior distribution on phylogenetic trees (Huelsenbeck et al., 2001) . In addition to uncertainty quantification, Bayesian methods enable integrating out tree uncertainty in order to get more confident estimates of parameters of interest, such as factors in the transmission of Ebolavirus (Dudas et al., 2017) . Bayesian methods also allow complex substitution models (Lartillot & Philippe, 2004) , which are important in elucidating deep phylogenetic relationships (Feuda et al., 2017) .

Ever since its introduction to the phylogenetic community in the 1990s, Bayesian phylogenetic inference has been dominated by random-walk Markov chain Monte Carlo (MCMC) approaches (Yang & Rannala, 1997; Mau et al., 1999; Huelsenbeck & Ronquist, 2001) . However, this approach is fundamentally limited by the complexities of tree space. A typical MCMC method for phylogenetic inference involves two steps in each iteration: first, a new tree is proposed by randomly perturbing the current tree, and second, the tree is accepted or rejected according to the Metropolis-Hastings acceptance probability. Any such random walk algorithm faces obstacles in the phylogenetic case, in which the high-posterior trees are a tiny fraction of the combinatorially exploding number of trees. Thus, major modifications of trees are likely to be rejected, restricting MCMC tree movement to local modifications that may have difficulty moving between multiple peaks in the posterior distribution (Whidden & Matsen IV, 2015) . Although recent MCMC methods for distributions on Euclidean space use intelligent proposal mechanisms such as Hamiltonian Monte Carlo (Neal, 2011) , it is not straightforward to extend such algorithms to the composite structure of tree space, which includes both tree topology (discrete object) and branch lengths (continuous positive vector) (Dinh et al., 2017) .

Variational inference (VI) is an alternative approximate inference method for Bayesian analysis which is gaining in popularity (Jordan et al., 1999; Wainwright & Jordan, 2008; Blei et al., 2017) . Unlike MCMC methods that sample from the posterior, VI selects the best candidate from a family of tractable distributions to minimize a statistical distance measure to the target posterior, usually the Kullback-Leibler (KL) divergence. By reformulating the inference problem into an optimization problem, VI tends to be faster and easier to scale to large data (via stochastic gradient descent) (Blei et al., 2017) . However, VI can also introduce a large bias if the variational distribution is insufficiently flexible. The success of variational methods, therefore, relies on having appropriate tractable variational distributions and efficient training procedures.

To our knowledge, there have been no previous variational formulations of Bayesian phylogenetic inference. This has been due to the lack of an appropriate family of approximating distributions on phylogenetic trees. However the prospects for variational inference have changed recently with the introduction of subsplit Bayesian networks (SBNs) (Zhang & Matsen IV, 2018) , which provide a family of flexible distributions on tree topologies (i.e. trees without branch lengths). SBNs build on previous work (Höhna & Drummond, 2012; Larget, 2013) , but in contrast to these previous efforts, SBNs are sufficiently flexible for real Bayesian phylogenetic posteriors (Zhang & Matsen IV, 2018) .

In this paper, we develop a general variational inference framework for Bayesian phylogenetics. We show that SBNs, when combined with appropriate approximations for the branch length distribution, can provide flexible variational approximations over the joint latent space of phylogenetic trees with branch lengths. We use recently-proposed unbiased gradient estimators for the discrete and continuous components separately to enable efficient stochastic gradient ascent. We also leverage the similarity of local structures among trees to reduce the complexity of the variational parameterization for the branch length distributions and provide an extension to better capture the between-tree variation. Finally, we demonstrate the effectiveness and efficiency of our methods on both synthetic data and a benchmark of challenging real data Bayesian phylogenetic inference problems.

2. BACKGROUND:

Phylogenetic Posterior A phylogenetic tree is described by a tree topology τ and associated nonnegative branch lengths q. The tree topology τ represents the evolutionary diversification of the species. It is a bifurcating tree with N leaves, each of which has a label corresponding to one of the observed species. The internal nodes of τ represent the unobserved characters (e.g. DNA bases) of the ancestral species. A continuous-time Markov model is often used to describe the transition probabilities of the characters along the branches of the tree. Let Y = {Y 1 , Y 2 , . . . , Y M } ∈ Ω N ×M be the observed sequences (with characters in Ω) of length M over N species. The probability of each site observation Y i is defined as the marginal distribution over the leaves

p(Y i |τ, q) = a i η(a i ρ ) (u,v)∈E(τ ) P a i u a i v (q uv ) (1)

where ρ is the root node (or any internal node if the tree is unrooted and the Markov model is time reversible), a i ranges over all extensions of Y i to the internal nodes with a i u being the assigned character of node u, E(τ ) denotes the set of edges of τ , P ij (t) denotes the transition probability from character i to character j across an edge of length t and η is the stationary distribution of the Markov model. Assuming different sites are identically distributed and evolve independently, the likelihood of observing the entire sequence set Y is p(Y |τ, q) = M i=1 p(Y i |τ, q). The phylogenetic likelihood for each site in equation 1 can be evaluated efficiently through the pruning algorithm (Felsenstein, 2003) , also known as the sum-product algorithm in probabilistic graphical models (Strimmer & Moulton, 2000; Koller & Friedman, 2009; Höhna et al., 2014) . Given a proper prior distribution with density p(τ, q) imposed on the tree topologies and the branch lengths, the phylogenetic posterior p(τ, q|Y ) is proportional to the joint density

p(τ, q|Y ) = p(Y |τ, q)p(τ, q) p(Y ) ∝ p(Y |τ, q)p(τ, q)

where p(Y ) is the intractable normalizing constant.

Subsplit Bayesian Networks:

We now review subsplit Bayesian networks (Zhang & Matsen IV, 2018) and the flexible distributions on tree topologies they provide. Let X be the set of leaf labels. The corresponding SBN assignments for the trees. For ease of illustration, subsplit (W, Z) is represented as W Z in the graph. The dashed gray subgraphs represent fake splitting processes where splits are deterministically assigned, and are used purely to complement the networks such that the overall network has a fixed structure. Right: The SBN for these examples.

ATGAAC • • • ATGCAC • • • ATGCAT • • • ATGCAT • • •

We call a nonempty subset of X a clade. Let be a total order on clades (e.g., lexicographical order). A subsplit (W, Z) of a clade X is an ordered pair of disjoint subclades of X such that W ∪ Z = X and W Z. A subsplit Bayesian network B X on a leaf set X of size N is a Bayesian network whose nodes take on subsplit or singleton clade values that represent the local topological structure of trees (Figure 1 ). Following the splitting processes (see the solid dark subgraphs in Figure 1 , middle right), rooted trees have unique subsplit decompositions and hence can be uniquely represented as compatible SBN assignments. Given the subsplit decomposition of a rooted tree τ = {s 1 , s 2 , . . .}, where s 1 is the root subsplit and {s i } i>1 are other subsplits, the SBN tree probability is

p sbn (T = τ ) = p(S 1 = s 1 ) i>1 p(S i = s i |S πi = s πi )

where S i denotes the subsplit-or singleton-clade-valued random variables at node i and π i is the index set of the parents of S i . The Bayesian network formulation of SBNs enjoys many benefits: i) flexibility. The expressiveness of SBNs is freely adjustable by changing the dependency structures between nodes, allowing for a wide range of flexible distributions; ii) normality. SBN-induced distributions are all naturally normalized if the associated conditional probability tables (CPTs) are consistent, which is a common property of Bayesian networks. The SBN framework also generalizes to unrooted trees, which are the most common type of trees in phylogenetics. Concretely, unrooted trees can be viewed as rooted trees with unobserved roots. Marginalizing out the unobserved root node S 1 , we have the SBN probability estimates for unrooted trees

p sbn (T u = τ ) = s1∼τ p(S 1 = s 1 ) i>1 p(S i = s i |S πi = s πi )

where ∼ means all root subsplits that are compatible with τ .

To reduce model complexity and encourage generalization, the same set of CPTs for parent-child subsplit pairs is shared across the SBN network, regardless of their locations. Similar to weight sharing used in convolutional networks (LeCun et al., 1998) for detecting translationally-invariant structure of images (e.g., edges, corners), this heuristic parameter sharing used in SBNs is for identifying conditional splitting patterns of phylogenetic trees. See Zhang & Matsen IV (2018) for more detailed discussion on SBNs.

3. VARIATIONAL PHYLOGENETIC INFERENCE VIA SBNS:

The flexible and tractable tree topology distributions provided by SBNs serve as an essential building block to perform variational inference (Jordan et al., 1999) for phylogenetics. Suppose that we have a family of approximate distributions Q φ (τ ) (e.g., SBNs) over phylogenetic tree topologies, where φ denotes the corresponding variational parameters (e.g., CPTs for SBNs). For each tree τ , we posit another family of densities Q ψ (q|τ ) over the branch lengths, where ψ is the branch length variational parameters. We then combine these distributions and use the product

Q φ,ψ (τ, q) = Q φ (τ )Q ψ (q|τ )

as our variational approximation. Inference now amounts to finding the member of this family that minimizes the Kullback-Leibler (KL) divergence to the exact posterior,

EQUATION

which is equivalent to maximizing the evidence lower bound (ELBO),

L(φ, ψ) = E Q φ,ψ (τ,q) log p(Y |τ, q)p(τ, q) Q φ (τ )Q ψ (q|τ ) ≤ log p(Y ).

As the ELBO is based on a single-sample estimate of the evidence, it heavily penalizes samples that fail to explain the observed sequences. As a result, the variational approximation tends to cover only the high-probability areas of the true posterior. This effect can be minimized by averaging over K > 1 samples when estimating the evidence (Burda et al., 2016; Mnih & Rezende, 2016) , which leads to tighter lower bounds

EQUATION

where

Q φ,ψ (τ 1:K , q 1:K ) ≡ K i=1 Q φ,ψ (τ i , q i )

; the tightness of the lower bounds improves as the number of samples K increases (Burda et al., 2016) . We will use multi-sample lower bounds in the sequel and refer to them as lower bounds for short.

3.1. VARIATIONAL PARAMETERIZATION:

The CPTs in SBNs are, in general, associated with all possible parent-child subsplit pairs. Therefore, in principle a full parameterization requires an exponentially increasing number of parameters. In practice, however, we can find a sufficiently large subsplit support of CPTs (i.e. where the associated conditional probabilities are allowed to be nonzero) that covers favorable subsplit pairs from trees in the high-probability areas of the true posterior. In this paper, we will mostly focus on the variational approach and assume the support of CPTs is available, although in our experiments we find that a simple bootstrap-based approach does provide a reasonable CPT support estimate for real data. We leave the development of more sophisticated methods for finding the support of CPTs to future work. Now denote the set of root subsplits in the support as S r and the set of parent-child subsplit pairs in the support as S ch|pa . The CPTs are defined according to the following equations

p(S 1 = s 1 ) = exp(φ s1 ) sr∈Sr exp(φ sr ) , p(S i = s|S πi = t) = exp(φ s|t ) s∈S •|t exp(φ s|t )

where S •|t denotes the set of child subsplits for parent subsplit t.

We use the Log-normal distribution Lognormal(µ, σ 2 ) as our variational approximation for branch lengths to accommodate their non-negative nature in phylogenetic models. Instead of a naive parameterization for each edge on each tree (which would require a large number of parameters when the high-probability areas of the posterior are diffuse), we use an amortized set of parameters over the shared local structures among trees. A simple choice of such local structures is the split, a bipartition (X 1 , X 2 ) of the leaf labels X (i.e. X 1 ∪ X 2 = X , X 1 ∩ X 2 = ∅), and each edge of a phylogenetic tree naturally corresponds to a split, the bipartition that consists of the leaf labels from both sides of the edge. Note that a split can be viewed as a root subsplit. We then assign µ(•, •), σ(•, •) for each split (•, •) in S r . We denote the corresponding split of edge e of tree τ as e/τ .

A Simple Independent Approximation Given a phylogenetic tree τ , we start with a simple model that assumes the branch lengths for the edges of the tree are independently distributed. The approximate density Q ψ (q|τ ), therefore, has the form

Q ψ (q|τ ) = e∈E(τ ) p Lognormal (q e | µ(e, τ ), σ(e, τ )) , µ(e, τ ) = ψ µ e/τ , σ(e, τ ) = ψ σ e/τ . (4) e ψ µ (W, Z) W 1 W 2 Z 1 Z 2 W ψ µ ( W 1 , W 2 |W , Z ) Z ψ µ ( Z 1 , Z 2 |W , Z ) + µ(e, τ )

Figure 2 : Branch length parameterization using primary subsplit pairs, which is the sum of parameters for a split and its neighboring subsplit pairs. Edge e represents a split (W, Z). Parameterization for the variance is the same as for the mean.

Capturing Between-Tree Branch Length Variation:

The above approximation equation 4 implicitly assumes that the branch lengths in different trees have the same distribution if they correspond to the same split, which fails to account for between-tree variation. To capture this variation, one can use a more sophisticated parameterization that allows other tree-dependent terms for the variational parameters µ and σ. Specifically, we use additional local structure associated with each edge as follows:

Definition 1 (primary subsplit pair) Let e be an edge of a phylogenetic tree τ which represents a split e/τ = (W, Z). Assume that at least one of W or Z, say W , contains more than one leaf label and denote its subsplit as (W 1 , W 2 ). We call the parent-child subsplit pair (W 1 , W 2 )|(W, Z) a primary subsplit pair.

We assign additional parameters for each primary subsplit pair. Denoting the primary subsplit pair(s) of edge e in tree τ as e/ /τ , we then simply sum all variational parameters associated with e to form the mean and variance parameters for the corresponding branch length (Figure 2 ):

µ(e, τ ) = ψ µ e/τ + s∈e/ /τ ψ µ s , σ(e, τ ) = ψ σ e/τ + s∈e/ /τ ψ σ s .

This modifies the density in equation 4 by adding contributions from primary subsplit pairs and hence allows for more flexible between-tree approximations. Note that the above structured parameterizations of branch length distributions also enable joint learning across tree topologies.

3.2. STOCHASTIC GRADIENT ESTIMATORS AND THE VBPI ALGORITHM:

In practice, the lower bound is usually maximized via stochastic gradient ascent (SGA). However, the naive stochastic gradient estimator obtained by differentiating the lower bound has very large variance and is impractical for our purpose. Fortunately, various variance reduction techniques have been introduced in recent years including the control variate (Paisley et al., 2012; Ranganath et al., 2014; Mnih & Gregor, 2014; Mnih & Rezende, 2016) for general latent variables and the reparameterization trick (Kingma & Welling, 2014) for continuous latent variables. In the following, we apply these techniques to different components of our latent variables and derive efficient gradient estimators with much lower variance, respectively. In addition, we also consider a stable gradient estimator based on an alternative variational objective. See Appendix A for derivations.

The VIMCO Estimator Let f φ,ψ (τ, q) = p(Y |τ,q)p(τ,q) Q φ (τ )Q ψ (q|τ ) . The stochastic lower bound with K samples is LK (φ, ψ) = log 1 K K i=1 f φ,ψ (τ i , q i ) . Mnih & Rezende (2016) propose a localized learning signal strategy that significantly reduces the variance of the naive gradient estimator by utilizing the independence between the multiple samples and the regularity of the learning signal, which estimates the gradient as follows

EQUATION

where

LK j|-j (φ, ψ) := LK (φ, ψ) -log 1 K   i =j f φ,ψ (τ i , q i ) + fφ,ψ (τ -j , q -j )  

is the per-sample local learning signal, with fφ,ψ (τ -j , q -j ) being some estimate of f φ,ψ (τ j , q j ) for sample j using the rest of samples (e.g., the geometric mean), and wj = f φ,ψ (τ j ,q j ) K i=1 f φ,ψ (τ i ,q i ) is the self-normalized importance weight. This gives the following VIMCO estimator

∇ φ L K (φ, ψ) K j=1 LK j|-j (φ, ψ) -wj ∇ φ log Q φ (τ j ) with τ j , q j iid ∼ Q φ,ψ (τ, q). (6)

The Reparameterization Trick The VIMCO estimator also works for the branch length gradient. However, as branch lengths are continuous latent variables, we can use the reparameterization trick to estimate the gradient. Because the Log-normal distribution has a simple reparameterization, q ∼ Lognormal(µ, σ 2 ) ⇔ q = exp(µ + σ ), ∼ N (0, 1), we can rewrite the lower bound:

L K (φ, ψ) = E Q φ, (τ 1:K , 1:K ) log   1 K K j=1 p(Y |τ j , g ψ ( j |τ j ))p(τ j , g ψ ( j |τ j )) Q φ (τ j )Q ψ (g ψ ( j |τ j )|τ j )   .

where g ψ ( |τ ) = exp(µ ψ,τ + σ ψ,τ

). Then the gradient of the lower bound w.r.t. ψ is

EQUATION

where wj =

f φ,ψ (τ j ,g ψ ( j |τ j )) K i=1 f φ,ψ (τ i ,g ψ ( i |τ i ))

is the same normalized importance weight as in equation equation 5. Therefore, we can form the Monte Carlo estimator of the gradient

EQUATION

Self-normalized Importance Sampling Estimator In addition to the standard variational formulation equation 2, one can reformulate the optimization problem by minimizing the reversed KL divergence, which is equivalent to maximizing the likelihood of the variational approximation

EQUATION

We can use an importance sampling estimator to compute the gradient of the objective

EQUATION

with the same importance weights wj as in equation 5. This can be viewed as a multi-sample generalization of the wake-sleep algorithm (Hinton et al., 1995) and was first used in the reweighted wake-sleep algorithm (Bornschein & Bengio, 2015) for training deep generative models. We therefore call the gradient estimator in equation 10 the RWS estimator. Like the VIMCO estimator, the RWS estimator also provides gradients for branch lengths. However, we find in practice that equation 8 that uses the reparameterization trick is more useful and often leads to faster convergence, although it uses a different optimization objective. A better understanding of this phenomenon would be an interesting subject of future research.

All stochastic gradient estimators introduced above can be used in conjunction with stochastic optimization methods such as SGA or some of its adaptive variants (e.g. Adam Kingma & Ba, 2015) to maximize the lower bounds. See algorithm 1 in Appendix B for a basic variational Bayesian phylogenetic inference (VBPI) approach.

4. EXPERIMENTS:

Throughout this section we evaluate the effectiveness and efficiency of our variational framework for inference over phylogenetic trees. The simplest SBN (the one with a full and complete binary tree structure) is used for the phylogenetic tree topology variational distribution; we have found it to provide sufficiently accurate approximation. For real datasets, we estimate the CPT supports from ultrafast maximum likelihood phylogenetic bootstrap trees using UFBoot (Minh et al., 2013) , which is a fast approximate bootstrap method based on efficient heuristics. We compare the performance of the VIMCO estimator and the RWS estimator with different variational parameterizations for the branch length distributions, while varying the number of samples in the training objective to see how these affect the quality of the variational approximations. For VIMCO, we use Adam for stochastic gradient ascent with learning rate 0.001 (Kingma & Ba, 2015) . For RWS, we also use AMSGrad (Sashank et al., 2018) , a recent variant of Adam, when Adam is unstable. Results were collected after 200,000 parameter updates. The KL divergences reported are over the discrete collection of phylogenetic tree structures, from trained SBN distribution to the ground truth, and a low KL divergence means a high quality approximation of the distribution of trees.

4.1. SIMULATED SCENARIOS:

To empirically investigate the representative power of SBNs to approximate distributions on phylogenetic trees under the variational framework, we first conduct experiments on a simulated setup. We use the space of unrooted phylogenetic trees with 8 leaves, which contains 10395 unique trees in total. Given an arbitrary order of trees, we generate a target distribution p 0 (τ ) by drawing a sample from the symmetric Dirichlet distributions Dir(β1) of order 10395, where β is the concentration parameter. The target distribution becomes more diffuse as β increases; we used β = 0.008 to provide enough information for inference while allowing for adequate diffusion in the target. Note that there are no branch lengths in this simulated model and the lower bound is

L K (φ) = E Q φ (τ 1:K ) log 1 K K i=1 p 0 (τ i ) Q φ (τ i ) ≤ 0

with the exact evidence being log(1) = 0. We then use both the VIMCO and RWS estimators to optimize the above lower bound based on 20 and 50 samples (K). We use a slightly larger learning rate (0.002) in AMSGrad for RWS.

Figure 3 shows the empirical performance of different methods. From the left plot, we see that the lower bounds converge rapidly and the gaps between lower bounds and the exact model evidence are close to zero, demonstrating the expressive power of SBNs on phylogenetic tree probability estimations. The evolution of KL divergences (middle plot) is consistent with the lower bounds. All methods benefit from using more samples, with VIMCO performing better in the end and RWS learning slightly faster at the beginning. 1 The slower start of VIMCO is partly due to the regularization term in the lower bounds, which turns out to be beneficial for the overall performance since the regularization encourages the diversity of the variational approximation and leads to more sufficient exploration in the starting phase, similar to the exploring starts (ES) strategy in reinforcement learning (Sutton & Barto, 1998) . The right plot compares the variational approximations obtained by VIMCO and RWS, both with 50 samples, to the ground truth p 0 (τ ). We see that VIMCO slightly underestimates trees in high-probability areas as a result of the regularization effect. While RWS provides better approximations for trees in high-probability areas, it tends to underestimate trees in low-probability areas which deteriorates the overall performance. We expect the biases in both approaches to be alleviated with more samples.

4.2. REAL DATA PHYLOGENETIC POSTERIOR ESTIMATION:

In the second set of experiments we evaluate the proposed variational Bayesian phylogenetic inference (VBPI) algorithms at estimating unrooted phylogenetic tree posteriors on 8 real datasets commonly used to benchmark phylogenetic MCMC methods (Lakner et al., 2008; Höhna & Drummond, 2012; Larget, 2013; Whidden & Matsen IV, 2015) (Table 1 ). We concentrate on the most challenging part of the phylogenetic model: joint learning of the tree topologies and the branch lengths. We assume a uniform prior on the tree topology, an i.i.d. exponential prior (Exp(10)) for the branch lengths and the simple Jukes & Cantor (1969) substitution model. We consider two different variational parameterizations for the branch length distributions as introduced in section 3.1. In the first case, we use the simple split-based parameterization that assigns parameters to the splits associated with the edges of the trees. In the second case, we assign additional parameters for the primary subsplit pairs (PSP) to better capture the between-tree variation. We form our ground truth posterior from an extremely long MCMC run of 10 billion iterations (sampled each 1000 iterations with the first 25% discarded as burn-in) using MrBayes (Ronquist et al., 2012) , and gather the support of CPTs from 10 replicates of 10000 ultrafast maximum likelihood bootstrap trees (Minh et al., 2013) . Following Rezende & Mohamed (2015) , we use a simple annealed version of the lower bound which was found to provide better results. The modified bound is:

L K βt (φ, ψ) = E Q φ,ψ (τ 1:K , q 1:K ) log 1 K K i=1 [p(Y |τ i , q i )] βt p(τ i , q i ) Q φ (τ i )Q ψ (q i |τ i )

where β t ∈ [0, 1] is an inverse temperature that follows a schedule β t = min(0.001, t/100000), going from 0.001 to 1 after 100000 iterations. We use Adam with learning rate 0.001 to train the variational approximations using VIMCO and RWS estimators with 10 and 20 samples.

Figure 4 (left and middle plots) shows the resulting KL divergence to the ground truth on DS1 as a function of the number of parameter updates. The results for methods that adopt the simple splitbased parameterization of variational branch length distributions are shown in the left plot. We see that the performance of all methods improves significantly as the number of samples is increased. The middle plot, containing the results using PSP for variational parameterization, clearly indicates that a better modeling of between-tree variation of the branch length distributions is beneficial for all method / number of samples combinations. Specifically, PSP enables more flexible branch length distributions across trees which makes the learning task much easier, as shown by the considerably smaller gaps between the methods. To benchmark the learning efficiency of VBPI, we also compare to MrBayes 3.2.5 (Ronquist et al., 2012) , a standard MCMC implementation. We run MrBayes with 4 chains and 10 runs for two million iterations, sampling every 100 iterations. For each run, we compute the KL divergence to the ground truth every 50000 iterations with the first 25% discarded Table 1 : Data sets used for variational phylogenetic posterior estimation, and marginal likelihood estimates of different methods across datasets. The marginal likelihood estimates of all variational methods are obtained by importance sampling using 1000 samples. We run stepping-stone in Mr-Bayes using default settings with 4 chains for 10,000,000 iterations and sampled every 100 iterations. The results are averaged over 10 independent runs with standard deviation in brackets. (e.g., SBNs) which in turn allows guided exploration in the tree topology space. VBPI also provides the same majority-rule consensus tree as the ground truth MCMC run (Figure 5 in Appendix D).

The variational approximations provided by VBPI can be readily used to perform importance sampling for phylogenetic inference (more details in Appendix C). The right plot of Figure 4 compares VBPI using VIMCO with 20 samples and PSP to the state-of-the-art generalized stepping-stone (GSS) (Fan et al., 2011) algorithm for estimating the marginal likelihood of trees in the 95% credible set of DS1. For GSS, we use 50 power posteriors and for each power posterior we run 1,000,000 MCMC iterations, sampling every 1000 iterations with the first 10% discarded as burn-in. The reference distribution for GSS was obtained from an independent Gamma approximation using the maximum a posterior estimate. Table 1 shows the estimates of the marginal likelihood of the data (i.e., model evidence) using different VIMCO approximations and one of the state-of-the-art methods, the stepping-stone (SS) algorithm (Xie et al., 2011) . For each data set, all methods provide estimates for the same marginal likelihood, with better approximation leading to lower variance. We see that VBPI using 1000 samples is already competitive with SS using 100000 samples and provides estimates with much less variance (hence more reproducible and reliable). Again, the extra flexibility enabled by PSP alleviates the demand for larger number of samples used in the training objective, making it possible to achieve high quality variational approximations with less samples.

5. DISCUSSION:

In this work we introduced VBPI, a general variational framework for Bayesian phylogenetic inference. By combining subsplit Bayesian networks, a recent framework that provides flexible distributions of trees, and efficient structured parameterizations for branch length distributions, VBPI exhibits guided exploration (enabled by SBNs) in tree space and provides competitive performance to MCMC methods with less computation. Moreover, variational approximations provided by VBPI can be readily used for further statistical analysis such as marginal likelihood estimation for model comparison via importance sampling, which, compared to MCMC based methods, dramatically reduces the cost at test time. We report promising numerical results demonstrating the effectiveness and efficiency of VBPI on a benchmark of real data Bayesian phylogenetic inference problems.

When the data are weak and posteriors are diffuse, support estimation of CPTs becomes challenging. However, compared to classical MCMC approaches in phylogenetics that need to traverse the enormous support of posteriors on complete trees to accurately evaluate the posterior probabilities, the SBN parameterization in VBPI has a natural advantage in that it alleviates this issue by factorizing the uncertainty of complete tree topologies into local structures.

Many topics remain for future work: constructing more flexible approximations for the branch length distributions (e.g., using normalizing flow (Rezende & Mohamed, 2015) for within-tree approximation and deep networks for the modeling of between-tree variation), deeper investigation of support estimation approaches in different data regimes, and efficient training algorithms for general variational inference on discrete / structured latent variables.

A GRADIENT DERIVATION FOR THE MULTI-SAMPLE OBJECTIVES:

In this section we will derive the gradient for the multi-sample objectives introduced in section 3. We start with the lower bound

L K (φ, ψ) = E Q φ,ψ (τ 1:K , q 1:K ) log   1 K K j=1 p(Y |τ j , q j )p(τ j , q j ) Q φ (τ j )Q ψ (q j |τ j )   = E Q φ,ψ (τ 1:K , q 1:K ) log   1 K K j=1 f φ,ψ (τ j , q j )   .

Using the product rule and noting that ∇ φ log f φ,ψ (τ j , q

j ) = -∇ φ log Q φ (τ j ), ∇ φ L k (φ, ψ) = E Q φ,ψ (τ 1:K , q 1:K ) ∇ φ log   1 K K j=1 f φ,ψ (τ j , q j )   + E Q φ,ψ (τ 1:K , q 1:K ) K j=1 ∇ φ Q φ (τ j ) Q φ (τ j ) log 1 K K i=1 f φ,ψ (τ i , q i ) = E Q φ,ψ (τ 1:K , q 1:K ) K j=1 f φ,ψ (τ j , q j ) K i=1 f φ,ψ (τ i , q i ) ∇ φ log f φ,ψ (τ j , q j )+ E Q φ,ψ (τ 1:K , q 1:K ) K j=1 log 1 K K i=1 f φ,ψ (τ i , q i ) ∇ φ log Q φ (τ j ) = E Q φ,ψ (τ 1:K , q 1:K ) K j=1 LK (φ, ψ) -wj ∇ φ log Q φ (τ j ).

This gives the naive gradient of the lower bound w.r.t. φ.

Using the reparameterization trick, the lower bound has the form

L K (φ, ψ) = E Q φ, (τ 1:K , 1:K ) log   1 K K j=1 p(Y |τ j , g ψ ( j |τ j ))p(τ j , g ψ ( j |τ j )) Q φ (τ j )Q ψ (g ψ ( j |τ j )|τ j )   = E Q φ, (τ 1:K , 1:K ) log   1 K K j=1 f φ,ψ (τ j , g ψ ( j |τ j ))  

Since ψ is not involved in the distribution with respect to which we take expectation,

∇ ψ L K (φ, ψ) = E Q φ, (τ 1:K , 1:K ) ∇ ψ log   1 K K j=1 f φ,ψ (τ j , g ψ ( j |τ j ))   = E Q φ, (τ 1:K , 1:K ) K j=1 f φ,ψ (τ j , g ψ ( j |τ j )) K i=1 f φ,ψ (τ i , g ψ ( i |τ i )) ∇ ψ log f φ,ψ (τ j , g ψ ( j |τ j )) = E Q φ, (τ 1:K , 1:K ) K j=1 wj ∇ ψ log f φ,ψ (τ j , g ψ ( j |τ j )).

Next, we derive the gradient of the multi-sample likelihood objective used in RWS

L(φ, ψ) = E p(τ,q|Y ) log Q φ,ψ (τ, q).

Again, p(τ, q|Y ) is independent of φ, ψ, and we have

∇ φ L(φ, ψ) = E p(τ,q|Y ) ∇ φ log Q φ,ψ (τ, q) = E Q φ,ψ (τ,q) p(τ, q|Y ) Q φ (τ )Q ψ (q|τ ) ∇ φ log Q φ,ψ (τ, q) = 1 p(Y ) E Q φ,ψ (τ,q) f φ,ψ (τ, q)∇ φ log Q φ (τ ) K j=1 f φ,ψ (τ j , q j ) k i=1 f φ,ψ (τ i , q i ) ∇ φ log Q φ (τ j ) with τ j , q j iid ∼ Q φ,ψ (τ, q). = K j=1 wj ∇ φ log Q φ (τ j )

The second to last step uses self-normalized importance sampling with K samples. ∇ ψ L(φ, ψ) can be computed in a similar way.

B THE VARIATIONAL BAYESIAN PHYLOGENETIC INFERENCE ALGORITHM:

Algorithm 1 The variational Bayesian phylogenetic inference (VBPI) algorithm. 1: φ, ψ ← Initialize parameters 2: while not converged do 3:

τ 1 , . . . , τ K ← Random samples from the current approximating tree distribution Q φ (τ )

4:

1 , . . . , K ← Random samples from the multivariate standard normal distribution N (0, I) 5:

g ← ∇ φ,ψ L K (φ, ψ; τ 1:K , 1:K ) (Use any gradient estimator from section 3.2) 6:

φ, ψ ← Update parameters using gradients g (e.g. SGA) 7: end while 8: return φ, ψ

C IMPORTANCE SAMPLING FOR PHYLOGENETIC INFERENCE VIA VARIATIONAL APPROXIMATIONS:

In this section, we provide a detailed importance sampling procedure for marginal likelihood estimation for phylogenetic inference based on the variational approximations provided by VBPI.

C.1 ESTIMATING MARGINAL LIKELIHOOD OF TREES:

For each tree τ that is covered by the subsplit support,

Q ψ (q|τ ) = e∈E(τ )

p Lognormal (q e | µ(e, τ ), σ(e, τ ))

can provide accurate approximation to the posterior of branch lengths on τ , where the mean and variance parameters µ(e, τ ), σ(e, τ ) are gathered from the structured variational parameters ψ as introduced in section 3.1. Therefore, we can estimate the marginal likelihood of τ using importance sampling with Q ψ (q|τ ) being the importance distribution as follows

p(Y |τ ) = E Q ψ (q|τ ) p(Y |τ, q)p(q) Q ψ (q|τ ) 1 M M j=1

p(Y |τ, q j )p(q j ) Q ψ (q j |τ ) with q j iid ∼ Q ψ (q|τ )

C.2 ESTIMATING MODEL EVIDENCE:

Similarly, we can estimate the marginal likelihood of the data as follows p(Y ) = E Q φ,ψ (τ, q) p(Y |τ, q)p(τ, q) Q φ (τ )Q ψ (q|τ )

1 K K j=1

p(Y |τ j , q j )p(τ j , q j ) Q φ (τ j )Q ψ (q j |τ j ) with τ j , q j iid ∼ Q φ,ψ (τ, q).

In our experiments, we use K = 1000. When taking a log transformation, the above Monte Carlo estimate is no longer unbiased (for the evidence log p(Y )). Instead, it can be viewed as one sample Monte Carlo estimate of the lower bound

EQUATION

whose tightness improves as the number of samples K increases. Therefore, with a sufficiently large K, we can use the lower bound estimate as a proxy for Bayesian model selection.

D CONSENSUS TREE COMPARISON ON DS1:

Alligator mississippiensis 

:

Although we use larger learning rates for RWS in this experiment, we found RWS generally learns slightly faster than VIMCO at the beginning. See Figure4for the real data phylogenetic inference problems in section 4.2 where we use Adam with learning rate 0.001 for both methods.

the extra factor of 2/4 is because the likelihood and the gradient can be computed together in twice the time of a likelihood(Schadt et al., 1998) and we run MCMC with 4 chains.

Figures and Tables:

Figure1: A simple subsplit Bayesian network for a leaf set that contains 4 species. Left: A leaf label set X of 4 species, each label corresponds to a DNA sequence. Middle (left): Examples of (rooted) phylogenetic trees that are hypothesized to model the evolutionary history of the species. Middle (right): The corresponding SBN assignments for the trees. For ease of illustration, subsplit (W, Z) is represented as W Z in the graph. The dashed gray subgraphs represent fake splitting processes where splits are deterministically assigned, and are used purely to complement the networks such that the overall network has a fixed structure. Right: The SBN for these examples.

Figure 3: Comparison of multi-sample objective on approximating a challenging distribution over unrooted phylogenetic trees with 8 leaves using VIMCO and RWS gradient estimators. Left: Evidence lower bound. Middle: KL divergence. Right: Variational approximations vs ground truth probabilities. The number in brackets specifies the number of samples used in the training objective.

Figure 4: Performance on DS1. Left: KL divergence for methods that use the simple split-based parameterization for the branch length distributions. Middle: KL divergence for methods that use PSP. Right: Per-tree marginal likelihood estimation (in nats): VBPI vs GSS. The number in brackets specifies the number of samples used in the training objective. MCMC results are averaged over 10 independent runs. The results for VBPI were obtained using 1000 samples and the error bar shows one standard deviation over 100 independent runs.

Figure 5: A comparison of majority-rule consensus trees obtained from VBPI and ground truth MCMC run on DS1. Left: Ground truth MCMC. Right: VBPI (10000 sampled trees). The plot is created using the treespace (Jombart et al., 2017) R package.

For a relatively fair comparison (in terms of the number of likelihood evaluations), we compare 10 (i.e. 2•20/4) times the number of MCMC iterations with the number of 20-sample objective VBPI iterations. 2 Although MCMC converges faster at the start, we see that VBPI methods (especially those with PSP) quickly surpass MCMC and arrive at good approximations with much less computation. This is because VBPI iteratively updates the approximate distribution of trees

Table content: <table><tr><td>DATA SET</td><td>REFERENCE</td><td>(#TAXA, #SITES)</td><td>VIMCO(10)</td><td>VIMCO(20)</td><td colspan="3">MARGINAL LIKELIHOOD (NATS) VIMCO(10) + PSP VIMCO(20) + PSP</td><td>SS</td></tr><tr><td>DS1</td><td>HEDGES ET AL. (1990)</td><td>(27, 1949)</td><td>-7108.43(0.26)</td><td colspan="2">-7108.35(0.21)</td><td>-7108.41(0.16)</td><td>-7108.42(0.10)</td><td>-7108.42(0.18)</td></tr><tr><td>DS2</td><td>GAREY ET AL. (1996)</td><td colspan="4">(29, 2520) -26367.70(0.12) -26367.71(0.09)</td><td>-26367.72(0.08)</td><td>-26367.70(0.10)</td><td>-26367.57(0.48)</td></tr><tr><td>DS3</td><td>YANG &amp; YODER (2003)</td><td colspan="4">(36, 1812) -33735.08(0.11) -33735.11(0.11)</td><td>-33735.10(0.09)</td><td>-33735.07(0.11)</td><td>-33735.44(0.50)</td></tr><tr><td>DS4</td><td>HENK ET AL. (2003)</td><td colspan="4">(41, 1137) -13329.90(0.31) -13329.98(0.20)</td><td>-13329.94(0.18)</td><td>-13329.93(0.22)</td><td>-13330.06(0.54)</td></tr><tr><td>DS5</td><td>LAKNER ET AL. (2008)</td><td>(50, 378)</td><td>-8214.36(0.67)</td><td colspan="2">-8214.74(0.38)</td><td>-8214.61(0.38)</td><td>-8214.55(0.43)</td><td>-8214.51(0.28)</td></tr><tr><td>DS6</td><td colspan="2">ZHANG &amp; BLACKWELL (2001) (50, 1133)</td><td>-6723.75(0.68)</td><td colspan="2">-6723.71(0.65)</td><td>-6724.09(0.55)</td><td>-6724.34(0.45)</td><td>-6724.07(0.86)</td></tr><tr><td>DS7</td><td>YODER &amp; YANG (2004)</td><td colspan="4">(59, 1824) -37332.03(0.43) -37331.90(0.49)</td><td>-37331.90(0.32)</td><td>-37332.03(0.23)</td><td>-37332.76(2.42)</td></tr><tr><td>DS8</td><td>ROSSMAN ET AL. (2001)</td><td>(64, 1008)</td><td>-8653.34(0.55)</td><td colspan="2">-8651.54(0.80)</td><td>-8650.63(0.42)</td><td>-8650.55(0.46)</td><td>-8649.88(1.75)</td></tr><tr><td colspan="2">as burn-in.</td><td/><td/><td/><td/><td/><td/></tr></table>

Peer review (iclr2019):
 This paper is well written, appears to be well executed, and the results look good. I am not particularly well informed about the area, but the work appears to be novel. MCMC for phylogenetic inference is hugely expensive, and anything we can do to reduce that cost would be beneficial (the computational expense is not given, or I*ve missed it, for the variational approach - presumably it is relatively small compared to MCMC?). My main criticism is that I found the details of subsplit Bayesian networks difficult to follow. Googling them suggests they are a relatively new model, which has not been well studied or used (there are no citations of the paper that introduces them for example!). The paper would be stronger if it discussed these in more detail - how close can they come to approximating the models usually used in phylogenetic analyses? Often the inferred phylogeny is itself of interest - how similar are the trees inferred here to those found from MrBayes?