ARTreeFormer: A Faster Attention-based Autoregressive Model for Phylogenetic Inference

Authors: Tianyu Xie, Cheng Zhang
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)
Keywords: phylogenetic inference, autoregressive model, attention mechanism

Abstract:
Probabilistic modeling of the combinatorially explosive tree topology space has posed a significant challenge in phylogenetic inference. Previous approaches often necessitate pre-sampled tree topologies, limiting their modeling capability to a subset of the entire tree space. A recent advancement is ARTree, a deep autoregressive model that offers unrestricted distributions for tree topologies. However, the repetitive computations of topological node embeddings via Dirichlet energy minimization and the message passing over all the nodes can be expensive, which may hinder its application to data sets with many species. This paper proposes ARTreeFormer, a novel approach that harnesses attention mechanisms to accelerate ARTree. By introducing attention-based recurrent node embeddings, ARTreeFormer allows the reuse of node embeddings from preceding ordinal tree topologies and fast vectorized computation as well. This, together with a local message passing scheme, significantly improves the computation speed of ARTree while maintaining great approximation performance. We demonstrate the effectiveness and efficiency of our method on a benchmark of challenging real data phylogenetic inference problems.

Decision: reject
Reviewer scores: 8, 8, 3, 5, 5, 6

Reviews and discussion:
<thread>
creation date:
1730384450133

thread_type:
official_review

<thread_object>
creator: reviewer_2sl2

<content>
summary:
This submission extends the existing ARTree model for inferring phylogenetic posterior distributions by using attention mechanisms to compute the node embeddings. The resulting algorithm, ARTreeFormer, furthermore employs a constrained message-passage scheme which in total improves the run times of the algorithm in contrast to ARTree. The algorithm is evaluated on popular datasets commonly used in the VI-based phylogenetic inference literature.

strengths:
The paper is well-written, and the proposed method tackles a problem which is of interest to the ML community and, importantly, reduces run times by significant factors. The use of attention mechanisms to tackle important problems is also aligned with the interests of the ICLR community. I think the submission is well-suited for publication at ICLR and that others will find it interesting, however I have some (actionable) concerns, given below.

weaknesses:
Apart from the impressive improvements in run times, the results in terms of posterior accuracy (Table 1) and ML estimates (Table 2) are modest. For the ML estimates, this is expected as the scores are highly saturated, why most of the other works referenced in this submission instead value the size of the standard deviation (smaller is better). W.r.t. std, ARTReeFormer never outperforms the other VBPI methods. Regarding this I have one question and one comment:

* Question: Why do you think the std is larger for ARTreeFormer than the other VBPI methods? It would be nice to see a short discussion about this. The run times are already very good, so there is nothing to loose from transparently discussing this limitation of ARTreeFormer. 
* Comment: There are other VBPI models available, such as [1,2,3,4], all of which (I believe) use pre-computed candidate trees. Notably [1] and [2] report better ML scores than those reported in Table 2 in this submission. See for instance DS8 in Table 2 in [1]. As such, the following statement in the current submission does not seem to hold: "Compared to other VBPI variants, the superiority of ARTreeFormer is mainly reflected by its unconfined support[...]". (This point weights down my soundness score).

Apart from my comment above, I think [1] and [2] should be referenced in the submission to demonstrate the growing interest in VI-based phylogenetic inference in the ML community. (This point weights down my presentation score).

In [5] they show that more importance samples during training can harm the inference of the variational approximation. In Eq. 3 in the submission it states that the models are trained using the importance weighted ELBO. How does the choice of $K$ affect the quality of your inferred approximations? Can you say something about the expected differences in scores for say $K=1$ or $K=100$?

If my concerns above are handled/discussed, I am willing to raise my overall rating to an acceptance level.

[1] https://arxiv.org/abs/2310.00941

[2] https://arxiv.org/abs/2406.07083

[3] https://pmc.ncbi.nlm.nih.gov/articles/PMC9949155/pdf/nihpp-2302.08840v1.pdf

[4] https://proceedings.neurips.cc/paper/2020/hash/d96409bf894217686ba124d7356686c9-Abstract.html

[5] https://arxiv.org/abs/1802.04537

questions:
I am aware that there are many papers now in ML conferences on VI-based phylogenetics. Something that I am missing in these papers is a strong motivation of why better phylogenetic posterior inference algorithms are needed? As I have stated above, most of the results on the considered experiments appear saturated, and so it would really increase the significance of these papers if there were motivations that capture the necessity of these models and in which sense biologists would be interested in them. Do you think it is possible to provide such a motivation, potentially with references to works were similar types of models are used to draw biological conclusions?

Note that my question above is not a critique of your submission. It would however strengthen the paper, and potentially make it impactful outside of the ML community, if such a motivation was to be provided.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730473029390

thread_type:
official_review

<thread_object>
creator: reviewer_frst

