PoNet: Pooling Network for Efficient Token Mixing in Long Sequences

Abstract

Transformer-based models have achieved great success in various NLP, vision, and speech tasks. However, the core of Transformer, the self-attention mechanism, has a quadratic time and memory complexity with respect to the sequence length, which hinders applications of Transformer-based models to long sequences. Many approaches have been proposed to mitigate this problem, such as sparse attention mechanisms, low-rank matrix approximations and scalable kernels, and token mixing alternatives to self-attention. We propose a novel Pooling Network (PoNet) for token mixing in long sequences with linear complexity. We design multi-granularity pooling and pooling fusion to capture different levels of contextual information and combine their interactions with tokens. On the Long Range Arena benchmark, PoNet significantly outperforms Transformer and achieves competitive accuracy, while being only slightly slower than the fastest model, FNet, across all sequence lengths measured on GPUs. We also conduct systematic studies on the transfer learning capability of PoNet and observe that PoNet achieves 95.7% of the accuracy of BERT on the GLUE benchmark, outperforming FNet by 4.5% relative. Comprehensive ablation analysis demonstrates effectiveness of the designed multi-granularity pooling and pooling fusion for token mixing in long sequences and efficacy of the designed pre-training tasks for PoNet to learn transferable contextualized language representations.

Accepted by ICLR 2022. Codes and checkpoints are also available on huggingface hub: https://huggingface.co/chtan/ponet-base-uncased

References in corpus (11)

• MLP-Mixer: An all-MLP Architecture for Vision
• Linformer: Self-Attention with Linear Complexity
• Generating Long Sequences with Sparse Transformers
• Reformer: The Efficient Transformer
• Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention
• Focal Self-attention for Local-Global Interactions in Vision Transformers
• Random Feature Attention
• Beyond Self-attention: External Attention using Two Linear Layers for Visual Tasks
• Fastformer: Additive Attention Can Be All You Need
• Nyströmformer: A Nyström-Based Algorithm for Approximating Self-Attention
• Masked Language Modeling for Proteins via Linearly Scalable Long-Context Transformers