[Paper Review] PoNet: Pooling Network for Efficient Token Mixing in Long Sequences

Chao-Hong Tan, Qian Chen+4 more|arXiv (Cornell University)|Oct 6, 2021

Topic Modeling37 references4 citations

TL;DR

PoNet proposes a linear-complexity pooling network that replaces self-attention in Transformers to enable efficient token mixing in long sequences. By combining multi-granularity pooling (global, segment, and local max-pooling) with feature fusion, PoNet achieves SOTA performance on Long Range Arena and 95.7% of BERT's GLUE accuracy—outperforming FNet by 4.5% relative—while being up to 9× faster and 10× more memory-efficient than standard Transformers.

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.

Motivation & Objective

• To address the quadratic complexity bottleneck of self-attention in Transformers for long sequence modeling.
• To design an efficient, scalable alternative to self-attention that maintains strong contextual modeling capacity.
• To improve transfer learning performance while achieving linear time and memory complexity.
• To explore the full potential of pooling mechanisms for capturing long-range dependencies in sequence modeling.

Proposed method

• Introduces a multi-granularity pooling block that combines global aggregation, segment max-pooling, and local max-pooling to capture diverse contextual levels.
• Fuses pooled representations through learnable projection matrices to produce a unified output feature per token.
• Applies residual connections to inject pooled features back into each input token for enhanced representation.
• Employs shared linear projections for query, key, and value heads across pooling types to maintain parameter efficiency.
• Uses a contextual decision mechanism in the pooling fusion layer to dynamically weight contributions from different pooling types.
• Supports causal attention via recursive mean and max-pooling updates, preserving O(N) complexity for streaming or autoregressive tasks.

Experimental results

Research questions

• RQ1 Can pooling mechanisms effectively replace self-attention in modeling long-range dependencies with linear complexity?
• RQ2 How does multi-granularity pooling compare to self-attention in capturing hierarchical contextual information across sequence lengths?
• RQ3 To what extent can a pooling-based architecture match the transfer learning performance of BERT on standard NLP benchmarks?
• RQ4 Does pooling fusion with learnable weighting improve representation quality compared to simple averaging or concatenation?
• RQ5 Can PoNet maintain efficiency and performance in causal attention settings without sacrificing O(N) complexity?

Key findings

• PoNet achieves 95.7% of BERT’s GLUE benchmark accuracy, outperforming FNet by 4.5% relative, demonstrating strong transfer learning capability.
• On the Long Range Arena benchmark, PoNet improves accuracy by +2.28 absolute (3.9% relative) over standard Transformers, with up to 9× faster inference and 10× less memory usage.
• Ablation studies confirm that both multi-granularity pooling and pooling fusion are essential for performance, with ablation reducing accuracy by up to 5.2%.
• The model maintains competitive performance on both single-sentence and long-sequence tasks, including CoLA and MLM, indicating robust sentence encoding.
• Removing the self-attention-like global aggregation (SS-GA) degrades performance minimally, suggesting that a simpler, recursive mean-based global pooling can preserve efficiency and performance.
• PoNet supports causal attention with O(N) complexity by leveraging recursive updates for mean and max operations, enabling streaming applications.