<content>
summary:
The ARTreeFormer model introduces an innovative attention-based approach to accelerate phylogenetic inference, addressing computational efficiency challenges in large-scale phylogenetic studies. Through recursive node embedding and local message-passing mechanisms, ARTreeFormer significantly reduces computational time while maintaining high performance across various tasks. Key improvements include the effective use of Transformer-based attention to prioritize meaningful interactions and the capability for parallel processing, enhancing its suitability for handling complex tree structures.

strengths:
1. The field has made significant contributions and solved important challenges in phylogenetic inference: The introduction of the ARTreeFormer model has important academic and practical significance. Phylogenetic inference is extremely valuable in biological and ecological research, but existing methods face efficiency bottlenecks when processing large-scale data. Through innovative design, ARTreeFormer accelerates the inference process without significantly sacrificing accuracy, filling the gap in efficient phylogenetic inference models and promoting research and practical applications in the field.
2. Significantly improves efficiency and is suitable for large-scale data analysis: ARTreeFormer effectively reduces the computational time of the model through a recursive node embedding and local message passing strategy using an attention mechanism. This design allows the model to quickly generate and update tree topologies, which is especially advantageous when dealing with large numbers of nodes and complex tree structures. Experimental results show that ARTreeFormer is significantly superior to traditional models in terms of running time, making it highly competitive in processing large-scale phylogenetic data and promising for use in diverse biological data that requires efficient inference.

weaknesses:
1. Limitations of performance improvement: Although ARTreeFormer shows significant improvement in computational efficiency, the improvement in inference accuracy is limited. In the tasks of maximum simplicity and tree topology density estimation, ARTreeFormer is only slightly better than existing models such as SBN-SGA and ARTree on some datasets (see the KL divergence results in Table 1). This result shows that the main advantage of ARTreeFormer is concentrated on computational speed, and the accuracy bottleneck has not yet been completely broken through. If the accuracy can be further improved, especially in terms of generalization ability and performance in complex scenarios, the model value and applicability will be greatly improved, making it more competitive for practical applications.
2. Model complexity increases and implementation difficulty: ARTreeFormer introduces a multi-layer attention mechanism and a local message passing module, which not only accelerates the calculation but also makes the model structure more complex. In “Section 3.2 Attention-Based Edge Decision Distribution”, the use of multi-head attention mechanisms increases the implementation complexity of the model and may place higher requirements on technical implementation. This complexity may limit the promotion and application of the model, especially in scenarios where resources are limited or efficient parallelization is difficult. Therefore, a better balance between structural optimization and implementation difficulty will help to improve the potential for widespread application of the model.
3. Lack of generalization verification on diverse datasets: The experimental evaluation mainly focused on eight specific phylogenetic datasets (“DS1-8”), which, although representative, have limited coverage. The applicability and robustness of the model to other types of biological data (such as viral evolutionary data or ecosystem data) have not been fully considered in the article, making it difficult to verify its generalization ability in diverse phylogenetic scenarios. Extending the experiment to a wider range of datasets would help to more comprehensively evaluate the applicability of ARTreeFormer, thereby enhancing the credibility of the model's application in the fields of bioinformatics and ecology.
4. Potential information loss in local message passing: Although the local message passing mechanism significantly reduces the computational burden, it may lead to potential global information loss. As mentioned in “Section 3.2 Local Message Passing Module”, only the neighborhood of the added node is updated, and the potential information loss of this mechanism in long-distance dependencies or complex tree structures is not discussed. Although the local update strategy improves efficiency, there may be a trade-off in terms of information integrity, which may affect model performance, especially in cases where node associations are complex or global consistency requirements are high. In the future, finding a better balance between efficient local updates and global information access will help improve the overall performance of the model.
5. Theoretical improvement: The depth of the explanation is insufficient. The article fails to provide sufficient theoretical support for several innovative points. For example, “Equation 12” only describes the calculation method of the local Dirichlet energy, without discussing in detail the theoretical basis of this energy minimization strategy and its relative advantages over global minimization. In addition, the specific optimization strategy of the local message passing mechanism is only briefly described. A more in-depth theoretical analysis would help to enhance the persuasiveness of the model design, so that readers can better understand its applicability and rationality in complex inference tasks.
6. Limitations of parallel processing in practical applications: ARTreeFormer supports the parallelization of node embedding and message passing on different tree topologies, which improves processing speed. However, “Figure 2 (right)” shows the significant acceleration effect of parallelization compared to non-parallelization, but does not discuss the feasibility in resource-constrained or single-machine environments. Although parallelization significantly improves running efficiency, its feasibility and practicality in environments with limited hardware resources still needs to be further verified. Optimizing model efficiency in single-machine or low-resource environments will be more conducive to the popularization of this model in practical applications.

