Variational Phylogenetic Inference with Products over Bipartitions

Thank you for your insightful review: see our responses below.

For inference accuracy, VIPR's trees' likelihoods lagged behind VBPI. For inference speed, no results report computation time.

The aspect where VIPR shines is in time (or number of parameter updates) to attain an approximation error, as shown in Figure 5 of Appendix B. For example, on DS1 VIPR-LOOR achieves an estimated marginal log-likelihood (MLL) of -7175 after ~100 iterations, while VBPI10 takes ~1000. For DS5, VIPR-LOOR achieves an MLL of -8300 after ~200 iterations, while VBPI10 takes ~5000.

Figure 5 compares iterations, and we agree that we should also report computation time. Thus, we ran all VI methods on simulated datasets with varying numbers of taxa for 1,000 iterations and reported computation time (see our response to Reviewer 1 for the simulation procedure).

Seconds/1,000 iterations:

One iteration is one parameter update. Our method is ~ twice as slow as VBPI per iteration for 8 taxa, but it scales better and outperforms VBPI for 512 taxa. We will add the results to Figure 4 and the Appendix.

We improved our code since submission (see our response to Reviewer 1). VIPR's primary computational bottleneck is now the phylogenetic likelihood, which takes $\mathcal{O}(NM)$ time.

How is VBPI used for ultrametric trees?

Zhang and Matsen IV [ICLR 2019] applies to general additive phylogenetic trees, but the follow-up paper [Zhang and Matsen IV, JMLR2024] extends the approach to ultrametric trees in Sections 6 and 7. The github repository (https://github.com/zcrabbit/vbpi-torch) contains code for ultrametric trees (in the directory "rooted") that we use in our experiments. We will make this more clear in the manuscript by providing a Github repository link and referencing the sections within the JMLR paper.

GeoPhy should be considered

GeoPhy is similar to VIPR in that it uses a tree construction algorithm on a distance matrix, but it is used for unrooted trees while we focus on ultrametric trees.

Figure 3(c) (VBPI is better) contradicts with Table 2 (VIPR-VIMCO is better).

Figure 3(c) reports the marginal log-likelihood, but Table 2 reports the ELBO. For the COVID dataset, VBPI is better in marginal log-likelihood, and VIPR-VIMCO is better in ELBO.

ARTree should be discussed.

We will change Section 4.2 (line 264, column 2): "We compare VIPR to the VBPI algorithm as implemented by Zhang and Matsen IV (2024), which uses MCMC runs to determine likely subsplits in an SBN."

We will also add ARTree to the introduction (line 26, column 2): "For example, ViaPhy (Koptagel et al., 2022) uses a gradient-free variational inference approach and directly sample from the Jukes and Cantor (1969) model, GeoPhy (Mimori and Hamada, 2023) uses a distance-based metric in hyperbolic space to construct unrooted phylogenetic trees, and ARTree (Xie and Zhang, 2023) uses graph neural networks to construct a deep autoregressive model for variational inference over phylogenetic tree structures."

Phyloformer is similar to VIPR

We will mention Phyloformer in the introduction: see our response to Reviewer 1.

The title page violates the ICML format

Thank you for catching the title formatting error. This arose due to a copy/paste mistake. As Reviewer 2 noted, our Figures are not too space-hungry, and we could resolve this by cutting down the length of the Results and Discussion as you suggest, and improving the location of the Figures.

Ultrametirc trees should be (re)defined

We will adopt the suggested definition (line 74, column 1):

"We focus on ultrametric trees, in which the leaves of the trees are all equidistant from the root. We denote our ultrametric trees with a rooted, binary tree topology tau and a set of coalescent times ..."

The $N_e$ in the prior distribution is unclear

We expanded section 2.3:

"We use the Kingman coalescent (Kingman 1982) as the prior distribution on the trees. This coalescent process proceeds backward in time with exponentially distributed inter-event intervals, and coalescent events occurring at rate $\lambda_k = \binom{k}{2}/N_e$, where $k$ is the number of taxa and $N_e$ is the effective population size, a parameter which governs the rate at which species coalesce. We fix $N_e = 5$ in our experiments. At each coalescent event, a pair of taxa are chosen to coalesce into a single taxon uniformly at random over all pairs of taxa. ... "

We fixed the "PhylogGFN" typo, thank you for catching this.

We appreciate the thoughtful suggestions below, and hope that we have addressed your comments sufficiently.

One minor concern is that a similar tree-structured matrix representation has been studied independently in another paper [Bouckaert2024] very recently.

As we discuss in our literature review, we have two primary significant improvements compared to [Bouckaert2024]: 1) we use optimization-based VI to maximize the ELBO, and 2) we provide a closed-form density in Proposition 1. Using this density means we do not have to restrict the tree space. In addition, our framework allows us to use a relatively general class of variational distributions for the pairwise distances, whereas Bouckaert require log-normal distributions in order to incorporate a covariance matrix $\Sigma$.

