自然言語処理で2倍のトレーニングスピードアップのためのPackedBERTの紹介
より効率的なトレーニングのために、この新しい BERT パッキング アルゴリズムを確認してください。
By マリオ・マイケル・クレル博士、Graphcore の主任機械学習リード、および マテイ・コセッチ、Graphcore の AI アプリケーション スペシャリスト
著者による画像。
新しいパッキング アルゴリズムを使用することで、BERT-Large のトレーニング中に自然言語処理を 2 倍以上高速化しました。 新しいパッキング技術によりパディングが削除され、大幅に効率的な計算が可能になります。
これは、ゲノミクスやタンパク質の折り畳みモデル、その他の歪んだ長さ分布を持つモデルにも適用でき、さまざまな業界や用途にさらに広範な影響を与えることができるのではないかと考えています。
私たちは、Graphcore の高効率な非負最小二乗ヒストグラム パッキング アルゴリズム (NNLSHP) と、パックされたシーケンスに適用される BERT アルゴリズムを新しい論文 [1] で紹介しました。
シーケンスパディングによる NLP における計算の無駄
私たちは、最近の取り組みに取り組みながら、BERT トレーニングを最適化する新しい方法の調査を開始しました。 MLPerf™ へのベンチマークの提出。 目標は、現実世界のアプリケーションに簡単に採用できる便利な最適化を開発することでした。 BERT は業界や多くの顧客によって広く使用されているため、これらの最適化に焦点を当てるモデルの XNUMX つとして自然な選択でした。
Wikipedia データセットを使用した独自の BERT-Large トレーニング アプリケーションでは、データセット内のトークンの 50% がパディングされており、その結果、多くのコンピューティングが無駄になっていたことを知って、私たちは本当に驚きました。
シーケンスをパディングしてすべてを同じ長さに揃えるのは、GPU で使用される一般的なアプローチですが、別のアプローチを試してみる価値はあると考えました。
シーケンスの長さは次の XNUMX つの理由で大きく異なります。
- 基礎となる Wikipedia データは、文書の長さに大きなばらつきがあることを示しています
- BERT の前処理自体は、トレーニング シーケンスを生成するために結合される抽出されたドキュメントのサイズをランダムに縮小します。
長さを最大長の 512 まで埋めると、すべてのトークンの 50% がパディング トークンになります。 パディングの 50% を実際のデータに置き換えると、同じ計算量で 50% 多くのデータが処理されるため、最適な条件下では 2 倍の速度が向上します。
図 1: Wikipedia データセットの分布。 著者による画像。
これはウィキペディアに特有のものですか? いいえ。
では、それは言語に特有のものなのでしょうか？ いいえ。
実際、歪んだ長さの分布は、言語、ゲノミクス、タンパク質の折り畳みなど、あらゆる場所で見られます。 図 2 と 3 は、SQuAD 1.1 データセットと GLUE データセットの分布を示しています。
図 2: 最大シーケンス長 1.1 の SQuAD 384 BERT 事前トレーニング データセットシーケンス長ヒストグラム。著者による画像。
図 3: 最大シーケンス長 128 の GLUE データセットシーケンス長ヒストグラム。著者による画像。
計算の無駄を避けながら、さまざまな長さを処理するにはどうすればよいでしょうか?
現在のアプローチでは、長さに応じて異なる計算カーネルを使用するか、エンジニアがプログラムでパディングを削除し、アテンション ブロックと損失の計算ごとにパディングを繰り返し追加し直す必要があります。 コードを大幅に増やしてより複雑にすることで計算量を節約することは魅力的ではなかったので、より良い方法を探しました。 複数のシーケンスを最大長のパックにまとめてまとめて処理することはできないでしょうか? 実は、それができるのです！
このアプローチには、次の XNUMX つの重要な要素が必要です。
- 残りのパディングをできるだけ少なくするためにどのサンプルをまとめるかを決定する効率的なアルゴリズム
- シーケンスではなくパックを処理するように BERT モデルを調整する
- そしてハイパーパラメータを調整する
梱包
当初は、Wikipedia のような大規模なデータセットを非常に効率的にパックできるとは考えられませんでした。 この問題は一般にビンパッキングとして知られています。 パッキングが XNUMX シーケンス以下に制限されている場合でも、結果として生じる問題は依然として強力な NP 完全性を持ち、効率的なアルゴリズムによる解決策が欠けています。 既存のヒューリスティック パッキング アルゴリズムは、少なくとも次のような複雑さがあるため、有望ではありませんでした。 O(ログ(n））、 どこ n はシーケンスの数です (Wikipedia の場合は約 16M)。 私たちは、数百万のシーケンスにうまく拡張できるアプローチに興味を持っていました。
次の XNUMX つのトリックにより、複雑さを大幅に軽減できました。
- パック内のシーケンス数を XNUMX に制限する (最初のソリューション アプローチの場合)
- 出現する長さごとに XNUMX つのビンを使用してシーケンス長のヒストグラムのみを操作する
最大シーケンス長は 512 でした。そのため、ヒストグラムに移行することで、次元と複雑さが 16 万シーケンスから 512 個の長さまで削減されました。 パック内で最大 22 つのシーケンスを許可すると、許可される長さの組み合わせの数が 4K に減りました。 これには、パック内のシーケンスを長さによってソートすることを要求するトリックがすでに含まれています。 では、22つのシーケンスを試してみてはいかがでしょうか? これにより、組み合わせの数が 940K から 3K に増加しましたが、これは最初のモデリング アプローチには多すぎました。 さらに、深さ XNUMX ではすでに驚くほど高い充填効率が達成されています。
当初、私たちは XNUMX つのパック内で XNUMX つを超えるシーケンスを使用すると、計算オーバーヘッドが増加し、トレーニング中の収束動作に影響を与えると考えていました。 ただし、さらに高速なリアルタイム パッキングを必要とする推論などのアプリケーションをサポートするために、非常に効率的な非負最小二乗ヒストグラム パッキング (NNLSHP) アルゴリズムを開発しました。
非負最小二乗ヒストグラム パッキング (NNLSHP)
ビン パッキングは、数学的な最適化問題として定式化されることがよくあります。 ただし、16 万 (またはそれ以上) のシーケンスでは、これは現実的ではありません。 問題の変数だけでも、ほとんどのマシンのメモリを超えてしまいます。 ヒストグラムベースのアプローチの数学的プログラムは非常に優れています。 簡単にするために、最小二乗法アプローチを採用することにしました(斧=b) ヒストグラムベクトル付き b。 戦略ベクトルをリクエストすることでそれを拡張しました x 負でないことを確認し、軽度のパディングを可能にするために重みを追加します。
難しい部分は戦略マトリックスでした。 各列は最大合計 512 を持ち、希望する全長に正確に一致するようにどのシーケンスをまとめてパックするかをエンコードします。 私たちの場合はXNUMXです。 行は、合計の長さに達する可能性のある組み合わせのそれぞれをエンコードします。 戦略のベクトル x これは私たちが探していたもので、20 通りの組み合わせのうちいずれかを選択する頻度を表します。 興味深いことに、最終的に選ばれた組み合わせはわずか約 600 個でした。 正確な解決策を得るには戦略が重要です x は正の整数である必要がありますが、非負の値だけで近似的に丸められた解が得られることに気付きました。 x 十分でした。 近似解の場合は、すぐに使えるシンプルなソルバーを使用して 30 秒以内に結果を得ることができます。
図 4: シーケンス長 8、パッキング深さ 3 の戦略行列の例。行は一緒にパックされる長さ 1 ～ 8 のシーケンスを表し、列は特定の順序を持たないパック内のすべての可能な長さの組み合わせを表します。 著者による画像。
最後に、戦略が割り当てられなかったいくつかのサンプルを修正する必要がありましたが、それらは最小限でした。 また、各シーケンスがパディングを使用してパックされることを強制し、パディングに応じた重み付けを行うバリアント ソルバーも開発しました。 はるかに長い時間がかかり、解決策もそれほど改善されませんでした。
最短パック優先ヒストグラム パッキング
NNLSHP は私たちに十分な梱包アプローチを提供してくれました。 しかし、私たちは理論的には、より高速なオンライン対応アプローチを実現し、3 つのシーケンスのみを組み合わせるという制限を取り除くことができないか考えていました。
したがって、既存のパッキング アルゴリズムからインスピレーションを得ましたが、それでもヒストグラムに焦点を当てました。
最初のアルゴリズムである Shortest-pack-first ヒストグラム パッキング (SPFHP) には XNUMX つの要素があります。
- 最長シーケンスから最短シーケンスまでのヒストグラムのカウントを操作します。
- 現在のシーケンスの長さがどのパックにも収まらない場合は、新しいパックのセットを開始します
- 複数の適合がある場合は、シーケンスの長さの合計が最も短いパックを取得し、それぞれのカウントを変更します
- 残りのカウントが一致しているかどうかを再度確認します
このアプローチは実装が最も簡単で、所要時間はわずか 0.02 秒でした。
変形例は、より完全に適合させるために、最短数と分割数の代わりにシーケンス長の最大合計を取得することでした。 全体として、これによって効率はあまり変わりませんでしたが、コードの複雑さは大幅に増加しました。
最短パック優先ヒストグラム パッキングの仕組み。 作者によるアニメーション。
Wikipedia、SQuAD 1.1、GLUE のパッキング結果
表 1、2、および 3 は、提案された XNUMX つのアルゴリズムのパッキング結果を示しています。 梱包深さ パックされたシーケンスの最大数を示します。 パッキング深さ 1 は、ベースライン BERT 実装です。 制限が設定されていない場合、発生する最大充填深さは追加の「最大」で示されます。 の パック数 は、新しいパックされたデータセットの長さを示します。 効率化 パックされたデータセット内の実際のトークンの割合です。 の パッキングファクター は、パッキング深さ 1 と比較して、結果として生じる潜在的な速度向上を示しています。
主な観察結果は次の XNUMX つでした。
- 分布の偏りが大きいほど、パッキングの利点が大きくなります。
- すべてのデータセットはパッキングの恩恵を受けます。 2倍を超えるものもあります。
- 充填深さが制限されていない場合、SPFHP はより効率的になります。
- 最大数が 3 つのパックされたシーケンスの場合、NNLSHP が複雑になればなるほど、効率が高くなります (99.75 対 89.44)。
表 1: Wikipedia で提案されているパッキング アルゴリズム (SPFHP および NNLSHP) の主なパフォーマンス結果。 著者による画像。
表 2: SQUaD 1.1 BERT 事前トレーニング用に提案されたパッキング アルゴリズムのパフォーマンス結果。 著者による画像。
表 3: GLUE データセットに対して提案されたパッキング アルゴリズムのパフォーマンス結果。 保圧深さを制限せずに、ベースラインと SPFHP 保圧結果のみが表示されます。 著者による画像。
BERT処理調整
BERT アーキテクチャの興味深い点は、ほとんどの処理がトークン レベルで行われることです。つまり、パッキングは妨げられません。 調整が必要なコンポーネントは、アテンション マスク、MLM 損失、NSP 損失、精度の XNUMX つだけです。
異なる数のシーケンスを処理する XNUMX つのアプローチすべてで重要なのは、ベクトル化と連結できるシーケンスの最大数を使用することでした。 注意点として、パディングに対処するためのマスクがすでに用意されています。 次の TensorFlow 疑似コードに見られるように、これを複数のシーケンスに拡張するのは簡単でした。 コンセプトは、注意が個別のシーケンスに限定され、それを超えて拡張できないようにするということです。
アテンションマスクのコードサンプル。
図 5: ゼロ XNUMX マスクの例
損失の計算では、原則としてシーケンスをアンパックして個別の損失を計算し、最終的に (パックではなく) シーケンス全体にわたる損失の平均を取得します。
MLM 損失の場合、コードは次のようになります。
損失計算のコードサンプル。
NSP 損失と精度についても、原理は同じです。 公開されているサンプルでは、社内のコードを使用してそれぞれのコードを見つけることができます。 PopART フレームワーク.
Wikipedia のオーバーヘッドと高速化の見積もり
BERT を修正したことで、次の XNUMX つの疑問が生じました。
- それに伴うオーバーヘッドはどれくらいですか?
- オーバーヘッドは、パックにまとめられるシーケンスの最大数にどの程度依存しますか?
BERT でのデータ準備は面倒な場合があるため、ショートカットを使用して、複数の異なる充填深さのコードをコンパイルし、それぞれの (測定された) サイクルを比較しました。 結果を表 4 に示します。 オーバーヘッド、パッキングを有効にするためのモデルへの変更 (注意のためのマスキング スキームや損失計算の変更など) によるスループットの減少率を示します。 の 高速化を実現 パッキングによる高速化の組み合わせです ( パッキングファクター)、およびそれに起因するスループットの低下 オーバーヘッド.
表 4: Wikipedia で提案されているパッキング アルゴリズム (SPFHP と NNLSHP) の推定高速化比較。 著者による画像。
ベクトル化技術のおかげで、オーバーヘッドは驚くほど小さく、多くのシーケンスを一緒にパックすることによる不都合はありません。
ハイパーパラメータの調整
梱包により、有効なバッチ サイズが (平均で) XNUMX 倍になります。 これは、トレーニングのハイパーパラメータを調整する必要があることを意味します。 簡単なトリックは、トレーニング前と同じ有効平均バッチ サイズを維持するために、勾配の累積数を半分に減らすことです。 事前トレーニングされたチェックポイントを備えたベンチマーク設定を使用すると、精度曲線が完全に一致していることがわかります。
図 6: パック処理とアンパック処理の学習曲線の比較 バッチサイズの縮小 パックされたアプローチに。 著者による画像。
精度は一致しています。MLM トレーニングの損失は、最初はわずかに異なる場合がありますが、すぐに追いつきます。 この最初の違いは、以前のトレーニングで短いシーケンスに偏っていた可能性のあるアテンション レイヤーのわずかな調整に起因する可能性があります。
速度の低下を回避するには、元のバッチ サイズを同じに保ち、増加した有効バッチ サイズ (XNUMX 倍) に合わせてハイパーパラメータを調整すると役立つ場合があります。 考慮すべき主なハイパーパラメータは、ベータ パラメータと学習率です。 一般的なアプローチの XNUMX つはバッチ サイズを XNUMX 倍にすることですが、この場合はパフォーマンスが低下します。 LAMB オプティマイザーの統計を見ると、ベータ パラメーターをパッキング係数で乗じることは、運動量と速度を同等に保つために複数のバッチを連続してトレーニングすることに相当することが証明できます。
図 7: パック処理とアンパック処理の学習曲線の比較 ヒューリスティック 適用済み。 著者による画像。
私たちの実験では、ベータの XNUMX 乗が優れたヒューリスティックであることがわかりました。 このシナリオでは、バッチ サイズを増やすと通常、目標精度に達するまでサンプル/エポックの意味での収束速度が低下するため、曲線が一致するとは予想されません。
ここで問題となるのは、実際のシナリオで、本当に予想通りの高速化が得られるのかということです。
図 8: パック処理とアンパック処理の学習曲線の比較 最適化されたセットアップ。 著者による画像。
はい、そうです！ データ転送を圧縮したため、さらに速度が向上しました。
まとめ
文をまとめて圧縮すると、計算量と環境を節約できます。 この手法は、PyTorch や TensorFlow を含む任意のフレームワークで実装できます。 私たちは明らかに 2 倍の速度向上を達成し、その過程でパッキング アルゴリズムの最先端を拡張しました。
私たちが興味を持っている他のアプリケーションは、同様のデータ分布が観察できるゲノミクスとタンパク質のフォールディングです。 ビジョン トランスフォーマーも、異なるサイズのパックされた画像を適用する興味深い分野になる可能性があります。 どのアプリケーションがうまく機能すると思いますか? ぜひご連絡ください。
ありがとうございました
この作業に貢献してくれた Graphcore のアプリケーション エンジニアリング チームの同僚である Sheng Fu と Mrinal Iyer に感謝します。また、貴重なフィードバックをくれた Graphcore のリサーチ チームの Douglas Orr にも感謝します。
参考文献
[1] M. コーセック、S. フー、MM クレル、 パッキング: NLP BERT の 2 倍の加速に向けて (2021)、arXiv
マリオ・マイケル・クレル博士 Graphcore の主任機械学習リードです。 マリオは 12 年以上にわたって機械学習アルゴリズムの研究と開発を行っており、ロボット工学、自動車、電気通信、ヘルスケアなどのさまざまな業界向けのソフトウェアを作成しています。 Graphcore では、彼は私たちの素晴らしい成果に貢献しました。 MLPerf の提出 また、統計的な COVID-19 データ分析のための近似ベイズ計算などの新しい非標準モデルを加速することに情熱を持っています。
マテイ・コセッチ パロアルトの Graphcore の AI アプリケーション スペシャリストです。 彼は以前、サンノゼの NIO で自動運転に関する AI 科学者として働いており、スタンフォード大学で航空学と宇宙航学の修士号を取得しています。
元の。 許可を得て転載。
関連する
PlatoAi。 Web3の再考。 増幅されたデータインテリジェンス。
アクセスするには、ここをクリックしてください。