questions:
1. Extensiveness of experimental data sets and the problem of updating: Current experiments focus on the DS1 to DS8 data sets, which are relatively old and have a limited sequence length, making it difficult to fully reflect the diversity and complexity of modern biological data. It is recommended that tests be conducted on updated and diverse data sets (such as longer sequence data or more complex phylogenetic tree structures) to more comprehensively evaluate the robustness and generalization ability of ARTreeFormer.
2. How to use Transformer to increase the interpretability of the model and open the black box model: The Transformer's attention mechanism is introduced in ARTreeFormer, which, despite enhancing performance and efficiency, still has limited interpretability. The attention weights of the Transformer provide a potential interpretable pathway. For example, the key nodes and paths that the model focuses on when generating tree topology can be revealed by analyzing the attention distribution. However, the current model lacks a systematic interpretative mechanism. Future research can further explore how to use the attention weights to reveal the decision-making process of ARTreeFormer, helping to open the black box and make it more transparent and understandable in the task of biological evolution inference.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730688022686

thread_type:
official_review

<thread_object>
creator: reviewer_3onl

<content>
summary:
The paper proposes ARTreeFormer, which aims to improve the computational efficiency in phylogenetic inference tasks by improving the time complexity of the global multi-head attention mechanism (MHA). The motivation of this paper focuses on solving the computational efficiency problems caused by the minimisation of Dirichlet energy and message passing to all nodes. The innovation of the paper lies in the use of the Transformer architecture to achieve faster ancestor sampling and probability evaluation than the baseline ARTree with the improvement of reducing complexity from O(N3) to O(N2). The experimental part verifies the time efficiency and performance of the model and demonstrates its performance in specific benchmark tests.

strengths:
* (S1) ARTreeFormer achieves a better computational complexity with the recurrent node embedding module using local operations in comparison to the traditional multi-head attention mechanism, which is innovative in terms of computational efficiency. It can be supported by comparison experiments with the improvement of the performance and efficiency on some specific tasks.

* (S2) The overall presentation is well-originaized and easy to follow. It is clear that the authors provide step-by-step designs to improve the attention mechanism for better performance and efficiency.

weaknesses:
* (W1) This paper emphasizes the time complexity of improving the MHA module but without showing some ablation experiments that remove the MHA module, which might not be a fair verification to support the contribution of the MHA module. The novelty and motivation of this paper are somewhat unclear (refer to Q1).

* (W2) The proposed graph pooling function with the MHA and the message-passing process all learn local structures and node information, which might lack the aggregation of global information. The authors should verify whether these designs are reasonable for phylogenetic inference. Are there any intrinsic drawbacks or necessities of capturing long-distance dependencies or global topologies in these tasks?

* (W3) Since computational efficiency is an important contribution, the authors should clarify the discussion of the time complexity formula in line 213. But it is only briefly mentioned in the main text. I recommended the authors provide a more detailed derivation to show the complexity reduction from $O(N^3)$ to $O(N^2)$. Some detailed explanation of each module is unclear, especially the simplification process of the MHA.

questions:
* (Q1) It seems that the motivation of this work is to address the computational efficiency of Dirichlet energy minimization and message passing to all nodes. Although the proposed module can achieve faster ancestor sampling and probability evaluation than ARTree without using the global MHA, the overall innovation still seems to be insufficient and needs clarification.

* (Q2) The practical value of the improvement results is unclear. The paper improves time efficiency but does not demonstrate the significance of this improvement in practical applications, nor does it provide a specific dataset or scenario to verify the actual necessity of such efficiency improvements in this field. The reduction of complexity from $O(N^3)$ to $O(N^2)$ is interesting, but it may be more suitable and important to improve the multi-head attention mechanism itself (e.g., using linear attentions and Mamba variants).
</content>
</thread_object>

</thread>
<thread>
creation date:
1730695798308

thread_type:
official_review

<thread_object>
creator: reviewer_3ipi

<content>
summary:
This paper introduces ARTreeFormer, a novel probabilistic model over tree topologies. ARTreeFormer is a variant of the previously proposed ARTree model, with the advantage that ARTreeFormer is faster while showing comparable accuracy. The speed improvement is achieved by replacing the topological node embedding module and the GNN edge-decision layer with (1) attention-based recurrent node embeddings and edge-decisions, (2) a local Dirichlet energy minimization procedure, and (3) a local message passing scheme which only updates the node embeddings in the neighbourhood of the newly added node. On three standard benchmarks (the large parsimony problem, TDE, and VBPI), the paper demonstrates that ARTreeFormer is ~3x faster on learning tasks while showing comparable accuracy to ARTree.

strengths:
This paper demonstrates that topological node embeddings derived from Dirichlet energy minimization are not needed to achieve near-SoTA results. Instead, learnt attention-based recurrent node embeddings can provide similar performance (if not slightly worse - see Table 1). The use of attention-based edge decisions and the local message passing scheme are also new ideas wrt ARTree which provide speedups to the model.

