[论文解读] PoNet: Pooling Network for Efficient Token Mixing in Long Sequences

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

Topic Modeling参考文献 37被引用 4

一句话总结

PoNet 提出了一种线性复杂度的池化网络，用以替代 Transformer 中的自注意力机制，从而在长序列中实现高效的 token 混合。通过结合多粒度池化（全局、分段和局部最大池化）与特征融合，PoNet 在 Long Range Arena 上达到 SOTA 性能，并在 GLUE 上达到 BERT 的 95.7%，相比 FNet 提升了 4.5% 的相对准确率，同时推理速度最快提升 9 倍，内存效率提升 10 倍，优于标准 Transformer。

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.

研究动机与目标

• 解决 Transformer 在长序列建模中自注意力机制带来的二次方复杂度瓶颈。
• 设计一种高效且可扩展的自注意力替代方案，同时保持强大的上下文建模能力。
• 在实现线性时间与内存复杂度的同时，提升迁移学习性能。
• 探索池化机制在序列建模中捕捉长距离依赖关系的全部潜力。

提出的方法

• 引入一种多粒度池化模块，结合全局聚合、分段最大池化和局部最大池化，以捕捉多样化的上下文层次。
• 通过可学习的投影矩阵融合池化表示，为每个 token 生成统一的输出特征。
• 应用残差连接，将池化特征重新注入每个输入 token，以增强表示能力。
• 在不同池化类型之间共享用于查询、键和值头的线性投影，以保持参数效率。
• 在池化融合层中引入上下文决策机制，动态加权不同池化类型对输出的贡献。
• 通过递归均值和最大池化更新支持因果注意力，保持 O(N) 复杂度，适用于流式或自回归任务。

实验结果

研究问题

• RQ1 池化机制是否能以线性复杂度有效替代自注意力，用于建模长距离依赖？
• RQ2 多粒度池化在捕捉不同序列长度下的分层上下文信息方面，与自注意力相比表现如何？
• RQ3 基于池化的架构在标准 NLP 基准测试中，能否在迁移学习性能上与 BERT 相当？
• RQ4 与简单平均或拼接相比，采用可学习加权的池化融合是否能提升表示质量？
• RQ5 PoNet 是否能在因果注意力设置中保持效率与性能，同时不牺牲 O(N) 复杂度？

主要发现

• PoNet 在 GLUE 基准测试中达到 BERT 准确率的 95.7%，相比 FNet 提升 4.5% 的相对准确率，展现出强大的迁移学习能力。
• 在 Long Range Arena 基准测试中，PoNet 相比标准 Transformer 提升了 +2.28% 的绝对准确率（相对提升 3.9%），推理速度最快提升 9 倍，内存使用减少 10 倍。
• 消融实验表明，多粒度池化与池化融合均对性能至关重要，其缺失可导致准确率下降最高达 5.2%。
• 模型在单句任务与长序列任务（包括 CoLA 和 MLM）中均保持竞争力，表明其具备稳健的句子编码能力。
• 移除类似自注意力的全局聚合（SS-GA）后性能下降极小，表明仅使用递归均值的全局池化即可在保持效率与性能的同时实现良好效果。
• PoNet 通过递归更新均值与最大值操作，实现因果注意力的 O(N) 复杂度，支持流式应用。