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

Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing, pages 3801–3815 August 1–6, 2021. ©2021 Association for Computational Linguistics 3801 H-Transformer-1D: Fast One-Dimensional Hierarchical Attention for Sequences Zhenhai Zhu Google Research zhenhai@google.com Radu Soricut Google Research rsoricut@google.com Abstract We describe an efficient hierarchical method to compute attention in the Transformer architec- ture. The proposed attention mechanism ex- ploits a matrix structure similar to the Hier- archical Matrix (H-Matrix) developed by the numerical analysis community, and has linear run time and memory complexity. We per- form extensive experiments to show that the inductive bias embodied by our hierarchical at- tention is effective in capturing the hierarchi- cal structure in the sequences typical for nat- ural language and vision tasks. Our method is superior to alternative sub-quadratic propos- als 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. 1 Introduction Linearly combining information using content- based weights, a method generically known as at- tention, is a key building block in many deep neu- ral networks such as recurrent neural networks (RNN) (Luong et al., 2015), convolutional neu- ral 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 nat- ural language processing (NLP) tasks such as ma- chine translation, document classification, entail- ment, 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 perfor- mance 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 Ma- trix (H-Matrix) (Hackbusch, 1999, 2000) and Multigrid method (Briggs et al., 2000). We pro- pose a hierarchical attention that has linear com- plexity in run time and memory, and only uti- lizes dense linear algebra operations optimized for GPUs or TPUs. We hypothesize that the inductive bias embod- ied by the proposed hierarchical structure for …

Content truncated to 500 words