H -Transformer-1 D : Fast One-Dimensional Hierarchical Attention for Sequences (2021)

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.

## introduction linearly combining information using contentbased weights, a method generically known as attention, is a key building block in many deep neural networks such as recurrent neural networks (rnn) @xcite , convolutional neural networks (cnn) @xcite and graph convolutional networks (gcn) @xcite . 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 @xcite , which has been shown to push the state-of-theart (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 @xcite @xcite , as well as music generation @xcite , image generation @xcite and genomics @xcite . the transformer is also the backbone architecture for models such as bert @xcite (and its numerous relatives) and gpt3 @xcite , 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 @xcite for a survey of such work. in this paper, we draw inspiration from two branches in numerical analysis: hierarchical matrix (h-matrix) @xcite and multigrid method @xcite . 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 @xcite , 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 big-bird algorithm @xcite , while pushing sota performance higher in 4 of the 5 successful tasks. furthermore, using this new atten-tion, a transformer-based language model trained on the one-billion word dataset @xcite sets a new sota performance record by reducing the test perplexity by 1.55 points comparing to the previous-best transformer-xl @xcite 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. ## related works it is well established in the nlp literature that the embeddings of nearby tokens tend to be more similar than the distant ones @xcite . 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 foot_0 . this motivates the sliding-window local attention @xcite @xcite 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 nonzeros per row in the attention matrix. the nonzero entry selection is either content-based @xcite @xcite , hand-crafted @xcite @xcite or simply random @xcite . it is also well known in the nlp literature that long-range contextual information is necessary for many nlp tasks @xcite . so a set of global tokens are also considered. this adds o(1) number of dense rows and columns to the attention matrix @xcite @xcite . a third approach is to approximate the attention matrix with a low-rank factored form @xcite @xcite . 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 small set of samples of these algorithms and their engineering applications include fast multipole method @xcite @xcite , pre-corrected fft @xcite , hierarchical singular value decomposition (svd) @xcite and hierarchical matrix (h-matrix) @xcite @xcite . these are generally called multilevel methods @xcite . 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. ## definition and notation 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 @xcite is defined as here, a, @xmath0, 1 @xmath1, 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 singlehead attention in the exposition of the proposed algorithm. extension to the multi-head case is straightforward since each attention head is computed independently @xcite . ## intuition for hierarchical attention 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 @xcite 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 interactions), 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 efficiently 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. ## algorithm and computational complexity 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 level-l. 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: it is important to note that only the diagonal blocks at level-0 and the super-diagonal and subdiagonal blocks at level-l 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 level-l. 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 ã(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. ( 25 ) and coarsen k using eq. ( 26 ) compute diagonal blocks in ã(0) and superdiagonal and sub-diagonal blocks in ã(l) using eq. ( 28 ) coarsen v using eq. ( 27 ) compute @xmath2. ( 2 ) using eq. ( 29 ) compute d in eq. ( 5 ) using eq. ( 29 ) compute @xmath3: ## experiments and results we have implemented the proposed hierarchical attention using jax, an open source library foot_3 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 @xcite 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 foot_4 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. ## conclusions and future work 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 onedimensional hierarchical attention to 2d images.