近年来，在机器学习领域Transformer模型已成为最先进的(SOTA) 序列建模模型，包括在自然语言处理 (NLP)、计算机视觉、语音处理、基因组数据等都有着广泛的应用。
Transformer 成功的关键原因在于它的自我注意（self-attention）机制，计算输入表征的每个位置之间的点积。Transformer被证明在学习上下文表征方面非常有效，它成为最主要的主干模型，例如 BERT和 RoBERTa。这些预训练语言模型展示了强大的迁移学习能力，并在广泛NLP任务中实现了 SOTA。
然而，由于Transformer模型中的self-attention机制相对于句子长度的复杂度是二次的(O(N^2))，因此在计算速度和显存空间方面都制约了它在长序列中的应用。我们提出了一种具有线性复杂度 (O(N)) 的计算模型 PoNet ，使用 pooling 网络替代 self-attention 机制对句子词汇进行混合，从而捕捉上下文信息。
实验表明，PoNet 在长文本测试 Long Range Arena (LRA) 榜[1] 上在准确率上比 Transformer 高 2.28 个点，在GPU上运行速度是Transformer的 9 倍，显存占用只有 1/10。此外，实验也展示了 PoNet 的迁移学习能力，PoNet-Base 在 GLUE 基准上达到了 BERT-Base 的 95.7% 的准确性。
|| 模型
受到用于视觉任务的外部注意EA[2]的启发，我们将其简化为 多层感知器 和 softmax，并观察到 softmax 通过分母项将序列信息融入到 token 中提供了上下文建模能力。然而，softmax 涉及到指数的计算，这仍然是很慢的。因此，我们考虑使用池化法作为替代方法，以显著降低的复杂度来捕捉语境信息。
模型主要由三个不同粒度的 pooling 组成，一个全局的pooling模块(GA)，分段的segment max-pooling模块(SMP)，和局部的max-pooling模块(LMP)，对应捕捉不同粒度的序列信息：
① 在第一阶段，GA沿着序列长度进行平均得到句子的全局表征g。为了加强对全局信息的捕捉，GA在第二阶段对g和输入训练计算cross-attention。由于g的长度为1，因此总的计算复杂度仍为O(N)。
② SMP按每个分段求取最大值，以捕获中等颗粒度的信息。
③ LMP沿着序列长度的方向计算滑动窗口max-pooling。
④ 然后通过池化融合(PF)将这些池化特征聚合起来。由于GA的特征在整个token序列是共享的，SMP的特征在segment内部也是共享的，直接将这些特征加到原始token上会使得token趋同（向量加法），而这种token表征同质化的影响将会降低诸如句子对分类任务的性能。因此，我们在PF层将原始的token于对应的GA，SMP特征计算元素乘法得到新的特征，使得不同的token对应了不同的特征。
|| 实验结果
长序列任务
Long Range Arena(LRA) 是用来评估捕捉长距离依赖关系的基准测试。在LRA上，PoNet取得了比Transformer更好的分数。
在速度和显存方面，仅次于FNet[3]，显著优于Transformer。
迁移学习
我们用大规模语料库对PoNet进行预训练，然后测试它在下游任务上的性能。下图是预训练的 MLM[4] 和 SSO[5] 两个子任务的训练曲线，可以看到，我们的模型在 MLM 上略微弱于 BERT ，在 SSO 上与 BERT 还有一定的差距，两个任务上都明显要优于 FNet 。
GLUE
PoNet取得了76.80的AVG分数，达到了 BERT 在 GLUE 上的准确率（80.21）的95.7%，相对来说比 FNet 要好4.5%。这些性能比较与图2中显示的预训练准确率一致。
长文本任务
我们还评估了预训练的 PoNet 在四个长文本分类数据集上的性能。从表4可以看出，PoNet-Base 在 HND 和 Arxiv 上优于 BERT-Base，在 IMDb 和 Yelp-5 上的F1分数达到了 BERT-Base 的99%。
消融分析
下面的消融实验也证明了每个组件的重要性。同时与 L_MN(MLM+NSP)，L_OM(MLM) 也说明了预训练任务使用 MLM+SSO 的必要性。
|| 总结
我们提出了一个使用多粒度的 Pooling 结构来替代 attention 的网络（PoNet），它可以捕捉到不同层次的上下文信息，让序列的 token 之间能够得到有效的交互。实验表明，PoNet 既实现了有竞争力的长距离依赖性建模能力，又实现了强大的迁移学习能力，而且具有线性的时间和显存复杂度。
|| Future Work
未来的工作包括进一步优化模型结构和预训练，以及将 PoNet 应用于包括生成任务在内的更广泛的任务。我们希望PoNet模型可以对探索更高效的序列建模模型提供一些启发。
参考资料：
[1] Yi Tay, Mostafa Dehghani, Samira Abnar, Yikang Shen, Dara Bahri, Philip Pham, Jinfeng Rao, Liu Yang, Sebastian Ruder, and Donald Metzler. Long Range Arena : A benchmark for efficient transformers. ICLR 2021. URL
https://openreview.net/forum?id=qVyeW-grC2k.
[2] Meng-Hao Guo, Zheng-Ning Liu, Tai-Jiang Mu, and Shi-Min Hu. Beyond self-attention: External attention using two linear layers for visual tasks. CoRR, abs/2105.02358, 2021. URL
https://arxiv.org/abs/2105.02358.
[3] James Lee-Thorp, Joshua Ainslie, Ilya Eckstein, and Santiago Ontanon. FNet: mixing tokens with fourier transforms. CoRR, abs/2105.03824, 2021. URL
https://arxiv.org/abs/2105.03824.
[4] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. NAACL-HLT 2019. URL
https://doi.org/10.18653/v1/n19-1423.
[5] Wei Wang, Bin Bi, Ming Yan, Chen Wu, Jiangnan Xia, Zuyi Bao, Liwei Peng, and Luo Si. StructBERT: incorporating language structures into pre-training for deep language understanding. ICLR 2020. URL
https://openreview.net/forum?id=BJgQ4lSFPH.
论文地址：
https://openreview.net/forum?id=9jInD9JjicF
https://arxiv.org/abs/2110.02442
代码已在GitHub上开源：