T REEFORMER : D ENSE G RADIENT T REES FOR E FFICIENT A TTENTION C OMPUTATION
A BSTRACT
Standard inference and training with transformer based architectures scale quadratically with input sequence length.
This is prohibitively large for a variety of applications especially in web-page translation, query-answering etc.
Consequently, several approaches have been developed recently to speedup attention computation by enforcing different attention structures such as sparsity (Zaheer et al., 2020), low-rank (Wang et al., 2020), approximating attention using kernels (Choromanski et al., 2021).
In this work, we view attention computation as that of nearest neighbor retrieval, and use decision tree based hierarchical navigation to reduce the retrieval cost per query token from linear in sequence length to nearly logarithmic.
Based on such hierarchical navigation, we design Treeformer which can use one of two efﬁcient attention layers – TF-A TTENTION and TC-A TTENTION .
TF-A TTENTION computes the attention in a ﬁne-grained style, while TC-A TTENTION is a coarse attention layer which also ensures that the gradients are “dense”.
To optimize such challenging discrete layers, we propose a two-level bootstrapped training method.
Using extensive experiments on standard NLP benchmarks, especially for long-sequences, we demonstrate that our T REEFORMER architecture can be almost as accurate as baseline Transformer while using 30x lesser FLOPs in the attention layer.
Compared to Linformer, the accuracy can be as much as 12% higher while using similar FLOPs in the attention layer.
I NTRODUCTION
Self attention layer is the key component of Transformers (Vaswani et al., 2017), enabling them to achieve state of the art performance across tasks in Natural Language Processing (Devlin et al., 2019;
Radford et al., 2019; Raffel et al., 2019) and Computer Vision
(Dosovitskiy et al., 2021).
Attention computation scales quadratically ( n 2 ) with the input sequence length ( n ), making it a key bottleneck in scaling Transformers to long inputs.
This has resulted in a lot of alternate proposals to efﬁciently compute attention using different approximations.
However, as shown in Tay et al.
(2021), methods that offer dramatic speedups (Wang et al., 2020; Choromanski et al., 2021) suffer in accuracy, and methods that match accuracy of Transformer (Zaheer et al., 2020) don’t offer signiﬁcant speedups.
In this paper we develop a novel efﬁcient attention mechanism using input dependent sparsity structure in attention computation, motivated by the analysis showing the sparse nature of attention (Shi et al., 2021; Gupta et al., 2021).
We propose to use decision trees to efﬁciently compute attention by only retrieving the top nearest neighboring keys for a given query.
Decision trees are low cost, hierarchical non-linear models that have been popular in both supervised (Chen & Guestrin, 2016; Quinlan, 2014) and unsupervised (Duda et al., 1973)
learning.
Given a vector, each node of the decision tree decides whether to navigate to the left or the right child depending on the sign of a simple linear projection.
We learn decision trees in each attention layer of a Transformer model so that a query pays attention only to keys that map to the same leaf nodes 1 of the decision trees.
We refer to this as TF-A TTENTION .
Such sparse decision trees that retrieve only a single leaf node are typically hard to train due to their discrete nature.
Further, queries that pay strong attention to multiple keys result in bunching many keys together in a single leaf node creating imbalance.
The key computational advantage of decision trees is realized only when they are balanced.
To address these issues we propose another attention mechanism TC-A TTENTION , that in addition to keys in the leaf node, pays attention to keys at all the intermediate nodes traversed by the query.
We further combine the retrieved value vectors with query independent, but trainable weights.
This allows us to learn a decision tree that clusters keys in a hierarchical and balanced manner.
In both these methods we train the decision tree in an end-to-end manner with the rest of the components of the Transformer.
Since the decision trees themselves are inexpensive to evaluate, this allows us to reduce attention computation cost dramatically, requiring typically only a few dot products per query.
In our goal to train these decision trees in an end to end manner with the rest of the Transformer components we faced some challenges.
Naively using the decision trees to restrict attention computation results in poor optimization with training getting stuck at high loss.
To solve this we introduce a novel bootstrapping method that allows us to gradually restrict attention computation in layers using decision trees.
While this keeps training expensive initially, it still dramatically reduces inference costs.
We further use Dense Gradient Trees (Karthikeyan et al., 2021), allowing us to use simple gradient based optimization.
We evaluate the proposed architecture on popular NLP models such as BERT (Devlin et al., 2019), and on the Long Range Arena benchmark (LRA) (Tay et al., 2021), developed speciﬁcally to compare efﬁcient attention methods on long sequence length tasks.
We show that T REEFORMER is able to match/improve the performance over other popular models while achieving lower computational cost, for e.g. - T REEFORMER has 8-9x lesser FLOPs than models like BigBird (Zaheer et al., 2020) and
Performer (Choromanski et al., 2021).
Our contributions in this paper are — (i)
We pose attention computation as a nearest neighbor retrieval problem, and propose a novel architecture T REEFORMER − using dense gradient trees for retrieving top keys to pay attention to for a given query.
We develop two novel attention mechanisms using decision trees - that, for a given query, pay attention to only keys in the matching leaf node (TF-A TTENTION ) or to all the keys along the traversed path (TC-A TTENTION ).
(ii) We develop a novel bootstrapping method to gradually sparsify attention computation thereby showing that decision trees can be trained as a part of neural networks using back propagation.
(iii) We experimentally show the effectiveness of the proposed architecture using BERT models and on LRA (Tay et al., 2021).
In particular, both TF-A TTENTION and TC-A TTENTION have matching accuracy to baseline
Transformer architecture, but with up to 30 × cheaper (in FLOPs) attention layer compared to standard
Transformers and up to 9 × cheaper attention layer compared to SOTA BigBird (Zaheer et al., 2020).
P RELIMINARIES : A TTENTION AND D ECISION TREES 2. A TTENTION
Let Q , K , V ∈ R n × d be the input query, key, and value matrices, where n is the sequence length and d is the embedding dimension, and W Q , W K ,
W V ∈ R
d × d are their respective projection matrices.
The standard attention in each Transformer layer can be deﬁned as
[Equation]
Note that for the self attention layer, all the three matrices Q , K , V are same as the input to the layer.
Computing both attention A and its product with the projected value matrix V W V is an
O ( n 2 d + d 2 n ) operation.
In particular, assuming cpqr cost for multiplying two matrices of size p × q and q × r , the main cost of this layer is: c (2 n 2 d + 2 d 2 n ) .
This computation becomes a major bottleneck in applications which have large sequence lengths n .
To mitigate this, several efﬁcient attention mechanisms have been proposed (Child et al., 2019; Wang et al., 2020; Choromanski et al., 2021; Kitaev et al., 2020; Sun et al., 2021; Wang et al., 2022).
Child et al.
(2019) proposed sparse attention ( A ) with different sparsity patterns that reduce attention cost from O ( n 2 ) to O ( n 1 . 5 ) .
Later Zaheer et al.
Yun et al.
(2020) proved universal approximation power of Transformers with sparse attention, improving the cost to O ( ns ) , with constant sparsity s per query.
Alternately Wang et al.
(2020) considered a low rank approximation to attention, reducing computation cost to O ( nk ) for rank k approximation.
Choromanski et al.
(2021) considered a kernel based approach to approximate attention computation reducing the cost to O ( nr ) , where r is the dimension of the random features used to compute the kernel.
We refer readers to Tay et al.
(2020) for a detailed survey of efﬁcient attention methods.
In this paper, we propose a decision tree based mechanism to speed up the attention computation to O ( nh ) , where h is the height of the decision tree.
We show that with T REEFORMER even a h as small as 8 is sufﬁcient for most applications, leading to the least amount of FLOPs required, while maintaining the performance.
[Table]
2. D ECISION TREES
Decision trees are a classical technique that enables partitioning of the input space in several oblique blocks that can be navigated efﬁciently using the hierarchical structure.
Decision trees and similar hierarchical structures are highly popular in supervised learning (XgBoost (Chen & Guestrin, 2016), C4.
(Quinlan, 2014)), unsupervised learning (hierarchical clustering (Duda et al., 1973)), and retrieval (KD tree (Bentley, 1975)).
Generally, hierarchical retrieval techniques are nominally data-dependent, e.g., PCA-tree (Sproull, 1991) and Cover Trees (Beygelzimer et al., 2006).
Our TC-A TTENTION and TF-A TTENTION methods combine hierarchical retrieval methods with supervised tree learning techniques to enable learning of the hierarchical attention layer in end-to-end differentiable manner.
See Figure 1 for overall representation of the hierarchical navigation based attention layer.
Each internal node of the tree has a classiﬁer that determines which children a given input should navigate to.
While the classiﬁer makes hard decision, we use technique from straight-through estimators (Bengio et al., 2013), and in particular, Dense Gradient Trees from Karthikeyan et al.
(2021) to train the tree parameters using gradient style methods.
Construction Formally, let T ( θ ) represent a binary tree with height h and 2 h nodes.
Let θ { l,j } denote the tree parameters at { l, j } -th node where l ∈ { 0 , 1 , · · · , h } is the node level and j ∈ { 0 , 1 , · · · , 2 l − 1 } is the node index at level l .
Let decision at each internal node to send the point q to left or the right child is based on SIGN ( f ( θ { l,j } ; q )) , where,
[Equation]
and θ { l,j } = { w { l,j } ; b { l ; j } } consists of a oblique hyperplane and a bias.
Let P T ( l, q ) denote the node-index of T that was traversed by q at level l .
That is, P T ( h, q ) denote the leaf node to which q is mapped.
Also, let S l,j ( K ) denote the indices of vectors from the key matrix K that traverses node { l, j } .
That is, i ∈ S l,j ( K ) iff j = P T ( l, k i ) .
T REEFORMER A TTENTION V ARIANTS
In this section, we describe the decision tree attention mechanism which replaces the standard dot product attention (Vaswani et al., 2017) in our T REEFORMER architecture.
Recall that for each query q
i , the attention is given by:
[Equation]
That is, attention computation for q i is equivalent to applying a linear scan over all keys, and computing value as the weighted average of value of most similar keys to q i .
Recent works have shown that we do not need to attend to all tokens in a given input sequence (Shi et al., 2021), and demonstrate that top- k attention is a good approximation for standard attention (Gupta et al., 2021).
We propose to leverage this observation by using a decision tree based hierarchical navigation instead of linear scan to ﬁnd the top-most similar keys.
Naturally, accuracy of the method would depend on accuracy of retrieval of most similar keys by the hierarchical navigation method.
So, we propose to learn the decision tree parameters so as to explicitly optimize the loss function.
We develop two approaches to compute the value for a query q i based on its traversal through the tree.
In the following subsections we present our two different tree based attention layers.
3. TF-A TTENTION ( Tree Fine-grained Attention )
In this variant, we compute the value for a query q as a weighted sum of value vectors of keys that lie in the same leaf node as q .
More formally,
[Equation]
where ¯ S = S h,P T ( h, ( W Q ) T q ) ( K W K ) is the set of keys K that are mapped to same leaf node (i.e.
P T ( h, ( W Q ) T q ) ) as ( W Q ) T q .
Note that this is a sparse attention variant as we compute dot products only between those queries and keys belonging to the same leaf node, and the attention score between a query and keys lying in different leaf nodes is 0.
Intuitively, this attention layer can be thought of as top- k attention where the tree structure is responsible for computing the k tokens which are the most important for a given query token.
Computational Complexity Assuming a uniform distribution of the input sequence tokens across the leaf nodes of the decision tree, the amortized computational cost for TF-A TTENTION is O ( n 2 d 2 h ) + 2 cd 2 n ≡ O ( nkd ) + 2 cd 2 n , where k is generally a small factor assuming the height of the decision tree is of the order O ( log ( n k )) , and c is as deﬁned in Section 2.1.
If the decision tree parameters are such that all the queries and keys get clustered into a single leaf node of the tree, we get the original full attention in this case.
Therefore the worst case time complexity of TF-A TTENTION is O ( n 2 d + d 2 n ) .
Furthermore, this also illustrates that TF-A TTENTION can be as expressive as standard attention.
We state this claim formally below.
Lemma 1.
There exists a tree T s.t.
TF-A TTENTION ( Q , K , V ; T ) =
Attention ( Q , K , V ) .
The lemma follows immediately from the above observation; for completeness, see Appendix for a proof of the above lemma.
k -ary TF-A TTENTION We extend the binary tree construction of the TF-A variant to construct a k -ary tree attention model for height h , where each node at level l has a branching factor of b l , l ∈ { 0 , 1 , · · · , h − 1 } .
The total number of leaf nodes in this case are (cid:81) l = h − 1 l = b l .
This is a general version of the TF-A TTENTION variant, where we can scale the number of leaf nodes even for smaller heights, and thus reducing the amount of queries and keys in each leaf node.
In all our experiments on k -ary TF-A TTENTION models, h = 2 and b = [100 , 10] , resulting in 1000 leaf nodes.
Although the proposed TF-A TTENTION T REEFORMER variant works well in principle, there are a couple of shortcomings of the same — (i)
In case the distribution of queries and keys is skewed in even one of the layers of the model, the computational cost approaches the worst case time complexity of the standard attention model.
So we might need to add explicit loss terms to encourage more uniform representation across various leaf nodes/subtrees.
(ii) It might be possible that for a given query, some key tokens important to this query lie in a different leaf node of the decision tree to minimize the overall loss.
But due to the sparse nature of the model, the attention score between the query and these keys will be 0.
Using multiple decision trees (forest structure) might help alleviate this but we leave it for future exploration.
We address the above problems using a novel variant that computes the attention values in a coarse manner and hence has denser gradients.
3. TC-A TTENTION ( Tree Coarse Attention )
In this variant, rather than computing attention as a weighted sum of values of keys in the leaf node, we take a coarse unweighted sum of the value vectors.
That is, the query is only used to navigate to a leaf node, but with-in a leaf node, the value vector is ﬁxed .
So given a leaf node, the attention value is independent of the query itself.
Note that this requirement completely ﬂips the tree’s tendency to map multiple keys to a single leaf node.
That is, to optimize the loss function, TF-A TTENTION would prefer all the keys to map to one leaf node,
TC-A TTENTION ’s incentives are completely reversed.
That is, if TC-A TTENTION maps all the keys to the same leaf node then it would give an unweighted average of the value vectors of all the keys, irrespective of query which is average attention - naturally a weaker model and would have large loss.
Instead, TC-A TTENTION would prefer to spread out keys as much as possible in leaf nodes so that it can compute more powerful query dependent attention.
We present key distributions of learned T REEFORMER models in Figure 4 supporting this argument.
We further generalize this technique to ensure contributions from other nodes to the value vector.
That is, we add contributions to a query’s value vector from not only the leaf node but also from the internal nodes that were traversed by the query.
Formally, let ν l,j be the value vector associated with a node.
Then,
[Equation]
i.e., it is an average of all value vectors of keys such that l, j lies in their paths.
Now, given ν l,j for each node and the tree T (including it’s parameters θ ), we get the attention for a query q as:
[Equation]
where α l is the learnable attention weight at each level of the decision tree.
See Figure 1 for a visual representation of TC-A TTENTION .
Note that as TC-A TTENTION applies query independent averaging, it is easy to design counterexamples where the standard attention layer is strictly more powerful than TC-A TTENTION .
Now, we can make TC-A TTENTION as powerful as a standard attention layer by allowing traversal of multiple paths and α l to be query dependent.
However, our results indicate that TC-A TTENTION is sufﬁciently powerful, and might not require more modeling power in the attention layer.
We leave further exploration of multiple traversal paths for future work.
Computational Complexity
The cost of TC-A TTENTION can be broken down into three parts: 1. projection of query, key, and value tensors - note that we can embed projection matrices for both query and key in the parameters of the tree (with two different sets of tree parameters).
So the cost is half of standard attention: cd 2 n 2.
storing projected key vectors in the decision tree -
O ((2 h
+1 − 1) d ) , and 3.
computing the path and attention layer transformation for all queries Q - O ( ndh ) .
Compared to the full attention mechanism - O ( n 2 d + d 2 n ) , the computational complexity for TCA TTENTION is linear in sequence length n : O ( ndh + (2 h +1 − 1)
d + d 2 n ) .
In all our experiments, decision tree height h is generally ≤ 10 .
Also, note that compared to TF-A TTENTION , the computational complexity is guaranteed to be linear in both n and h , irrespective of the trained trees/model.
3. T RAINING OF T REEFORMER MODELS
Pre-training It is well-known that learning tree is a challenging problem due to the discrete structure.
We use recent techniques (Karthikeyan et al., 2021) to help with the optimization, and ensure better ﬂow of gradients throughout the network.
Despite this, we found in our experiments that pretraining complete T REEFORMER models with tree attentions in all L layers is challenging and often leads to poor sub-optimal solutions.
To overcome this challenge issue and avoid inaccurate saddle points, we use a method of incremental bootstrapping for training T REEFORMER (see Algorithm 1).
Speciﬁcally, we ﬁrst pre-train an existing model (say a transformer or BigBird).
We then introduce tree attention (either TF-A TTENTION or TC-A TTENTION ) only in the last few layers (layer j to
L ).
The new architecture – denoted by M { j, ··· ,L } – is trained for a ﬁxed number of steps using the standard self-supervised learning objective.
Now, all the parameters of M { j, ··· ,L } are used to initialize the M { i, ··· ,j, ··· ,L } model.
We do this bootstrapping until we have the complete M { 1 , ··· ,L } T REEFORMER model and it is also trained for a given number of steps.
In a similar fashion, we can do this bootstrapping process to initialize T REEFORMER models for a given height using a
T REEFORMER model with a smaller tree height.
Based on TF-A TTENTION or TC-A TTENTION model, multiple bootstrapping techniques for tree levels can be explored: random assignments in next level, assignment to one node only etc.
In this work, we mostly use random assignments based approach where the tree parameters in layers { i, · · · , j − 1 } are randomly initialized.
Fine-tuning We don’t do any bootstrapping during the downstream ﬁne-tuning tasks and tree attention exists in all layers of the models.
But, we tried two variants of keeping the pre-trained decision tree weights ﬁxed and allowing them to be trained during ﬁne-tuning, and the corresponding results are reported in Section 4.
Overall, bootstrapping in ﬁnetuning phase leads to similar results as standard ﬁne-tuning with all the model and tree parameters.
E XPERIMENTS
In this section, we present experimental results comparing the T REEFORMER models with Transformers and with popular efﬁcient Transformer architectures – BigBird (Zaheer et al., 2020),
Linformer (Wang et al., 2020) and Performer (Choromanski et al., 2021).
We ﬁrst evaluate our
T REEFORMER on the standard NLP benchmark of BERT-based pre-training and GLUE+SQuAD ﬁnetuning (Devlin et al., 2019).
Next, to study effectiveness of T REEFORMER on long sequence length tasks, we present results on the standard benchmark for methods in long-sequence regime:
Long
Range Arena (LRA) (Tay et al., 2021).
We also provide ablation studies to understand challenges in optimizing T REEFORMER and the utility of the bootstrapping approach (Algorithm 1).
The complete experimental setup along with the hyper-parameters for each of the tasks and hardware details is provided in Appendix A. Please note that we will be using a shorter notation of TF-A for
TF-A TTENTION and TC-A for TC-A TTENTION .
4. BERT P RE - TRAINING / F INE - TUNING
Following Devlin et al.
we use a masked LM (MLM) objective to pretrain our T REEFORMER
BERT-Base models on Wikipedia, Books (Zhu et al., 2015), CC-News (Guu et al., 2020), and
Stories (Trinh & Le, 2018) datasets.
[Table]
Fine-tuning - GLUE/SQuAD
We ﬁne-tune the pre-trained BERT-Base models on 7 datasets from the GLUE benchmark (Wang et al., 2018) including MNLI and the SQuAD (Rajpurkar et al., 2016) dataset for question answering.
Table 2 shows the comparison between various architectures on these tasks.
We observe that the T REEFORMER model with TF-A TTENTION and k -ary
TF-A TTENTION work well - beating the BERT baseline in MNLI and SQuAD while using fewer FLOPs.
The bracketed values in Table 2 for the T REEFORMER models represent ﬁne-tuning while the decision tree parameters are ﬁxed and obtained from the corresponding pre-trained T REEFORMER models.
We would like to make multiple observations here: a).
TF-A TTENTION is competitive with both baseline BERT model as well as the BigBird model.
In fact, for MNLI task, it is 1% more accurate than BERT, and on SQuAD, it is 2% more accurate than BigBird, despite requiring 3x fewer FLOPs in the attention layer.
b).
TF-A TTENTION is more accurate than Linformer as well, while still being cheaper c).
TC-A TTENTION provides accuracy values competitive with baseline Linformer, BigBird only when TC-A TTENTION layer is applied in 6 of the 12 transformer blocks.
With a larger number of TC-A TTENTION layers, the accuracy drops signiﬁcantly, conﬁrming challenges in optimizing a difﬁcult discrete layer across multiple layers.
We leave further investigation into better optimization of TC-A TTENTION in deep architectures for future work.
Beneﬁts of Bootstrapping Figure [-
-]2 shows the training progression of the T REEFORMER models with bootstrapping (Algorithm 1).
The sudden decrease (increase) in MLM accuracy (loss) happens at the bootstrapping stage, when a model with more decision tree layers is initialized with a model with less decision tree layers.
We observe that while bootstrapping causes a momentary drop, it eventually leads to a higher ﬁnal accuracy for TF-A. For TC-A we notice improvement in bootstrapping uptodecision tree layers (till ≈ 1 e 6 steps), but going to 12 decision tree layers causes training instabilities.
We next compare the training progress for T REEFORMER models with different decision tree heights in Figure 3.
We notice that TF-A performs better with bootstrapping for smaller tree heights, but suffers as tree height goes to 7.
On the other hand we ﬁnd TC-A to be much more challenging in bootstrapping to more that 6(out of 12) layers.
For larger heights the training is hard due to vanishing gradients and MLM accuracy approaches 0 at the last stage of the bootstrapping process.
Leaf distribution Figure 4 shows the distribution of key vectors in the leaf nodes of the T REEFORMER models in speciﬁc layers.
Lighter lines represent the initial stages of the pre-training process and darker lines represent the distribution at later stages in the pre-training – showing the evolution of distribution as the training progresses.
We observe that the distribution is more skewed in later layers compared to the initial layers.
Moreover, the distribution is much more skewed in the TF-A variant compared to TC-A.
This is because TF-A, being a sparse attention mechanism, will try to cluster more keys in a given leaf node to attend to all relevant keys for a given query (see § 3).
4. L ONG R ANGE A RENA (LRA)
We next perform experiments on the Long Range Arena Benchmark (LRA) (Tay et al., 2021).
It comprises of 6 tasks ( ListOps , Text , Retrieval , Image , Pathﬁnder , and Path-X ) with inputs having long
[Figure]
sequence lengths ranging from 1024 to 16384 .
Three of the tasks are text-based and the remaining three comprise of visual inputs.
In our comparison, we skip the Pathﬁnder task in the LRA benchmark because of the very high variance in the accuracy for all models, and also skip Path-X because all model architectures fail in this task due to the memory or computational bottlenecks at sequence length 16384.
The results of various model architectures across the remaining 4 tasks are presented in
Table 3.
Along with the performance, we also show the computational FLOPs relative to the standard
Transformer averaged over the 4 tasks for each of the models in the self attention layer of an encoder block in the model architecture.
On average, both TF-A TTENTION and TC-A TTENTION are almost as accurate as the baseline
Transformer as well as the BigBird.
Further, attention layer FLOPs of TF-A TTENTION is almost 3x lower than that of Transformer.
These factors are further improved by using the k -ary TF-A TTENTION
[Table] [Table]
variant.
TC-A TTENTION on the other hand is almost 33x cheaper than Transformer and 9x cheaper than BigBird despite achieving almost the same accuracy.
Only Linformer has similar computational cost as TC-A TTENTION , but TC-A TTENTION is almost 5% more accurate than Linformer.
In terms of total FLOPs for the transformer encoder block, the self attention comprises of 50.84% of the total computational FLOPs in the ListOps task, and the rest is used in dropouts and the two feed forward layers.
For Text , the distribution is 73.3% and 26.7% respectively.
So, we can get a total speedup of 2-3x by using T REEFORMER models instead of the standard Transformer.
Inference times In Table 4, we ﬁnally present a comparison of inference wall time between the
T REEFORMER TF-A TTENTION model and the standard attention on a Intel Xeon Platinum PCPU using one thread and 100GB of RAM.
We measure the inference time for the self-attention module using 8 attention heads and hidden dimension 768 for a batch size of 1.
We use a decision tree of height 6 for TF-A TTENTION .
Our implementation of T REEFORMER TF-A TTENTION model iterates over the leaf nodes and gathers the subset of queries and keys that lie in a given leaf node, and only computes attention scores by doing matrix multiplication between these gathered queries and their relevant keys.
The ﬁnal projections for the gathered keys are then scattered into the original query matrix.
For sequence lengths 2048 - 8192, we get a speedup ranging from 1.8x - 6.7x .
C ONCLUSION
In this work we proposed T REEFORMER with a novel efﬁcient attention mechanism that leverages the idea of ﬁnding top nearest neighbors using decision trees and applying it to attention computation.
We proposed two attention variants TF-A TTENTION and TC-A TTENTION with different sparsity levels.
We showed that using techniques like dense gradient trees and bootstrapping allows us to train the T REEFORMER models in an end to end manner with simple back propagation.
Using extensive experiments, we showed that T REEFORMER models achieve comparable or better accuracy than other efﬁcient attention models, but with signiﬁcantly lower computational complexity.
While decision trees are popular in other applications (XgBoost (Chen & Guestrin, 2016), C4.
(Quinlan, 2014), KD tree (Bentley, 1975)), there have been few works integrating and jointly training them with neural networks.
Our work takes a ﬁrst step in that direction and shows that while optimizing with decision trees can be challenging, it is also rewarding.
Some interesting future research directions include developing better optimization techniques for discrete models, efﬁciently combining decision trees with different types of neural network architectures, and using multiple decision trees to reduce discretization.
A E XPERIMENTAL S ETUP
In this section we give a detailed description of our experimental setup including various hyperparameters.
A. BERT P RE - TRAINING /F INE - TUNING
Table 5 detail the four public datasets we used for pre-training the T REEFORMER and baseline models.
For all datasets, we truncate documents with length greater than 512 to multiple sequences of length 512.
We follow the recommendations of Devlin et al.
for MLM pre-training
[Table]
and mask 15% tokens in each input sequence.
We report the MLM performance in terms of accuracy of prediction on the masked positions.
We list the hyper-parameters used for pre-training in Table 6.
We use Adam optimizer with decoupled weight decay ( AdamW ).
We also use a learning rate warmup for the ﬁrst 10k steps, with linear decay of the learning rate afterwards.
We train the T REEFORMER models for a total of 1.8M steps - 450000 steps during each of the four bootstrapped processes shown in Figuresand 3.
The layer-width for each bootstrapping is 3, meaning 3 tree attention layers are added in every subsequent bootstrapping process.
[Table]
[Table]
For ﬁne-tuning
[Table]
we use the same pre-trained models and ﬁne-tune on the downstream tasks with most of the model parameters borrowed from the pre-training process.
The parameters that are different for various datasets in the GLUE benchmark and SQuAD are detailed in Table 7.
For MNLI, we report the matched accuracy, Matthews Correlation Coefﬁcient (MCC) for COLA, F1 Scores for MRPC, and accuracies for the remaining tasks - QQP, QNLI, SST-
[Table]
and RTE.
On the SQuAD dataset, we report the F1 scores on v1.1 version of the dataset
[Table] [Table]
A. L ONG R ANGE A RENA (LRA)
The task speciﬁc hyper-parameters used for the LRA Benchmark (Tay et al., 2021) are presented in
Table 8.
We use Adam optimizer for all the experiments but the weight decay factor is different for
Image task as highlighted in Table 8.
We start the bootstrapping for TC-A TTENTION variant of the
T REEFORMER from tree height h
= 2 up to all the way to h = 10 .
We notice that h = 6 performs the best on most of the LRA tasks.
So we report the performance numbers for h = 6 .
We also include the number of training steps during each bootstrapping process in the table.
[Table]
B E XPERIMENTAL R ESULTS
In this section
[Table]
we include some extra results in addition to those present in the main paper.
B. BERT P RE - TRAINING /F INE - TUNING
We ﬁrst present the task-wise performance of the different models on the GLUE benchmark in Table9.
In this section, we also demonstrate the importance of bootstrapping in the training process with the help of some empirical evidence using the same set of experiments as in the paper.
We track the gradient norms during the pre-training process for T REEFORMER models with bootstrapping and T REEFORMER models trained from scratch without any bootstrapping to see how they compare during the pre-training stage.
Figure 5 shows the corresponding plots for BERT-based T REEFORMER models in layers 1 , 4 , 8 , 12 of the models.
We observe that gradient norms for the bootstrapped
T REEFORMER model are higher compared to the non-bootstrapped version of the same model.
This shows the challenges in training decision trees directly with neural networks and the advantage of the proposed bootstrapping approach.
B. L ONG R ANGE A RENA (LRA)
Table 10 extends Table 3 to include more models like Kitaev et al.
(2020), Sun et al.
(2021), and
Gu et al.
Although Gu et al.
(2021) is the state of the art in the LRA benchmark, the paper itself mentions that S4 cannot be scaled to sequence modeling using very large models because of its compute and memory requirements.
S4 experiments in the paper are on smaller models and it cannot scale to the level of Transformer-based models involving billions of parameters.
Sun et al.
(2021) performs better than T REEFORMER models on LRA, but it loses performance on GLUE tasks compared to the baseline Transformer models.
In Table 3, we presented the average speedup (FLOPs column) over the four tasks in the LRA
Benchmark for each of the models.
Table 11 breaks down that average and details the speedup for each representative task in LRA.
[Table] [Table]
[Table] [Table]
We also perform experiments on sequence-to-sequence (seq2seq) tasks with encoder-decoder models. Following Zaheer et al. (2020), we only introduce our attention mechanism in the encoder and use standard full attention in the decoder, as the number of output tokens is usually small ( ≤ 512 ) compared to the input.
It has been well established that pre-training also helps in seq2seq tasks like summarization, machine translation, and other generative tasks Lewis et al.
Liu et al.
So, we also pre-train our encoder-decoder T REEFORMER models using the BART objective. We use the pre-trained encoder-decoder models and ﬁnetune on the abstractive summarization task on the PubMed and Arxiv datasets consisting of long documents with a maximum sequence length of 3072 for the input.
The results on this task are presented in Table 12.
Our TF-A TTENTION models outperform the baseline Transformer and BART models, and are competitive with the BigBird models with signiﬁcant reduction in FLOPs in the attention layer. The FLOPs in TF-A TTENTION models are higher compared to BigBird because of the skewness in leaf nodes of the decision tree. We leave further optimization of TF-A TTENTION to enable less leaf-node skew for future work.
C P ROOFS
We prove Lemma 1 by construction - we show that there exists a choice of the decision tree weights for all layers of the T REEFORMER TF-A TTENTION model, such that all the queries and keys cluster to the same leaf node.
From Equation 2, we have the tree parameters ( w { l,j } ; b { l,j } ) for tree T .
We assign the parameters as follows:
[Equation]
where (cid:15) is any small positive value.
We denote the tree with the above parameters as T