[論文レビュー] 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における自己注意機構を置き換える線形計算量のプーリングネットワークを提案する。マルチスケールプーリング（グローバル、セグメント、ローカルの最大プーリング）と特徴統合を組み合わせることで、Long Range ArenaでSOTA性能を達成し、BERTのGLUE精度の95.7％を達成（FNetより相対的に4.5％優れる）する。標準Transformerに比べて最大9倍速く、10倍メモリ効率が良い。

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の自己注意機構の2次計算量ボトルネックを解消すること。
• 強力な文脈モデリング能力を維持しつつ、効率的かつスケーラブルな自己注意機構の代替手段を設計すること。
• 線形時間およびメモリ計算量を達成しながら、転移学習性能を向上させること。
• プーリング機構の潜在的限界を解き放ち、系列モデリングにおける長距離依存関係を捉える能力を最大限に引き出すこと。

提案手法

• グローバル集約、セグメント最大プーリング、ローカル最大プーリングを組み合わせたマルチスケールプーリングブロックを導入し、多様な文脈レベルを捉える。
• 学習可能な投影行列を用いてプールド表現を統合し、各トークンごとに統一された出力特徴を生成する。
• 残差接続を適用して、プールド特徴を各入力トークンに再統合し、表現を強化する。
• プールングタイプ間でクエリ、キー、バリューのヘッドに共通の線形投影を採用し、パラメータ効率を維持する。
• プールフィュージョン層に文脈に応じた意思決定メカニズムを導入し、異なるプールタイプからの寄与を動的に重み付ける。
• 再帰的平均および最大プールング更新を活用することで、ストリーミングまたは自己回帰的タスクでもO(N)の計算量を維持する因果的注意機構をサポートする。

実験結果

リサーチクエスチョン

• RQ1 プールリング機構は、線形計算量のもとで長距離依存関係を効果的にモデル化できるか？
• RQ2 マルチスケールプールリングは、さまざまな系列長における階層的文脈情報の捉え方において、自己注意機構と比べてどのように差をつけるか？
• RQ3 プールリングベースのアーキテクチャは、標準的なNLPベンチマークでBERTの転移学習性能にどの程度近づけるか？
• RQ4 学習可能な重み付けを用いたプールフィュージョンは、単純な平均化や連結処理と比べて表現品質を向上させるか？
• RQ5 PoNetは、O(N)計算量を損なわず、因果的注意機構の設定でも効率性と性能を維持できるか？

主な発見

• PoNetはBERTのGLUEベンチマーク精度の95.7％を達成し、FNetより相対的に4.5％優れている。これは強力な転移学習能力を示している。
• Long Range Arenaベンチマークでは、標準Transformerに比べて精度が+2.28（相対的に3.9％）向上し、推論速度が最大9倍速く、メモリ使用量が10倍少ない。
• アブレーションスタディの結果、マルチスケールプールリングとプールフィュージョンの両方が性能向上に不可欠であることが確認され、それぞれの除去により精度が最大5.2％低下した。
• 単一文および長時間系列タスクの両方で競争力ある性能を維持しており、CoLAやMLMを含むタスクでも、強固な文書エンコーディング能力を示している。
• 自己注意機構に類似したグローバル集約（SS-GA）を削除しても性能低下は最小限に抑えられる。これは、単純な再帰的平均ベースのグローバルプールリングが、効率性と性能を両立できる可能性を示している。
• 再帰的更新を活用することで、平均および最大操作にO(N)の計算量を維持しつつ、因果的注意機構を実現しており、ストリーミング応用に適している。