PoNet: Pooling Network for Efficient Token Mixing in Long Sequences

1 INTRODUCTION . Transformer ( Vaswani et al. , 2017 ) has become the state-of-the-art ( SOTA ) architecture for sequence modeling in a wide variety of fields , including natural language processing ( NLP ) , computer vision , speech processing , applications to genomics data , etc . The key reason for Transformer ’ s success is its self-attention mechanism , which computes dot-product between input representations for each pair of positions in the full sequence . Proved to be greatly effective in learning contextualized representations , Transformer becomes the backbone for dominant pre-trained language models ( PLM ) in NLP , such as BERT ( Devlin et al. , 2019 ) and RoBERTa ( Liu et al. , 2019 ) . These PLMs demonstrate strong transfer learning capabilities and have achieved SOTA widely on NLP tasks . However , self-attention has quadratic time and memory complexity to the input sequence length ( Vaswani et al. , 2017 ) , which becomes the bottleneck for applying the vanilla Transformer to long sequence modeling tasks and for scaling up Transformer-based PLMs . A broad spectrum of approaches have been proposed to address the efficiency problem of selfattention , as summarized in ( Tay et al. , 2020 ) . One major direction approximates the dense full self-attention using techniques such as introducing sparsity ( Child et al. , 2019 ; Beltagy et al. , 2020 ; Zaheer et al. , 2020 ; Zhang et al. , 2021a ; Yang et al. , 2021 ; Zhu & Soricut , 2021 ; Zhang et al. , 2021b ) , low-rank approximations of the softmax attention matrix ( Katharopoulos et al. , 2020 ; Wang et al. , 2020a ; Zhu & Soricut , 2021 ; Xiong et al. , 2021 ; Peng et al. , 2021 ; Choromanski et al. , 2021 ) , and locality sensitive hashing ( Kitaev et al. , 2020 ) . These approximation approaches exploit observations that token interactions have strong locality , hence the importance and in turn the attention should decrease with the increase of the distance between a query token and a key token . Several of these works achieve O ( N ) theoretical complexity . However , these models often require selecting relatively large local regions for fine-grained attention in order to approximate full 1We will release the source code upon the publication of the paper . self-attention , hence the scaling constants hidden by O ( N ) are often large and hinder significant improvements in speed and memory usage . It is observed that performance of the approximation approaches is usually inversely related to their speed ( Beltagy et al. , 2020 ; Zaheer et al. , 2020 ) . Another major direction replaces the self-attention structure with more efficient structures , such as MLP-Mixer ( Tolstikhin et al. , 2021 ) , FNet ( Lee-Thorp et al. , 2021 ) , AFT ( Zhai et al. , 2021 ) , and Fastformer ( Wu et al. , 2021 ) . Despite significant accelerations gained by efficient transformers , they are rarely evaluated both on effectiveness of the inductive bias and the transfer learning capability , except BigBird , FNet , and Nyströmformer2 as we understand . BigBird-Base ( Zaheer et al. , 2020 ) outperforms BERT-Base on the GLUE benchmark ( Wang et al. , 2019 ) without significant speedup ( slowdown for input length < 3K ) ( Tay et al. , 2021 ) . FNet-Base trains 80 % faster than BERT on GPUs at 512 input lengths but only achieves 92 % of BERT-Base ’ s accuracy ( Lee-Thorp et al. , 2021 ) . In this work , we propose a novel Pooling Network ( PoNet ) , aiming to simultaneously advance long sequence modeling capacity and transfer learning capabilities while improving speed and memory efficiency . PoNet replaces the O ( N2 ) complexity self-attention with a O ( N ) complexity multi-granularity pooling block . We design multi-granularity pooling and pooling fusion to model different levels of token interactions . Multi-granularity pooling incorporates three types of pooling , from coarse-grained to fine-grained , in each sublayer . Global aggregation aggregates information of the entire sequence into a single token . Segment max-pooling captures the paragraph or sentence level information . Local max-pooling captures the more important local information . These three poolings are fused to produce the output feature of the multi-granularity pooling block . Then through the residual connection , this output feature is further aggregated into each token . The contributions of this paper are summarized as follows : • We propose a novel PoNet architecture as a drop-in replacement for self-attention in Transformer , achieving linear time and memory complexity . We propose multi-granularity pooling and pooling fusion to capture different levels of contextual information and comprehensively model token interactions . To the best of our knowledge , our work is the first to explore the full potential of the simple pooling mechanism for token mixing and modeling long-range dependencies . • Extensive evaluations show that PoNet achieves competitive performance on the Long Range Arena benchmark ( Tay et al. , 2021 ) and significantly outperforms Transformer by +2.28 absolute ( +3.9 % relative ) on accuracy , with efficiency up to 9 times faster and 10 times smaller than Transformer on GPU . Also , PoNet demonstrates competitive transfer learning capabilities , with PoNet-Base reaching 95.7 % of the accuracy of BERT-Base on the GLUE benchmark . Ablation analysis further proves effectiveness of designed multi-granularity pooling and pre-training tasks . 2 RELATED WORK . Efficient Transformer Variants Among the models to approximate full self-attention , Longformer ( Beltagy et al. , 2020 ) ( O ( N ) ) sparsifies the full self-attention into three attention patterns of sliding window , dilated sliding window , and global attention . BigBird ( Zaheer et al. , 2020 ) ( O ( N ) ) combines global attention , local attention , and random attention . Poolingformer ( Zhang et al. , 2021a ) ( O ( N ) ) uses a two-level attention schema , with the first level using a smaller sliding window to aggregate local information and the second level using a larger window with pooling attention to reduce time and memory cost . Focal Transformer ( Yang et al. , 2021 ) ( O ( N ) ) uses both fine-grained local interactions and coarse-grained global interactions to balance efficiency and effectiveness of capturing short- and long-range dependencies . Transformer-LS ( Zhu et al. , 2021 ) ( O ( N ) ) approximates the full attention by aggregating long-range attention via dynamic projections and short-term attention via segment-wise sliding window . H-Transformer-1D ( Zhu & Soricut , 2021 ) ( O ( N ) ) exploits a matrix structure similar to Hierarchical Matrix . AdaMRA ( Zhang et al. , 2021b ) ( O ( N ) ) leverages a multi-resolution multi-head attention mechanism and kernel attention . Luna ( Ma et al. , 2021 ) ( O ( N ) ) introduces an additional fixed length sequence served as query to attend to the original input while the output is served as key and value to attend to the original input . Apart from the sparse attention models , other approximation approaches explore locality sensitive hashing and matrix approximation methods . Reformer ( Kitaev et al. , 2020 ) O ( NlogN ) replaces self-attention with locality sensitive hashing . Performer ( Choromanski et al. , 2020 ; 2021 ) ( O ( N ) ) approximates softmax attention by leveraging random features . Linformer ( Wang et al. , 2020a ) 2Nyströmformer was evaluated only on subsets of GLUE ( Xiong et al. , 2021 ) . approximates the self-attention matrix with a low-rank factorization . Nyströmformer ( Xiong et al. , 2021 ) ( O ( N ) ) approximates the softmax attention with the Nyström method by sampling a subset of columns and rows . However , these approximation methods have strengths on certain tasks and may cause accuracy degradation on many other tasks . Our work is in another line of research on replacing self-attention with more efficient token mixing mechanisms . MLP-Mixer ( Tolstikhin et al. , 2021 ) ( O ( N2 ) ) applies two separate linear transformations on the hidden state dimension and the sequence dimension . FNet ( Lee-Thorp et al. , 2021 ) ( O ( NlogN ) ) replaces the self-attention sublayer with 2D-FFT mixing sublayer . AFTlocal/conv ( Zhai et al. , 2021 ) ( O ( sN ) , s < N ) first combines the key and value with a set of learned position biases and then combines the query with this result via element-wise multiplication . Fastformer ( Wu et al. , 2021 ) ( O ( N ) ) first models global context via additive attention then models interactions between global context and input representations through element-wise product . Pre-training Tasks It has been observed that both underlying model architecture and pre-training are crucial to performance of PLMs . BERT ( Devlin et al. , 2019 ) with a Transformer encoder is pretrained with masked language modeling ( MLM ) and next sentence prediction ( NSP ) tasks on largescale unlabeled text corpora including the English Wikipedia and BooksCorpus . MLM predicts the masked token from context . NSP predicts whether a sentence pair is contiguous or not in the original source . Many approaches are proposed to improve these two tasks and show that more challenging pre-training tasks may help PLMs learn better and more transferable language representations . Whole word masking ( WWM ) ( Devlin et al. , 2019 ; Cui et al. , 2019 ) and SpanBERT ( Joshi et al. , 2020 ) outperform BERT on many tasks . WWM simultaneously masks all WordPiece tokens belonging to the same word and forces the model to recover a complete whole word . SpanBERT randomly samples contiguous spans inside of individual tokens and augments MLM with a new task to predict the entire masked span . RoBERTa ( Liu et al. , 2019 ) reports ineffectiveness of NSP and removes it from pre-training . ALBERT ( Lan et al. , 2020 ) replaces NSP with a sentence-order prediction ( SOP ) task to predict whether two consecutive sentences are in the right order or not , for learning fine-grained inter-sentence coherence . StructBERT ( Wang et al. , 2020b ) extends SOP to a new sentence structural objective ( SSO ) as a ternary classification on two sentences ( S1 , S2 ) to decide whether S1 precedes or follows S2 or the two sentences are noncontiguous . More challenging tasks for learning inter-sentence relations and document/discourse structures ( Iter et al. , 2020 ; Lee et al. , 2020 ; Ding et al. , 2021 ) show promising performance improvements on PLMs . 3 MODEL . Our work is inspired by the External Attention ( EA ) approach proposed in ( Guo et al. , 2021 ) for visual tasks . An input sequence of tokens x = { x1 , ... , xN } are mapped to an embedding matrix denoted by H ( ∈ RN×d ) = { h1 , ... , hN } , where N is the sequence length and d is the hidden dimension . EA uses two linear layers to implement external and shared memories , which facilitates learning correlations across all samples and hence serves strong regularization to and improves generalization of the attention mechanism with linear complexity . We simplify EA into multi-layer perceptron and softmax , and observe that by infusing the sequence-level information into each token through the denominator term ∑N n=1 e hn , softmax provides context modeling capabilities . However , softmax involves calculations of exponents , which is still slow . Consequently , we consider using pooling as an alternative to capture contextual information with significantly reduced complexity . We propose a Pooling Network ( PoNet ) as a drop-in replacement for the self-attention sublayer in Transformer , as shown in Figure 1 . PoNet models different levels of contextual information through a multi-granularity pooling ( MP ) block consisting of three components , namely , global aggregation ( GA ) , segment max-pooling ( SMP ) , and local maxpooling ( LMP ) . These pooling features are then aggregated through pooling fusion .

Summary: PoNet addresses the quadratic time and memory complexity of Transformer with a new attention mechanism that has only linear complexity. Theoretical (Section 3.2) and experimental results (Table 2) bear this out. Experimental results on standard datasets like Long Range Arena and GLUE are very competitive with SOTA. Ablations were done on the 3 components of the attention mechanism, too.