weaknesses:
For a work whose main contribution is a more computationally efficient model (in the author's words, line 338: "It should be emphasized that we mainly pay attention to the computational efficiency improvement of ARTreeFormer and only expect it to attain similar accuracy with baseline methods"), it is surprising that there is no analysis whatsoever of the computational complexity of ARTree and ARTreeFormer. My expectation coming into the paper was that (1) the paper would discuss the computational complexity of each main step of the ARTree model using Big-O/Theta/Omega notation, deriving the runtime as a function of key quantities such as the number of nodes in the tree ($n$), the node feature dimension ($d$), etc., and discussing the degree of vectorization/parallelism achievable in each step, (2) the paper would then go on to profile each step of ARTree to confirm that the theoretical runtimes of these step align with the empirical runtimes, thereby establishing the bottlenecks, and (3) the alternative ARTreeFormer model would be proposed and its runtime again analyzed theoretically and empirically. The paper does not attempt to perform any computational complexity analysis of this kind (with the exception of the comments on line 213 and line 266). This leaves me wondering which are the real computational bottlenecks and why. Or - are they the result of inefficient implementation?

In particular, the paper claims that the node embedding module is a big computational bottleneck: "As N increases, the total run time of ARTree grows rapidly and the node embedding module dominates the total time (≈ 65%), which makes ARTree prohibitive when the number of leaf nodes is large. The reason behind this is that compared to other modules, the node embedding module can not be easily vectorized and parallelized w.r.t. different tree topologies and different nodes, resulting in great computational inefficiency.". I am skeptical about this claim. By looking at the code provided by the authors, as well as at the original ARTree code, the implementation of Dirichlet energy minimization used by the node embedding module seems to be a highly inefficient Python implementation. Here is what I found:

```
def node_embedding(self, tree, level):
    name_dict = {}
    j = level
    rel_pos = np.arange(max(4,2*level-4))
    for node in tree.traverse('postorder'):
        if node.is_leaf():
            node.c = 0
            node.d = self.leaf_features[node.name]
        else:
            child_c, child_d = 0., 0.
            for child in node.children:
                child_c += child.c
                child_d += child.d
            node.c = 1./(3. - child_c) 
            node.d = node.c * child_d
            if node.name != '':
                rel_pos[node.name] = j
            node.name, j = j, j+1
        name_dict[node.name] = node
    
    node_features, node_idx_list, edge_index = [], [], []            
    for node in tree.traverse('preorder'):
        neigh_idx_list = []
        if not node.is_root():
            node.d = node.c * node.up.d + node.d
            neigh_idx_list.append(node.up.name)
            
            if not node.is_leaf():
                neigh_idx_list.extend([child.name for child in node.children])
            else:
                neigh_idx_list.extend([-1, -1])              
        else:
            neigh_idx_list.extend([child.name for child in node.children])
        
        edge_index.append(neigh_idx_list)                
        node_features.append(node.d)
        node_idx_list.append(node.name)
    
    branch_idx_map = torch.sort(torch.tensor(node_idx_list,device=self.device).long(), dim=0, descending=False)[1]
    edge_index = torch.tensor(edge_index,device=self.device).long() 
    
    return torch.index_select(torch.stack(node_features), 0, branch_idx_map), edge_index[branch_idx_map], torch.from_numpy(rel_pos).to(self.device), name_dict
```

Notice not only the use of for loops which are notoriously slow in Python, but also the ete3 tree object which is traversed as part of the algorithm. An efficient implementation of Dirichlet energy minimization, say in C++, would likely be a lot faster (say >=10x). If so, ARTree might catch up to the speed of ARTreeFormer while giving slightly better results (see Table 1). Dirichlet energy minimization being a bottleneck is also suspicious given that it is a linear-time algorithm ($O(nk)$ where $n$ is the number of leaves and $k$ is the embedding size).

There are other natural options beyond an efficient C++ (or numba, or cython) implementation to potentially speed up the Dirichlet energy minimization, for example:
- The linear system $Ax=b$ involved in Dirichlet energy minimization has very particular, sparse structure. For example, one can perform fixed point iteration via $f^{t+1}(u) = 1/N(u) \sum_{w \in N(u)} f^{t}(w)$ until convergence. This is a (sparse) matrix-vector multiply $f^{t+1} = A'f^t$ for a suitable matrix $A'$. One can use repeated squaring for fast convergence ($A'^{2t} = (A'^{t})^2$). This is essentially just computing the eigenvector of the matrix $A'$ with repeated squaring. This might be fast and it is naturally vectorizable (just batch all the $A'$ matrices from different trees together).
- Using low-dimensional embeddings: to each leaf node assign a fixed random embedding on the $n'$-dimensional simplex, with $n' < n$, then set the internal node embeddings via Dirichlet energy minimization. This should be faster than using $n$-dimensional embeddings. Random embeddings usually have good statistical properties. Furthermore, since these embeddings get projected into dimension $d$ when computing the node features with the GNN in ARTree (as mentioned in line 166), it would seem reasonable to take e.g. $n'=d$. I.e why not just start in dimension $d$ to begin with?