Would it be difficult to see the reduction in performance when restricted to only restricted trees within the framework of the proposed method?

We believe that this comment refers to selecting an ordering of taxa similarly to [Bouckaert2024], and then re-running VIPR while restricted to the "cube space" from [Bouckaert2024] consistent with the ordering. This is an excellent idea to investigate the effect of restricting tree space.

We have now constructed the maximum clade credibility (MCC) tree from BEAST using our gold standard MCMC run, selected an order from the MCC tree, and then calculated the percentage of tree topologies from the BEAST gold standard that are within the "cube space" implied by this ordering. This process estimates the percentage of the posterior that is impossible to reach using the restricted tree space from [Bouckaert2024]:

These results are striking, but [Bouckaert2024] mentions that CubeVB may struggle on high-entropy posteriors in their discussion. We will add this Table and discussion to the Appendix of the camera-ready.

It is not easy to intuitively understand what improvement VIPR has over the variational distribution of the standard mean-field approximation.

One of the key challenges for variational inference over phylogenetic trees is that using a mean-field approximation is not straightforward. There are two main reasons for this.

First, we can only apply a mean-field approximation after decomposing the distribution of the tree as a product over cliques of random variables. There is no standard way of doing this for trees. In Matsen IV (2024) this is done by forming a subsplit Bayesian network with one node per subtree appearing in the MCMC samples used to initialize the support. Our novel decomposition (Proposition 1) is another way of decomposing the distribution of the tree as a product over coalescent times. This results in $\mathcal{O}(N^2)$ parameters, improving upon the worst-case performance of Matsen IV (2024), in which the number of parameters could be super-exponential in the worst case. We then use our novel decomposition for the mean-field approximation (using state-of-the-art techniques for gradient evaluation and optimization: autograd, VIMCO, REINFORCE and the Reparameterization Trick).

Second, for ultrametric trees we cannot assume independence between coalescent times, lest the resulting tree violate the ultrametric constraint. To overcome this challenge, we form our our variational family to approximate the matrix of pairwise coalescent times (the matrix bold T). We then map T to ultrametric trees using single-linkage clustering.

A quantitative comparison with another related study might have strengthened its persuasiveness.

We chose VBPI for a baseline comparison because it was the only VI-based method for ultrametric trees that we are aware of in the literature. We did not include [Bouckaert2024] because it does not rely on optimization, so we could not include it in our trace plots of marginal log-likelihood vs iteration number. As you have suggested, restricting our method (and the BEAST gold-standard) to the same restricted tree space as [Bouckaert2024] is a valuable experiment to isolate the effects of optimization versus unrestricted tree space, and we aim to complete this experiment in a follow-up paper.

Thank you for catching the title formatting error—we have now corrected it.

Thank you for your thoughtful review: please find detailed responses below.

