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

Thank you for your careful review and constructive feedback. Here are our responses to your concerns.

(W1) Question: Why do you think the std is larger for ARTreeFormer than the other VBPI methods? (...)

Response: The standard deviations of ARTreeFormer are larger because ARTreeFormer provides a slightly less accurate approximation than ARTree. However, the correct MLL means of ARTreeFormer suggests it works fairly well in terms of mode coverage. We have clarified this in our revision.

(W2) Comment: There are other VBPI models available, such as [1,2,3,4], all of which (I believe) use pre-computed candidate trees. (...)

Response: Thank you for providing these related works. We have cited VBPI-mixture [1][2] and discuss the potential improvement by employing mixtures in VBPI in Section 4.3. Note that although the above methods can provide better ML scores, they all require pre-computed candidate trees (which can be difficult to obtain) and hence would suffer from restricted support of tree topologies, especially when the posterior is diffuse [6]. On the other hand, ARTree and ARTreeFormer enjoy unconfined support due to the autoregressive models for tree topologies that naturally cover the entire tree space. This is why we say compared to SBN-based VBPI variants, the superiority of ARTreeFormer is mainly reflected by its unconfined support. We'd like to clarify that what we want to emphasize in this sentence is the methodological improvement instead of the accuracy gain. We will modify this sentence to be more concise about ARTreeFormer's advantages over other baseline methods in our revision.

(W3) Apart from my comment above, I think [1] and [2] should be referenced (...)

Response: Thanks for providing these related works. We have cited [1,2] in our revision.

(W4) In [5] they show that more importance samples during training can harm the inference of the variational approximation. (...)

Response: Thanks for your question! We expect a moderate K to perform well in practice. That is because a small K would have limited exploration efficiency and hence hinder mode discovery, while a large K may lead to lower signal-to-noise ratio and hence deteriorate the training of the variational approximations as shown in [5]. It is hard to tell exactly the differences between $K=1$ and $K=100$. Roughly speaking, $K=1$ may lead to a variational approximation that gets stuck at limited modes and $K=100$ may have good mode coverage but worse overall approximation quality.

(Q1) I am aware that there are many papers now in ML conferences on VI-based phylogenetics. (...)

Response: Thanks for your question! The main motivation is to develop efficient Bayesian phylogenetic inference algorithms that are capable of delivering accurate posterior estimates in a timely manner. The commonly used MCMC tends to be slow and often requires long runs to generate high quality samples. As an alternative, variational inference can provide comparable posterior estimation to MCMC, while requiring less computation due to a more efficient tree exploration mechanism and its optimization nature. Therefore, these VI-based phylogenetic inference algorithms have the potential to scale up Bayesian phylogenetic inference. As phylogenetic inference can provide extremely helpful information when making important public health decisions (especially during pandemics), it would be invaluable (with the potential to save thousands of lives) if significant computational efficiency improvement of Bayesian phylogenetic inference methods can be made in practice. For example, VI methods have been used for rapid analysis of pandemic-scale data (e.g., SARS-CoV-2 genomes) to provide accurate estimates of epidemiologically relevant quantities that are verifiable via alternative sources of public health data [7]. We have added some related discussion to the conclusion section in our revision.

[6] Cheng Zhang and Frederick A Matsen IV. A variational approach to Bayesian phylogenetic inference. Journal of Machine Learning Research, 25(145):1–56, 2024.
[7] Caleb Ki and Jonathan Terhorst. Variational phylodynamic inference using pandemic-scale data. Mol. Biol. Evol., July 2022.

Thank you for your careful review and constructive feedback. Here are our responses to your concerns.

(W1) Limitations of performance improvement

Response: We would like to clarify that there is no "accuracy bottleneck"; SBN, ARTree, and ARTreeFormer are all the SOTA methods for VBPI, and work much better than other models (see Table 3). The bottleneck of SBN is that it requires support estimation and hence suffers from a confined support and the bottleneck of ARTree is the large computation cost. These two bottlenecks are both overcome by ARTreeFormer. We admit there is an accuracy loss of ARTreeFormer compared to ARTree, but this is negligible for practical applications.

Table: KL divergence difference between different methods on TDE

(W2) Model complexity increases and implementation difficulty