Another comment regarding runtime analysis: runtimes for (1) generating tree topologies vs (2) learning the model parameters may differ because training requires gradients while generating does not. This means that while the node embedding module represents 65% of the time for tree generation (Figure 2) it may be less of a bottleneck for model training. Without proper profiling, it is hard to tell what the computational bottlenecks really are and whether the paper is solving any real computational bottlenecks at all, or just implicitly overcoming inefficient implementations. It is important to recall that ARTreeFormer performs slightly worse than ARTree on several tasks (see e.g. Table 1 where ARTree beats ARTreeFormer on all but 1 dataset). Thus, the architectural changes proposed by ARTreeFormer should be well justified. If speeding up ARTree's Dirichlet energy minimization step as above makes ARTree just as fast as ARTreeFormer, the value of ARTreeFormer drops significantly.

Architecturally speaking, ARTree (Algorithm 3 / Figure 5) is simpler and more elegant that ARTreeFormer (Algorithm 1 / Figure 1), since ARTree simply uses the Dirichlet node embeddings (which are interpretable and have some nice theoretical properties [Zhang 2023, "Learnable Topological Features for Phylogenetic Inference via Graph Neural Networks"]) into a GNN to obtain edge decisions. ARTreeFormer, in contrast, uses an ad-hoc arquitecture consisting of attention-based learnt recurrent node embeddings, local Dirichlet edge minimization, and a local message passing scheme (each of which might explain the reduced performance of ARTreeFormer in some cases, as acknowledged by the paper and seen in Table 1 - without proper ablations, which are missing, it is hard to know which module causes performance to be lost). I generally found it hard to understand the ARTreeFormer model (Section 3.2). It suffices to compare the schematics for the two models (ARTree - Figure 5, ARTreeFormer - Figure 1) to notice the increase in architecture  complexity. In particular, note the presence of cycles in ARTreeFormer in Figure 1. There seems to be unnecessary complexity in ARTreeFormer, which is hard to justify without proper runtime analysis.

The paper is plagued with vague qualitative statements such as:
- abstract: "significantly improves the computation speed"
- line 69: "much more computationally efficient"
- line 177: "require much computational cost".
- line 184: "resulting in great computational inefficiency".
- line 307: "greatly improved computational efficiency".
- line 313: "runtime [...] are significantly reduced".
- line 323: "considerably reduced computation cost".
- etc.

Please provide quantitative claims instead of qualitative ones. For example, say "3x faster".

The words "vectorization" and "parallelization" seem to be used interchangeably throughout the paper. I believe in standard use, "parallelization" is used to describe the process of running code on several processes in parallel, while "vectorization" is used to describe speeding up code via use of vector operations (such as matrix-matrix products) which can be efficiently implemented on GPUs or CPUs (e.g. via SIMD). The terms "vectorization" and "parallelization" seem to be used interchangeably in the paper, which leads to confusion. For example, in Fig. 2 right the paper shows the "runtime of ARTreeFormer for generating 100 tree topologies with and without parallelization". Is this parallelization as in multiprocessing, or as in vectorization? (e.g. via batching the trees). Since the figure caption says "All tests are run on a single 2.4 GHz CPU" I believe this is vectorization via batching, not parallelization.

Overall, the lack of proper profiling and computational complexity analysis, the suspicous Dirichlet energy minimization code which seems highly inefficient, the increased complexity of the model architecture which makes it hard to understand compared to ARTree, all without an increase in accuracy (as seen in Table 1), lead me to believe this work requires a major revision before being ready for publication, so I recommend rejection with a score of 3.

I believe if pursued in the right direction, the ideas in this paper could lead to not just faster but also more elegant and performant probabilistic models for tree topologies. The idea of learnt node embeddings and use of attention instead of global message passing all are appealing.

questions:
1. What implementation of Dirichlet energy minimization are you benchmarking? Is it the Python implementation I cited above?
2. As given by Big-O/Theta/Omega notation, what is the computational complexity of each step of ARTree and ARTreeFormer, as a function of key parameters such as the number of nodes ($n$), the feature dimension ($d$), etc.?
3. Does the empirical runtime of each step align with the theoretical runtimes from question 2? For example, if the Dirichlet energy minimization step has total runtime of $\Theta(n^3)$ for sampling one tree, does it align with the empirical runtime? (A back-of-the-envelope calculation used frequently is that in a performant CPU implementation, e.g. in C++, $10^8$ to $10^9$ operations will be executed per second).
4. How costly is it to compute the gradients of the model parameters, which are needed for learning? This is not revealed by Figure 2, which only considers generation.
5. Line 283: "with or without parallelization" - how was the code "parallelized"? Provide details.
6. line 946: "two peaks of the posterior distribution": just curious, how do you know there are "two peaks"? I am not familiar with this dataset.
7. Have you considered other simple options to speed up ARTree's Dirichlet energy minimization step? I already mentioned two above.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730731335095

thread_type:
official_review

<thread_object>
creator: reviewer_ogwq

<content>
summary:
The authors present **ARTreeFormer**, an improvement on a previous method: **ARTree**.  
 These methods are autoregressive GNN-based models and allow for a flexible way to model the entire tree space. Both **ARTree** and **ARTreeFormer** can be used in Bayesian phylogenetic tasks, including variational Bayesian phylogenetic inference.    

