[論文レビュー] Blockwise Self-Attention for Long Document Understanding

Jiezhong Qiu, Hao Ma+4 more|arXiv (Cornell University)|Nov 7, 2019

Topic Modeling参考文献 44被引用数 16

ひとこと要約

BlockBERTは、自己注意機構におけるスパースブロック構造を導入することで、全系列ではなく局所的ブロックに限定された注意にすることで、メモリと計算コストを削減する。RoBERTaと比較して、メモリ使用量が18.7–36.1%低減され、学習が12.0–25.1%高速化され、推論が27.8%高速化され、長文処理タスクにおける下流性能は同等または向上を達成している。

ABSTRACT

We present BlockBERT, a lightweight and efficient BERT model for better modeling long-distance dependencies. Our model extends BERT by introducing sparse block structures into the attention matrix to reduce both memory consumption and training/inference time, which also enables attention heads to capture either short- or long-range contextual information. We conduct experiments on language model pre-training and several benchmark question answering datasets with various paragraph lengths. BlockBERT uses 18.7-36.1% less memory and 12.0-25.1% less time to learn the model. During testing, BlockBERT saves 27.8% inference time, while having comparable and sometimes better prediction accuracy, compared to an advanced BERT-based model, RoBERTa.

研究の動機と目的

• 長文シーケンス上でBERTを学習・推論する際の高いメモリと計算コストを解決すること。これは、モデルのスケーラビリティと展開性を制限する要因である。
• 自己注意機構におけるドット積注意が、シーケンス長に対して2次関数的複雑性を示すため、BERTにおける主なメモリボトルネックであることを同定すること。
• 長期間依存関係を扱うタスクにおいて、強力な性能を維持しながら、顕著にメモリと推論時間を削減する軽量なBERT変種を設計すること。
• モデルの正確性を損なわず、構造的スパarsityを用いた注意行列を通じて、長距離依存関係の効率的モデリングを可能にすること。
• カスタムカーネルや限定的ハードウェア支援に依存する複雑なスパース注意メカニズムの代替として、シンプルで実装可能な代替案を提供すること。

提案手法

• シーケンスを固定サイズのブロックに分割し、自己注意をブロック内および選択されたブロック間接続に制限することで、ブロックスパース注意を導入する。
• BERTにおける完全で密な注意行列を、スパarsityパターンに比例してFLOPsとメモリ使用量を削減できるブロックスパース行列構造に置き換える。
• マルチヘッド注意を維持するが、各ヘッドが事前に定義されたブロック内またはブロック間でのみ注目できるように制限することで、局所的および長距離のモデリングを両立させる。
• 標準的なディープラーニングフレームワークと互換性があるシンプルで微分可能なスパarsityパターンを採用し、カスタムCUDAカーネルを回避する。
• 元のBERTアーキテクチャと事前学習目的を維持し、効率性向上のための注目機構の変更に焦点を当てる。
• すべてのTransformer層にブロックスパース注意を適用することで、既存の微調整パイプラインとのエンドツーエンドの互換性を確保する。

実験結果

リサーチクエスチョン

• RQ1 構造的スパarsityを注意行列に適用することで、長期間依存関係タスクにおける性能を劣化させることなく、BERTのメモリと計算コストを削減できるか？
• RQ2 特に長文処理において、ブロックスパース注意は密な注意と比較して、下流タスクの精度にどのように影響するか？
• RQ3 ブロックスパース注意は、長距離依存関係のモデル容量を維持しつつ、学習および推論時間をどの程度削減できるか？
• RQ4 ブロックスパース設計は、意図したように、短距離および長距離依存関係の両方を効果的に捉えることができるか？
• RQ5 特に実世界の展開シナリオにおいて、BlockBERTの効率性向上は、RoBERTaなどの最先端BERT変種と比較してどの程度顕著か？

主な発見

• BlockBERTは、RoBERTaと比較して、学習中におけるメモリ消費量を18.7–36.1%削減しており、主に注意行列のスパース化に起因する。
• FLOPsの削減と低いメモリ圧力により、より大きなバッチサイズが可能になったことから、学習時間が12.0–25.1%短縮された。
• 推論時間はRoBERTaと比較して27.8%短縮され、リアルタイムまたはリソース制限のあるアプリケーションに適している。
• さまざまな段落長さの質問応答ベンチマークにおいて、BlockBERTはRoBERTaと同等またはより高い性能を達成している。
• ブロックスパース注意メカニズムは、注意ヘッドが局所的および長距離依存関係を効果的にモデリングできることを成功裏に示している。
• この手法はシンプルで、標準的なディープラーニングフレームワークと互換性があり、カスタムカーネルや複雑なスパarsity管理の必要がない。