[論文レビュー] Packing: Towards 2x NLP BERT Acceleration

Matej Kosec, Sheng Fu+1 more|arXiv (Cornell University)|Jun 29, 2021

VLSI and FPGA Design Techniques参考文献 6被引用数 4

ひとこと要約

本稿では、BERT事前学習におけるパディングを排除するための決定的ヒストグラムベースのパッキングアルゴリズムを提案する。複数のシーケンスを1つのサンプルに統合することで、1秒あたりのシーケンス数を2倍に高速化した。SPFHPおよびNNLSHPという手法は、線形時間の効率性を実現するため、シーケンス長ヒストグラム上で動作し、モデル性能への影響を最小限に抑えつつ、ほぼ最適なパッキングを実現する。BERT-Largeの事前学習においても完全な収束が達成された。

ABSTRACT

We find that at sequence length 512 padding tokens represent in excess of 50% of the Wikipedia dataset used for pretraining BERT (Bidirectional Encoder Representations from Transformers). Therefore by removing all padding we achieve a 2x speed-up in terms of sequences/sec. To exploit this characteristic of the dataset, we develop and contrast two deterministic packing algorithms. Both algorithms rely on the assumption that sequences are interchangeable and therefore packing can be performed on the histogram of sequence lengths, rather than per sample. This transformation of the problem leads to algorithms which are fast and have linear complexity in dataset size. The shortest-pack-first histogram-packing (SPFHP) algorithm determines the packing order for the Wikipedia dataset of over 16M sequences in 0.02 seconds. The non-negative least-squares histogram-packing (NNLSHP) algorithm converges in 28.4 seconds but produces solutions which are more depth efficient, managing to get near optimal packing by combining a maximum of 3 sequences in one sample. Using the dataset with multiple sequences per sample requires additional masking in the attention layer and a modification of the MLM loss function. We demonstrate that both of these changes are straightforward to implement and have relatively little impact on the achievable performance gain on modern hardware. Finally, we pretrain BERT-Large using the packed dataset, demonstrating no loss of convergence and the desired 2x speed-up.

研究の動機と目的

• シーケンス長512でWikipediaデータセット上で50％を超える過剰なパディングが生じるBERT事前学習の非効率性を解消する。
• 個々のサンプルの処理を避けるために、シーケンス長ヒストグラム上で動作する、高速で決定的なパッキングアルゴリズムを開発する。
• モデルの収束性と性能を維持しつつ、ほぼ最適なパッキング効率を達成する。
• パッケージドシーケンス用に注意マスクと修正されたMLM損失が、現代のハードウェア上で軽量かつ効果的であることを実証する。

提案手法

• パッキング問題を個々のシーケンス処理からヒストグラムベースの最適化に変換することで、データセットサイズに比例する線形時間の複雑さに低減する。
• 短いシーケンスを優先する最短パック優先のヒストグラムパッキング（SPFHP）アルゴリズムを設計し、1600万シーケンスに対して0.02秒でパッキングを実現した。
• 深さ効率を最適化する非負の最小二乗ヒストグラムパッキング（NNLSHP）アルゴリズムを実装し、1パッケージドサンプルあたり最大3つのシーケンスを組み合わせることを可能にした。
• パッケージドサンプル内において、元の異なるシーケンスに属するトークン同士が注意計算を行わないように、自己注意層にシーケンス固有のマスクを導入した。
• パッキングによってマスクされる一部のトークンを考慮に入れることで、マスク言語モデル（MLM）損失関数を修正した。
• パッケージドデータセットを用いてBERT-Largeを事前学習し、性能が維持され、2倍の高速化が達成されることを検証した。

実験結果

リサーチクエスチョン

• RQ1 決定的ヒストグラムベースのパッキングが、収束性に影響を与えることなく、BERT事前学習における50％を超えるパディングを排除できるか？
• RQ2 個々のサンプル処理と比較して、ヒストグラムベースのパッキングアルゴリズムは実行時間およびパッキング深さの点でどの程度効率的か？
• RQ3 パッケージドBERT学習における注意マスクと修正されたMLM損失の性能への影響は何か？
• RQ4 下流タスクの性能に最小限の劣化で、1秒あたりのシーケンス数を2倍に高速化できるか？

主な発見

• パッキングにより、シーケンス長512における入力のパディング割合が50％未満に低下し、BERT事前学習における1秒あたりのシーケンス数が2倍に向上した。
• SPFHPアルゴリズムは、1600万シーケンスのWikipediaデータセットをわずか0.02秒でパッキングし、線形時間スケーラビリティを実証した。
• NNLSHPアルゴリズムは高いパッキング深さ効率を達成し、1パッケージドサンプルあたり最大3つのシーケンスを組み合わせることで、ほぼ最適なパッキングに近づいた。
• 修正された注意マスクとMLM損失関数は、モデル性能にほとんど影響を与えず、BERT-Largeの事前学習においても収束が保証された。
• パッケージドデータセットを用いたBERT-Largeの事前学習は、標準的な学習と同等の収束行動を示し、本手法の実用性を裏付けた。
• 本手法は現代のハードウェアでも実用的であり、両方のパッキングアルゴリズムと推論の修正は、実装が簡単である。