The core methodological aspects of both methods are the same: starting with the simplest tree topology $\tau_3$ with 3 nodes, they iteratively add leaf-nodes to the tree by selecting edges on which to attach them until all $N$ leaf nodes are contained in the tree. 
At a step $n<N$, a probability vector $q_n\in\mathbb{R}^{|E_n|}$ over the set of edges of the tree is computed using a GNN, and the edge on which to attach the $(n+1)^{th}$ leaf node is sampled from $E_n$ according to $q_n$. 
The probability vector $q_n$ is computed from edge embeddings which, in turn, are computed from tree *node embeddings* after several rounds of *message passing*.  
 
The latter two steps are where the two methods differ and allow **ARTreeFormer** to be faster than **ARTree**: 
 1. *Node embedding*: **ARTreeFormer** uses the very popular self-attention mechanism to learn topological embeddings for tree nodes. This eliminates a tree-traversal based algorithm used in **ARTree** to compute these embeddings allowing for **(1)** faster execution as self-attention can be efficiently implemented in a vectorized fashion and **(2)** batched parallel execution of this embedding step on several topologies at once. 
 2. *Message passing*:  **ARTreeFormer** uses *local* message-passing to update node embeddings after addign a new leaf-node, where only the nodes in the direct neighbourhood of the newly-created internal node are updated. This contrasts with the approach in **ARTree** where all the nodes in the graph would be updated at each step of the iterative algorithm. 

These two modifications allow **ARTreeFormer** to consistently run faster than **ARTree** without sacrificing performance.

strengths:
1. Overall the paper is fairly well-written and flows quite well
2. The learnable node embeddings using attention are a simple but effective way to **(1)** make the embeddings and the model more flexible **(2)** improve the runtime
3. There is a significant improvement in runtime shaving *hours* of CPU runtime off (e.g for the Maximum parsimony application). I believe this makes the method more applicable to modern phylogenetic datasets where there are several hundred/thousands of taxa.

weaknesses:
The paper, I believe, has no real "deal-breaking" weakness but would benefit from addressing the following points: 

1. In the VBPI experiment, **ARTreeFormer** is compared to other methods: $\phi$-CSMC and GeoPhy in terms of approximation accuracy. However, in terms of runtime, **ARTreeFormer** results are only compared to methods upon which it was built *(i.e. SBNs and ARTree)*. Since the focus of this paper is the fast runtime of **ARTreeFormer** I think a computation-speed comparison to all methods is warranted. 
2. While I understand the motivation behind choosing a one dimensional query vector $q_n$ in the node embedding phase *(a lowered time complexity and runtime)* , does that not limit the expressivity of the model ? It is not evident to me that each node benefit from querying the same node features. 
3. If the main benefit of **ARTreeFormer**  is a decreased runtime, why not explore alternatives to VIMCO to estimate gradients of the tree topology parameters ? As I understand it this approach requires $K$ samplings of the topology distribution to estimate gradients. It might be faster to compute gradients using methods like CONCRETE [1] or the Gumbel-softmax trick [2] to backpropagate directly through the discrete iterative tree building algorithm. 

[1] https://doi.org/10.48550/arXiv.1611.00712  
[2] https://doi.org/10.48550/arXiv.1611.01144  


*Minor:*
-  In (12), should it be $\ell(f_n,\tau_{n+1},w):=\sum_{(u,w)\in E_{n+1}} ||f_n(u) - f_n(w)||^2$ ? i.e. $f_n(w)$ instead of $f_n(v)$

questions:
1. Could the authors try and compare the runtime of **ARTreeFormer** to independent methods as mentioned in weakness 1 ?
2. The paper might benefit form another ablation study, (similar to what has already been done for $K$ and $d$ in appendix $E$), to quantify the impact of the choice of a 1D query vector as mentioned in weakness 2.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731129629946

thread_type:
official_review

<thread_object>
creator: reviewer_5tht

<content>
summary:
The paper introduces ARTreeFormer, which aims to improve computational efficiency in phylogenetic inference by leveraging attention mechanisms. ARTreeFormer is based on ARTree. For traditional ARTree, the total run time grows rapidly and the
node embedding module dominates the total time (the message passing layers are relatively short in time consumption). In comparison to ARTree, ARTreeFormer uses (1) the recurrent node embedding that is learnable through a deep graph neural network layers (2) a local updating scheme in the neighborhood of the newly added internal node, instead of a global one. In the experiment, ARTreeFormer is evaluated with tree topology density estimation (TDE) and variational Bayesian phylogenetic inference (VBPI) on DS1-8. Experiments also show that ARTreeFormer is significantly faster than ARTree in training and evaluation.

strengths:
(1) ARTreeFormer improves significantly over ARTree by recognizing the several design bottlenecks and running time complexity of ARTree. The ARTreeFormer proposes novel attention and GNN mechanism that is tailored to phylogenetic inference.

