Paper: PhyloGFN: Phylogenetic inference with generative flow networks (ICLR-2024)

Reviewer comment (Evaluation):
The biggest weakness looks to me like the results in Table 1. The log-likelihood scores look very similar to each other. For instance -7108.95 for PhyloGFN and 7108.95 VBPI-GNN. Similar for other data sets.

Author response:
## 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.

Reviewer comment (Evaluation):
The paper reports log-likelihood results, which measures model fit. It measures parsimony, which appears to be about the computational aspect. The log densities reported in Figure 2 is an indicator of diversity. It may be better to have a more direct and interpretable score of the discovery performance, for instance prediction accuracy of evolution tree links or ancestral relationship detection between pairs of species in cases where there is agreement on the ground truth. The current result landscape looks a bit too exploratory. While the results make intuitive sense, they still leave many gray areas in their detailed interpretation.

Author response:
## 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.

## Metric to show diversity
We do address the diversity of samples with Temperature-conditioned PhyloGFN. Recall that PhyloGFN-Parsimony learns to fit an energy distribution of tree topologies where the parsimony scores of trees are scaled by the temperature $T$.  At a small $T$, trees with low parsimony score would dominate the energy distribution, so the sample diversity would be low. For example, in dataset DS1, there is only 1 tree topology with the minimum parsimony score, and it would make up most of the sample as $T\to0$.

Therefore, we have proposed the temperature-conditioned PhyloGFN, a single sampler that accepts the temperature $T$ as input and is able to sample from the target energy distribution at that temperature. Figure 3 on page 9 shows the sampling results of our temperature-conditioned PhyloGFN on DS1. On the left, we plot the histograms of parsimony scores of trees sampled at different $T$. We see that as we increase $T$ gradually from 1 to 8, the sample diversity increases accordingly. The plot on the right shows that under different temperatures, PhyloGFN-Parsimony is still able to sample proportionally from the ground truth distribution. 

## Validate PhyloGFN against ground truth evolution data

For any given set of species, their ground truth evolutionary history and relationships are generally unknown, so we cannot directly evaluate the accuracy of any of the phylogenetic inference approaches. When this is of interest, researchers use simulated data, where the tree is known, to assess the ability of an approach to recover the correct tree. However our objective is slightly different: our goal is to sample from the posterior distribution over trees defined by a given set of sequences. This is a purely mathematical problem; the correctness of the sampling produced by a given method can thus be assessed without the knowledge of the evolutionarily correct tree. This is what is done in our paper and is consistent with a long line of prior work in phylogenetic inference.

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

Reviewer comment (Evaluation):
Are the ± values given in Table 1 standard deviation or standard error? What are these standard deviations/error over? Experiment repetitions? If standard deviation across repetitions, the results may be good enough. They may be alarming otherwise.

Author response:
## 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.

Reviewer comment (Evaluation):
I also find it somewhat surprising how much VBPI-GNN outperforms PhyloGFN on the MLL estimation task.

Author response:
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.

Reviewer comment (Presentation):
Are the axes flipped on Figure 2? If I understand correctly the x-axis should be the unnormalized posterior under the long run of MrBayes, and the y-axis is the unnormalized posterior for either GFN or VBPI. Why do the points fall on different x-axis ranges for the two columns but very similar y-axis ranges? (Similarly for Figure S1)

Author response:
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)$.

Reviewer comment (Writing):
"Given a set of observed sequence" --> "Given a set of observed sequences"

Author response:
Thank you for pointing out the typos in our manuscript. We have corrected them.

Reviewer comment (Writing):
"ground-trueth" --> "ground-truth"

Author response:
Thank you for pointing out the typos in our manuscript. We have corrected them.

Reviewer comment (Novelty):
A preprint that was posted shortly after the ICLR deadline is highly relevant: https://arxiv.org/abs/2310.09553. Since that was posted after the ICLR deadline, it did not influence my review of this paper. The method in that preprint, ARTree, is based on reinforcement learning and also seeks to learn a posterior via framing a tree as the outcome of sequential tree building. Yet, the training objectives and MDP formulation are substantially different, and so I do not see ARTree as reducing in any way the novelty of the present work. I just bring this up in the hopes that it is interesting/helpful to the authors.

