Venue: ICLR 2023 Conference

Reviews:

final_score: rating: 8: accept, good paper
review_content: summary_of_the_paper: Attention computation of Transformers scales quadratically (n^2) with the input sequence length (n), making it a key bottleneck
in scaling Transformers to long inputs. Targeting at such limitation, this paper proposes a new architecture called TREEFORMER to use decision trees to efficiently compute attention by only retrieving the top nearest neighboring keys for a given query. In addition, to make the decision trees to be trained as part of neural networks using back propagation, this paper proposes a novel bootstrapping method to gradually sparsify attention computation. Finally, this paper extensively evaluates TREEFORMER on both GLUE and LRA benchmarks to show that TREEFORMER architecture can be almost as accurate as baseline Transformer while using 30x lesser FLOPs in the attention
layer.
strength_and_weaknesses: Strength:
1. This paper proposes a new architecture called TREEFORMER to use decision trees to efficiently compute attention by only retrieving the top nearest neighboring keys for a given query. The TREEFORMER comes with two novel attention mechanisms named TF-ATTENTION and TC-ATTENTION.

2. TREEFORMER is evaluated extensively on both GLUE and LRA (long sequences data) benchmarks against many SOTA models such as Big Bird and Performer to show that TREEFORMER architecture can be almost as accurate as baseline Transformer while using 30x lesser FLOPs in the attention layer. Thus the result shown is convincing.

3. Ablation studies give good insight on why bootstrapping method works well by showing the gradient norms.

Weakness:
My main questions are in the boost strapping method.
1. With boost strapping and learning different heights of trees (specified in Algorithm 1), how much longer of the pre-training process for TREEFORMER require compared with standard transformer models? It mentions that "we first pre-train an existing model (say a transformer or BigBird)". Does it pre-train the existing model for the same number of steps as a regular transformer's pre-training process and then trigger boost strapping process? Knowing this will help readers to better understand if there is any cost / limitation for using TREEFORMER to achieve SOTA FLOPS reduction in the attention layer. 

2. Lack of justification on why tree attention is added backwardly starting from the last layer? what happen if you add the tree attention layer forwardly starting from the beginning layer? It would be interesting to see such comparison in the ablation study. 

3. What is the configuration of hyper-parameter layer width for boost strapping? (couldn't find it in the appendix as well) and the previously mentioned training steps of pre-training an existing model before boost strapping. Those missing information is required for researchers to reproduce the experimental results.

Some minor comments/questions:
1. It would be better to explicitly mention that there is no pre-training in the LRA. So that readers can better understand the setting of training TREEFORMER.

2. In the computational complexity -- point 2, in section 3.2, should there be n multiplied in front of the complexity value as there are n key vectors to store where n is sequence length?
clarity,_quality,_novelty_and_reproducibility: The idea presented in the paper is novel to address the quadratic complexity of the attention layer respect to the sequence length. Extensive evaluations on both GLUE and LRA and benchmarking against multiple SOTA models show a convincing result for the proposed method TREEFORMER. Overall, it is a high quality paper.

I do have minor concerns on the missing details of boost strapping method to reproduce the experimental results. 
summary_of_the_review: Attention computation of Transformers scales quadratically (n^2) with the input sequence length (n), making it a key bottleneck
in scaling Transformers to long inputs. Targeting at such limitation, this paper proposes a new architecture called TREEFORMER to use decision trees to efficiently compute attention by only retrieving the top nearest neighboring keys for a given query. In addition, to make the decision trees to be trained as part of neural networks using back propagation, this paper proposes a novel bootstrapping method to gradually sparsify attention computation. Finally, the paper conducts an extensive evaluations on both GLUE and LRA and benchmarking against multiple SOTA models, show a convincing result for the proposed method TREEFORMER. Overall, it is a high quality paper.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes a new sparse attention mechanism that uses decision trees to cluster the queries and keys such that the queries only attend to the keys that fall into the same leaf. To mitigate the collapsing problem (all keys and queries are allocated to one leaf node) in the proposed TF-Attention (Tree Fine-grained Attention), the paper proposes two mitigations: 1) Tree Coarse Attention, where the attention is summed unweighted for the keys in the leaf node. 2) a boot-strapping training strategy that gradually increases the tree height.

strength_and_weaknesses: 
Strength:

1. The idea of using decision trees to cluster the queries and keys is novel to my knowledge. The authors made essential technical contributions (TC-Attention & Bootstrapping) to make their model work.

2. In terms of FLOPs as computational efficiency metrics, the proposed TreeFormers achieve competitive pre-training & fine-tuning performance with significantly reduced FLOPs in attention layers.


Weaknesses:


1. The paper only evaluated their model in short-sequence (512) real NLP tasks and a synthetic long-sequence task. Compared to GLUE, evaluation on real long-sequence NLP tasks such as long-document summarization & long-document question answering [1] is more suitable for this paper.