(2) ARTreeFormer is evaluated on several standard phylogenetic inference benchmarks, while the authors have provided several useful metrics in analyzing the performance of the model. The author explains in details about what each metrics stand for, and make insightful comments on why ARTreeFormer is suprior to SOTA methods in certain metrics.

(3) The paper is written with a logical flow and clear exposition, and the authors provide an algorithm pseudo codes that makes the algorithm easy to understand.

(4) The ARTreeFormer's performance in FLOPs and time complexity seem to significantly improve over ARTree, which showcase the effectiveness of the proposed algorithm.

weaknesses:
(1) Table 1 reveals that ARTreeFormer consistently lags behind ARTree in terms of KL divergence. Why does ARTreeFormer not achieve competitive performance with ARTree, which it was designed to improve? What training components or procedures might contribute to ARTreeFormer’s reduced performance?

(2) It seems that the recurrent node embedding with simplified attention mechanism and local message passing updates module are built upon two separate components of ARTree, respectively. The authors should consider to bring an ablation study on only changing one of the modules for ARTree and observe the metrics (time,FLOP, KL divergence, etc.). Since attention mechanisms are a key part of ARTreeFormer’s design, the paper could also have a deeper ablation of different attentions. 

(3) There is not a theoretical analysis on how ARTreeFormer improves over ARTree in terms of running time. While the authors mention that ARTree builds topological node embeddings with two-pass algorithm which needs additional traversal over tree topology, the authors could mention how big-O complexity of the ARTree is for node-embedding and message passing module. Then, it would also be beneficial to include the recurrent node embedding and local message passing updates' complexity in terms of big-O and model sizes. Overall, the authors could make a more compelling argument than showing Figure 2 (and give a roughly estimate like 65%).

(4)  Small thing on the writing, but the authors spend two paragraphs explaining the VBPI score and parsimony score before introducing the ARTree model, which is a bit distracting. Line 199, Line 263 use two different notations $F$ to describe the multi-headed attention mechanism, while the attention is denoted as M in eq. 8b. The notations could be unified for readability.

(5) A lot of the terms in the experiments, for example DS1-8, VIMCO estimator, MrBayes, are wordings that are not familiar to general audience in the ML literatures. The authors should consider expand the appendix section to include some introduction on the technical components, instead of simply citing the terms from the literatures.

questions:
(1) Is there any explanations about why for in Figure 4 ARTreeFormer converges slower than SBN and ARTree at the beginning? What seem to be difficult to learn for ARTreeFormer, and can we attribute this phenomenon to simplification of attention-based architecture? 

(2) Similar to the (1) in the weakness section, can the authors provide some justifications of why ARTreeFormer under-perform the ARTree?

(3) The multi-sample lower bound and annealing schedule play significant roles in training. Have the authors tested different schedules or multi-sample configurations? What are the observed effects on the model’s performance?

(4) How sensitive is ARTreeFormer/ARTree ’s performance to different orders of pre-specified leaf nodes?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732436214438

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global response

comment:
We thank all reviewers for the constructive feedback. 
We have revised the paper, and have incorporated their suggestions with the following major changes:
- In section 2, we merge the introduction of phylogenetic likelihood and the parsimony score to keep a more compact writing.
- We add a discussion about the computational complexity in Table 1, Section 3.2, and Appendix B.2, which validate that the node embedding module dominates the computation time.
- We added an explanation of the standard deviations of MLL in Section 4.3 and discussed the possibility of using mixture techniques to further improve the results.
- We substitute the qualitative description of acceleration with quantitative ones if applicable. The word "parallelization" is changed to "vectorization" throughout this paper.
- We added the details about the VIMCO estimator in Appendix C.
- We added the time consumption of GeoPhy in Table 8.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732567503652

thread_type:
official_comment

<thread_object>
creator: area_chair_kdhe

<content>
title:
Last day for reviewers to ask questions to the authors!

comment:
Dear reviewers,

Tomorrow (Nov 26) is the last day for asking questions to the authors. With this in mind, please read the rebuttal provided by the authors, as well as the other reviews. If you have not already done so, please explicitly acknowledge that you have read the rebuttal and reviews, provide your updated view _accompanied by a motivation_, and raise any outstanding questions for the authors.

**Timeline**: As a reminder, the review timeline is as follows:
- November 26: Last day for reviewers to ask questions to authors.
- November 27: Last day for authors to respond to reviewers.
- November 28 - December 10: Reviewer and area chair discussion phase.

Thank you again for your hard work, 

Your AC
</content>
</thread_object>

</thread>
<thread>
creation date:
1733297346341

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Global Response

comment:
Dear AC and Reviewers,

During the discussion phase, the reviewers raised several concerns regarding the novelty of ARTreeFormer and the claim of improved computational efficiency, and we feel it is necessary to clarify some of them here.