Author response:
Thank you for sharing this paper with us. At the inception of the project, we had considered constructing phylogenetic trees in a similar top-down fashion, but in the end we designed the current bottom-up PhyloGFN, because this way we can obtain a more efficient subtree representation based on the Fitch/Felsenstein algorithms.  

One property of the ARTree construction process is that one phylogenetic tree is mapped uniquely to one construction path. For this, the order of input sequence needs to be predetermined. However, if an MDP is designed exactly this way, actions at each step may be less intuitive. Depending on the order of sequences, the agent may have difficulty identifying the optimal decision for inserting the new species to the current partial tree topology. Of course, through learning from a large set of examples, the agent will know that its current action will be beneficial at certain points in the future. However, let us now consider a more challenging problem: learning a generalized agent that can take any input set of sequences and perform phylogenetic inference. Given an agent and a set of sequences that the model has never seen before,  it may have difficulty estimating each action’s future impact without seeing any examples in advance. On the other hand, with our bottom up approach, we believe it may be easier to solve the generalized problem. For example in parsimony analysis, the best action in many scenarios is to select and join the two subtrees with the closest root Fitch feature, which is a relatively easy concept to learn for PhyloGFN.

Reviewer comment (Evaluation):
*Performance vs. Efficiency*: While GFlowNets might perform comparably to PAUP* in the parsimony-based inference setting, the real differentiator would be computational efficiency on a new inference task. Unfortunately, no wall-clock time data is provided, making it challenging to discern any advantages of GFlowNets in this scenario.

Author response:
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:

| DS1 | DS2 | DS3 | DS4 | DS5 | DS6 | DS7 | DS8 |
|-----|-----|-----|-----|-----|-----|-----|-----|
| 15\% | 20\% | 20\% | 45\% | 65\% | 75\% | 45\% | 45\% |

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.

Reviewer comment (Novelty):
*Methodological Novelty*: The paper does not seem to bring forth significant machine learning methodological advancements, with much of the methodology being straightforward applications without any notable novel ML innovations for the particular problem at hand. The biggest methodological contribution is the use of discretized bins for branch lengths, which is a compromised solution.

Author response:
## ML innovation
We acknowledge the reviewer's concern regarding the methodological contribution and novelty of PhyloGFN. We believe that our study is an innovative application of the GFlowNet framework to the problem of phylogenetic inference. Our approach outperforms two other state-of-the-art variational inference methods, VaiPhy and GeoPhy, both of which also sample from the entire phylogenetic tree space. The success of our method relies on several key methodological innovations:

First, we have proposed a Markov decision process that allows us to generate phylogenetic trees in a bottom-up fashion, through iteratively joining pairs of subtrees and creating new roots, until a complete tree is formed. To the best of our knowledge, this is the first bottom-up phylogenetic tree construction scheme proposed in the ML-for-phylogenetics literature. In VBPI-GNN and its more recent extension ARTree, trees are constructed in a top-down fashion, which requires learning representations of the ancestral nodes using Dirichlet energy minimization and graph neural networks. For ARTree in particular, the Dirichlet energy minimization and graph representation learning at each step raises questions of scalability to larger trees. 

Second, our bottom-up tree construction scheme allows us to employ two kinds of provably-powerful and yet easy-to-compute features for the ancestral nodes -- those produced by applying the Fitch and Felsenstein algorithms. After joining the pair of subtrees, the newly created root node is represented by these features. This offers several advantages: (1) it allows us to directly model the new subtree without introducing any learnable parameters such as a graph neural network; (2) this feature representation is very fast to compute exactly; (3) it is scalable to larger phylogenetic inference datasets. Again, to the best of our knowledge, we are the first to design and employ these features in the machine learning context, and we theoretically justify that these features provide sufficient information for any optimal sampling policy.

Third, the parsimony problem is traditionally treated as one of discrete optimization. In this work, we frame it as a hierarchical inference problem: we train a PhyloGFN to sample from an energy distribution over tree topologies, and by manipulating the temperature on which the model is conditioned we can control the trade-off between the sampled tree parsimony scores and diversity. At inference time, we can sample from a series of energy distributions at different temperatures. 

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

Reviewer comment (Evaluation):
*Utility of Results*: While the GFlowNets surpass VBPI methods in estimating probabilities for suboptimal structures, the practical significance of this in the context of phylogenetic applications remains unexplained. The experiments focus solely on the accuracy of probability estimates and fail to provide insights into their relevance for the broader application.

