Mr Bayes benchmark

The base- lines we compare to are the MCMC-based MrBayes combined with the stepping-stone sampling tech- nique (Ronquist et al., 2012) (page 8)

How do you obtain that and are you sure of how it is computed?

Thank you for pointing out that our description for this benchmark result is unclear. Stepping stone is the MCMC-based sampling algorithm to generate phylogenetic trees and estimate MLL scores. MrBayes is the software that implements the stepping stone algorithm. In table 1, the MR-Bayes results for MLL estimation are obtained by using MrBayes to run the Stepping stone algorithm. We will update the manuscript to resolve this confusion.

In table 1, the displayed results are the MrBayes results recorded in GeoPhy, VBPI-GNN. To perform the computation, they follow the same procedure described in Zhang & Matsen IV (2018b). In appendix section F1, we provide the MrBayes script used in Zhang & Matsen IV (2018b).

Question on model setup

To select pairs of trees to join, we evaluate tree-pair features for every pair of trees ... (Page 16)

We didn't see a question or comment for the last section in your questions. Please let us know if you have any question on the model setup.

Question about continuous branch length variables

continuous vari- ables that capture the level of sequence divergence along each branch of the tree. (Page 1)

Also the ml version has those.

We are not sure what you mean by this comment. We are happy to answer if you could clarify the question.

On rooted/unrooted trees

Each action chooses a pair of trees and join them at the root, thus creating a new tree (page 5)

This suggests binary trees. “At the root” is a poor formulation.

its two children (page 5)

This should close the case!

We are not quite sure what you mean here, but we can revise the sentence to end "each action chooses a pair of trees and joins their roots by a common parent node".

Suggestions for edits

high complexity of tree space (Page 1)

It is large! complex is less clear.

Thank you for pointing out this issue. We will make the change “high complexity of tree space” to “extremely large tree space” in the manuscript.

I don’t know, but the better fit for lower part of the posterior is worth mentioning. It may be key.

We agree with you. We will update the manuscript to highlight the significance of modeling suboptimal trees.

Coming from the intersection of variational inference and reinforcement learning is Reformulate.

We will reformulate this to make it more clear.