- **Our novelty and contribution**: By leveraging learnable topological node embedding and local message passing, ARTreeFormer overcomes the large time consumption bottleneck of ARTree while maintaining ARTree's favorable approximation accuracy. Its main methodological contribution can be summarized in the following aspects: (i) the learnable node embedding with attention mechanism in ARTreeFormer can be vectorized over different nodes and tree topologies, which reduces the time complexity by one order and has been verified empirically. (ii) Global message passing is substituted by an attention-based local message passing in ARTreeFormer, which can also reduce the time complexity of an order without sacrificing approximation accuracy. (iii) Compared to the GNNs in ARTree, the attention mechanism in the message-passing step of ARTreeFormer is capable of capturing the long-distance dependencies. The following ablation study shows the advantage of the attention mechanism. Attention mechanism for phylogenetic tree modeling itself is an interesting topic and we take the first step in this direction.

*Table: KL divergence obtained by different message passing modules on the TDE task.*
| Model | DS1 | DS2 | DS3 | DS4 |
| --- | --- | --- | --- | --- |
| Learnable Node Embedding + GNN-based global message passing | **0.0044** | 0.0123 | 0.0916 | 0.0346 | 
| Learnable Node Embedding + Attention-base local message passing (= ARTreeFormer)| 0.0067 | **0.0102** | **0.0777** | **0.0320** |

- **Complexity analysis**. We provide the time complexity analysis for ARTree and ARTreeFormer, for both cases (with and without vectorized computation). It should be emphasized that with vectorization, the time complexity of ARTree's node embedding module is always 1-order higher than that of ARTreeFormer.
The main reason is that ARTreeFormer enables efficient vectorization across different nodes on a tree, while ARTree does not due to Dirichlet energy minimization for topological node embedding.
Note that although Dirichlet energy minimization can be done in linear time ($O(n)$) for each subtree, this operation needs to be repeated for a sequence of subtrees during the autoregressive generation process of the tree topology.
On the other hand, in ARTreeFormer, for each of the subtrees, we can reduce the complexity to $O(1)$ using learnable node embedding. 

- **Computation time baseline**. As suggested by Reviewer 3iPi, the topological node embedding can be calculated via the fix-point algorithm with the square trick.
We agree that this can reduce the computation time of ARTree's node embedding module, but is still much slower (around 6x computation time in our experiments) than that of ARTreeFormer.
We will add this important baseline in our revision.
We also want to point out that although the adjacency matrix is sparse, its $n$-power becomes more dense as $n$ increases, so the acceleration technique for sparse matrix multiplication may not apply for large $n$.

- **About the C++ implementation**. It is undoubted that C++ runs much faster than Python, and we apologize for not being able to deliver a runtime comparison in C++ (this requires rewriting our code base in C++) between ARTree and ARTreeFormer given this short time budget for rebuttal. However, we think that the current comparison in Python is reasonable as explained below.
Based on the profiling of ARTree's node embedding module, a large amount of time (around $55\\%$) is spent on the PyTorch-based numerical computation, which should have a similar consumption in C++.
This is to say, although the overhead of the pure Python part makes the current implementation of ARTree slow, the main part of the node embedding module in ARTree (which is written using Pytorch that has C++ under the hood) can not be sped up by using C++.
Therefore, the current results in our manuscript (Figure 2, Figure 3, Figure 4) are sufficient to show the superior efficiency of ARTreeFormer.
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper proposed a faster model than previous work, called ARTree, for phylogenetic inference. Several reviewers agreed that decreasing runtimes for Bayesian phylogenetics is an important objective by itself. 

During the author-reviewer discussions, several reviewers raised the importance of including a complexity analysis of the proposed work, which should explain why the proposed model is faster than the original ARTree. While the authors have responded with a complexity analysis in their rebuttal, these contained notations to include the effect of vectorized operations that were unfamiliar to all reviewers and the AC, and unfortunately further explanations by the authors did not clarify this sufficiently. 

Furthermore, reviewer 3iPi had extensive and lively discussions with the authors on the fairness of comparison against previously (supposed inefficient) implementations in python, profiling of the bottlenecks in speed of the proposed work and the original ARTree work, and the subsequent impact of the proposed improvement. Reviewer 3iPi suggested that the improved performance in speed based on a very inefficient implementation in python of the prior ARTree work is not fair. After a significant back and forth between authors and the reviewer on supposed factors of speed up by moving code from python/PyTorch  to pure C++ implementations, the authors conclude that they estimate that a C++ implementation will give only a very limited speed up the main part of the node embedding module in ARTree. However, reviewer 3iPi has provided several arguments against this, and has taken the extensive effort to provide a C++ implementation that runs around 30x faster than the author's Python implementation. 
During the AC-reviewer discussions, all reviewers who engaged with the discussion and read the review of 3iPI have agreed that reviewer 3iPi raised valid concerns with ample evidence to back up their claims, and indicated to want to lower their scores. 

Given this strong consensus I recommend to reject this paper. I'd like to thank the authors for actively engaging in the discussion phase and for their effort in trying to address all concerns. Furthermore, thank you to all reviewers for engaging in discussion, and in particular thank you to reviewer 3iPi for going far above and beyond what is normally expected and observed from a reviewer.