How do parameter numbers influence the time complexity of VBPI?

In [Zhang and Matsen, JMLR2024], as the number of taxa grows, the number of parameters grows with the number of trees in the SBN. There is no closed form for the number of trees; it depends on the MCMC algorithm and posterior concentration. We empirically calculated the number of parameters in VBPI vs the number of taxa on datasets simulated with "ms" [Hudson 2002] with 1,000 sites.

# of tree structure parameters:

Computing the variational density of VBPI is linear in the number of taxa, but normalizing the SBN scales with the number of parameters. We will add this experiment to the Appendix of the camera-ready copy.

Why should VIPR have a lower computational complexity?

In our experiments, VIPR attains accurate marginal log-likelihoods estimates in fewer parameter updates than VBPI (Appendix B, Figure 5). VIPR has O($N^2$) parameters, and the number of parameters in VBPI can be larger than that if the SBN support is large (see Table above).

We improved our code using scipy.cluster.hierarchy.linkage and streamlined our phylogenetic likelihood function. We performed new speed comparisons for all methods on simulated datasets with varying numbers of taxa. See our response to Reviewer 3 for results.

What is the impact of log-normal branch-lengths? Any way to relax this?

After running BEAST, we plotted histograms of pairwise log-coalescent times across sampled trees for some of the datasets. In most cases these histograms looked normal, motivating our log-normal branch lengths. We will include a some of these histograms as a supplement. VIPR can incorporate any branch length distribution with continuously differentiable density. We will consider flexible branch-length distributions in future work.

Sequence divergence is not included.

We calculated pairwise Hamming distances between each taxa for each dataset (dropping sites with missingness). Values in parentheses are standard deviations:

We will add this Table to Appendix B in the camera-ready.

Include simulated datasets with better control on tree depth, sequence divergence, mutation rates, etc.

This simulation is an excellent idea. We aim to do this in a follow-up paper.

Jukes-Cantor may be over-simplified vs. complex evolutionary models

We agree that Jukes-Cantor is a simplified assumption. We aim to include K2P [Kimura 1980] and GTR [Rodriguez 1990] in a follow-up paper. (Note that Zhang and Matsen JMLR2024 only consider JC.)

Add more baseline methods like RAxML, neighbour-joining, and Vaiphy? The experiment does not cover uncertainty estimation.

These methods do not apply specifically to variational inference over ultrametric trees. RAxML and Neighbour-joining do not provide estimates of the marginal likelihood. Vaiphy is for multifurcating trees. We have already mentioned VaiPhy and we will add more about non-Bayesian methods in the introduction:

"Phylogenetic inference can also be performed using non-Bayesian methods, including RAxML, Neighbour-joining, and Phyloformer. Phyloformer uses deep learning to construct pairwise representations of evolutionary distances between taxa. Phyloformer then uses pairwise distances to construct a tree using a neighbor-joining algorithm similar to the method described here. Non-Bayesian methods do not provide estimates of marginal likelihood, which are useful for model selection."

Regarding uncertainty estimation, we aim to add posterior predictive checks of tree length and clade support on simulated data to the Appendix for the camera-ready.

How would VIPR handle non-ultrametric trees?

A natural approach for non-ultrametric trees with our framework is to extend our strict clock models to relaxed clock models. Preliminary calculations suggest it is possible to do so at the cost of roughly twice as many variational parameters compared to the VIPR variational family.

How robust is VIPR when handling noisy or high divergent dataset?

VIPR struggled most with the COVID-19 dataset, where genomes varied relatively little (see Hamming distance Table above). We assume that highly divergent datasets would also be challenging. Future studies can quantify how VI methods such as VIPR are affected by noisy or divergent datasets.

What is the limitation of VIPR on dataset with 100+ taxa?

VIPR's empirical computation time per iteration is approximately linear in the number of taxa (see response to Reviewer 3). Future work can apply VIPR to larger datasets.