Author response:
Dear reviewer tiKV,

Thank you for your comments and questions. We address your specific questions and comments below.

## Applying PhyloGFN to zero-shot learning 

> the best utility of a trained GFN might be zero-shot learning on a new inference task. Have you tried that

We agree with the reviewer that extending PhyloGFN to zero-short learning is a valuable direction, but as of currently we have not tested PhyloGFN on new inference tasks with unseen sequences. We must point out that zero-shot learning is a highly challenging task of phylogenetic inference which has yet to be explored in any existing machine learning literature, but we do intend to answer this question with our PhyloGFN in a follow-up work. 

In order to enable zero-shot learning, that is to apply a pre-trained PhyloGFN on any arbitrary set of sequences and inferring their phylogeny, further modifications on the model architecture and training is required. Just to give a few examples:
- as is mentioned in appendix section D, we currently apply a linear transformation on the input sequences to obtain lower-dimensional embeddings, before they are input to the Transformer encoder. However, this linear transformation is unable to process sequences of varying length. We are currently investigating more flexible models to model to embed sequences of arbitrary length.
- modification to PhyloGFN training is also needed. At the moment, each model is trained on only one dataset with up to 64 sequences. For the zero-shot learning, it would be imperative for our PhyloGFN to be trained on a large amount of sequences across many datasets, in order to ensure generalizability.

Once again, we thank our reviewer for bringing up this research direction, and as part of our future plans, we have highlighted this direction in the conclusion of the manuscript. 

## Possible confusions on advantage of PhyloGFN 

> The approach not only matches up to VBPI-GNN-based methods in Bayesian posterior inference but offers capabilities beyond them, like generating from arbitrarily filled-in subtrees which preceding methods could not tackle
>
> 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


As in our response to Reviewer SMRb, We would like to re-emphasize two significant issues of VBPI-GNN:

1. VBPI-GNN 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.

We can summarize the advantages of PhyloGFN as follows:
- Comparing with VI algorithms that model the entire tree space, PhyloGFN improves significantly the performance in term of MLL estimation. Comparing with VBPI-GNN, the state-of-the-art VI algorithm, PhyloGFN models the entire tree space and achieves competitive MLL estimation without the need of a pre-generated high-quality tree set. Moreover, even if we consider limited modeling space of VBPI-GNN, PhyloGFN is superior at modeling sub-optimal trees. 
- Comparing with MCMC, the advantages of PhyloGFN are (i). PhyloGFN is an amortized sampler, able to very quickly generate a large number of independent tree samples once it has been trained, (ii). scalability to high-dimensional distribution

Reviewer comment (Evaluation):
*Time Efficiency*: Given the comparable performance of GFlowNets and PAUP*, can the authors provide information on the computational efficiency (wall-clock time) of GFlowNets to discern its advantages or disadvantages?

Author response:
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:

| DS1 | DS2 | DS3 | DS4 | DS5 | DS6 | DS7 | DS8 |
|-----|-----|-----|-----|-----|-----|-----|-----|
| 15\% | 20\% | 20\% | 45\% | 65\% | 75\% | 45\% | 45\% |

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.

Reviewer comment (Evaluation):
*Utility of Suboptimal Structures*: Can the authors elucidate the practical implications of having estimates for suboptimal structures in the context of phylogenetics? Is there a tangible benefit in real-world applications?

Author response:
## On suboptimal trees
We appreciate the reviewer’s suggestion that there could be further interesting experiments concerning the use of PhyloGFN to analyze suboptimal phylogenetic trees. We note that in our response to all reviewers we have discussed some of these application areas where PhyloGFN could be of high interest. For example, understanding cancer and virus evolution where the capability of modeling the full tree space, identifying common substructures of trees and quantifying uncertainties of phylogenetic hypotheses are particular advantages of our PhyloGFN. While these are exciting directions for future work, unfortunately we would not be able to include them during the rebuttal period due to time constraints.  We respectfully note that the breadth and thoroughness of our experimentation is quite in line with related works exploring novel methods for phylogenetic inference, such as VBPI-GNN, VaiPhy and the previously mentioned concurrent work (GeoPhy and ARTree).

We thank our reviewer for bringing up these research directions. We hope that our responses have addressed your questions and concerns sufficiently. Please let us know if we can provide any more clarifications.