2. Due to the non-parallelization of the decision-tree design, the paper mainly uses FLOPs and CPU run-time as efficiency metrics. While this is acceptable, when comparing with baselines on LRA, the paper should also include other state-of-the-art non-Transformer models such as S4 [2].

3. The paper misses citation and discussion with other learning-based sparse attention methods [3-5], to which the proposed method belongs.

4. The paper claims the computational complexity for TC-ATTENTION is linear in sequence length n, which is not correct. Notice that the real complexity of TC-Attention is $O(nkdh + (2^{h+1} − 1)d)$, where $k$ is the number of key tokens in the same leaf node as queries. When $h$ is fixed, $k$ increases linearly with $n$, so the complexity is quadratic. When $k$ is fixed, $h$ increases in the linearly with $\log n$, so the complexity is $n \log n$.



[1] Guo, Mandy, Joshua Ainslie, David Uthus, Santiago Ontañón, Jianmo Ni, Yun-Hsuan Sung, and Yinfei Yang. "LongT5: Efficient Text-To-Text Transformer for Long Sequences."

[2] Gu, Albert, Karan Goel, and Christopher Re. "Efficiently Modeling Long Sequences with Structured State Spaces." In International Conference on Learning Representations. 2021.

[3] Kitaev, Nikita, Lukasz Kaiser, and Anselm Levskaya. "Reformer: The Efficient Transformer." In International Conference on Learning Representations. 2019.

[4] Sun, Zhiqing, Yiming Yang, and Shinjae Yoo. "Sparse Attention with Learning to Hash." In International Conference on Learning Representations. 2021.

[5] Wang, Ningning, Guobing Gan, Peng Zhang, Shuai Zhang, Junqiu Wei, Qun Liu, and Xin Jiang. "ClusterFormer: Neural Clustering Attention for Efficient and Effective Transformer." In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 2390-2402. 2022.
clarity,_quality,_novelty_and_reproducibility: The paper is written clearly and the idea is novel. But the evaluation (in terms of FLOPs and CPU latency) and baseline comparisons (on LRA) are not consistent.
summary_of_the_review: The paper proposes a new sparse attention mechanism that uses decision trees to cluster the queries and keys. However, there are several problems in the paper (See Weaknesses). Therefore, I believe the paper is below the acceptance threshold.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper proposes Treeformer, a modified Transformer that replaces the multi-head self-attention layers with tree-based attention mechanisms. There are three variants studied in this paper: TF-attention (Tree Fine-grained attention), $k$-ary TF-attention, and TC-attention (Tree coarse attention). The authors claim that Treeformer outperforms Transformers while being more efficient (linear time complexity with respect to sequence length). Experimental results on MLM, MNLI, SQuAD, GLUE, and LRA are provided.
strength_and_weaknesses: ## Strengths
- The methods are novel.
- The authors conduct experiments on a lot of benchmark datasets. The results on SQuAD and GLUE are strong.
- The proposed method is novel and it is interesting that the proposed tree-based attention can outperform conventional self-attention

## Weaknesses
- It would be better if the models can include results on generative tasks where casual attention is used such as causal language modeling, 
- The author omits stronger baselines on LRA and only includes old inferior models. For example, recent papers including Luna, H-Transformer-1D, S4, and DSS achieve better performance while being efficient.
- Despite saving a significant number of FLOPs, the method cannot be implemented efficiently. As shown in Table 4, Treeformer TF-A starts to be faster than Transformer only after the sequence length gets to 4096. Moreover, I think the authors should include the inference time of Treeformer TF-C, and Treeformer $k$-ary TF-A. Moreover, it would be better to see the inference time comparison between Treeformers and other efficient Transformer variants.
clarity,_quality,_novelty_and_reproducibility: - Clarity: the paper is written clearly.
- Quality: the results on MNLI, GLUE and SQuAD are convincing, but the experiments on LRA is weak.
- Novelty: the proposed method is novel.
- Reproducibility: the code is not provided, but the authors provide hyperparameters for all experiments.
summary_of_the_review: In a nutshell, this paper proposes a novel idea and the results on SQuAD and GLUE benchmarks are strong and convincing. Despite its weaknesses in LRA experiments and actual inference time, I believe that the paper is still slightly above the borderline and give it a weak accept.
review_title: null

Meta-review: metareview:_summary,_strengths_and_weaknesses: The paper proposes TREEFORMER to learn and use decision trees to efficiently compute attention by only retrieving the top nearest neighboring keys for a given query. It uses a bootstrapping method to gradually sparsify attention computation. Reviewers generally likes the idea and the experiments. The authors also clarified some of their (and my) concerns during the rebuttal. I am in favour of accepting this paper.  
justification_for_why_not_higher_score: good paper
justification_for_why_not_lower_score: na
Decision: Accept: poster