Response: We have provided the Python implementation of ARTreeFormer which can be readily used for further applications. The message passing steps, attention mechanisms, etc, are common techniques in graph learning methods. The model architecture of ARTreeFormer is not complex -- instead, we make a careful design that enables vectorized computation. Benefiting from this, the actual implementation of ARTreeFormer is highly organized and batched, which also greatly improves the computational speed and does not show additional complexity than other graph models. See Figure 2 for a detailed runtime comparison of different modules. In contrast, it is the inefficient design (although easy to describe) of ARTree, e.g., the multiple and irreducible loops (repetitive dirichlet energy minimization for topological node embeddings), that causes the highly inefficient computations. Therefore, a more efficient and faster design is very necessary. See a more detailed discussion on page 6 (from line 301 to line 323).

(W3) Lack of generalization verification on diverse datasets

Response: We would like to clarify that DS1-8 are commonly used benchmarks in the field of Bayesian phylogenetic inference [1,2,3,4,5]. This is because DS1-8 covers all sorts of duplicated sub-structures and multimodal posteriors, as shown in [6]. Therefore, we can expect ARTreeFormer to work well for other data sets if it works for DS1-8.

[1] Zhang, C. and Matsen IV, F. A. "Variational Bayesian phylogenetic inference." International Conference on Learning Representations (2019).
[2] Koptagel, Hazal, et al. "Vaiphy: a variational inference based algorithm for phylogeny." Advances in Neural Information Processing Systems 35 (2022): 14758-14770.
[3] Mimori, Takahiro, and Michiaki Hamada. "GeoPhy: differentiable phylogenetic inference via geometric gradients of tree topologies." Advances in Neural Information Processing Systems 36 (2023).
[4] Xie, T. and Zhang, C. "ARTree: A deep autoregressive model for phylogenetic inference." Advances in Neural Information Processing Systems 36 (2023).
[5] Zhou, Mingyang, et al. "PhyloGFN: Phylogenetic inference with generative flow networks." arXiv preprint arXiv:2310.08774 (2023).
[6] Chris Whidden and Frederick A Matsen IV. Quantifying MCMC exploration of phylogenetic tree space. Syst. Biol., 64(3):472–491, May 2015.

(W4) Potential information loss in local message passing

Response: To inspect which component contributes to the reduced performance of ARTreeFormer, we conducted an ablation study. We see that the strong representation power of topological node embedding is the most crucial part of a model's performance. Importantly, the attention-based local message passing improves both accuracy and efficiency.

Table: KL divergence obtained by various model architectures on the TDE task.

(W5) Theoretical improvement

Response: The main contribution of this paper lies in the model architecture improvement and the superior acceleration results. Exploring the in-depth theoretical analysis is an interesting topic and we leave this for future work.

Thank you for your careful review and constructive feedback. Here are our responses to your concerns.

(W1) This paper emphasizes the time complexity of improving the MHA module but without showing some ablation experiments that remove the MHA module (...)

Response: Thanks for this constructive question. To inspect which component contributes to the reduced performance of ARTreeFormer, we conducted an ablation study. We see that whatever the node embedding module is, the attention-based local message passing always improves both accuracy (the following table) and efficiency (Figure 2).

Table: KL divergence obtained by various model architectures on the TDE task.