This is called dynamic programming (for Felsenstein's algorithm)

Yes, the recursive computation in Felsenstein's algorithm is an instance of dynamic programming. We can use the term "dynamic programming" if the text if you find it suitable.

a state $s$ consists of a set of rooted trees

Disjoint rooted …

Indeed, $s$ is a disjoint set of rooted trees (that is, a forest). We will update the manuscript accordingly to include the word "disjoint".

PhyloGFN explores and samples from the entire phylogenetic tree space, achieving a balance between exploration in this vast space and high-fidelity modeling of the modes. (Page 2)

No it has capacity or potential to do this.

Thanks for the suggestion, we will update the manuscript to "PhyloGFN has the capacity to explore ...".

PhyloGFN leverages a novel tree representation inspired by Fitch and Felsenstein’s algorithms to represent rooted trees without introducing additional learnable parameters (Page 2)

Make this clearer. In particular, in addition to what?

Thanks for pointing this out. We will update the manuscript to "without introducing additional learnable parameters to the model" .

We hope we have addressed all of your concerns. Please let us know if something else needs further clarification or if we missed something.

We appreciate your feedback and detailed comments. We may not have fully comprehended some of your questions. If this is the case, please let us know, and we will be happy to clarify.

GFlowNets introduction

It is not really trying to explain the GFN approach to the reader, but rather enumerating required items

This section could provide more insight (background section on GFlowNets)

Space limitations prevent us from giving a complete introduction to GFlowNets, but we believe that our introduction is self-contained and sufficient for understanding the application to phylogenetic inference. However, we have tried to improve the presentation. We point the reader to ["Trajectory balance...", NeurIPS 2022] as an accessible and self-contained introduction to the GFlowNet framework as relevant to our setting.

Running time and computation resources

The running time and performance in general require further explanation

The comparison must be made relative to resources since you use mcmc as gold standard

What is the computational resources required for these methods?

We describe the computation resources for PhyloGFN and compare the running time with other benchmark algorithms in our answer to all reviewers. We hope this answers your question.

Move the parsimony analysis to appendix

The parsimony case is clearly less interesting and could partially be move to the appendix

Bayesian inference and parsimony analysis are two types of problems that are usually solved by different families of algorithms. We believe that the fact that PhyloGFN is robust enough to solve both of them under the same framework is an important contribution worth highlighting. In particular, the following technical contributions related to parsimony-based inference are noteworthy:

• We solve the parsimony score minimization problem with a sampling approach by designing an energy distribution with the parsimony score as the energy. With sufficiently small temperature $T$, the distribution will be dominated by optimal solutions, and the learned GFlowNet can retrieve all parsimonious trees.
• We propose a tree representation based on Fitch node features. While it has similarity to the Felsenstein-based tree features for Bayesian inference, the proof of its sufficiency for optimal policy is different.
• We propose a temperature-conditioned PhyloGFN that takes $T$ as input and can thus trade off between diversity of tree topologies and low parsimony score at inference time.

GFlowNet SOTA

With their theoretical foundations laid out in Bengio et al. (2023); (Page 5)

Is this the gfn sota? Then point out this fact.

We are not sure what you mean by this question. GFlowNets do indeed achieve state-of-the-art sampling performance for many problems, but they never been applied to phylogenetic inference. The adaptation of GFlowNets to phylogenetic inference is a major contribution of our work.

On multifurcating trees

The tree topology can be either a rooted binary tree or a bifurcating unrooted tree. (Page 2)

This is a potential weakness. Can you restrict you method to binary trees? Does this mean that the posterior support of a subsplit consist of both binary and multifurcating trees? On an earlier reading i got the impression that you considered binary trees. Please clarify.

We believe there is a misunderstanding here. Like almost all approaches in phylogenetics (including all the recent work addressing the problem with machine learning), we do not consider multifurcating trees. The two types of trees considered are rooted and unrooted binary trees.

Question on trajectories/states of GFlowNets

With polynomial length trajectories, which you need, exponentially many trajectories implies exponentially many states. The ladder also yields your optimization infeasible, in worst case.

We respectfully disagree that this is a limitation. All machine learning-based phylogenetic inference algorithms learn to sample a distribution over an exponentially large space of of tree topologies, but they use the power of amortization and deep neural nets to generalize from a much smaller number of topologies seen in training.

To be precise, given $n$ sequences, there are indeed $(2n-5)!!$ tree topologies, which is superexponential. You are correct that it is unrealistic to let the model visit every single tree to learn the full distribution. This is exactly the reason we use deep neural networks to parameterize the action policy. With deep neural networks, the algorithm can learn the underlying statistical patterns of the trees and generalize to those unseen in training.

Dear Reviewer 7zzf,

We have responded to your questions and comments above. Could you please let us know if you have any further questions before the end of the rebuttal period and if they have affected your assessment of the paper?

We have also just posted a revised pdf. Please see the comment to all reviewers for details of the changes.

Thank you,

The authors

Tree topology construction and transformer model

do you have so consider all pairs of roots among the trees?

Yes, given a state consisting of $N$ disjoint trees, we need to choose a pair of trees out of all $N \choose 2$ possible pairs. We do not use any heuristic to limit the set of pairs. In this way, PhyloGFN models the entire tree space.

Can you give a more extensive account of the transformer, which is crucial since it allows you to, in contrast to VBPI, potentially can consider any possible tree?

We want to emphasize that the capability of modeling the entire tree space does not depend on the specific architecture of the transformer model. As we explained above, PhyloGFN can sample all possible trees because the action policy covers all possible pairs of trees at every state. In the PhyloGFN architecture, this is achieved by applying an MLP to each pair of condensed tree features, yielding logits for all possible actions of joining two trees by a common parent node. The transformer encoder is utilized to generate condensed feature representations of trees. While a transformer architecture is suitable for our input, which consists of a set of disjoint trees, it is not the determining factor enabling PhyloGFN to model the entire tree space. This is also why we relegated the architecture details to the appendix in our paper.

That being said, we are happy to provide more details regarding the transformer architecture itself: the hyperparameters used for the transformer encoder are recorded in Table S3 in the appendix. For all datasets, the transformer encoder consists of 6 self-attention blocks with 4 heads in each block. The feature embedding size is 128 and the activation function is GELU. Each self-attention block's layout is similar to the self-attention block of the ViT model (Dosovitskiy et al., 2020). It consists of a layer normalization, a self-attention layer, another layer normalization, and finally an MLP block. For a problem with $N$ input sequences, the time complexity of an $L$-layer transformer is $O(N^2 L)$.

We hope that our responses have addressed your questions and concerns sufficiently. Please let us know if we can provide any more clarifications.

Thank you for your feedback and interesting questions. We address them in turn below.

MLL estimation similar among different algorithms

While VBPI-GNN demonstrates MLL performance closely aligned with the state-of-the-art MrBayes, it is essential to re-emphasize two significant issues with VBPI-GNN:

1. The algorithm requires a pre-generated high-quality tree set to constrain the action space for tree construction. The performance of VBPI-GNN is highly contingent on the quality of this pre-generated tree set.

2. Due to the limitations on the action space imposed by the pre-generated tree set, VBPI-GNN cannot model the entire phylogenetic tree space. In other words, VBPI-GNN does not support the majority of phylogenetic trees. We are happy to offer a more detailed explanation of how VBPI-GNN constrains its model space. (In case there is confusion, this issue stands apart from our analysis in Table 2, Figure 2, and Figure S1, which highlight VBPI-GNN's limitations in learning suboptimal trees.)

Prior to August 2023, VaiPhy was the leading VI algorithm capable of modeling the entire tree space. However, as indicated in Table 1 (column $\phi$-CSMC), its MLL estimation is significantly inferior to the state-of-the-art. This observation serves as a key motivation for our work: the design of a phylogenetic inference algorithm leveraging GFlowNets, with the objective of modeling the entire tree space and improving MLL estimation.

Please also see the response to all reviewers, where we discuss the importance of modeling the entirety of the full phylogenetic tree space and VBPI-GNN's inferiority to PhyloGFN in this regard.

Experiment specification in main text

The $\pm$ are standard deviation and they are obtained from a single PhyloGFN model estimating the marginal log likelihood 10 times, using 1000 trajectories in each round. We are currently repeating the experiments for several runs to ensure significance. As described in the response to all reviewers, we are also repeating 3 sets of experiments with reduced training examples. We will update the MLL results table with new runs once we have obtained all the results and detail the setup in the text.

Explanation of tree representation is unclear

Perhaps it was not made clear how the node representations are derived: in Section 3.1, we describe Felsenstein’s algorithm and Fitch’s algorithm, which compute vectors of features for each node. For a given subtree with root node $r$, these features are exactly the vectors computed by these two dynamic programming algorithms for node $r$ (used for full Bayesian and parsimony-based inference, respectively).

Section 3.1 describes in detail how these two features can be calculated recursively using the features of the child nodes. We will update the manuscript to make the reference more clear in the tree representation section.

Possible misunderstanding regarding Table 2 and Figure 2

The analysis presented in Table 2 and Figure 2 is not intended to show diversity. Recall that we have obtained similar MLL compared to VBPI-GNN (e.g., -7018.6 for VBPI-GNN and -7018.9 for PhyloGFN on DS1). However, MLL mostly measures how well the algorithms model the high-probability trees. Therefore, in Table 2 and Figure 2, we are trying to answer the question of how well the algorithms model the rest of the tree space.

In Table 2, Figure 2, and Figure S1, we show that on PhyloGFN is better at modeling the entire tree space compared to VBPI-GNN. For a given set of trees, we compare the model-learned sampling probability against their ground truth posterior density. For each data set, we performed the analysis on three sets of data with high/medium/low posterior densities. The analysis shows that while both algorithms model high-probability trees well, PhyloGFN is much better at modeling the rest of tree space.

Please also see the responses to all reviewers regarding the importance of modeling the full space.

Dear Reviewer SMRb,

We have responded to your questions and comments above. Could you please let us know if you have any further questions before the end of the rebuttal period and if they have affected your assessment of the paper?

We have also just posted a revised pdf. Please see the comment to all reviewers for details of the changes.

Thank you,

The authors

Dear Reviewer SMRb,

We have completed the repeated experiments for all datasets and we have updated the manuscript. We believed we have also addressed your questions and concerns in our previous responses.

Could you please let us know if you have any further questions before the end of the rebuttal period and if they have affected your assessment of the paper?

Thank you,

The authors

Thank you for your feedback. We hope the answers below clarify your concerns.

Comparing running time with PAUP*

For parsimony analysis, our models are trained with 25.6 million examples on all datasets, with training duration of 3-8 days. The heuristic search algorithm of PAUP*; is able to discover optimal trees in less than 30 CPU seconds for all datasets. This is largely due to the optimization efforts that have been put into the software over the past two decades.

In this context, it is interesting to compare the number of trees evaluated by each method. For example in DS8, PAUP*; has attempted a total of 8,113,732 tree shape rearrangements (TBR branch swapping) for finding the 21 most parsimonious trees, while PhyloGFN is trained with 25,600,000 trajectories.

It is worth bearing in mind that, in contrast to PAUP*, (1) Bayesian phylogenetic inference algorithms, such as PhyloGFN, tackle the far more challenging problem of sampling from a distribution over the entire phylogenetic tree space where the most parsimonious trees are the global minima, and (2) the amortization technique allows us to control the distribution of tree samples using a single model conditioned on the temperature variable.

Despite all models undergoing training with 25.6 million trajectories, they exhibit the capability to produce parsimonious trees at significantly earlier stages in the training process. At intervals of every 5% of the training process for each model, we generate 10,000 samples and evaluate the parsimony scores of the sampled trees. This table shows the percentage of training time required for each model to sample at least one of the most-parsimonious trees:

We haven't fine-tuned the training setup specifically for fast parsimonious tree retrieval. By training with a more aggressive temperature annealing schedule, integrating an early stopping mechanism, and various code and hardware optimizations, it could become feasible to significantly reduce the number of training examples needed to generate the most parsimonious trees. However, we emphasize again that the goal of PhyloGFN and other Bayesian phylogenetic inference algorithms is to model the full distribution over trees, not simply to find the most parsimonious ones.

For running time for Bayesian inference and comparison with other variational inference algorithms, please refer to the response to all reviewers.

Utility of suboptimal structures

This is a good question; please refer to our answer to all reviewers.

Thank you for your response. "PAUP*; is able to discover optimal trees in less than 30 CPU seconds for all datasets." It seems much faster than training a GFN. In this case, the best utility of a trained GFN might be zero-shot learning on a new inference task. Have you tried that and if so, how dos the test inference time compare with PAUP*?

On modeling the distribution of suboptimal trees, if this is what PhyloGFN is better at compared with MCMC and VI, i.e. provide probability estimate or identify substructures, I really hope to see experiments that illustrate this in experiments. Especially given that this is a ML-application paper, I would hope to see how this proposed method unlock interesting applications that previous methods fail to achieve.

Dear Reviewer tiKV,

We have responded to your questions and comments above. Could you please let us know if you have any further questions before the end of the rebuttal period and if they have affected your assessment of the paper?

We have also just posted a revised pdf. Please see the comment to all reviewers for details of the changes.

Thank you,

The authors

Dear Reviewer tiKV,

Thank you for your effort in reviewing our work and so thoroughly justifying your assessment. We respect your decision.

Nevertheless, we still want to re-emphasize the following, not to convince you to change the score, but to clear up any possible misunderstandings for you or other reviewers: while modeling suboptimal trees well is a highlight for PhyloGFN, it is only one consequence of our powerful approach to Bayesian inference via sequential construction, rather than the main focus. To elaborate:

• Our improved modeling of the entire tree space without the need for pre-generated trees is an important contribution. Modeling the entire tree space is not quite the same as "modeling well suboptimal trees" or "generating from arbitrarily filled-in subtrees" as you described in your original comment. Without prior knowledge (for example, pre-generated trees), the algorithm would have no idea what trees are suboptimal or arbitrary unless it can properly explore and model the tree space.

• Modeling the full tree space faithfully is perhaps the primary issue that researchers in Bayesian phylogenetics are currently addressing: the technical challenges intensify across multiple scales when the algorithm does not depend on pre-generated tree sets to constrain the model space. Past and concurrent works in the last two years (VaiPhy, GeoPhy, and the concurrent ARTree) have proposed innovative solutions to tackle these challenges, and our current work builds and improves upon this line of research, showing for the first time that deep reinforcement learning approaches bring unique advantages to this task.

Thank you for your consideration.

The authors

Thank you for the detailed response. I fully agree with the importance of sampling suboptimal trees and believe this is what the model should focus on. I also agree that current experiments show that GFN is promising in sampling from the tree distribution. But just that I wish to see something more interesting that can be achieved with GFN (in the tree sampling setting) beyond log-likelihood since it is an application paper. The time comparison with MCMC is helpful in understanding how GFN trains. I don't have further questions at the moment. My current standpoint is neural towards accepting so I will not change my score, although 5.5 better reflects where I stand.

Dear Reviewer tiKV,

We have updated the manuscript to add more references delineating the importance of sampling suboptimal trees (section appendix H). We thank you again for suggesting interesting future work directions. We believed we have addressed your questions and concerns in our previous responses.

Could you please let us know if you have any further questions before the end of the rebuttal period and if they have affected your assessment of the paper?

Thank you,

The authors

Thank you for feedback. We answer your questions and concerns below.

On axes in Figure 2

In Table 2 and Figure 2, we compute log sampling probability of PhyloGFN and VBPI-GNN and compare them to the ground-truth unnormalized posterior density for optimal trees and suboptimal trees. Log sampling probability (X-axis) is calculated using the algorithms in comparison (PhyloGFN and VBPI-GNN). Ground-truth unnormalized density (Y-axis) of tree $(z,b)$ is directly calculated as $P(Y,z,b) = P(Y|z,b)P(z)P(b)$. We assume a uniform distribution over tree space, hence $P(z) = \frac{1}{(2n-5)!!}$ and $P(b) = \prod_e P(b(e)), \ \ b(e) \sim \text{Exp}(10)$.

For the reviewer’s question "why do points fall in similar Y ranges but different X ranges", we see two interpretations.

If we interpret it as "comparing the VBPI-GNN plot to the corresponding PhyloGFN plot, why do the points fall in similar Y ranges but different X ranges": to ensure fair comparison, we are using the same set of trees. As the Y-axis corresponds to the ground truth joint density, they indeed should have the same range. The X-axis labels the sampling probability learned by the two models. In the VBPI-GNN plot, the points falling in a narrower X range indicate that VBPI-GNN's sampling probabilities for these suboptimal trees are learned incorrectly: VBPI-GNN has overestimated the posterior probability of these suboptimal trees.

If we interpret it as "why do the X-axis and Y-axis ranges within each plot differ", this is because the Y-axis shows the joint log-density $\log P(Y,z,b)$, not the posterior log-density $\log P(z,b|Y)$. For a perfectly trained model, all points in the graph would lie on a diagonal line with slope 1 and intercept equal to the MLL $\log P(Y,z,b)-\log P(z,b|Y)=\log P(Y)$.

Interpretable analysis for Bayesian inference algorithms

We agree with the reviewer’s comment. Using the MLL alone, it can be difficult to assess and interpret the relative performance of different algorithms, particularly when the estimated MLLs are similar or have high variance. This is precisely what motivates our quantitative and qualitative analysis that compares the learned sampling probability of PhyloGFN and VBPI-GNN to the ground truth unnormalized posterior density (Table 2, Figure 2, and Figure S1). These analyses show that:

1. For all datasets except DS7, although PhyloGFN has slightly lower MLL than VBPI-GNN, both algorithms are able to model the high-probability trees well. Therefore, from a practical standpoint, PhyloGFN is at least as good as VBPI-GNN in its tree sampling capability.

2. For suboptimal trees, both the correlation scores (Table 2) and scatter plots (Figure 2, Figure S1) show that PhyloGFN is actually superior to VBPI-GNN at modeling suboptimal trees.

PhyloGFN'S MLL estimation is inferior to VBPI-GNN's

We have identified three potential reasons for VBPI-GNN's better MLL estimation:

1. PhyloGFN models branch lengths using discrete multinomial distributions. When estimating MLL, the branch lengths model is converted into a piecewise-constant continuous form, inducing a small quantization error. In contrast, VBPI-GNN models branch lengths with a continuous distribution directly.

2. VBPI-GNN utilizes a set of pre-generated high-likelihood trees to restrict the search space. Since most trees with high posterior density also have high likelihood, it is easier for VBPI-GNN to explore the high posterior density region. (In this sense, the comparison is not entirely fair, since VBPI-GNN uses information that the training of PhyloGFN does not have access to.)

3. There is room for improvement in model fitting for DS7. Both the MLL study and the study on sampling probability against ground truth indicate that PhyloGFN poorly models the high-probability trees for DS7.

In future work, we intend to address these limitations and improve PhyloGFN's MLL estimation.

Significance of modeling the entire tree space well

Please see the response to all reviewers, where we describe important biological questions where properly sampling suboptimal trees is important.

Dear Reviewer iXKc,

Thank you for taking the effort of reviewing our paper.

The authors

Significance of modeling suboptimal trees well

Reviewers tiKV and iXKc asked about the significance of modeling suboptimal trees well. Several approaches that rely on phylogenetic inference require not only the identification of optimal (maximum likelihood or most parsimonious) trees, but also the proper sampling of suitably weighted suboptimal trees. In evolutionary studies, this includes the computation of branch length support (i.e., the probability that a given subtree is present in the true tree), as well as the estimation of confidence intervals for the timing of specific evolutionary events. These are particularly important in the very common situations where the data only weakly define the posterior on tree topologies (e.g. small amount of data per species, or high degree of divergence between sequences). In such cases, because suboptimal trees vastly outnumber optimal trees, they can contribute to a non-negligible extent to the probability mass of specific marginal probabilities. Full modeling of posterior distributions on trees is also critical in studying tumor or virus evolution within a patient, e.g., to properly assess the probability of the ordered sequence of driver or passenger mutations that may have led to metastasis or drug resistance.