zhenhai@google.com Radu Soricut Google Research rsoricut@google.com Abstract We describe an efficient hierarchical method to compute attention in the Transformer architecture. The proposed attention mechanism exploits a matrix structure similar to the Hierarchical Matrix (H-Matrix) developed by the numerical analysis community, and has linear run time and memory complexity. We perform extensive experiments to show that the inductive bias embodied by our hierarchical attention is effective in capturing the hierarchical structure in the sequences typical for natural language and vision tasks. Our method is superior to alternative sub-quadratic proposals by over +6 points on average on the Long Range Arena benchmark. It also sets a new SOTA test perplexity on One-Billion Word dataset with 5x fewer model parameters than that of the previous-best Transformer-based models. Linearly combining information using content-based weights, a method generically known as attention, is a key building block in many deep neural networks such as recurrent neural networks (RNN) (Luong et al., 2015), convolutional neural networks (CNN) (Bello et al., 2019) and graph convolutional networks (GCN) (Velickovic et al., 2018). One particular type of such attention, called multi-head scaled dot-product attention, is one of the main components of the Transformer architecture proposed by Vaswani et al. (2017), which has been shown to push the state-of-the-art (SOTA) performance for various understanding and generation tasks. These include standard natural language processing (NLP) tasks such as machine translation, document classification, entailment, summarization and question answering (Za-heer et al., 2020; Dai et al., 2019; Baevski and Auli, 2019), as well as music generation (Huang et al., 2018), image generation (Parmar et al., 2018; Chen et al., 2020) and genomics (Zaheer et al., 2020; Choromanski et al., 2020). The Transformer is also the backbone architecture for models such as BERT (Devlin et al., 2019) (and its numerous relatives) and GPT3 (Brown et al., 2020), which have delivered impressive performance across many NLP tasks. However, the standard attention mechanism of the Transformer has a run time and memory usage that scales quadratically with sequence length. Therefore, this quadratic complexity has become a critical bottleneck in processing long sequences (over 1,000 tokens), and has since motivated many new attention algorithms, see (Tay et al., 2020d) for a survey of such work. In this paper, we draw inspiration from two branches in numerical analysis: Hierarchical Matrix (H-Matrix) (Hackbusch, 1999, 2000) and Multigrid method (Briggs et al., 2000). We propose a hierarchical attention that has linear complexity in run time and memory, and only utilizes dense linear algebra operations optimized for GPUs or TPUs. We hypothesize that the inductive bias embodied by the proposed hierarchical structure for the attention matrix is effective in capturing the hierarchical structure in the sequences typically seen in many natural language processing and computer vision tasks. The main benchmark we use in this paper is the Long Range Arena (LRA) benchmark (Tay et al., 2020c), which has been specifically designed to evaluate and compare various sub-quadratic attention algorithms. Our new hierarchical attention mechanism achieves best average performance to-date on the LRA benchmark by more than 6 points over the previous-best BigBird algorithm (Zaheer et al., 2020), while pushing SOTA performance higher in 4 of the 5 successful tasks. Furthermore, using this new attention, a Transformer-based language model trained on the One-Billion Word dataset (Chelba et al., 2014) sets a new SOTA performance record by reducing the test perplexity by 1 . 55 points comparing to the previous-best Transformer-XL (Dai et al., 2019) with 5x more parameters. Overall, these empirical results both validate the soundness of our approximation method for computing attention weights, as well as the the appropriateness of the inductive bias present in the proposed hierarchical attention. It is well established in the NLP literature that the embeddings of nearby tokens tend to be more similar than the distant ones (Manning and Schutze, 1999). This leads to the intuition that token similarity and hence the attention should decrease with the sequence distance between a query token and a key token 1 . This motivates the sliding-window local attention (Parmar et al., 2018; Ramachandran et al., 2019; Qiu et al., 2019) which amounts to truncating off-diagonal entries in the attention matrix beyond a user-specified sequence distance. A second approach is to keep O (1) number of nonze-ros per row in the attention matrix. The nonzero entry selection is either content-based (Kitaev et al., 2020; Roy et al., 2020; Tay et al., 2020b; Zhou et al., 2020), hand-crafted (Beltagy et al., 2020; Brown et al., 2020; Child et al., 2019; Ho et al., 2019) or simply random (Zaheer et al., 2020). It is also well known in the NLP literature that long-range contextual information is necessary for many NLP tasks (Khandelwal et al., 2018; Liu and Lapata, 2019). So a set of global tokens are also considered. This adds O (1) number of dense rows and columns to the attention matrix (Zaheer et al., 2020; Ainslie et al., 2020; Belt-agy et al., 2020). A third approach is to approximate the attention matrix with a low-rank factored form (Choromanski et al., 2020; Wang et al., 2020; Tay et al., 2020a). The first two approaches are based on the premise that one needs to explicitly zero out entries in the attention matrix in order to reduce the quadratic complexity. Decades of research by the scientific computing and numerical analysis community has resulted in more sophisticated algorithms to sparsify matrices. A 1 Eq. small set of samples of these algorithms and their engineering applications include Fast Multipole Method (Greengard and Rokhlin, 1987; Greengard, 1994; Nabors et al., 1994; Shi et al., 1998), Pre-corrected FFT (Phillips and White, 1997; Zhu et al., 2005), Hierarchical Singular Value Decomposition (SVD) (Kapur and Long, 1997) and Hierarchical Matrix (H-Matrix) (Hackbusch, 1999, 2000; Zhu and White, 2005). These are generally called Multilevel Methods (Brandt and Lubrecht, 1990). The hierarchical attention proposed in this paper is inspired by these Multilevel Methods in general and the H-Matrix in particular. The hierarchical matrix structure allows a linear complexity in both constructing and applying the attention matrix. Given matrices Q , K and V , with rows representing sequences of token embedding or feature vectors for query, key and value respectively, the output weighted by the scaled dot-product attention in the Transformer (Vaswani et al., 2017) is defined as where Z, Q, K, V  RL  d , L is the length of the sequences, and d is the embedding or feature size. In a more compact matrix form, Eq. (1) can be written as Z = D  1 AV (2) where A = e S (3) S i,j = Q i K Tj  d (4) D = diag { A  1 L } (5) 1 L = [1 , 1 , ..., 1] T . Here, A, S  RL  L , 1 L  RL is a vector with all ones, and S i,j represents the unnormalized cosine similarity between query embedding Q i (the i -th row in Q ) and key embedding K j (the j -th row in K ). For the sake of clarity, we focus on the single-head attention in the exposition of the proposed algorithm. Extension to the multi-head case is straightforward since each attention head is computed independently (Vaswani et al., 2017). Computing the similarity matrix S in Eq. (4) and the attention matrix A in Eq. (3) takes O ( L 2 d ) time and O ( L 2 ) memory. Similarly, computing AV in Eq. (2) takes O ( L 2 d ) time, and computing A  1 L in Eq. (5) takes O ( L 2 ) time. The O ( L 2 d ) and O ( L 2 ) complexities are the bottlenecks for applying the attention mechanism over very long sequences. 4.1 H-Matrix The singular-value decomposition of the attention matrix A in Eq. (3) is A = U  VT (7) where  = diag {  1 ,  2 , ...,  L } and  i is the i -th singular value. The numerical rank of matrix A is r if (cid:80) L i = r +1  i < (cid:15) for a given tolerance (cid:15) (Tre-fethen and Bau, 1997). The standard rankr approximation to matrix A is A   U   VT =  U  VT (8) where  = diag {  1 ,  2 , ...,  r } ,  U,  V  RL  r have the first r columns of U and V , and  V =  V  . This is the low-rank approximation used in (Choromanski et al., 2020; Wang et al., 2020; Tay et al., 2020a). This approximation compresses L 2 entries in A to 2 rL entries in  U and  VT . So the compression rate is L 2 r . The H-Matrix generalizes this low-rank approximation by using matrix block hierarchy. Consider a two-level H-Matrix with 4  4 and 2  2 block partition at level-0 and level-1, respectively. Matrix A is partitioned as A =    A (0)11 A (0)12 A (0)21 A (0)22 A (1)12 A (1)21 A (0)33 A (0)34 A (0)43 A (0)44    . The low-rank approximation in Eq. (8) is applied to the off-diagonal blocks at each level. For example, A ( l ) 12   U ( l ) 12  V ( l ) 12 T (10) where l = 0 , 1 . To give a concrete example, suppose each entry in matrix A has the analytical form A i,j = e S i,j (11) S i,j = 2 e  ( i  j ) 2  1 (12) where i, j = 0 , 1 , 2 , ..., 15 2 . With the block hierarchy defined in Eq. (9), the size of the matrix block at level-1 and level-0 is 8  8 and 4  4 , respectively. For tolerance (cid:15) = 10  3 , one can verify that the numerical rank map of matrix A is   4 2 2 4 2 2 4 2 2 4   (13) where the number in each block is the numerical rank of the corresponding block in Eq. (9). Note that matrix A still has full numerical rank of 16 at a looser tolerance 10  1 . So the standard low-rank approximation is ineffective in this case. But even this simple two-level H-matrix already offers a compression rate of 43 since storing an H-matrix with the rank map in Eq. (13) takes 192 entries 3 . In addition, one can verify that no entry A i,j in Eq. (11) is very small, since S i,j  [  1 , 1] in Eq. (12). Therefore, truncating off-diagonal entries of matrix A , as proposed in (Parmar et al., 2018), would produce a poor approximation. In practice, the number of levels is adapted to the underlining governing equations that result in matrix A and it can easily be over 10 (Kapur and Long, 1997; Hackbusch, 2000; Zhu and White, 2005). In turn, this can substantially increase the compression rate. In general, the computation complexity of the H-Matrix is either O ( L ) or O ( L log L ) , depending on the underlining physics (Hackbusch, 1999, 2000). Multigrid Method is a multi-level nested iterative method for solving large-scale sparse matrices resulting from discretized partial-differential equations (PDEs) (Briggs et al., 2000; Trottenberg et al., 2000). At its core are two simple but powerfully complementary ideas: relaxation and correction. Our proposed hierarchical attention only uses the correction scheme as a building block since there is no sparse matrix to relax on. The correction scheme has two components: restriction or coarsening, and interpolation or pro-2 Matrix A in Eq.(11) is a symmetric Toeplitz matrix (Golub and Loan, 1996) and hence only has 16 unique entries. But we ignore this fact and treat A as a general matrix here. 3 Each one of four diagonal blocks at level-0 takes 16 entries. Each one of four off-diagonal blocks at level-0 takes 16 entries. Each one of two off-diagonal blocks at level-1 takes 32 entries. longation. Consider a vector  v h of scalar values defined on a set of N grids with uniform interval h . The simplest coarsening is to take the average of the scalar values on each pair of grids, i.e.,  v 2 hj = 1 2( v h 2 j +  v h 2 j +1 ) (14) where j = 0 , 1 , 2 , ...N/ 2  1 . The superscript in Eq. (14) indicates that the grid interval at these two levels is h and 2 h , respectively. The simplest interpolation is to duplicate the value on each coarse grid to values on a pair of fine grids, i.e.,  v h 2 j =  v 2 hj ,  v h 2 j +1 =  v 2 hj (15) where j = 0 , 1 , 2 , ...N/ 2  1 . The hierarchical low-rank structure like Eq. (13) turns out to be pervasive in many if not all physics phenomena. Much of the theoretical analysis by (Greengard and Rokhlin, 1987; Hackbusch, 1999) is concerned with quantifying such aspects. The key insight into these Multilevel Methods can be summarized as follows: perform no approximation for near interactions, and apply progressively lower-precision approximation for progressively longer distance interactions . The simple case shown in Eq. (9)-(13) is a good example. To satisfy the tolerance of 10  3 , we need full rank (no approximation) for the diagonal blocks (near inter-actions), higher precision approximation (rank-2 vs full-rank of 4) for the 4  4 off-diagonal blocks at level-0 (mid-distance) and lower precision approximation (rank-2 vs full-rank of 8) for the 8  8 off-diagonal blocks at level-1 (long-distance). In this section, we present some intuition to answer two important questions: 1) Does the hierarchical low-rank structure hold for the attention matrix A in Eq. (3)? 2) What is the algorithm to ef-ficiently compute the hierarchical low-rank structure? We only give an informal exposition of the hierarchical attention. The formal mathematical derivation is deferred to the Appendix. The error analysis in (Greengard and Rokhlin, 1987; Hackbusch, 1999) offers little direct insight since the attention matrix A in Eq. (3) is data dependent by definition and hence its analytical form like Eq. (11) and (12) is generally unknown. So gathering empirical evidences seems the only viable path to answer the first question listed above. The ablation studies by (Khandelwal et al., 2018) examine the effect of context words on a language model. Within the context range of about 200 tokens, word order is only relevant within the 20 most recent tokens or about a sentence. In the long-range context, order has almost no effect on performance, suggesting that the model maintains a high-level, rough semantic representation of faraway words. The observation is succinctly summarized by the title of the paper sharp nearby, fuzzy far away. Remarkably, this is in spirit very close to the key insight into the Multilevel Methods. A few recent attention-related studies have explored this direction with some success, such as word-level and sentence-level attentions in (Miculicich et al., 2018; Abreu et al., 2019), and sentence-level and paragraph-level attentions in (Liu and Lapata, 2019). Even though the proposed hierarchical attention in these studies only has two levels, as opposed to ten or more levels typically used by the Multilevel Methods, the reported positive results are quite suggestive. We therefore hypothesize that the same hierarchical low-rank structure as shown in Eq (13) might also hold for the attention matrix in many NLP tasks. And we treat it as the inductive bias in the hierarchical attention mechanism proposed in this paper. As pointed out in (Goyal and Ben-gio, 2020), inductive biases encourage the learning algorithm to prioritise solutions with certain properties. Hence good benchmark performance delivered by a Transformer-based model with proposed hierarchical attention can be regarded as a positive evidence to support the hierarchical low-rank structure hypothesis. In the standard definition of attention in Eq. (3) and (4), there is no preference given to any keys based on the sequence distance between a query and keys. The observation in (Khandelwal et al., 2018) clearly suggests that a distance-dependent attention mechanism should be a better alternative. We will take three steps to informally explain the hierarchical attention mechanism. First, the attention matrix blocks for nearby, mid-distance and long-distance attention are separated in section 5.2.1. This is the first step toward the distance-dependent attention mentioned above. Second, a token hierarchy is established in section 5.2.2. Third, the hierarchical attention is constructed in section 5.2.3 5.2.1 Attention Partition Consider a 16-word sentence in Fig. 1. The sentence is partitioned at three segment granularity. This induces a three-level partition of the attention matrix A for the original sequence: A = A (2) + A (1) + A (0) (16) where A (2) = (cid:34) 0 A (2)12 A (2)21 0 (cid:35) (17) A (1) =     A (1)12 A (1)21 A (1)23 A (1) 32 A (1) 34 A (1)43     (18) A (0) =   A (0)11 A (0)12 A (0)21 A (0)22 A (0)23 ... ... ... (19) Note that the nonzero entries in A (0) , A (1) and A (2) are the same as the corresponding entries of matrix A in Eq. (3). Matrix block size of A (0) ij , A (1) ij and A (2) ij is 2  2 , 4  4 and 8  8 , respectively. Following the key insight into Multilevel Methods, we perform no approximation to any level-0 matrix block A (0) ij and apply a low-rank approximation to off-diagonal matrix blocks in A (1) and A (2) . If we set the numerical rank of all these blocks to 2, then we can assemble the three rank maps into a single rank map as 4    2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2    . this sentence is to illustrate how to set up token hierarchy level by level with aggregation a) Level-0: 16 tokens partitioned into 8 segments b) Level-1: 16 tokens partitioned in 4 segments c) Level-2: 16 tokens partitioned in 2 segments this sentence is to to set up token illustrate how hierarchy level by level with aggregation this sentence is to illustrate how to set up token hierarchy level by level with aggregation Figure 1: Token sequence partitions in three segment granularity. The hierarchical structure embodied by the predetermined rank map in Eq. (20) represents the inductive bias for the attention matrix A in Eq. (16). But this construction step is inefficient because we need to form the original attention matrix and then perform SVD to discover the low-rank approximation. To illustrate the notion of token hierarchy, consider the same 16-word sentence in Fig. 2. A simple 3-level binary-tree hierarchy can be set up by following the simple coarsening defined in Eq. (14): 1) At level-0, each one of the 16 words is mapped to its word embedding; 2) At level-1, each token (parent node) corresponds to a pair of adjacent words at level-0 (child nodes), which are shown inside each box. The embedding of each parent token is simply the average of its child token embeddings; 3) At level-2, each token (parent node) corresponds to one pair of adjacent tokens at level-1 (child nodes) or 4 adjacent words at level-0 (grand child nodes), which are shown inside each box. The embedding of each parent token is simply the average of its child token embeddings. In general, the height of the binary tree is O ( log 2 ( L ) and the total number of tree nodes is O (2 L ) , where L is the sequence length. We only need word embeddings for the leaf nodes since the embeddings of all other tree nodes can be recursively computed. The formal definition and notations of the recursion for query and key are detailed in section 6.1. It is clear from Fig. 2 that the embeddings of higher level tokens represent a coarser level representation of a larger chunk of the text. The tokens at different levels can be understood as multi-scale snapshots of the original token sequence at level-0. Hence this token hierarchy naturally induces a set of multi-scale attention matrices. Let  A ( i ) be the attention matrix induced by the tokens at leveli . It is clear from Fig. 2 that the size of  A (0) ,  A (1) and  A (2) is 16  16 , 8  8 and 4  4 , respectively. This multi-scale viewpoint does not directly lead to a useful algorithm since matrix  A (0) contains all the information and there is little additional information from  A (1) and  A (2) . A key step to arrive at the hierarchical attention is to apply the contextual sliding window at each hierarchy level. The tokens at each level are partitioned into segments of size 2 in Fig. 2. One way to implement the local attention is to allow each query token segment to attend only two adjacent key token segments, one to its left and another to its right. At level-0, each query token segment also attends to the collocated key token segment. The token segment partition and local attention lead to a tri-diagonal block sparse matrix structure for  A (0) and bi-diagonal block sparse matrix structure for  A (1) and  A (2) . Their sparsity patterns are  A (0)     2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2 2    (21)  A (1)     2 2 2 2 2 2    (22)  A (2)  (cid:20) 2 2 (cid:21) (23) where the 2 in the nonzero blocks indicates that these are dense blocks of size 2  2 . It is clear that  A (0) is identical to A (0) in Eq. (19). The efficiency gain comes from  A (2) and  A (1) . Each nonzero entry in  A (2) and  A (1) captures the aggregated or coarse attention between two disjoint chunk of four and two tokens, respectively. Progressively larger token chunks lead to progressively lower-precision approximation to the original attention blocks. This is precisely the intention of the rank map in Eq. (20). We can now see that  A (2) and  A (1) provide an efficient way to approximate A (2) in Eq. (17) and A (1) in Eq. (18), respectively. The simple example in Fig. 2 can be easily generalized. Eq. (14) is used to coarsen or merge rows in matrices Q , K and V in Eq. (1). For sequence length L = 2 M +1 , the coarsening establishes a binary tree of depth M for Q , K and V , respectively. Each tree node represents a matrix row and there are 2 M +1  l nodes or rows at levell . To facilitate the discussion, we define a few hierarchy related notations here. Let  Q ( l ) ,  K ( l ) and  V ( l ) be coarsened versions of Q , K and V at levell in the binary tree. We note that l = 0 is a special case, which is defined as  Q (0) = Q,  K (0) = K,  V (0) = V. (24) Following Eq. (14), the recursion to coarsen Q , K and V is:  Q ( l +1) j = 1 2(  Q ( l ) 2 j +  Q ( l ) 2 j +1 ) (25)  K ( l +1) j = 1 2(  K ( l ) 2 j +  K ( l ) 2 j +1 ) (26)  V ( l +1) j = (  V ( l ) 2 j +  V ( l ) 2 j +1 ) (27) where l = 0 , 1 , ..., M  2 and j = 0 , 1 , 2 , ..., 2 M  l . It should be noted that the coarsening of V in Eq. (27) does not have the averaging factor 12 . We defer more details on coarsening to Appendix Section A.1. Now we are ready to compute the nonzero entries in Eq. (21), (22) and (23) and construct hierarchical attention matrix  A ( l ) . Substituting Eq. (25) and (26) into (4) and then into (3), we obtain  A ( l ) ij = e  S ( l ) ij = e  Q ( l ) i (  K ( l ) j ) T  d (28) Again, we note that l = 0 is a special case because  A (0) ij = A ij . The hierarchical matrix structure in Eq. (17), (18) and (19) naturally leads to a hierarchical approach to the matrix-matrix multiplication in Eq. (2) and the matrix-vector multiplication in Eq. (5). We use the matrix-matrix multiplication as an example since matrix-vector multiplication is just a special case of the matrix-matrix multiplication. In view of Eq. (17), (18) and (19), we write the matrix-matrix multiplication in Eq. (2) as Y = AV  A (0) V (0) +  A (1)  V (1) +  A (2)  V (2) = Y (0) + P (0) (cid:16)  Y (1) + P (1)  Y (2) (cid:17) (29) where  Y ( l ) =  A ( l )  V ( l ) , l = 1 , 2 (30) We defer the detailed derivation of Eq. (29) to Appendix Section A.5 and A.6. To facilitate the description and the complexity analysis of the algorithm, we define a few more hierarchy-related notations. In addition to sequence length L , number of hierarchy levels M and embedding or feature size d in Eq. (1), the new notations include: 1) N r : numerical rank of the off-diagonal blocks (for instance, 2 in Eq. (20)). This is also the diagonal block size at level-0; 2) N ( l ) b : number of blocks at levell . Note that L and d are usually data-dependent hyper-parameters, while N r is the only model hyper-parameter responsible for our method's inductive bias. In turn, N ( l ) b and M are derived parameters, computed as: N (0) b = LN r , N ( l +1) b = N ( l ) b 2 (31) M = log 2 ( N (0) b ) . It is important to note that only the diagonal blocks at level-0 and the super-diagonal and sub-diagonal blocks at levell are needed in applying the hierarchical attention matrix. This is clearly shown in Eq. (21)(23). This means that only N ( l ) b  1 super-diagonal and sub-diagonal blocks are computed at levell . This is crucial to the overall linear complexity in run time and memory. We should also note that all matrix blocks in coarse attention matrix  A ( l ) have the same size N r  N r . This is due to the rank map in Eq. (20). This is crucial for efficiency reason since the single-instruction-multiple-data (SIMD) programming style supported by the dense linear algebra libraries for GPU and TPU encourages uniform tensor shapes. We summarize the main steps to construct and apply the hierarchical attention in Algorithm 1. Algorithm 1 H-Transformer-1D Input: Q (query), K (key), V (value) Output: Z Coarsen Q using Eq. The computational cost for Algorithm 1 has two parts: (a) diagonal blocks at level0 : dN 2 r N (0) b (b) Superand sub-diagonal blocks at level-l : 4 dN 2 r ( N ( l ) b  1) (c) total: 5 dLN r = O ( dL ) 2. Computing matrix-matrix (MM) multiplication in Eq. (2) and matrix-vector (MV) multiplication in Eq. (5): (a) MM: 5 dLN r (b) MV: 5 LN r (c) total: 5( d + 1) LN r = O ( dL ) So the overall run time complexity of the hierarchical attention algorithm is O ( dL ) . Likewise, the memory complexity can be shown to be O ( dL ) as well. We defer the detailed analysis to appendix Section A.5 and A.6. We have implemented the proposed hierarchical attention using Jax, an open source library 5 for automatic gradient computation and linear algebra operations on GPUs and TPUs. All numerical operations in our algorithm use the Numpy native linear algebra functions supported by Jax. In all our experiments in this section, we use the standard Transformer architecture described in (Vaswani et al., 2017) as the backbone for our H-Transformer-1D model. Unless specified otherwise, the model parameters are: number of layers is 6, number of heads is 8, word embedding size is 512 and the feed-forward module (FFN) size is 2048. We follow the API for the standard multihead scaled dot-product attention implementation 6 so that we can perform a simple drop-in replacement of the standard multihead attention with our hierarchical attention implementation. This allows for an easy and fair comparison. The open-source Long-Range Arena (LRA) benchmark 7 has been proposed as a standard way to probe and quantify the capabilities of various xformer (long-range Transformer) architectures (Tay et al., 2020c). In our case, it also serves to highlight the effectiveness of the inductive bias 5 https://github.com/google/jax 6 https://github.com/google/flax/blob/master/flax/nn 7 https://github.com/google-research/long-range-arena inspired by the H-Matrix method, as well as the capability of our hierarchical attention to handle long sequences. The LRA has several desirable qualities that made us focus on it as a primary evaluation benchmark: generality (restricted to encoder-only tasks to accommodate most proposals); simplicity (no pretraining, no data augmentation allowed); diffi-culty (large headroom with existing approaches); long-input focus (so that modeling improvements in this area are visible); diverse (6 tasks, covering math, language, image, and spatial modeling); and lightweight (so that modeling improvements are measurable independently of the ability to train and run high-capacity models). The tasks that comprise LRA are: ListOps (sequences of arithmetical expressions of lengths of up to 2K that tests the ability to reason hierarchically while handling long context); Text (byte/character-level text classification at document level, which both simulates longer input sequences  max length 4K  and increases the diffi-culty level); Retrieval (byte/character-level document retrieval, which simulates the ability to model document similarity as a score between two independently-encoded long input sequences  max length 4K + 4K = 8K); Image (image classification based on the CIFAR-10 dataset, where an NxN image is flattened to a sequence of length N 2 pixels); Pathfinder (long-range spatial dependency task, with images consisting of two small Model perplexity parameters (Dai et al., 2019) 21.8 800M (Baevski and Auli, 2019) 23.02 1000M (Dai et al., 2019) 23.5 465M (Baevski and Auli, 2019) 23.91 465M (Shazeer et al., 2018) 24.0 4900M Transformer baseline 30.04 53M Transformer baseline 24.8 144M H-Transformer-1D N r = 16 23.95 53M H-Transformer-1D N r = 16 20.25 144M Table 2: Experimental results on one-billion word benchmark. circles and dash-line paths that either connect the two circles or not  image dimensions of 32x32 for a pixel sequence of length 1,024); Path-X (same as Pathfinder, but for image dimensions of 128x128 for a total pixel sequence of length 16,384). The default Transformer model parameters such as number of layers and number of heads etc are pre-determined by the benchmark configu-ration for each task. The results obtained by our H-Transformer-1D model on the LRA benchmark are given in Table 1. Overall, the H-Transformer-1D model achieves 61.41 average accuracy, a +6.4 points improvement over the previous-best average performance from BigBird (Zaheer et al., 2020). We want to highlight ListOps, Text and Retrieval because they all involve long sequences and H-Transformer-1D model improves SOTA performance by relatively large margins. These should be strong evidences to support our hypothesis in section 5.1 and validate the inductive bias due to the hierarchical attention. We have used Flax, an open-source library 8 to train neural networks, as the code base for the model training. Our H-Transformer-1D model uses the standard Transformer decoder implementation in Flax as the backbone. Only the attention is replaced with our hierarchical attention. We trained both the Transformer baseline and H-Transformer-1D on the One-Billion Word benchmark (Chelba et al., 2014). We tried different N r 8 https://github.com/google/flax (numerical rank) in our H-Transformer-1D model. These represent different inductive bias. We found that H-Transformer-1D with N r = 16 generated text with quality comparable to that of the baseline Transformer. For both Transformer baseline and H-Transformer-1D, we also tried two sets of model parameters: 1) embedding size is 512 and feed-forward module size is 2048 and hence the parameter count is 53M; 2) embedding size is 1024 and feed-forward module size is 4096 and hence the parameter count is 144M. The test perplexity results of these four models and various SOTA models are shown in table 2. H-Transformer-1D delivers the lowest perplexity to-date while using 5  smaller model capacity than that of the previous SOTA model Transformer-XL (Dai et al., 2019). This is another strong evidence to support our hypothesis in section 5.1 and validate the inductive bias due to the hierarchical attention. We have proposed a new Transformer attention using the inductive bias inspired by the H-Matrix. The new algorithm has linear complexity in run time and memory usage and is fully compatible with dense linear algebra libraries on GPU and TPU. The effectiveness of this new attention is demonstrated by the empirical evidences from long-range arena benchmark and One-Billion word language modeling. Future work include applying the new attention to music and genomics, developing proper inductive bias for cross-attention and extending the one-dimensional hierarchical attention to 2D images.
abstain abstain abstain abstain method abstain result abstain abstain abstain abstain other abstain abstain abstain abstain objective objective method abstain abstain abstain objective other other other other other other other other other other other other other other objective other abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain other abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain method abstain method other abstain abstain abstain method abstain abstain abstain abstain abstain abstain abstain method abstain abstain other abstain method abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain abstain other other other abstain abstain result result abstain abstain abstain abstain abstain other abstain result abstain abstain abstain objective abstain abstain abstain