(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.

Response: We want to clarify that ARTreeFormer does not discard global information. First, in the node embedding module, MHA extracts all the node features to form the features of the new nodes. Second, we considered a global representation vector $r_n$ in forming edge decision distributions. ARTreeFormer is local only in the attention-based message passing module; however, the local attention can be more powerful than the global GNN, as suggested by the above table.

(W3) Since computational efficiency is an important contribution, the authors should clarify the discussion of the time complexity formula in line 213.

Response: Thanks for this suggestion. The following table compares the computational complexity of the two major modules in ARTree/ARTreeFormer, where $\alpha\in(0,1)$ refers to the accelerated complexity order due to vectorized operations in PyTorch.

We see that $\alpha$ can be small in practice (i.e., fast computation of batched tensors in PyTorch). Even in such a case, the complexity of ARTree's node embedding module is still higher than $O(N^2)$, while those of other modules are reduced to approximately equal to or less than $O(N)$. This validates the observation that the topological node embedding dominates the computation time. We have added a detailed computation and explanation on this in our revision.

(Q1) It seems that the motivation of this work (...)

Response: Thanks for this question. We have clarified our innovation in the last paragraph in Section 3, and put it here for your convenience.

• First, the learnable node embedding based on the attention mechanism is novel and overcomes the non-vectorizable bottleneck of ARTree.
• Second, we incorporate message passing and local Dirichlet energy minimization within the neighborhood structure, tailored for phylogenetic trees.
• Third, adapting graph techniques to phylogenetic trees is not straightforward and requires careful design, and we are the first to show that this simplified attention-based architecture exhibits strong approximation capacity with considerably reduced computational cost.

(Q2) The practical value of the improvement results is unclear. (...)

Response: As we claimed in our paper, ARTree suffers from extremely large computational costs. For example, it takes 12 days for training ARTree on DS8 (just multiply the training time in Figure 4 (right) by 40000.) Given that timely Bayesian phylogenetic analysis is crucial for making important public health decisions (especially during pandemics), it would be invaluable (with the potential to save thousands of lives) if significant computational efficiency improvement of Bayesian phylogenetic inference methods can be made in practice.

We have also tried other types of attention mechanisms (e.g., linear attention) and report the results in the following table. We find that linear attention works fairly well but is inferior to softmax attention.

Table: KL divergence obtained by different attention types on the TDE task.

Thanks for your careful review and helpful suggestions! Here are our responses to your concerns.

Time complexity analysis Thanks for your constructive suggestion! The following table compares the computational complexity of the two major modules in ARTree/ARTreeFormer, where $\alpha\in(0,1)$ refers to the accelerated complexity order of linear (1-order) vectorized operations in PyTorch. The results are interpreted as follows.

• For the node embedding module of ARTree, the cubic order of $N$ can be split into three aspects: (i) Iterating $N$ subtrees when autoregressively adding leaves; (ii) Iterating all $N$ internal nodes when computing the embeddings of a subtree; (iii) Summation and scalar multiplication of $N$-dimension vectors. We note that only (iii) can be accelerated with the vectorized operation, and (i) & (ii) always lead to two for-loops even if implemented in C++ or the fix-point algorithm. (Note that the number of fix-point iterations until convergence is $O(N)$.) This gives a complexity (with vectorization) of $O(N^{2+\alpha})=O(N)\cdot O(N)\cdot O(N^\alpha)$.
• For the message passing module of ARTree, there is only one for-loop: Iterating $N$ subtrees when autoregressively adding leaves. Other computations in $N$ nodes and $d$-dimension features can be vectorized. Therefore, this gives an accelerated complexity of $O(N^{1+\alpha}d^{2\alpha})=O(N)\cdot O(N^\alpha d^{2\alpha})$.
• For the node embedding module and message passing module in ARTreeFormer, all computations can be vectorized except for "Iterating $N$ subtrees when autoregressively adding leaves". Therefore, the complexity with vectorization is $O(N^{1+\alpha}d^{\alpha}+Nd^{2\alpha})$ and $O(Nd^{2\alpha})$.

We would like to emphasize that $\alpha$ can be small in practice (i.e., fast computation of batched tensors in PyTorch). Even in such a case, the complexity of ARTree's node embedding module is still higher than $O(N^2)$, while those of other modules are reduced to approximately equal to or less than $O(N)$. This validates the observation that the topological node embedding dominates the computation time.

Topological embedding with fix-point algorithm We thank the reviewer for the suggestion of using a fix-point algorithm. We try this method and report the time consumption in the following Table, where we see there is no evident improvement of fix-point algorithm. The reasons are analyzed as follows:

• There are three for-loops in the topological node embedding within the two-pass algorithm of ARTree: (a) a traversal over the nodes on a tree; (b) a traversal over all the trees in a batch; (c) apply (a)(b) to each ordinal subtree as the tree autoregressively grows. All these (a)(b)(c) loops cannot be vectorized in the current implementation.
• For the topological node embedding with fix-point algorithm, (b) is vectorized and (a)(c) are kept. Note that the (a) is replaced by an iteration process in the fix-point algorithm, and the number of iterations until convergence is in the same order as the number of taxa (because the information should be propagated from the leaf nodes with initial features to the internal nodes). Finally, as the corresponding linear system is not homogeneous, the square trick $(A')^2$ cannot be applied.
• For the learnable topological embedding in ARTreeFormer, (a)(b) are vectorized and (c) is kept.

Table: Time consumption (seconds) of node embedding module in generating 100 trees. "n'=n" means setting embedding dimension $n'$ as the number of leaf nodes $n$ in a subtree.

We admit that a C++ implementation is a valuable suggestion for the improvement of computational efficiency and will leave this for future work.

Thanks for your follow-up questions. Before addressing your concerns, we think we should first establish a consensus that two models should be compared in the same programming language, the same device, and possibly, the same embedding dimension.

• Firstly, we want to clarify your potential misunderstanding of our acceleration claim. The '3$\times$' is the acceleration of training (training includes running the whole deep model, backpropagation, and likelihood computation), and '10$\times$' is the acceleration of generating (including inference of the whole deep model). In both cases, topological node embedding is only a part of the whole process. If you are talking about the whole process and the acceleration of the C++ implemented ARTree is 10x, then 10x acceleration should be the same for C++ implemented ARTreeFormer. Now we only talk about the implementation of the topological node embedding module.

(a) We assume the topological node embedding of ARTree is computed by the two-pass algorithm, and defer the discussion on other choices in the following points. In the Python implementation, the acceleration of ARTreeFormer upon the current two-pass algorithm of ARTree is $20\times$. A potential overhead in Python is the tree traversal implementation, which we reported below. Compared to the full node embedding module, we see that the pure traversal of a tree is negligible in ARTree in the Python implementation, so it is the numerical computation that dominates the runtime, not the tree traversal.

Table: Computation time (seconds)

(b) In a C++ implementation, the two for-loops (iteration over all immediate subtrees; iteration over all nodes) still cannot be reduced. We apologize for not being able to deliver a C++ runtime comparison given the limited time budget. However, as PyTorch is based on C++ backend, we expect the ratio of runtimes of ARTreeFormer and ARTree implemented in C++ to be similar to what we observed on the Python side, that is (runtime of ARTreeFormer in C++)/(runtime of ARTree in C++) $\approx$ (runtime of ARTreeFormer in Python)/(runtime of ARTree in Python).

Table R1: Time consumption (seconds) of node embedding module in generating 100 trees. ARTree's default setting means the embedding dimension is $N$.

• We use the $\alpha$ notation here to reflect the acceleration achievable by vectorized computation. As different programming languages or computation module may lead to different vectorization efficiency, we use a flexible $\alpha$ here. If one wants to see how the model can be "perfectly vectorized", just let $\alpha=0$.

• Thanks for this example. We would like to give the following explanations for the complexity of ARTree and ARTreeFormer. For generating one tree with $N$ leaves:

  • The time complexity of the node embedding module of ARTree is $O(N^3)$, where the third dimension can be vectorized; The time complexity of the node embedding module in ARTreeFormer is $O(N^2d+Nd^2)$ where for both terms, the second and third dimensions can be vectorized.
  • The time complexity of the message passing module of ARTree is $O(N^2d^2)$, where the 2-nd, 3-rd, and 4-th dimensions can be vectorized. The time complexity of the message passing module of ARTreeFormer module is $O(Nd^2)$, where the 2-nd and 3-rd dimensions can be vectorized.

  We see that the main difference lies in the node embedding module: ARTree has two dimensions that cannot be vectorized; ARTreeFormer only has one dimension that cannot be vectorized. When vectorization is enabled, the complexity of node embedding module would be dominant for both ARTreeFormer and ARTree.

def fixpoint(edge_index, X, tol=1e-5, max_iters=10000):
    # edge_index describes the index of the adjacency nodes of a node. Its shape is (number of trees, number of nodes, 3).
    # X is the initialization of the features of the nodes.
    # (each row of X is initialized as [1/feature_dim]*feature_dim)

    bs, dim, nf = X.shape
    identity = torch.eye(dim+2, nf).unsqueeze(0).repeat(bs, 1,1)

    for it in range(1, max_iters + 1):
        X_old = X.clone()
        neigh = torch.gather(torch.concat([identity, X], dim=1), dim=1, index=edge_index.reshape(bs, -1).unsqueeze(-1).repeat(1,1,nf))
        neigh = neigh.reshape(bs, dim, 3, nf)
        X = torch.mean(neigh, dim=2)
        Lnorm = torch.mean(torch.abs(X-X_old), dim=(1,2))

        if not torch.any(Lnorm > tol):
            break
    return X

We thank you for your follow-up questions. Let us focus on several notable points for a more precise understanding of the implementation and time consumption of ARTree/ARTreeFormer.

Your comment "ARTree's node embedding module in Python" is inaccurate. All modules of ARTree, including the node embedding module, are implemented in PyTorch, not pure Python. One can tell this by simply looking at the torch.Tensor type of the leaf_features in the following code block. We have also profiled the node embedding module of ARTree, and the results have been shown below. We see that 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, ARTree's node embedding module does not have much room for speedup by using C++ (at most 2 times faster based on our profiling).

Regarding your example, we would like to point out that all the runtime is reported with vectorized computation enabled in pytorch. Therefore, your simple calculation of time complexity is misleading. With vectorization, the time complexity of topological node embedding step (1) is roughly $O(N^2)$, and message passing step (2) is $O(N)$. Therefore, the node embedding step of ARTree can be a bottleneck in practice as demonstrated in our experiments. In contrast, with vectorization, the node embedding module of ARTreeFormer can be much faster. This is the main reason behind the results in Table R1. For your back envelop time calculation, we'd like to remind you that the absolute speed depends on the computing hardware and the big $O$ notation omits the constant term. Therefore, only the relative speed can be fairly compared.

Your table of benchmarked runtime is also misleading. As mentioned above, the node embedding module of ARTree also uses Pytorch. The overhead of Python is at most half the computation. The numerical computation part (coded in Pytorch) takes at least half of the runtime reported in Table R1, and the acceleration of the learnable node embedding module in ARTreeFormer remains significant.

• (runtime of ARTree's node embedding module in Pytorch (C++))
• (runtime of ARTree's message passing module in PyTorch (C++))
• (runtime of ARTreeFormer's node embedding module PyTorch (C++))
• (runtime of ARTreeFormer's message passing module in PyTorch (C++))

Regarding the fixed point iteration method: We'd like to point out that the topological embeddings do not decay exponentially fast with distance, as there is also a constant term $d_u$ there for the update equation. In fact, all the topological node embeddings after convergence lie in the simplex, that is the topological node embedding $x_u$ of node $u$ satisfies $\sum_{n=1}^N x_u[n]=1$, where $x_u[n]$ is the $n$-th element of $x_u$, as suggested by Theorem 1 in [Zhang, 2023]. Therefore, for graphs with diameter $\Omega(n)$, fixed point iteration with $O(\log n)$ iterations can still be far away from convergence. And those nodes with embedding coordinates of near-zero would be incapable of representing local geometric structures of the tree topology, leading to a potential drop in performance.

Regarding your example: With perfect vectorization, the time complexity of fix-point algorithm should be $O(N * \log N)$ (the best case of a balanced tree), while that of the two-pass algorithm should be $O(N^2)$. Now consider one tree with $N=100$ leaf nodes, and 20 iterations are required for the convergence of the fixed point algorithm. With vectorization, the time complexity of fix-point algorithm should be $C_1\cdot 100\cdot 20$, while that of the two-pass algorithm should be $C_2 \cdot 100^2$. This complexity ratio aligns with our empirical results.

To summarize: We have done an additional profiling of ARTree's node embedding module. We found that the overhead of Python takes at most half the runtime. This means that the numerical computation in Pytorch takes at least one half of the runtime, which can not be reduced much when using C++. Therefore, the claim that the overhead of Python is >=10x is inaccurate. Moreover, for the runtime analysis, we want to point out that those runtimes are all obtained with vectorized computation enabled through Pytorch. Therefore, a more appropriate analysis should use the corresponding vectorized complexities instead of the standard ones (which were used by the reviewer). As shown in the discussion above, our runtime results align with the vectorized complexities of ARTree and ARTreeFormer. We sincerely hope these clarifications can address your concerns and help understand our time consumption better.

Line #      Hits         Time  Per Hit   % Time  Line Contents
==============================================================
    35                                           @profile
    36                                           def node_embedding(tree, level):
    37     10000       5253.0      0.5      0.0      name_dict = {}
    38     10000       4410.0      0.4      0.0      j = level
    39     10000      46496.0      4.6      0.1      rel_pos = np.arange(max(4,2*level-4))
    40   1990000    6115957.0      3.1      9.5      for node in tree.traverse('postorder'):
    41   1980000    1536086.0      0.8      2.4          if node.is_leaf():
    42   1000000     291708.0      0.3      0.5              node.c = 0
    43   1000000    2545394.0      2.5      3.9              node.d = leaf_features[node.name]
    44                                                   else:
    45    980000     947045.0      1.0      1.5              child_c, child_d = 0., 0.
    46   2950000    1344661.0      0.5      2.1              for child in node.children:
    47   1970000     670931.0      0.3      1.0                  child_c += child.c
    48   1970000    7407605.0      3.8     11.5                  child_d += child.d
    49    980000     499189.0      0.5      0.8              node.c = 1./(3. - child_c) 
    50    980000    5303259.0      5.4      8.2              node.d = node.c * child_d
    51    980000     697394.0      0.7      1.1              if int(node.name) < len(rel_pos):
    52    960000     563908.0      0.6      0.9                  rel_pos[node.name] = j
    53    980000     485237.0      0.5      0.8              node.name, j = j, j+1
    54   1980000     988536.0      0.5      1.5          name_dict[node.name] = node
    55                                               
    56     10000       4767.0      0.5      0.0      node_features, node_idx_list, edge_index = [], [], []            
    57   1990000    4912675.0      2.5      7.6      for node in tree.traverse('preorder'):
    58   1980000     443753.0      0.2      0.7          neigh_idx_list = []
    59   1980000    1597575.0      0.8      2.5          if not node.is_root():
    60   1970000   17219507.0      8.7     26.6              node.d = node.c * node.up.d + node.d
    61   1970000    1596294.0      0.8      2.5              neigh_idx_list.append(node.up.name)
    62                                                       
    63   1970000    1758935.0      0.9      2.7              if not node.is_leaf():
    64    970000    2501343.0      2.6      3.9                  neigh_idx_list.extend([child.name for child in node.children])
    65                                                       else:
    66   1000000     542745.0      0.5      0.8                  neigh_idx_list.extend([-1, -1])              
    67                                                   else:
    68     10000      23428.0      2.3      0.0              neigh_idx_list.extend([child.name for child in node.children])
    69                                                   
    70   1980000     894422.0      0.5      1.4          edge_index.append(neigh_idx_list)                
    71   1980000     762492.0      0.4      1.2          node_features.append(node.d)
    72   1980000     770827.0      0.4      1.2          node_idx_list.append(node.name)
    73                                               
    74     10000     522331.0     52.2      0.8      branch_idx_map = torch.sort(torch.tensor(node_idx_list).long(), dim=0, descending=False)[1]
    75     10000    1082825.0    108.3      1.7      edge_index = torch.tensor(edge_index).long() 
    76                                               
    77     10000     539947.0     54.0      0.8      return torch.index_select(torch.stack(node_features), 0, branch_idx_map), edge_index[branch_idx_map], torch.from_numpy(rel_pos), name_dict

Thanks very much for your constructive review and positive rating of our paper. Here are our responses:

Weaknesses

(1): In the VBPI experiment, ARTreeFormer is compared to other methods (...)

Response: We have supplemented the training time of GeoPhy in Table 8, and also put it here for your convenience. $\phi$-CSMC is an SMC-based method that does not require training but may cost a large amount of time for inference (more than 10 minutes for sampling 100 trees), and therefore is not considered here. We'd like to emphasize that although ARTreeFormer can cost more time than GeoPhy, it achieves much better approximation accuracy.

Table: Training time (seconds) per passing 100 trees

(2): While I understand the motivation behind choosing a one dimensional query vector in the node embedding phase (...)

Response: Thanks for your question! Note that we want to compute a one-dimensional feature vector for the newly added nodes, therefore one dimensional query is enough in terms of model design. This choice also follows [Han et al, 2023] who also considers one-dimensional query vector for forming node features. Employing a multi-dimensional query certainly also works for this problem, and we reported the results in the following Table, where we see that the effect of the additional query vectors is negligible.

(3): If the main benefit of ARTreeFormer is a decreased runtime, why not explore alternatives to VIMCO (...)

Response: Thanks for your suggestion! As can be seen in Table 8 in our revision (or in our response to (1)), VIMCO is not the bottleneck for the computation efficiency of ARTree, as it is also used in SBN which is much faster than the other methods. The inefficiency of ARTree mainly comes from the repetitive non-vectorizable topological node embedding and the global message passing scheme, as discussed in Section 3.1. Alternative gradient estimation methods like CONCRETE [1] or Gumbel-softmax trick [2] have also been used for discrete models. However, the main difficulty of applying these gradient estimators to phylogenetic models is that the phylogenetic likelihood function depends on the tree topology in a rather subtle way, which makes it hard to differentiate through the discrete iterative tree building algorithm. Thanks again for this interesting suggestion and we will explore this in future work.

Minor: In (12), should it be (...)

Response: Thanks! We have fixed it in our revision.

Questions

(1): Could the authors try and compare the runtime of ARTreeFormer to independent methods as mentioned in weakness 1?

Response: Thanks for this question. Please refer to our response to weakness 1.

(2): The paper might benefit form another ablation study, (similar to what has already been done for $K$ and $d$ in appendix), to quantify the impact of the choice of a 1D query vector as mentioned in weakness 2.

Response: Please refer to our response to weakness 2.

We highly appreciate your constructive review and positive rating of our paper. Here are our responses:

Weaknesses

(1): Table 1 reveals that ARTreeFormer consistently lags behind ARTree in terms of KL divergence. (...)

Response to (1): Thanks for this constructive question. First, we want to clarify that ARTreeFormer is mainly designed to improve the computational efficiency of ARTree, rather than the approximation accuracy. To inspect which component contributes to the reduced performance of ARTreeFormer, we conducted an ablation study. We see that the strong representation power of topological node embedding is the most crucial part of a model's performance. Notably, the attention-based local message passing improves both accuracy and efficiency compared to GNN used in ARTree, showing the additional benefit of attention-based modeling for tree topologies.

Table: KL divergence obtained by various model architectures on the TDE task.

(2): It seems that the recurrent node embedding with simplified attention mechanism and (...)

Response to (2): Thanks for the suggestion! For the ablation study on only changing one of the modules of ARTree, please check our response to Weakness (1). For the ablation study on different attention mechanisms, we investigate an additional attention mechanism, i.e., linear attention, and report the KL divergences in the following table. We see that linear attention can deteriorate the inference accuracy, although it may offer more efficient computation.

Table: KL divergence obtained by different attention types on the TDE task.

(3): There is not a theoretical analysis on how ARTreeFormer improves over ARTree in terms of running time.

Response to (3): Thanks for this constructive suggestion. The following table compares the computational complexity of the two major modules in ARTree/ARTreeFormer, where $\alpha\in(0,1)$ refers to the accelerated complexity order of linear (1-order) vectorized operations in PyTorch. We have added this discussion in our revision.

We would like to emphasize that $\alpha$ can be small in practice (i.e., fast computation of batched tensors in PyTorch). Even in such a case, the complexity of ARTree's node embedding module is still higher than $O(N^2)$, while those of other modules are reduced to approximately equal to or less than $O(N)$. This validates the observation that the topological node embedding dominates the computation time.

(4): Small thing on the writing, but the authors spend two paragraphs explaining (...)

Response to (4): Thanks for your suggestion! We have now merged the paragraphs for VBPI score and parsimony score together in our revision. We would like to clarify that Line 196 and Line 256 are deliberately distinguished by us. 'graph' is for graph-level pooling and 'message' is for local message passing. They represent two distinct modules in ARTreeFormer (although they are both implemented by MHA).

(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.

Response to (5): We apologize for the frequent usage of abbreviations.

• DS1-8 refers to eight benchmark datasets that are used to evaluate phylogenetic inference problems. These data sets consist of collections of benchmark biological sequences, and we aim at inferring trees from them. These data sets have been cited in the second paragraph of Section 4.
• MrBayes is an MCMC software for Bayesian phylogenetic inference and is often considered as a golden criterion in this field. We added a citation in the second paragraph of Section 4.
• VIMCO is a gradient estimator for multi-sample lower bounds and is not related to phylogenetic inference. We added a description of VIMCO in Appendix C in our revision.