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: In this paper, the authors aim to resolve the quadratic time and memory complexity of the standard attention mechanism. They introduce a multi-granularity pooling and pooling fusion that captures contextual information from various levels (global, segment, and token).  They conduct experiments on long sequence tasks and large-scale pre-training and fine-tuning. As a result, their model achieves great performance with speed-up.